package web
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"craftdesign.group/quarry/internal/api"
"craftdesign.group/quarry/internal/ci"
"craftdesign.group/quarry/internal/ops"
"craftdesign.group/quarry/internal/streamcap"
)
// deadlineRecorder is a response writer with a connection's write
// deadline, recording every deadline a handler sets on it.
type deadlineRecorder struct {
*httptest.ResponseRecorder
mu sync.Mutex
deadlines []time.Time
flushed chan struct{}
once sync.Once
}
func newDeadlineRecorder() *deadlineRecorder {
return &deadlineRecorder{ResponseRecorder: httptest.NewRecorder(), flushed: make(chan struct{})}
}
func (d *deadlineRecorder) SetWriteDeadline(t time.Time) error {
d.mu.Lock()
defer d.mu.Unlock()
d.deadlines = append(d.deadlines, t)
return nil
}
func (d *deadlineRecorder) Flush() {
d.ResponseRecorder.Flush()
d.once.Do(func() { close(d.flushed) })
}
func (d *deadlineRecorder) set() []time.Time {
d.mu.Lock()
defer d.mu.Unlock()
return append([]time.Time(nil), d.deadlines...)
}
// A frame written through a live stream carries a write deadline in the
// future when it is written and again when it is flushed, so a reader
// that stops reading is cut off rather than holding the handler open.
func TestALiveFrameCarriesAWriteDeadline(t *testing.T) {
rec := newDeadlineRecorder()
w, fl, ok := liveStream(rec)
if !ok {
t.Fatal("FAIL: a flushing writer was refused")
}
before := time.Now()
for i := 0; i < 5; i++ {
fmt.Fprintf(w, "data: line %d\n", i)
}
fmt.Fprint(w, "\n")
fl.Flush()
got := rec.set()
if len(got) != 2 {
t.Fatalf("FAIL: a five-line frame set %d deadlines, want one at its first write and one at its flush", len(got))
}
for _, d := range got {
if !d.After(before) {
t.Fatalf("FAIL: deadline %v is not in the future", d)
}
}
t.Log("ok: one frame arms the deadline at its first write and at its flush")
// The next frame re-arms at its first write, so a deadline left
// from an earlier frame never expires under a new one.
fmt.Fprint(w, ": ping\n\n")
if n := len(rec.set()); n != 3 {
t.Fatalf("FAIL: the next frame's first write did not re-arm (deadlines %d, want 3)", n)
}
t.Log("ok: the next frame re-arms at its first write")
}
// Compression wraps every response, so a handler asking for a write
// deadline asks through it. A wrapper that hides the connection makes
// the deadline silently unsupported, which is the same as having none.
func TestADeadlineReachesTheConnectionThroughCompression(t *testing.T) {
errc := make(chan error, 1)
srv := httptest.NewServer(Compress(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
errc <- http.NewResponseController(w).SetWriteDeadline(time.Now().Add(time.Minute))
})))
defer srv.Close()
req, _ := http.NewRequest("GET", srv.URL, nil)
req.Header.Set("Accept-Encoding", "gzip")
resp, err := srv.Client().Do(req)
if err != nil {
t.Fatal(err)
}
resp.Body.Close()
if err := <-errc; err != nil {
t.Fatalf("FAIL: the deadline did not reach the connection: %v", err)
}
t.Log("ok: a deadline set under compression reaches the connection")
}
// A real stream route arms the deadline: the run page's stream is the
// representative, since every live handler opens through liveStream.
func TestARunStreamArmsAWriteDeadline(t *testing.T) {
eng := &countingEngine{first: ci.Run{RunID: "run-5", Status: "running",
Steps: map[string]ci.Step{"build": {Status: "running"}}}}
ix := landingIndex(t, false)
r := ops.NewRegistry()
deps := api.Deps{Version: "v",
Repos: api.NewRepos(map[string]*api.Repo{"alpha": {Index: ix}}),
CI: fakeCI{runs: []api.CIRun{{RunID: "run-5", Checkin: "c1", Workflow: "w",
Pipeline: "p", Status: "running", Started: "2026-08-29T01:00:00Z"}}}}
if err := api.RegisterAll(r, deps); err != nil {
t.Fatal(err)
}
h := HandlerWith(r, nil, nil, []Option{
WithSessions(fakeSessions{user: "dan"}), WithCaps(stubCaps{letters: "i8"}), WithEngine(eng)})
ctx, cancel := context.WithCancel(context.Background())
req := httptest.NewRequest("GET", "/r/alpha/ci/run.sse/run-5", nil).WithContext(ctx)
req.AddCookie(&http.Cookie{Name: sessionCookie, Value: "good"})
rec := newDeadlineRecorder()
done := make(chan struct{})
go func() {
defer close(done)
h.ServeHTTP(rec, req)
}()
select {
case <-rec.flushed:
case <-time.After(3 * time.Second):
t.Fatal("FAIL: the run stream never flushed")
}
cancel()
<-done
if len(rec.set()) == 0 {
t.Fatal("FAIL: the run stream flushed without a write deadline")
}
t.Log("ok: the run stream arms a write deadline")
}
// The chat page's own stream tails the same channel the API's
// chat.stream does, and is just as open to anonymous readers, so it
// draws on the same capped budget: past the per-address cap it answers
// 429, and an ended stream gives its slot back.
func TestAnonymousChatPageStreamsAreCappedPerAddress(t *testing.T) {
r := chatRegistry(t, &webChatStub{}, &webForumReadStub{},
func(context.Context, string) bool { return true })
live := newFakeChatLive()
// The fake makes its activity channel on first use without a lock,
// and this test is the one that opens many streams at once, so the
// channel is made here, before any of them can race to make it.
live.Activity()
srv := httptest.NewServer(Handler(r, nil, live))
defer srv.Close()
open := func() (int, func()) {
ctx, cancel := context.WithCancel(context.Background())
req, _ := http.NewRequestWithContext(ctx, "GET", srv.URL+"/r/test/chat.sse", nil)
resp, err := srv.Client().Do(req)
if err != nil {
cancel()
t.Fatal(err)
}
return resp.StatusCode, func() { cancel(); resp.Body.Close() }
}
var closers []func()
defer func() {
for _, c := range closers {
c()
}
}()
for i := 0; i < streamcap.PerAddress; i++ {
code, done := open()
closers = append(closers, done)
if code != http.StatusOK {
t.Fatalf("FAIL: stream %d refused below the cap: %d", i+1, code)
}
}
code, done := open()
done()
if code != http.StatusTooManyRequests {
t.Fatalf("FAIL: stream %d answered %d, want 429", streamcap.PerAddress+1, code)
}
t.Logf("ok: stream %d is refused with 429", streamcap.PerAddress+1)
closers[0]()
closers = closers[1:]
deadline := time.Now().Add(3 * time.Second)
for {
code, done := open()
if code == http.StatusOK {
closers = append(closers, done)
break
}
done()
if time.Now().After(deadline) {
t.Fatalf("FAIL: an ended stream never gave its slot back (last answer %d)", code)
}
time.Sleep(20 * time.Millisecond)
}
t.Log("ok: an ended stream gives its slot back")
}