palette/ratelimit_test.go

package main

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func newTestServer(t *testing.T) *apiServer {
	t.Helper()
	globalLimiter = newLimiter() // fresh buckets per test
	store, err := OpenStore(t.TempDir() + "/test.db")
	if err != nil {
		t.Fatal(err)
	}
	if webUIInstance == nil {
		ui, err := NewWebUI()
		if err != nil {
			t.Fatal(err)
		}
		webUIInstance = ui
	}
	return &apiServer{store: store, cfg: Config{MaxTextBytes: 1024 * 1024}}
}

func postJSON(t *testing.T, h http.Handler, path string, body any) *httptest.ResponseRecorder {
	t.Helper()
	b, _ := json.Marshal(body)
	req := httptest.NewRequest("POST", path, bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()
	h.ServeHTTP(rr, req)
	return rr
}

// TestRateLimitCreateBurst: burst of 5 creates allowed, then 429.
func TestRateLimitCreateBurst(t *testing.T) {
	srv := newTestServer(t)
	h := srv.routes()
	// unique IP per test run so tests don't share buckets
	reqIP := "10.9.9.1:1234"
	for i := 0; i < 5; i++ {
		req := httptest.NewRequest("POST", "/api/pastes", bytes.NewReader([]byte(`{"content":"hi"}`)))
		req.RemoteAddr = reqIP
		rr := httptest.NewRecorder()
		h.ServeHTTP(rr, req)
		if rr.Code != 201 {
			t.Fatalf("req %d: want 201, got %d: %s", i, rr.Code, rr.Body.String())
		}
	}
	req := httptest.NewRequest("POST", "/api/pastes", bytes.NewReader([]byte(`{"content":"hi"}`)))
	req.RemoteAddr = reqIP
	rr := httptest.NewRecorder()
	h.ServeHTTP(rr, req)
	if rr.Code != 429 {
		t.Fatalf("6th req: want 429, got %d", rr.Code)
	}
	if ra := rr.Header().Get("Retry-After"); ra == "" {
		t.Fatal("missing Retry-After header")
	}
	if ra := rr.Header().Get("X-RateLimit-Limit"); ra == "" {
		t.Fatal("missing X-RateLimit-Limit header")
	}
}

// TestRateLimitRefill: after waiting >1s a token refills and a create succeeds.
func TestRateLimitRefill(t *testing.T) {
	srv := newTestServer(t)
	h := srv.routes()
	reqIP := "10.9.9.2:1234"
	for i := 0; i < 6; i++ {
		req := httptest.NewRequest("POST", "/api/pastes", bytes.NewReader([]byte(`{"content":"hi"}`)))
		req.RemoteAddr = reqIP
		rr := httptest.NewRecorder()
		h.ServeHTTP(rr, req)
	}
	time.Sleep(1100 * time.Millisecond)
	req := httptest.NewRequest("POST", "/api/pastes", bytes.NewReader([]byte(`{"content":"hi"}`)))
	req.RemoteAddr = reqIP
	rr := httptest.NewRecorder()
	h.ServeHTTP(rr, req)
	if rr.Code != 201 {
		t.Fatalf("after refill: want 201, got %d", rr.Code)
	}
}

// TestRateLimitGuess: guess-language endpoint is limited too.
func TestRateLimitGuess(t *testing.T) {
	srv := newTestServer(t)
	h := srv.routes()
	reqIP := "10.9.9.3:1234"
	for i := 0; i < 6; i++ {
		req := httptest.NewRequest("POST", "/api/guess-language", bytes.NewReader([]byte(`{"content":"def f(): pass"}`)))
		req.RemoteAddr = reqIP
		rr := httptest.NewRecorder()
		h.ServeHTTP(rr, req)
		if i < 5 && rr.Code != 200 {
			t.Fatalf("req %d: want 200, got %d", i, rr.Code)
		}
	}
	req := httptest.NewRequest("POST", "/api/guess-language", bytes.NewReader([]byte(`{"content":"x"}`)))
	req.RemoteAddr = reqIP
	rr := httptest.NewRecorder()
	h.ServeHTTP(rr, req)
	if rr.Code != 429 {
		t.Fatalf("want 429, got %d", rr.Code)
	}
}

// TestRateLimitUnlock: 5 unlock attempts per IP+paste per minute, then 429.
func TestRateLimitUnlock(t *testing.T) {
	srv := newTestServer(t)
	h := srv.routes()
	// create a password-protected paste
	rr := postJSON(t, h, "/api/pastes", map[string]any{"content": "secret", "password": "pw1", "visibility": "unlisted"})
	if rr.Code != 201 {
		t.Fatalf("create failed: %d", rr.Code)
	}
	var created map[string]any
	json.Unmarshal(rr.Body.Bytes(), &created)
	id := created["id"].(string)
	reqIP := "10.9.9.4:1234"
	for i := 0; i < 6; i++ {
		req := httptest.NewRequest("POST", "/"+id, bytes.NewReader([]byte("password=wrong")))
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		req.RemoteAddr = reqIP
		rr2 := httptest.NewRecorder()
		h.ServeHTTP(rr2, req)
		if i < 5 && rr2.Code == 429 {
			t.Fatalf("req %d: unexpected 429", i)
		}
	}
	req := httptest.NewRequest("POST", "/"+id, bytes.NewReader([]byte("password=wrong")))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.RemoteAddr = reqIP
	rr2 := httptest.NewRecorder()
	h.ServeHTTP(rr2, req)
	if rr2.Code != 429 {
		t.Fatalf("want 429, got %d", rr2.Code)
	}
}

// TestHighlightCode basic expectations.
func TestHighlightCode(t *testing.T) {
	in := "func main() {\n\t// comment\n\tfmt.Println(\"hello\")\n}\n"
	out := highlightCode(in, "go")
	if !bytes.Contains([]byte(out), []byte(`<span class="tok-kw">func</span>`)) {
		t.Fatalf("no keyword span: %s", out)
	}
	if !bytes.Contains([]byte(out), []byte(`<span class="tok-com">// comment</span>`)) {
		t.Fatalf("no comment span: %s", out)
	}
	if !bytes.Contains([]byte(out), []byte(`tok-str">&#34;hello&#34;</span>`)) {
		t.Fatalf("no string span: %s", out)
	}
	// unsupported language returns escaped plain text
	plain := highlightCode("<b>x</b>", "text")
	if plain != "&lt;b&gt;x&lt;/b&gt;" {
		t.Fatalf("plain escaping wrong: %q", plain)
	}
	// line count preserved (gutter alignment)
	if got := len(splitLines(highlightCode("a\nb\nc", "go"))); got != 3 {
		t.Fatalf("want 3 lines, got %d", got)
	}
}

func splitLines(s string) []string {
	var out []string
	start := 0
	for i := 0; i < len(s); i++ {
		if s[i] == '\n' {
			out = append(out, s[start:i])
			start = i + 1
		}
	}
	out = append(out, s[start:])
	return out
}

// TestCreatorAutoUnlock: create with password, then POST the password to
// /{id}, then GET /{id} with the cookie shows the paste (#26).
func TestCreatorAutoUnlock(t *testing.T) {
	srv := newTestServer(t)
	h := srv.routes()
	rr := postJSON(t, h, "/api/pastes", map[string]any{"content": "secret stuff", "password": "pw2", "visibility": "unlisted"})
	if rr.Code != 201 {
		t.Fatalf("create failed: %d", rr.Code)
	}
	var created map[string]any
	json.Unmarshal(rr.Body.Bytes(), &created)
	id := created["id"].(string)

	// locked GET shows unlock page
	req := httptest.NewRequest("GET", "/"+id, nil)
	rr2 := httptest.NewRecorder()
	h.ServeHTTP(rr2, req)
	if bytes.Contains(rr2.Body.Bytes(), []byte("secret stuff")) {
		t.Fatal("locked paste leaked content")
	}

	// unlock POST with ?next= should set cookie and redirect
	req = httptest.NewRequest("POST", "/"+id, bytes.NewReader([]byte("password=pw2&next=/"+id+"?created=1")))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	rr3 := httptest.NewRecorder()
	h.ServeHTTP(rr3, req)
	if rr3.Code != http.StatusSeeOther {
		t.Fatalf("unlock POST: want 303, got %d", rr3.Code)
	}
	var cookie *http.Cookie
	for _, c := range rr3.Result().Cookies() {
		if c.Name == "pw_"+id {
			cookie = c
		}
	}
	if cookie == nil {
		t.Fatal("no pw_ cookie set")
	}

	// GET with cookie shows content
	req = httptest.NewRequest("GET", "/"+id, nil)
	req.AddCookie(cookie)
	rr4 := httptest.NewRecorder()
	h.ServeHTTP(rr4, req)
	if !bytes.Contains(rr4.Body.Bytes(), []byte("secret stuff")) {
		t.Fatalf("cookie unlock failed: %d %s", rr4.Code, rr4.Body.String())
	}
}