Prevent shared caching of session cookie responses

AI.md

@@ -278,6 +278,10 @@ creation patterns:
 * Enable `Jaws.AutoSession` to create an anonymous session during a successful
   WebSocket upgrade when the Request has none.
 
+`SessionMiddleware` sends `Cache-Control: no-store` when it creates a Session;
+other responses keep the handler's cache policy. Callers using `NewSession`
+directly must prevent shared caching when they send its cookie.
+
 Create or retrieve the session before `NewRequest` when initial rendering or
 authentication depends on it. Later Requests with the same valid cookie and IP
 can access the Session. `Request.Get` returns nil and `Request.Set` is a no-op
@@ -351,8 +355,8 @@ entries and the final drain.
 
 Register `Jaws.ServeHTTP` for the `/jaws/` prefix. It owns these routes:
 
-* `/jaws/.jaws.<hash>.css` -- built-in stylesheet; cache indefinitely.
-* `/jaws/.jaws.<hash>.js` -- built-in client; cache indefinitely.
+* `/jaws/.jaws.<hash>.css` -- built-in stylesheet; publicly cacheable by default.
+* `/jaws/.jaws.<hash>.js` -- built-in client; publicly cacheable by default.
 * `/jaws/<key>` and `/jaws/<key>/noscript` -- single-use Request callback. The
   key must parse to a nonzero value through `key.Parse`; parsing is
   case-insensitive, while generated URLs use canonical lowercase base 32. A

session.go

@@ -1,6 +1,8 @@
 package jaws
 
 import (
+	"bufio"
+	"net"
 	"net/http"
 	"net/netip"
 	"net/textproto"
@@ -396,6 +398,10 @@ func (jw *Jaws) GetSession(r *http.Request) (sess *Session) {
 // of the same HTTP request. If a concurrent [Session.Close] wins first, neither
 // w nor r receives its live cookie.
 //
+// Callers that send the cookie must prevent shared caching of the response.
+// [Jaws.SessionMiddleware] applies Cache-Control: no-store when it creates the
+// Session itself.
+//
 // It returns nil without closing matching Sessions if r is nil, shutdown has
 // begun, [Jaws.MaxSessions] or [Jaws.MaxSessionsPerIP] is reached, or cookie
 // publication fails; w may be nil.
@@ -534,6 +540,55 @@ type sessioner struct {
 	h  http.Handler
 }
 
+// sessionCacheWriter protects a newly created Session's cookie at header commit.
+type sessionCacheWriter struct {
+	http.ResponseWriter
+	committed bool
+}
+
+func (w *sessionCacheWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
+
+func (w *sessionCacheWriter) Hijack() (conn net.Conn, rw *bufio.ReadWriter, err error) {
+	conn, rw, err = http.NewResponseController(w.ResponseWriter).Hijack()
+	if err == nil {
+		w.committed = true
+	}
+	return
+}
+
+func (w *sessionCacheWriter) commit() {
+	if !w.committed {
+		w.Header().Set("Cache-Control", headerCacheControlNoStore)
+		w.committed = true
+	}
+}
+
+func (w *sessionCacheWriter) WriteHeader(statusCode int) {
+	if statusCode == http.StatusSwitchingProtocols || statusCode >= http.StatusOK {
+		w.commit()
+	}
+	w.ResponseWriter.WriteHeader(statusCode)
+}
+
+func (w *sessionCacheWriter) Write(p []byte) (int, error) {
+	w.commit()
+	return w.ResponseWriter.Write(p)
+}
+
+func (w *sessionCacheWriter) FlushError() error {
+	w.commit()
+	return http.NewResponseController(w.ResponseWriter).Flush()
+}
+
+func (w *sessionCacheWriter) Flush() { _ = w.FlushError() }
+
+func (w *sessionCacheWriter) WriteHeaderNow() {
+	if whn, ok := w.ResponseWriter.(interface{ WriteHeaderNow() }); ok {
+		w.commit()
+		whn.WriteHeaderNow()
+	}
+}
+
 func (sess sessioner) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if sess.jw.GetSession(r) == nil {
 		created, limitIP := sess.jw.newSession(w, r)
@@ -545,6 +600,10 @@ func (sess sessioner) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			http.Error(w, http.StatusText(status), status)
 			return
 		}
+		sw := &sessionCacheWriter{ResponseWriter: w}
+		sess.h.ServeHTTP(sw, r)
+		sw.commit() // net/http can send an implicit 200 after the handler returns.
+		return
 	}
 	sess.h.ServeHTTP(w, r)
 }
@@ -554,6 +613,8 @@ func (sess sessioner) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 // Before invoking h, it creates a JaWS [Session] when the request has none.
 // A full per-client bucket returns HTTP 429 when global capacity remains;
 // other creation failures return HTTP 503. In either case it does not invoke h.
+// Responses for which it creates a Session receive Cache-Control: no-store;
+// other responses keep the wrapped handler's cache policy.
 //
 // It is distinct from the session accessors:
 // [Jaws.GetSession] and [Request.Session] look up an existing [Session], while

session_test.go

@@ -395,6 +395,215 @@ func TestSessionMiddleware_CloseDuringResponseHeader(t *testing.T) {
 	}
 }
 
+func TestSessionMiddlewareCookieAssetCacheControl(t *testing.T) {
+	jw, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(jw.Close)
+	h := jw.SessionMiddleware(jw.SecureHeadersMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Cache-Control", "no-store")
+		jw.ServeHTTP(w, r)
+	})))
+
+	for _, tc := range []struct{ name, path string }{
+		{"JavaScript", jw.serveJS.Name},
+		{"CSS", jw.serveCSS.Name},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			w := httptest.NewRecorder()
+			h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, tc.path, nil))
+			response := w.Result()
+			if response.StatusCode != http.StatusOK {
+				t.Fatalf("asset status = %d, want 200", response.StatusCode)
+			}
+			if got := response.Header.Get("Cache-Control"); got != "no-store" {
+				t.Fatalf("cookie response Cache-Control = %q, want no-store", got)
+			}
+			cookies := response.Cookies()
+			if len(cookies) != 1 {
+				t.Fatalf("cookie response has %d cookies, want 1", len(cookies))
+			}
+
+			r := httptest.NewRequest(http.MethodGet, tc.path, nil)
+			r.AddCookie(cookies[0])
+			w = httptest.NewRecorder()
+			h.ServeHTTP(w, r)
+			response = w.Result()
+			if got := response.Header.Get("Cache-Control"); !strings.Contains(got, "public") {
+				t.Fatalf("existing session asset Cache-Control = %q, want public", got)
+			}
+			if cookies := response.Cookies(); len(cookies) != 0 {
+				t.Fatalf("existing session asset has %d cookies, want none", len(cookies))
+			}
+		})
+	}
+}
+
+func TestSessionMiddlewareExistingSessionUsesOriginalWriter(t *testing.T) {
+	jw, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(jw.Close)
+	var want http.ResponseWriter
+	h := jw.SessionMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		if want != nil && w != want {
+			t.Fatal("existing Session received a wrapped ResponseWriter")
+		}
+	}))
+	w := httptest.NewRecorder()
+	h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/", nil))
+	cookies := w.Result().Cookies()
+	if len(cookies) != 1 {
+		t.Fatalf("new Session response has %d cookies, want 1", len(cookies))
+	}
+	w = httptest.NewRecorder()
+	want = w
+	r := httptest.NewRequest(http.MethodGet, "/", nil)
+	r.AddCookie(cookies[0])
+	h.ServeHTTP(w, r)
+}
+
+func BenchmarkSessionMiddlewareExistingSession(b *testing.B) {
+	jw, err := New()
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.Cleanup(jw.Close)
+	r := httptest.NewRequest(http.MethodGet, "/", nil)
+	if jw.NewSession(nil, r) == nil {
+		b.Fatal("session creation failed")
+	}
+	w := httptest.NewRecorder()
+	h := jw.SessionMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	b.ReportAllocs()
+	b.ResetTimer()
+	for b.Loop() {
+		h.ServeHTTP(w, r)
+	}
+}
+
+func TestSessionMiddlewareCookieCacheControlOnCommit(t *testing.T) {
+	for _, tc := range []struct {
+		name   string
+		commit func(http.ResponseWriter) error
+	}{
+		{"Write", func(w http.ResponseWriter) error {
+			_, err := w.Write([]byte("ok"))
+			return err
+		}},
+		{"Flush", func(w http.ResponseWriter) error {
+			return http.NewResponseController(w).Flush()
+		}},
+		{"Flusher", func(w http.ResponseWriter) error {
+			w.(http.Flusher).Flush()
+			return nil
+		}},
+		{"Return", func(http.ResponseWriter) error { return nil }},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			jw, err := New()
+			if err != nil {
+				t.Fatal(err)
+			}
+			t.Cleanup(jw.Close)
+			h := jw.SessionMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+				w.Header().Set("Cache-Control", "public, max-age=3600")
+				if err := tc.commit(w); err != nil {
+					t.Error(err)
+				}
+			}))
+			w := httptest.NewRecorder()
+			h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/", nil))
+			response := w.Result()
+			if got := response.Header.Get("Cache-Control"); got != "no-store" {
+				t.Fatalf("Cache-Control = %q, want no-store", got)
+			}
+			if cookies := response.Cookies(); len(cookies) != 1 {
+				t.Fatalf("response has %d cookies, want 1", len(cookies))
+			}
+			if tc.name == "Write" && response.Header.Get("Content-Type") != "text/plain; charset=utf-8" {
+				t.Fatalf("implicit Write lost content type: %q", response.Header.Get("Content-Type"))
+			}
+		})
+	}
+}
+
+type sessionHeaderNowRecorder struct {
+	*httptest.ResponseRecorder
+	cacheControlAtCall string
+}
+
+func (w *sessionHeaderNowRecorder) WriteHeaderNow() {
+	w.cacheControlAtCall = w.Header().Get("Cache-Control")
+}
+
+func TestSessionMiddlewareWriteHeaderNow(t *testing.T) {
+	jw, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(jw.Close)
+	h := jw.SessionMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Cache-Control", "public, max-age=3600")
+		w.(interface{ WriteHeaderNow() }).WriteHeaderNow()
+	}))
+	w := &sessionHeaderNowRecorder{ResponseRecorder: httptest.NewRecorder()}
+	h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/", nil))
+	if w.cacheControlAtCall != "no-store" {
+		t.Fatalf("Cache-Control at WriteHeaderNow = %q, want no-store", w.cacheControlAtCall)
+	}
+}
+
+func TestSessionMiddlewareWebSocketUpgrade(t *testing.T) {
+	jw, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(jw.Close)
+	release := make(chan struct{})
+	accepted := make(chan error, 1)
+	h := jw.SessionMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := http.NewResponseController(w).SetWriteDeadline(time.Time{}); err != nil {
+			accepted <- err
+			return
+		}
+		conn, err := websocket.Accept(w, r, nil)
+		accepted <- err
+		if err == nil {
+			<-release
+			if err := conn.CloseNow(); err != nil {
+				t.Error(err)
+			}
+		}
+	}))
+	server := httptest.NewServer(h)
+	t.Cleanup(server.Close)
+	defer close(release)
+
+	conn, response, err := websocket.Dial(t.Context(), server.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := conn.CloseNow(); err != nil {
+			t.Error(err)
+		}
+	}()
+	if err := <-accepted; err != nil {
+		t.Fatal(err)
+	}
+	if got := response.Header.Get("Cache-Control"); got != "no-store" {
+		t.Fatalf("upgrade Cache-Control = %q, want no-store", got)
+	}
+	if cookies := response.Cookies(); len(cookies) != 1 {
+		t.Fatalf("upgrade has %d cookies, want 1", len(cookies))
+	}
+}
+
 func TestSession_NewSessionReplacesDuplicateCookieSessions(t *testing.T) {
 	jw, err := New()
 	if err != nil {