Drain WebSocket connections on shutdown

cmd/meshtender/main.go

@@ -149,6 +149,15 @@ func run(logger *slog.Logger) error {
 		cancel()
 	}
 
+	// Shutdown doesn't close hijacked/WebSocket connections, so close the active
+	// console/confirm sockets and give their handlers a moment to finish (they do
+	// DB work) before we stop the flusher and close the pool.
+	wsCtx, wsCancel := context.WithTimeout(context.Background(), 5*time.Second)
+	if !srv.DrainWebSockets(wsCtx) {
+		logger.Warn("shutdown: some WebSocket handlers did not finish before the deadline")
+	}
+	wsCancel()
+
 	// Now stop the flusher and wait for its final flush to complete before the
 	// deferred st.Close() closes the pool underneath it.
 	stopAnalytics()

internal/core/confirm.go

@@ -97,14 +97,20 @@ func (s *Handlers) wsConfirm(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Track the socket so shutdown can drain it (http.Server.Shutdown doesn't close
+	// hijacked/WebSocket conns). Add before Accept so a shutdown racing the upgrade
+	// still waits for this handler.
+	s.wsWG.Add(1)
+	defer s.wsWG.Done()
+
 	ws, err := websocket.Accept(w, r, nil) // same-origin (request host) authorized by default
 	if err != nil {
 		return
 	}
 
-	// A connection-lifetime context, independent of the request context which
-	// is unsafe to use after Accept.
-	ctx, cancel := context.WithTimeout(context.Background(), confirmTimeout)
+	// A connection-lifetime context derived from the server's WS context (cancelled
+	// on shutdown), not the request context which is unsafe to use after Accept.
+	ctx, cancel := context.WithTimeout(s.wsCtx, confirmTimeout)
 	defer cancel()
 
 	bridge := wsbridge.New(ctx, ws)

internal/core/console.go

@@ -190,11 +190,18 @@ func (s *Handlers) wsConsole(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Track the socket so shutdown can drain it (http.Server.Shutdown doesn't close
+	// hijacked/WebSocket conns). Add before Accept so a shutdown racing the upgrade
+	// still waits for this handler.
+	s.wsWG.Add(1)
+	defer s.wsWG.Done()
+
 	ws, err := websocket.Accept(w, r, nil)
 	if err != nil {
 		return
 	}
-	ctx, cancel := context.WithCancel(context.Background())
+	// Derive from the server's WS context so shutdown cancels the session.
+	ctx, cancel := context.WithCancel(s.wsCtx)
 	defer cancel()
 	idle := time.AfterFunc(consoleIdleTimeout, cancel)
 	defer idle.Stop()

internal/core/web.go

@@ -8,6 +8,7 @@ import (
 	"embed"
 	"net"
 	"net/http"
+	"sync"
 
 	"github.com/go-chi/chi/v5"
 
@@ -27,14 +28,44 @@ var templatesFS embed.FS
 type Handlers struct {
 	*web.Env
 	Auth *auth.Service
+
+	// wsCtx is the base context for WebSocket handlers; cancelling it (on shutdown)
+	// signals active console/confirm sockets to close. wsWG tracks those handlers so
+	// shutdown can wait for them — http.Server.Shutdown does not close hijacked
+	// (WebSocket) connections.
+	wsCtx    context.Context
+	wsCancel context.CancelFunc
+	wsWG     sync.WaitGroup
 }
 
 // Server is the built HTTP entry point.
-type Server struct{ handler http.Handler }
+type Server struct {
+	handler http.Handler
+	app     *Handlers
+}
 
 // Handler returns the root HTTP handler.
 func (s *Server) Handler() http.Handler { return s.handler }
 
+// DrainWebSockets closes active WebSocket connections and waits for their handlers
+// to return, up to ctx's deadline. Call it after http.Server.Shutdown, which does
+// not close hijacked/WebSocket connections. Reports whether all handlers finished
+// before the deadline.
+func (s *Server) DrainWebSockets(ctx context.Context) bool {
+	s.app.wsCancel()
+	done := make(chan struct{})
+	go func() {
+		s.app.wsWG.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}
+
 // NewServer assembles every surface and the host dispatcher. The deployment is
 // always split-host (auth / root / app); the auth host serves sign-in + account,
 // the root host serves marketing + public org discovery, and the app host (plus
@@ -63,6 +94,7 @@ func NewServer(st *store.Store, authSvc *auth.Service, idSvc *identity.Service,
 		return nil, err
 	}
 	app := &Handlers{Env: coreEnv, Auth: authSvc}
+	app.wsCtx, app.wsCancel = context.WithCancel(context.Background())
 
 	authH, err := auth.NewWeb(deps, authSvc)
 	if err != nil {
@@ -78,7 +110,7 @@ func NewServer(st *store.Store, authSvc *auth.Service, idSvc *identity.Service,
 	// public org pages).
 	appHandler := mkH.CustomDomain(app.appRouter())
 	handler := web.Dispatcher(cfg, authH.Routes(), mkH.Routes(), appHandler)
-	return &Server{handler: handler}, nil
+	return &Server{handler: handler, app: app}, nil
 }
 
 // sessionMW loads and validates the SSO session (two DB touches). It's applied to

internal/core/ws_drain_test.go

@@ -0,0 +1,111 @@
+package core
+
+import (
+	"context"
+	"crypto/rand"
+	"errors"
+	"net/http"
+	"net/http/cookiejar"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+	meshcore "github.com/meshcore-go/meshcore-go"
+
+	"github.com/jleight/meshtender/internal/auth"
+	"github.com/jleight/meshtender/internal/identity"
+	"github.com/jleight/meshtender/internal/store"
+)
+
+// TestDrainWebSockets: an active console WebSocket is closed and its handler
+// returns when the server drains on shutdown — http.Server.Shutdown alone leaves
+// hijacked sockets running, which is what this covers.
+func TestDrainWebSockets(t *testing.T) {
+	t.Parallel()
+	st, ctx := coreStore(t)
+
+	var masterKey [32]byte
+	_, _ = rand.Read(masterKey[:])
+	idSvc, err := identity.LoadOrCreate(ctx, st, masterKey)
+	if err != nil {
+		t.Fatalf("identity: %v", err)
+	}
+	authSvc, err := auth.New(st, st.Pool(), testAuthConfig())
+	if err != nil {
+		t.Fatalf("auth: %v", err)
+	}
+	srv, err := NewServer(st, authSvc, idSvc, testConfig())
+	if err != nil {
+		t.Fatalf("server: %v", err)
+	}
+	ts := httptest.NewServer(srv.Handler())
+	defer ts.Close()
+
+	jar, _ := cookiejar.New(nil)
+	user := seedSession(t, ts, st, ctx, jar, "wsuser")
+	repeater, err := meshcore.GenerateLocalIdentity(rand.Reader)
+	if err != nil {
+		t.Fatalf("repeater identity: %v", err)
+	}
+	rep, err := st.CreateRepeater(ctx, &store.Repeater{
+		OwnerID: user.ID, Name: "Test", PublicKeyHex: repeater.String(),
+		RadioFreqHz: 869525000, RadioBwHz: 250000, RadioSF: 11, RadioCR: 5,
+	})
+	if err != nil {
+		t.Fatalf("create repeater: %v", err)
+	}
+
+	wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/repeaters/" + rep.PublicID + "/console/ws"
+	hdr := http.Header{}
+	if cs := jar.Cookies(mustURL(t, ts.URL)); len(cs) > 0 {
+		var parts []string
+		for _, c := range cs {
+			parts = append(parts, c.Name+"="+c.Value)
+		}
+		hdr.Set("Cookie", strings.Join(parts, "; "))
+	}
+	dctx, dcancel := context.WithTimeout(ctx, 5*time.Second)
+	defer dcancel()
+	ws, _, err := websocket.Dial(dctx, wsURL, &websocket.DialOptions{HTTPHeader: hdr})
+	if err != nil {
+		t.Fatalf("ws dial: %v", err)
+	}
+	defer ws.Close(websocket.StatusInternalError, "")
+
+	// Send ready and read one status frame so the handler is definitely past the
+	// upgrade/connect and running its session loop before we drain.
+	rw, rwcancel := context.WithTimeout(ctx, 10*time.Second)
+	defer rwcancel()
+	if err := ws.Write(rw, websocket.MessageText, []byte(`{"type":"ready"}`)); err != nil {
+		t.Fatalf("write ready: %v", err)
+	}
+	if _, _, err := ws.Read(rw); err != nil {
+		t.Fatalf("read first frame: %v", err)
+	}
+
+	// Drain: cancels the WS context and waits for the handler. It must finish well
+	// within the deadline (the handler is context-aware).
+	drainCtx, dcancel2 := context.WithTimeout(context.Background(), 5*time.Second)
+	defer dcancel2()
+	if !srv.DrainWebSockets(drainCtx) {
+		t.Fatal("DrainWebSockets timed out — a WebSocket handler did not exit on shutdown")
+	}
+
+	// The server should have closed the socket. Read past any buffered status frames
+	// until a read fails; a close (clean status or EOF) is expected — a read
+	// deadline would mean the socket was left open.
+	readCtx, rcancel := context.WithTimeout(ctx, 3*time.Second)
+	defer rcancel()
+	for {
+		_, _, err := ws.Read(readCtx)
+		if err == nil {
+			continue // a buffered status frame; keep reading toward the close
+		}
+		if errors.Is(err, context.DeadlineExceeded) {
+			t.Fatal("socket still open after drain (read deadline hit)")
+		}
+		break // socket closed
+	}
+}