feat(session): extend grace for bounded external flows

doc/sessions.md

@@ -64,6 +64,15 @@ grace deadline; detaching a claimed Request refreshes it, while detaching an
 unclaimed Request leaves it unchanged. Maintenance removes expired Sessions
 without attached Requests. Cookies have no explicit expiry or MaxAge.
 
+For a bounded external flow, such as an OAuth provider visit, call
+`Session.ExtendGrace(duration)` on the current Session before redirecting the
+browser. A later WebSocket disconnect does not shorten the extended deadline.
+The method returns false if the Session is nil, expired, closed or replaced, or
+the duration is not positive. Check that result before starting the flow.
+`ExtendGrace` does not verify identity or authorize the request; the application
+must do that first. The default one-minute grace still applies to Sessions that
+do not call it.
+
 `Session.Reload` asks its active pages to reload. `Session.Close` invalidates
 the Session and queues reloads for its associated pages, including those still
 waiting to connect. It leaves stored data intact; call `Session.Clear` to remove

serve.go

@@ -185,7 +185,7 @@ func (jw *Jaws) maintenance(requestTimeout time.Duration) {
 			_ = jw.retireNonRunningRequestLocked(rq, nil)
 		}
 	}
-	// Every tenth maintenance pass checks unattached Sessions against sessionGrace deadlines.
+	// Every tenth maintenance pass checks unattached Sessions against their deadlines.
 	jw.sessionSweep++
 	if jw.sessionSweep == 10 {
 		jw.sessionSweep = 0

session.go

@@ -23,9 +23,11 @@ const sessionGrace = time.Minute
 //
 // Sessions live in memory and are bound to the remote IP that created them.
 // Their cookies have no explicit expiry or MaxAge. A Session starts with a
-// one-minute grace deadline; detaching a claimed Request refreshes that deadline,
-// while detaching an unclaimed Request leaves it unchanged. It expires once the
-// deadline passes with no attached Requests, and maintenance removes it.
+// one-minute grace deadline; ExtendGrace can lengthen it for a bounded external
+// flow. Detaching a claimed Request refreshes the normal grace without shortening
+// an extended deadline, while detaching an unclaimed Request leaves it unchanged.
+// It expires once the deadline passes with no attached Requests, and
+// maintenance removes it.
 // [Jaws.Close] invalidates all Sessions and clears their data.
 //
 // Its exported methods are safe to call on a nil *Session; methods with results
@@ -100,7 +102,9 @@ func (sess *Session) delRequest(rq *Request, wasClaimed bool) {
 		// must fire even when other requests remain attached, otherwise an aged
 		// session whose last departing request is an unclaimed bootstrap render
 		// would be reaped with its stale deadline despite recent live activity.
-		sess.deadline = time.Now().Add(sessionGrace)
+		if deadline := time.Now().Add(sessionGrace); deadline.After(sess.deadline) {
+			sess.deadline = deadline
+		}
 	}
 	// For an unclaimed request (its bootstrap render finished before the
 	// WebSocket connected) leave the existing deadline intact: the creation-time
@@ -143,6 +147,29 @@ func (sess *Session) Set(key string, value any) {
 	}
 }
 
+// ExtendGrace keeps a current, live Session available for at least duration
+// from now when it has no attached Requests. It never shortens an existing
+// deadline and does not authorize the caller. It returns false for a nil,
+// expired, closed, or replaced Session, or a non-positive duration. A concurrent
+// Close can invalidate the Session immediately after this method returns.
+func (sess *Session) ExtendGrace(duration time.Duration) bool {
+	if sess == nil || duration <= 0 {
+		return false
+	}
+	jw := sess.jw
+	jw.mu.RLock()
+	defer jw.mu.RUnlock()
+	sess.mu.Lock()
+	defer sess.mu.Unlock()
+	if jw.sessions[sess.sessionID] != sess || sess.isDeadLocked() {
+		return false
+	}
+	if deadline := time.Now().Add(duration); deadline.After(sess.deadline) {
+		sess.deadline = deadline
+	}
+	return true
+}
+
 // ID returns the session ID, a 64-bit random value.
 func (sess *Session) ID() (id uint64) {
 	if sess != nil {

session_test.go

@@ -10,6 +10,7 @@ import (
 	"net/netip"
 	"slices"
 	"strings"
+	"sync"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -2097,6 +2098,138 @@ func TestSession_GetSessionExpiredBeforeCleanup(t *testing.T) {
 	}
 }
 
+func TestSession_ExtendGraceForHTTPAndWebSocketHandoff(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		jw, err := New()
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer jw.Close()
+		r := httptest.NewRequest(http.MethodGet, "/", nil)
+		sess := jw.NewSession(httptest.NewRecorder(), r)
+		if sess == nil || !sess.ExtendGrace(10*time.Minute) {
+			t.Fatal("failed to extend a new session")
+		}
+		otherRequest := httptest.NewRequest(http.MethodGet, "/", nil)
+		if other := jw.NewSession(httptest.NewRecorder(), otherRequest); other == nil || other == sess {
+			t.Fatal("failed to create an independent session")
+		}
+		time.Sleep(2 * time.Minute)
+		if jw.GetSession(otherRequest) != nil {
+			t.Fatal("unextended session survived the default grace")
+		}
+		if got := jw.GetSession(r); got != sess {
+			t.Fatalf("HTTP session expired during provider visit: got %p, want %p", got, sess)
+		}
+		if !sess.ExtendGrace(2 * time.Minute) {
+			t.Fatal("shorter extension rejected a live session")
+		}
+		time.Sleep(3 * time.Minute)
+		if got := jw.GetSession(r); got != sess {
+			t.Fatalf("shorter extension shortened grace: got %p, want %p", got, sess)
+		}
+		time.Sleep(5*time.Minute + time.Second)
+		if jw.GetSession(r) != nil || sess.ExtendGrace(time.Minute) {
+			t.Fatal("expired HTTP session was revived")
+		}
+	})
+
+	synctest.Test(t, func(t *testing.T) {
+		jw, err := New()
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer jw.Close()
+		r := httptest.NewRequest(http.MethodGet, "/", nil)
+		sess := jw.NewSession(httptest.NewRecorder(), r)
+		if sess == nil {
+			t.Fatal("expected session")
+		}
+		rq := jw.newRequest(r)
+		if rq.Session() != sess {
+			t.Fatal("request did not attach to session")
+		}
+		rq.storeState(reqClaimed)
+		time.Sleep(2 * time.Minute) // The attached request keeps an old deadline live.
+		if !sess.ExtendGrace(10 * time.Minute) {
+			t.Fatal("attached session with an old deadline was rejected")
+		}
+		rq.killSession(true) // Browser leaves for the provider.
+		time.Sleep(9 * time.Minute)
+		if got := jw.GetSession(r); got != sess {
+			t.Fatalf("WebSocket disconnect shortened OAuth grace: got %p, want %p", got, sess)
+		}
+		time.Sleep(time.Minute + time.Second)
+		if jw.GetSession(r) != nil {
+			t.Fatal("extended WebSocket session survived its deadline")
+		}
+	})
+}
+
+func TestSession_ExtendGraceRejectsInvalidAndReplaced(t *testing.T) {
+	jw, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer jw.Close()
+	var absent *Session
+	if absent.ExtendGrace(time.Minute) {
+		t.Fatal("nil session was extended")
+	}
+
+	const sessionID = key.Key(0x1122334455667788)
+	ip := netip.MustParseAddr("192.0.2.1")
+	stale := newSession(jw, sessionID, ip, false)
+	jw.mu.Lock()
+	jw.registerSessionLocked(stale, false)
+	jw.mu.Unlock()
+	if stale.ExtendGrace(0) || stale.ExtendGrace(-time.Second) {
+		t.Fatal("non-positive duration was accepted")
+	}
+	jw.deleteSessionIfCurrent(stale)
+	replacement := newSession(jw, sessionID, ip, false)
+	jw.mu.Lock()
+	jw.registerSessionLocked(replacement, false)
+	jw.mu.Unlock()
+	if stale.ExtendGrace(time.Minute) || !replacement.ExtendGrace(time.Minute) {
+		t.Fatal("replaced session was revived or current session was rejected")
+	}
+	replacement.Close()
+	if replacement.ExtendGrace(time.Minute) {
+		t.Fatal("closed session was revived")
+	}
+}
+
+func TestSession_ExtendGraceConcurrentClose(t *testing.T) {
+	jw, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer jw.Close()
+	r := httptest.NewRequest(http.MethodGet, "/", nil)
+	sess := jw.NewSession(httptest.NewRecorder(), r)
+	if sess == nil {
+		t.Fatal("expected session")
+	}
+	start := make(chan struct{})
+	var workers sync.WaitGroup
+	for range 32 {
+		workers.Go(func() {
+			<-start
+			sess.ExtendGrace(10 * time.Minute)
+		})
+	}
+	workers.Go(func() {
+		<-start
+		sess.Close()
+	})
+	close(start)
+	workers.Wait()
+	if sess.ExtendGrace(time.Minute) || jw.GetSession(r) != nil {
+		t.Fatal("concurrent extension revived a closed session")
+	}
+}
+
 func TestSession_GetSessionRejectsCookieValueWithTail(t *testing.T) {
 	jw, _ := New()
 	defer jw.Close()