Reject tail fetches for running WebSockets

AI.md

@@ -110,8 +110,8 @@ The normal page flow has two related HTTP requests:
 1. Before writing the response, a page handler calls `Jaws.NewRequest(w, r)`,
    which replaces `Cache-Control` with `no-store`. `HeadHTML` normally emits the
    configured resources and request-key metadata. `TailHTML` is optional;
-   placing it before `</body>` applies queued initial updates before the
-   WebSocket connects and can reduce flicker.
+   placing it before `</body>` can apply queued initial updates before the
+   WebSocket starts and reduce flicker.
 2. The bundled script connects to `/jaws/<key>`. `Jaws.ServeHTTP` decodes the
    key, claims the pending Request through `UseRequest`, upgrades the connection,
    and begins event and DOM-update processing.

jaws_test.go

@@ -3367,6 +3367,64 @@ func TestServeHTTP_TailScript_EndpointIsPerRequest(t *testing.T) {
 	is.Equal(w.Code, http.StatusNoContent)
 }
 
+func TestServeHTTP_TailScript_RunningRequestDoesNotDrain(t *testing.T) {
+	jw, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+	go jw.Serve()
+	defer jw.Close()
+	server := httptest.NewServer(jw)
+	defer server.Close()
+
+	initial := httptest.NewRequest(http.MethodGet, server.URL+"/", nil)
+	initial.RemoteAddr = "127.0.0.1:1"
+	rq := jw.NewRequest(httptest.NewRecorder(), initial)
+	elem := rq.NewElement(&testUi{})
+	ready := make(chan struct{})
+	resume := make(chan struct{})
+	rq.SetConnectFn(func(*Request) error {
+		close(ready)
+		<-resume
+		return nil
+	})
+	conn := dialJawsRequest(t, server.URL, rq)
+	defer func() {
+		close(resume)
+		if err := conn.CloseNow(); err != nil {
+			t.Error(err)
+		}
+	}()
+	select {
+	case <-ready:
+	case <-time.After(testTimeout):
+		t.Fatal("WebSocket did not reach ConnectFn")
+	}
+
+	elem.SetClass("cls")
+	res, err := server.Client().Get(server.URL + "/jaws/.tail/" + rq.JawsKeyString())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := res.Body.Close(); err != nil {
+		t.Fatal(err)
+	}
+	if res.StatusCode != http.StatusNotFound {
+		t.Fatalf("tail status = %d, want 404", res.StatusCode)
+	}
+	resume <- struct{}{}
+	ctx, cancel := context.WithTimeout(t.Context(), testTimeout)
+	defer cancel()
+	_, data, err := conn.Read(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := (&wire.WsMsg{Jid: elem.Jid(), What: what.SClass, Data: "cls"}).Format()
+	if !strings.Contains(string(data), want) {
+		t.Fatalf("WebSocket frame %q lacks %q", data, want)
+	}
+}
+
 // TestServeHTTP_TailScript_RejectsRecycledKey covers the finished-request behavior
 // of the /jaws/.tail endpoint: completion reserves the key with a nil tombstone in
 // jw.requests, so a tail fetch for the old key finds no live Request and returns

serve.go

@@ -588,10 +588,11 @@ func (*Request) writeTailResponse(w http.ResponseWriter, b []byte, sent bool) (e
 	return
 }
 
-// TailHTML writes optional HTML code at the end of the page's BODY section that
-// will immediately apply HTML attribute and class updates made during initial
-// rendering, which minimizes flicker without having to write the correct
-// value in templates or during [Renderer.JawsRender].
+// TailHTML writes optional HTML code at the end of the page's BODY section.
+//
+// It can apply queued attribute and class updates before the WebSocket starts,
+// reducing flicker without requiring their values in templates or
+// [Renderer.JawsRender].
 //
 // It also adds a <noscript> tag that warns of reduced functionality.
 func (rq *Request) TailHTML(w io.Writer) (err error) {
@@ -604,29 +605,20 @@ func (rq *Request) TailHTML(w io.Writer) (err error) {
 }
 
 // serveTailScript handles a GET /jaws/.tail/<key> fetch, draining the one-shot
-// attribute/class updates queued for the matching Request and writing them. It
+// attribute/class updates queued for a non-running Request and writing them. It
 // reports whether it produced a response; a false return means the path was not a
 // handled tail fetch and [Jaws.ServeHTTP] should keep dispatching.
 func (jw *Jaws) serveTailScript(w http.ResponseWriter, r *http.Request) (handled bool) {
 	if jawsKeyString, ok := strings.CutPrefix(r.URL.Path, "/jaws/.tail/"); ok {
 		if jawsKey, tail := key.Parse(jawsKeyString); tail == "" {
 			remoteIP := jw.clientIP(r)
-			// Hold jw.mu (read) across both the lookup and the drain: finishing needs
-			// the jw.mu write lock, so rq cannot be unregistered while we drain its
-			// queue. A stale key either misses the map (404) or drains its own genuine
-			// content. The network write is done after releasing jw.mu so a slow client
-			// cannot stall completion or the Serve loop.
+			// Hold jw.mu across the state check and drain; starting the WebSocket
+			// and finishing the Request both need the write lock.
 			jw.mu.RLock()
 			rq := jw.requests[jawsKey]
-			// Bind the tail fetch to the client like the WebSocket claim path
-			// (Request.claim): the one-shot tail is drained only when the fetch comes from
-			// the same client IP the initial request was issued to (loopback-aware, see
-			// equalIP). rq.remoteIP is stable here because finishing requires the jw.mu
-			// write lock. A mismatch is treated as not found, so a leaked key cannot drain
-			// (and thereby deny) another client's tail. The WebSocket carries all live
-			// data, so this only closes the cross-IP read of the already-rendered
-			// attribute/class fragments and the cross-IP one-shot race.
-			if rq != nil && !equalIP(remoteIP, rq.remoteIP) {
+			// Match the WebSocket's client-IP binding so a leaked key cannot
+			// consume another client's one-shot tail.
+			if rq != nil && (!equalIP(remoteIP, rq.remoteIP) || rq.loadState() == reqRunning) {
 				rq = nil
 			}
 			var b []byte