diff --git a/cmd/meshtender/main.go b/cmd/meshtender/main.go index b01fc13..69cf49d 100644 --- a/cmd/meshtender/main.go +++ b/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() diff --git a/internal/core/confirm.go b/internal/core/confirm.go index c012bee..f920486 100644 --- a/internal/core/confirm.go +++ b/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) diff --git a/internal/core/console.go b/internal/core/console.go index d7df1a5..17ec09b 100644 --- a/internal/core/console.go +++ b/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() diff --git a/internal/core/web.go b/internal/core/web.go index 50b1d1d..998632a 100644 --- a/internal/core/web.go +++ b/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 diff --git a/internal/core/ws_drain_test.go b/internal/core/ws_drain_test.go new file mode 100644 index 0000000..45339e5 --- /dev/null +++ b/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 + } +}