package api
import (
"bufio"
"context"
"fmt"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"craftdesign.group/quarry/internal/chat"
"craftdesign.group/quarry/internal/ops"
"craftdesign.group/quarry/internal/streamcap"
)
// openChat is a chat plane whose tail stays open until the stream ends,
// the way a real JetStream consumer does.
type openChat struct{ *fakeChat }
func (o openChat) Tail(string, uint64) (<-chan chat.Message, func()) {
ch := make(chan chat.Message)
var once sync.Once
return ch, func() { once.Do(func() { close(ch) }) }
}
func chatStreamServer(t *testing.T) *httptest.Server {
t.Helper()
return chatStreamServerWith(t, nil)
}
// chatStreamServerWith serves chat.stream with wrap applied around the
// router, the way the composition root applies the client-address rule.
func chatStreamServerWith(t *testing.T, wrap func(http.Handler) http.Handler) *httptest.Server {
t.Helper()
r := ops.NewRegistry()
repos := testRepos(wikiIndex(t))
mustRepo(t, repos, "test").Code = "codeXYZ"
if err := RegisterAll(r, Deps{Version: "v", Repos: repos, Chat: openChat{&fakeChat{}}}); err != nil {
t.Fatal(err)
}
h, err := Router(r, "v")
if err != nil {
t.Fatal(err)
}
if wrap != nil {
h = wrap(h)
}
srv := httptest.NewServer(h)
t.Cleanup(srv.Close)
return srv
}
// openStream opens chat.stream and waits for the status line; the
// returned cancel ends the stream from the reader's side.
func openStream(t *testing.T, srv *httptest.Server) (int, func()) {
t.Helper()
return openStreamWith(t, srv, nil)
}
func openStreamWith(t *testing.T, srv *httptest.Server, header map[string]string) (int, func()) {
t.Helper()
ctx, cancel := context.WithCancel(context.Background())
req, _ := http.NewRequestWithContext(ctx, "GET", srv.URL+"/api/v1/r/test/chat/stream", nil)
for k, v := range header {
req.Header.Set(k, v)
}
resp, err := srv.Client().Do(req)
if err != nil {
cancel()
t.Fatal(err)
}
if resp.StatusCode == http.StatusOK {
// Read the opening comment so the stream is known to be running.
bufio.NewReader(resp.Body).ReadString('\n')
}
return resp.StatusCode, func() { cancel(); resp.Body.Close() }
}
// Anonymous readers of chat.stream are capped per address: each one
// holds a goroutine and a broker consumer, and nothing else stops one
// client from opening them without end. The stream past the cap is
// refused with 429, and an ended stream gives its slot back.
func TestAnonymousChatStreamsAreCappedPerAddress(t *testing.T) {
srv := chatStreamServer(t)
var closers []func()
defer func() {
for _, c := range closers {
c()
}
}()
for i := 0; i < streamcap.PerAddress; i++ {
code, done := openStream(t, srv)
closers = append(closers, done)
if code != http.StatusOK {
t.Fatalf("FAIL: stream %d of %d refused below the cap: %d", i+1, streamcap.PerAddress, code)
}
}
code, done := openStream(t, srv)
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 := openStream(t, srv)
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")
}
// A header the client wrote never decides whose budget a stream draws
// on: a client that forges a new X-Forwarded-For on every request is
// still one address, and is refused at N+1 like any other.
func TestAForgedForwardedHeaderDoesNotChangeTheCapKey(t *testing.T) {
srv := chatStreamServer(t)
var closers []func()
defer func() {
for _, c := range closers {
c()
}
}()
forged := func(i int) map[string]string {
return map[string]string{"X-Forwarded-For": fmt.Sprintf("198.51.100.%d", i+1)}
}
for i := 0; i < streamcap.PerAddress; i++ {
code, done := openStreamWith(t, srv, forged(i))
closers = append(closers, done)
if code != http.StatusOK {
t.Fatalf("FAIL: stream %d refused below the cap: %d", i+1, code)
}
}
code, done := openStreamWith(t, srv, forged(streamcap.PerAddress))
done()
if 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")
}
// The composition root decides who the client is; the router counts
// against whatever address that rule gives, so two readers the rule
// tells apart do not share a cap.
func TestTheInjectedAddressRuleKeysTheCap(t *testing.T) {
rule := func(req *http.Request) string { return req.Header.Get("X-Test-Reader") }
srv := chatStreamServerWith(t, func(h http.Handler) http.Handler { return WithClientAddr(rule, h) })
var closers []func()
defer func() {
for _, c := range closers {
c()
}
}()
for i := 0; i < streamcap.PerAddress; i++ {
code, done := openStreamWith(t, srv, map[string]string{"X-Test-Reader": "a"})
closers = append(closers, done)
if code != http.StatusOK {
t.Fatalf("FAIL: reader a refused below the cap: %d", code)
}
}
code, done := openStreamWith(t, srv, map[string]string{"X-Test-Reader": "b"})
closers = append(closers, done)
if code != http.StatusOK {
t.Fatalf("FAIL: reader b was refused for reader a's streams: %d", code)
}
t.Log("ok: the injected rule decides the cap key")
}
// deadlineWriter is a response writer with a connection's write
// deadline, recording each one a handler sets.
type deadlineWriter struct {
*httptest.ResponseRecorder
mu sync.Mutex
n int
}
func (d *deadlineWriter) SetWriteDeadline(t time.Time) error {
d.mu.Lock()
defer d.mu.Unlock()
if t.After(time.Now()) {
d.n++
}
return nil
}
// Every frame the API's stream writes carries a write deadline, so a
// reader that stops reading is cut off at the next frame instead of
// holding the stream and its consumer open forever.
func TestEveryAPIStreamFrameCarriesAWriteDeadline(t *testing.T) {
w := &deadlineWriter{ResponseRecorder: httptest.NewRecorder()}
st := newSSEStreamer(w, w)
if err := st.Emit(1, "one"); err != nil {
t.Fatal(err)
}
if err := st.Event("gap", "two"); err != nil {
t.Fatal(err)
}
st.Ping()
if w.n != 3 {
t.Fatalf("FAIL: three frames set %d future deadlines, want 3", w.n)
}
t.Log("ok: each frame arms a write deadline")
}