package web
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"sync"
"testing"
"time"
"craftdesign.group/quarry/internal/api"
"craftdesign.group/quarry/internal/chat"
"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")
}
// The chat page's stream keys its cap on the trusted-proxy rule, never
// on a header the client wrote: forging a new X-Forwarded-For on every
// request does not buy a stream past the cap.
func TestAForgedForwardedHeaderDoesNotBuyAChatStream(t *testing.T) {
r := chatRegistry(t, &webChatStub{}, &webForumReadStub{},
func(context.Context, string) bool { return true })
live := newFakeChatLive()
live.Activity() // made before the streams race to make it; see above
srv := httptest.NewServer(Handler(r, nil, live))
defer srv.Close()
var closers []func()
defer func() {
for _, c := range closers {
c()
}
}()
open := func(i int) int {
ctx, cancel := context.WithCancel(context.Background())
req, _ := http.NewRequestWithContext(ctx, "GET", srv.URL+"/r/test/chat.sse", nil)
req.Header.Set("X-Forwarded-For", fmt.Sprintf("198.51.100.%d", i+1))
resp, err := srv.Client().Do(req)
if err != nil {
cancel()
t.Fatal(err)
}
closers = append(closers, func() { cancel(); resp.Body.Close() })
return resp.StatusCode
}
for i := 0; i < streamcap.PerAddress; i++ {
if code := open(i); code != http.StatusOK {
t.Fatalf("FAIL: stream %d refused below the cap: %d", i+1, code)
}
}
if code := open(streamcap.PerAddress); code != http.StatusTooManyRequests {
t.Fatalf("FAIL: a forged header bought stream %d past the cap (answered %d)", streamcap.PerAddress+1, code)
}
t.Log("ok: forged forwarded headers share one cap key")
}
// openCappedStream opens a stream at path with header set,
// closed when the test ends.
func openCappedStream(t *testing.T, srv *httptest.Server, path string, header map[string]string) (int, string) {
t.Helper()
ctx, cancel := context.WithCancel(context.Background())
req, _ := http.NewRequestWithContext(ctx, "GET", srv.URL+path, nil)
for k, v := range header {
req.Header.Set(k, v)
}
resp, err := srv.Client().Do(req)
if err != nil {
cancel()
t.Fatal(err)
}
t.Cleanup(func() { cancel(); resp.Body.Close() })
return resp.StatusCode, resp.Header.Get("Retry-After")
}
// assertProxyKeysTheCap drives one surface behind a configured proxy:
// every stream carries the proxy's header for one reader and a forged
// X-Forwarded-For whose leading hop changes each time. The proxy's
// address is the key, so the forgeries buy nothing; the refusal says
// when to retry; and a second reader behind the same proxy still gets in.
func assertProxyKeysTheCap(t *testing.T, srv *httptest.Server, path string) {
t.Helper()
as := func(reader string, i int) map[string]string {
return map[string]string{
"CF-Connecting-IP": reader,
"X-Forwarded-For": fmt.Sprintf("198.51.100.%d, 192.0.2.1", i+1),
}
}
for i := 0; i < streamcap.PerAddress; i++ {
if code, _ := openCappedStream(t, srv, path, as("203.0.113.7", i)); code != http.StatusOK {
t.Fatalf("FAIL: %s stream %d refused below the cap: %d", path, i+1, code)
}
}
code, retry := openCappedStream(t, srv, path, as("203.0.113.7", streamcap.PerAddress))
if code != http.StatusTooManyRequests {
t.Fatalf("FAIL: %s: a forged leading hop bought stream %d past the cap (answered %d)",
path, streamcap.PerAddress+1, code)
}
t.Logf("ok: %s: forged leading hops share the proxy's key", path)
if n, err := strconv.Atoi(retry); err != nil || n <= 0 {
t.Fatalf("FAIL: %s: the 429 carries no usable Retry-After (%q)", path, retry)
}
t.Logf("ok: %s: the 429 says when to retry", path)
if code, _ := openCappedStream(t, srv, path, as("203.0.113.8", 0)); code != http.StatusOK {
t.Fatalf("FAIL: %s: a second reader behind the proxy was refused: %d", path, code)
}
t.Logf("ok: %s: the proxy's header tells readers apart", path)
}
// The chat page's stream, behind a configured proxy.
func TestBehindAConfiguredProxyTheChatPageCapKeysOnItsHeader(t *testing.T) {
r := chatRegistry(t, &webChatStub{}, &webForumReadStub{},
func(context.Context, string) bool { return true })
live := newFakeChatLive()
live.Activity() // made before the streams race to make it; see above
srv := httptest.NewServer(HandlerWith(r, nil, live,
[]Option{WithTrustedProxy("CF-Connecting-IP")}))
// Registered before any stream opens, so it runs after they close:
// Close waits for open connections.
t.Cleanup(srv.Close)
assertProxyKeysTheCap(t, srv, "/r/test/chat.sse")
}
// openTailChat keeps chat.stream open the way a broker consumer does,
// where the stub's own tail closes at once.
type openTailChat struct{ *webChatStub }
func (openTailChat) Tail(string, uint64) (<-chan chat.Message, func()) {
ch := make(chan chat.Message)
var once sync.Once
return ch, func() { once.Do(func() { close(ch) }) }
}
// The API's chat.stream, behind a configured proxy, with the rule
// injected the way the composition root injects it.
func TestBehindAConfiguredProxyTheAPICapKeysOnItsHeader(t *testing.T) {
r := chatRegistry(t, openTailChat{&webChatStub{}}, &webForumReadStub{},
func(context.Context, string) bool { return true })
h, err := api.Router(r, "v")
if err != nil {
t.Fatal(err)
}
proxy := TrustedProxy("CF-Connecting-IP")
srv := httptest.NewServer(api.WithClientAddr(proxy.ClientAddr, h))
// Registered before any stream opens, so it runs after they close:
// Close waits for open connections.
t.Cleanup(srv.Close)
assertProxyKeysTheCap(t, srv, "/api/v1/r/test/chat/stream")
}