fix: serialize AutoSession creation with Session.Close

README.md

@@ -667,12 +667,12 @@ complete contract.
 
 ### Session handling
 
-JaWS has non-persistent session handling integrated. Sessions won't 
+JaWS has non-persistent session handling integrated. Sessions won't
 be persisted across restarts and must have an expiry time.
 
 Use one of these patterns:
 
-* Wrap page handlers with `Jaws.SessionMiddleware(handler)` to ensure a session exists.
+* Wrap page handlers with `Jaws.SessionMiddleware(handler)` to create a session when none exists.
 * Call `Jaws.NewSession(w, r)` explicitly to create and attach a fresh session cookie.
 * Set `Jaws.AutoSession` to lazily create an anonymous session during a
   successful WebSocket upgrade when a Request has none.

jaws.go

@@ -100,7 +100,7 @@ type Jid = jid.Jid // convenience alias
 // concurrency behavior and may be called concurrently when stated.
 type Jaws struct {
 	CookieName              string          // Name for session cookies; defaults to a name derived from the executable ([assets.DefaultCookieName]), falling back to "jaws"
-	AutoSession             bool            // Create a session during a successful WebSocket upgrade when a Request has none. Defaults to false.
+	AutoSession             bool            // Create and associate a session during a successful WebSocket upgrade when a Request has none. Defaults to false.
 	TrustForwardedHeaders   bool            // Trust X-Forwarded-* headers: governs the session cookie Secure flag (X-Forwarded-Proto) and the client IP used for session/request binding (X-Forwarded-For/X-Real-IP). Defaults to false; only enable behind a single reverse proxy you control that sets these headers.
 	Logger                  Logger          // Optional logger to use
 	Debug                   bool            // Set to true to enable debug info in generated HTML code. Call GenerateHeadHTML after changing it.

request.go

@@ -364,17 +364,32 @@ func (rq *Request) sessionDestKey(sess *Session) (k key.Key) {
 }
 
 func (rq *Request) ensureAutoSession(w http.ResponseWriter, r *http.Request) {
-	if rq.Jaws.AutoSession && rq.Session() == nil {
-		sess := rq.Jaws.newSession(w, r)
-		rq.mu.Lock()
-		if rq.session == nil {
-			rq.session = sess
-			sess.addRequest(rq)
+	if rq.Jaws.AutoSession {
+		if sess := rq.newAutoSession(r); sess != nil {
+			sess.addCookie(w, r)
 		}
-		rq.mu.Unlock()
 	}
 }
 
+// newAutoSession creates and associates an AutoSession before making it visible
+// through Jaws session lookups.
+func (rq *Request) newAutoSession(r *http.Request) (sess *Session) {
+	jw := rq.Jaws
+	secure := secureheaders.RequestIsSecure(r, jw.TrustForwardedHeaders)
+	remoteIP := jw.clientIP(r)
+	jw.mu.Lock()
+	defer jw.mu.Unlock()
+	rq.mu.Lock()
+	defer rq.mu.Unlock()
+	if rq.session == nil {
+		sess = jw.newSessionLocked(remoteIP, secure)
+		sess.addRequest(rq)
+		rq.session = sess
+		jw.sessions[sess.sessionID] = sess
+	}
+	return
+}
+
 // releaseBuffersLocked detaches the reusable storage from a finished Request and
 // returns it for the caller to return to [Jaws.requestBufferPool].
 //
@@ -1045,9 +1060,9 @@ func normalizedWebSocketAcceptRequest(r *http.Request) (normalized *http.Request
 // sets a session cookie.
 //
 // Accept writes the 101 through the [http.ResponseWriter] it was given before
-// hijacking the connection, so the Set-Cookie header is in the header map in
-// time. Should the hijack itself fail after the 101, the session already
-// exists and expires through the normal session grace period.
+// hijacking the connection, so a live session cookie can enter the header map
+// in time. A session that survives this commit but whose hijack fails expires
+// through the normal session grace period.
 type autoSessionWriter struct {
 	http.ResponseWriter
 	rq *Request

request_test.go

@@ -3189,6 +3189,120 @@ func TestWS_AutoSessionCreatesSession(t *testing.T) {
 	}
 }
 
+type autoSessionCloseResponseWriter struct {
+	http.ResponseWriter
+	jw     *Jaws
+	closed chan<- *Session
+	once   sync.Once
+	t      *testing.T
+}
+
+func (w *autoSessionCloseResponseWriter) Header() http.Header {
+	if sessions := w.jw.Sessions(); len(sessions) > 0 {
+		w.once.Do(func() {
+			done := make(chan struct{})
+			go func(sess *Session) {
+				sess.Close()
+				w.closed <- sess
+				close(done)
+			}(sessions[0])
+			select {
+			case <-done:
+			case <-time.After(testTimeout):
+				w.t.Error("Session.Close blocked while ResponseWriter.Header re-entered Jaws")
+			}
+		})
+	}
+	return w.ResponseWriter.Header()
+}
+
+func (w *autoSessionCloseResponseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+
+func TestWS_AutoSessionCloseAtPublication(t *testing.T) {
+	ts := newTestServerNoSession(t)
+	defer ts.Close()
+	ts.jw.AutoSession = true
+
+	closedCh := make(chan *Session, 1)
+	ts.srv.Close()
+	ts.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ts.ServeHTTP(&autoSessionCloseResponseWriter{
+			ResponseWriter: w,
+			jw:             ts.jw,
+			closed:         closedCh,
+			t:              t,
+		}, r)
+	}))
+	ts.setInitialRequestOrigin()
+
+	connectSessionCh := make(chan *Session, 1)
+	ts.rq.SetConnectFn(func(rq *Request) error {
+		connectSessionCh <- rq.Session()
+		return nil
+	})
+
+	dialCtx, cancelDial := context.WithTimeout(t.Context(), testTimeout*2)
+	defer cancelDial()
+	hdr := http.Header{}
+	hdr.Set("Origin", ts.origin())
+	conn, resp, err := websocket.Dial(dialCtx, ts.Url(), &websocket.DialOptions{HTTPHeader: hdr})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = conn.CloseNow() }()
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		t.Fatalf("WebSocket status = %d, want %d", resp.StatusCode, http.StatusSwitchingProtocols)
+	}
+
+	var closedSession *Session
+	select {
+	case closedSession = <-closedCh:
+	case <-time.After(testTimeout):
+		t.Fatal("ResponseWriter.Header did not close the published AutoSession")
+	}
+
+	connectSession := waitForConnectSession(t, connectSessionCh)
+
+	const marker = "post-connect marker"
+	ts.jw.Broadcast(wire.Message{Dest: ts.rq.JawsKey, What: what.Alert, Data: marker})
+	readCtx, cancelRead := context.WithTimeout(t.Context(), testTimeout)
+	defer cancelRead()
+	var messages strings.Builder
+	for !strings.Contains(messages.String(), marker) {
+		messageType, data, err := conn.Read(readCtx)
+		if err != nil {
+			t.Fatalf("reading WebSocket messages: %v (got %q)", err, messages.String())
+		}
+		if messageType != websocket.MessageText {
+			t.Fatalf("WebSocket message type = %v, want text", messageType)
+		}
+		messages.Write(data)
+	}
+
+	if connectSession != nil {
+		t.Errorf("ConnectFn Session() = %v, want nil after Session.Close", connectSession)
+	}
+	if got := ts.rq.Session(); got != nil {
+		t.Errorf("Request Session() = %v, want nil after Session.Close", got)
+	}
+	if requests := closedSession.Requests(); len(requests) != 0 {
+		t.Errorf("closed Session Requests() = %v, want none", requests)
+	}
+	if got := ts.jw.SessionCount(); got != 0 {
+		t.Errorf("SessionCount() = %d, want 0", got)
+	}
+	for _, cookie := range resp.Cookies() {
+		if cookie.Name == ts.jw.CookieName && cookie.MaxAge >= 0 {
+			t.Errorf("response contains live session cookie: %v", cookie)
+		}
+	}
+	if got := strings.Count(messages.String(), what.Reload.String()+"\t"); got != 1 {
+		t.Errorf("got %d Reload commands, want exactly 1: %q", got, messages.String())
+	}
+}
+
 func TestWS_AutoSessionKeepsExistingSession(t *testing.T) {
 	ts := newTestServer(t)
 	defer ts.Close()

requestpool.go

@@ -50,8 +50,8 @@ import (
 // while the retired Request is reachable; no deadline is guaranteed for later key
 // reuse.
 //
-// NewRequest panics if the system CSPRNG ([crypto/rand]) fails while generating
-// the request key, which does not happen on supported platforms.
+// It panics if the [crypto/rand.Reader] captured by [New] returns an error while
+// generating the request key. Go's default reader does not return errors.
 func (jw *Jaws) NewRequest(r *http.Request) (rq *Request) {
 	remoteIP := jw.clientIP(r)
 

session.go

@@ -175,6 +175,35 @@ func (sess *Session) Cookie() (cookie *http.Cookie) {
 	return
 }
 
+// addCookie adds sess's cookie to w and r while sess is current and live.
+func (sess *Session) addCookie(w http.ResponseWriter, r *http.Request) {
+	var h http.Header
+	if w != nil {
+		// ResponseWriter.Header is caller code and may re-enter Jaws, including
+		// by closing sess, so call it before taking either core lock below.
+		h = w.Header()
+	}
+	jw := sess.jw
+	jw.mu.RLock()
+	defer jw.mu.RUnlock()
+	sess.mu.RLock()
+	defer sess.mu.RUnlock()
+	// Close unregisters before marking dead, while an unattached Session can
+	// expire during Header; require both registry identity and liveness.
+	if jw.sessions[sess.sessionID] == sess && !sess.isDeadLocked() {
+		cookie := sess.cookie
+		if h != nil {
+			if v := cookie.String(); v != "" {
+				// Header.Add and Request.AddCookie mutate concrete header maps
+				// without invoking caller code. Keep both under these read locks so
+				// cookie publication precedes a losing Session.Close.
+				h.Add("Set-Cookie", v)
+			}
+		}
+		r.AddCookie(&cookie)
+	}
+}
+
 // Close invalidates and expires the [Session].
 // Future [Request] values won't be able to associate with it, and [Session.Cookie] will return a deletion cookie.
 //
@@ -269,15 +298,18 @@ func (sess *Session) Broadcast(msg wire.Message) {
 	}
 }
 
-// SessionCount returns the number of active sessions.
+// SessionCount returns the number of registered sessions.
 func (jw *Jaws) SessionCount() (n int) {
 	jw.mu.RLock()
 	n = len(jw.sessions)
 	jw.mu.RUnlock()
 	return
 }
 
-// Sessions returns a list of all active sessions, which may be nil.
+// Sessions returns a snapshot of all registered sessions, which may be nil.
+//
+// Auto-created [Session] values are registered only after their initiating
+// [Request] is associated.
 func (jw *Jaws) Sessions() (sessions []*Session) {
 	jw.mu.RLock()
 	if n := len(jw.sessions); n > 0 {
@@ -362,12 +394,16 @@ func (jw *Jaws) GetSession(r *http.Request) (sess *Session) {
 // match used everywhere else; see [Jaws.GetSession] and [Jaws.TrustForwardedHeaders]
 // for the reverse-proxy caveat.
 //
-// As a side effect, the session cookie is also added to r itself, so the new
-// [Session] is visible to [Jaws.GetSession] and [Jaws.NewRequest] for the
-// remainder of the same HTTP request.
+// If the new [Session] remains current and live during cookie publication, its
+// cookie is written to w when w is non-nil and added to r itself. This makes the
+// [Session] visible to [Jaws.GetSession] and [Jaws.NewRequest] for the remainder
+// of the same HTTP request. If a concurrent [Session.Close] wins first, neither
+// w nor r receives its live cookie.
 //
-// It panics if the system CSPRNG ([crypto/rand]) fails while generating the session
-// ID, which does not happen on supported platforms.
+// It returns nil and has no effect if r is nil; w may be nil.
+//
+// It panics if the [crypto/rand.Reader] captured by [New] returns an error while
+// generating the session ID. Go's default reader does not return errors.
 func (jw *Jaws) NewSession(w http.ResponseWriter, r *http.Request) (sess *Session) {
 	if r != nil {
 		if sessionIDs := getCookieSessionsIDs(r.Header, jw.CookieName); len(sessionIDs) > 0 {
@@ -389,26 +425,27 @@ func (jw *Jaws) NewSession(w http.ResponseWriter, r *http.Request) (sess *Sessio
 
 func (jw *Jaws) newSession(w http.ResponseWriter, r *http.Request) (sess *Session) {
 	secure := secureheaders.RequestIsSecure(r, jw.TrustForwardedHeaders)
-	var cookie http.Cookie
+	remoteIP := jw.clientIP(r)
 	func() {
 		jw.mu.Lock()
 		defer jw.mu.Unlock()
-		for sess == nil {
-			sessionID := jw.nonZeroRandomLocked()
-			if _, ok := jw.sessions[sessionID]; !ok {
-				sess = newSession(jw, sessionID, jw.clientIP(r), secure)
-				jw.sessions[sessionID] = sess
-				cookie = sess.cookie
-			}
-		}
+		sess = jw.newSessionLocked(remoteIP, secure)
+		jw.sessions[sess.sessionID] = sess
 	}()
+	sess.addCookie(w, r)
+	return
+}
 
-	// http.SetCookie calls the caller-provided ResponseWriter.Header, which may
-	// re-enter Jaws, so emit the cookie only after releasing jw.mu.
-	if w != nil {
-		http.SetCookie(w, &cookie)
+// newSessionLocked allocates a Session whose ID is absent from jw.sessions.
+//
+// The caller must hold jw.mu and publish the Session before releasing it.
+func (jw *Jaws) newSessionLocked(remoteIP netip.Addr, secure bool) (sess *Session) {
+	for sess == nil {
+		sessionID := jw.nonZeroRandomLocked()
+		if _, ok := jw.sessions[sessionID]; !ok {
+			sess = newSession(jw, sessionID, remoteIP, secure)
+		}
 	}
-	r.AddCookie(&cookie)
 	return
 }
 
@@ -430,10 +467,13 @@ func (sess sessioner) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	sess.h.ServeHTTP(w, r)
 }
 
-// SessionMiddleware returns an [http.Handler] that ensures a JaWS [Session]
-// exists before invoking h, creating one if the request has none.
+// SessionMiddleware returns a session-creating [http.Handler].
+//
+// Before invoking h, it creates a JaWS [Session] when the request has none. If
+// a concurrent [Session.Close] wins the new Session's cookie publication, h
+// runs without that Session or its live cookie.
 //
-// It is the session-ensuring middleware, distinct from the session accessors:
+// It is distinct from the session accessors:
 // [Jaws.GetSession] and [Request.Session] look up an existing [Session], while
 // this wraps a handler. It composes with [Jaws.SecureHeadersMiddleware].
 func (jw *Jaws) SessionMiddleware(h http.Handler) http.Handler {