chore: initial commit

2026-10-23 18:35:44 +08:00
commit f6921b7b23
69 changed files with 12148 additions and 0 deletions
@@ -0,0 +1,359 @@
+package server
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestClientIPTrust(t *testing.T) {
+	s := &Server{}
+	if err := s.SetTrustedProxies([]string{"10.0.0.0/8", "::1/128"}); err != nil {
+		t.Fatal(err)
+	}
+	if err := s.SetTrustedProxies([]string{"bad"}); err == nil {
+		t.Fatal("invalid CIDR accepted")
+	}
+	for _, tt := range []struct{ peer, xff, want string }{
+		{"192.0.2.1:42", "198.51.100.1", "192.0.2.1"},
+		{"10.0.0.1:42", "198.51.100.1, 10.0.0.2", "198.51.100.1"},
+		{"10.0.0.1:42", "198.51.100.1, 192.0.2.2", "192.0.2.2"},
+		{"10.0.0.1:42", "bad, 192.0.2.2", "10.0.0.1"},
+		{"10.0.0.1:42", "", "10.0.0.1"},
+		{"10.0.0.1:42", "10.0.0.2", "10.0.0.1"},
+		{"[::1]:42", "2001:db8::1", "2001:db8::1"},
+	} {
+		r := httptest.NewRequest("POST", "/", nil)
+		r.RemoteAddr = tt.peer
+		r.Header.Set("X-Forwarded-For", tt.xff)
+		r.Header.Set("X-Real-IP", "203.0.113.1")
+		if got := s.clientIP(r); got != tt.want {
+			t.Errorf("%+v: got %s", tt, got)
+		}
+	}
+}
+
+func TestLoginProxyIsolation(t *testing.T) {
+	s, err := New(":memory:", "password", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+	if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
+		t.Fatal(err)
+	}
+	login := func(peer, forwarded string) int {
+		r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"wrong"}`))
+		r.RemoteAddr = peer
+		r.Header.Set("X-Forwarded-For", forwarded)
+		w := httptest.NewRecorder()
+		s.login(w, r)
+		return w.Code
+	}
+	for i := 0; i < 10; i++ {
+		if got := login("10.0.0.1:1", "192.0.2.1"); got != 401 {
+			t.Fatal(got)
+		}
+	}
+	if got := login("10.0.0.1:1", "192.0.2.1"); got != 429 {
+		t.Fatal(got)
+	}
+	if got := login("10.0.0.1:1", "192.0.2.2"); got != 401 {
+		t.Fatal(got)
+	}
+	for i := 0; i < 10; i++ {
+		login("192.0.2.3:1", fmt.Sprintf("198.51.100.%d", i))
+	}
+	if got := login("192.0.2.3:1", "203.0.113.1"); got != 429 {
+		t.Fatal(got)
+	}
+}
+
+func TestBodyDeadline(t *testing.T) {
+	result := make(chan error, 1)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
+		result <- err
+		w.WriteHeader(400)
+	}))
+	defer srv.Close()
+	conn, err := net.Dial("tcp", srv.Listener.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 20\r\n\r\nx")
+	select {
+	case err := <-result:
+		if err == nil {
+			t.Fatal("slow body accepted")
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("body deadline did not fire")
+	}
+}
+
+func TestBodyDeadlineCleared(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
+		if err != nil {
+			http.Error(w, "read", 400)
+			return
+		}
+		time.Sleep(100 * time.Millisecond)
+		select {
+		case <-r.Context().Done():
+			http.Error(w, "canceled", 500)
+			return
+		default:
+		}
+		fmt.Fprint(w, "data: ok\n\n")
+	}))
+	defer srv.Close()
+	conn, err := net.Dial("tcp", srv.Listener.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	conn.SetDeadline(time.Now().Add(2 * time.Second))
+	fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 1\r\n\r\nx")
+	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != 200 {
+		t.Fatal(resp.StatusCode)
+	}
+}
+
+func TestLoginBodyLimit(t *testing.T) {
+	r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+strings.Repeat("x", 16<<10)+`"}`))
+	var v any
+	if decode(httptest.NewRecorder(), r, &v) == nil {
+		t.Fatal("oversize login accepted")
+	}
+}
+
+func TestLimiterReleaseAndCleanup(t *testing.T) {
+	now := time.Now()
+	l := newRequestLimiter(2, 1, 2)
+	a, _ := l.acquire("a", now)
+	if a == nil {
+		t.Fatal("denied")
+	}
+	if release, _ := l.acquire("a", now); release != nil {
+		t.Fatal("per-key concurrency bypass")
+	}
+	b, _ := l.acquire("b", now)
+	if b == nil {
+		t.Fatal("denied")
+	}
+	if release, _ := l.acquire("c", now); release != nil {
+		t.Fatal("global concurrency bypass")
+	}
+	a()
+	a()
+	b()
+	a, _ = l.acquire("a", now)
+	if a == nil {
+		t.Fatal("not released")
+	}
+	a()
+	if release, retry := l.acquire("a", now); release != nil || retry != "60" {
+		t.Fatal("rpm bypass")
+	}
+	c, _ := l.acquire("c", now.Add(time.Minute))
+	if c == nil {
+		t.Fatal("window not reset")
+	}
+	c()
+	if len(l.entries) != 1 {
+		t.Fatalf("stale entries: %d", len(l.entries))
+	}
+	for i := 0; i < maxLimitEntries; i++ {
+		l.entries[fmt.Sprint(i)] = &limitEntry{window: now.Add(time.Minute)}
+	}
+	if release, _ := l.acquire("overflow", now.Add(time.Minute)); release != nil {
+		t.Fatal("map cap bypass")
+	}
+}
+
+func TestLimiterConcurrent(t *testing.T) {
+	l := newRequestLimiter(10000, 8, 64)
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if release, _ := l.acquire("key", time.Now()); release != nil {
+				release()
+			}
+		}()
+	}
+	wg.Wait()
+	if l.active != 0 || l.entries["key"].active != 0 {
+		t.Fatal("permit leak")
+	}
+}
+
+func TestProxyLimitRejectsBeforeBodyAndLogging(t *testing.T) {
+	s, err := New(":memory:", "test-password", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+	if err := s.SetLimits(1, 1, 1); err != nil {
+		t.Fatal(err)
+	}
+	cookie := loginCookie(t, s)
+	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
+	key := integrationCreateKey(t, s, cookie, "limited", ch)
+	request := func(body io.Reader) *httptest.ResponseRecorder {
+		r := httptest.NewRequest("POST", "/v1/chat/completions", body)
+		r.Header.Set("Authorization", "Bearer "+key.raw)
+		w := httptest.NewRecorder()
+		s.ServeHTTP(w, r)
+		return w
+	}
+	if w := request(strings.NewReader("{}")); w.Code != 400 {
+		t.Fatal(w.Code)
+	}
+	if s.limiter.active != 0 {
+		t.Fatal("permit leaked after invalid request")
+	}
+	body := &observedReader{}
+	w := request(body)
+	if w.Code != 429 || w.Header().Get("Retry-After") == "" || body.read {
+		t.Fatalf("status=%d read=%t", w.Code, body.read)
+	}
+	if err := s.flushLogs(); err != nil {
+		t.Fatal(err)
+	}
+	var count int
+	if err := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&count); err != nil {
+		t.Fatal(err)
+	}
+	if count != 1 {
+		t.Fatalf("rejected request logged: %d logs", count)
+	}
+}
+
+type observedReader struct{ read bool }
+
+func (r *observedReader) Read(p []byte) (int, error) { r.read = true; return 0, io.EOF }
+
+func TestLoginFailuresOnlyAndGlobalCap(t *testing.T) {
+	s, err := New(":memory:", "password", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+	if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
+		t.Fatal(err)
+	}
+	login := func(peer, forwarded, password string) int {
+		r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+password+`"}`))
+		r.RemoteAddr = peer
+		if forwarded != "" {
+			r.Header.Set("X-Forwarded-For", forwarded)
+		}
+		w := httptest.NewRecorder()
+		s.login(w, r)
+		return w.Code
+	}
+	for i := 0; i < 15; i++ {
+		if got := login("10.0.0.1:1", "192.0.2.1", "password"); got != 200 {
+			t.Fatalf("correct login %d rejected: %d", i, got)
+		}
+	}
+	for i := 0; i < loginFailuresPerMinute; i++ {
+		if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 401 {
+			t.Fatalf("failure %d: %d", i, got)
+		}
+	}
+	if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 429 {
+		t.Fatal(got)
+	}
+	// The per-IP lockout is uniform: it must not confirm candidate passwords.
+	if got := login("10.0.0.1:1", "192.0.2.2", "password"); got != 429 {
+		t.Fatal("lockout leaks password validity:", got)
+	}
+	// A successful login clears its IP bucket, so fat-fingers do not accumulate.
+	if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
+		t.Fatal(got)
+	}
+	if got := login("10.0.0.1:1", "192.0.2.3", "password"); got != 200 {
+		t.Fatal(got)
+	}
+	for i := 0; i < loginFailuresPerMinute; i++ {
+		if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
+			t.Fatalf("post-clear failure %d: %d", i, got)
+		}
+	}
+	if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 429 {
+		t.Fatal(got)
+	}
+	for i := 0; i < loginGlobalFailuresPerMinute; i++ {
+		login("10.0.0.1:1", fmt.Sprintf("198.18.%d.%d", i/256, i%256), "wrong")
+	}
+	if got := login("10.0.0.1:1", "203.0.113.1", "wrong"); got != 429 {
+		t.Fatal(got)
+	}
+	// The global cap fails closed too: aggregate brute force locks every
+	// login out without revealing candidate-password validity.
+	if got := login("10.0.0.1:1", "203.0.113.1", "password"); got != 429 {
+		t.Fatal("global lockout leaks password validity:", got)
+	}
+}
+
+func TestAuthFailureGate(t *testing.T) {
+	s := testServer(t)
+	cookie := loginCookie(t, s)
+	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
+	key := integrationCreateKey(t, s, cookie, "gated", ch)
+	request := func(auth string) *httptest.ResponseRecorder {
+		r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(integrationRequest("openai-chat", "abc/gpt-5", false)))
+		r.Header.Set("Authorization", "Bearer "+auth)
+		w := httptest.NewRecorder()
+		s.ServeHTTP(w, r)
+		return w
+	}
+	if w := request(key.raw); w.Code != 502 {
+		t.Fatal("valid key setup:", w.Code, w.Body.String())
+	}
+	for i := 0; i < authFailuresPerMinute; i++ {
+		if w := request("wrong"); w.Code != 401 {
+			t.Fatalf("%d: %d", i, w.Code)
+		}
+	}
+	if w := request("wrong"); w.Code != 429 || w.Header().Get("Retry-After") == "" {
+		t.Fatal(w.Code, w.Header())
+	}
+	// The gate applies before the key lookup: a valid key from the same
+	// (attacker-shared) IP is throttled the same way.
+	if w := request(key.raw); w.Code != 429 {
+		t.Fatal(w.Code)
+	}
+}
+
+func TestProxyBodyTooLarge(t *testing.T) {
+	s := testServer(t)
+	cookie := loginCookie(t, s)
+	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
+	key := integrationCreateKey(t, s, cookie, "limit", ch)
+	r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"abc/gpt-5","messages":[],"pad":"`+strings.Repeat("x", maxBody)+`"}`))
+	r.Header.Set("Authorization", "Bearer "+key.raw)
+	w := httptest.NewRecorder()
+	s.ServeHTTP(w, r)
+	if w.Code != 413 || !strings.Contains(w.Body.String(), "too large") {
+		t.Fatal(w.Code, w.Body.String())
+	}
+}