go_blog/handlers/ratelimit_test.go

package handlers

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

	"github.com/gin-gonic/gin"

	"go_blog/models"
)

// TestWindowLimiterFixedWindow 单元测试固定窗口行为：窗口内超限拒绝、
// 窗口过期后计数重置、不同键互不影响、键有界（sweep 生效）。
func TestWindowLimiterFixedWindow(t *testing.T) {
	l := NewWindowLimiter(2, time.Minute)
	now := time.Unix(1_000_000, 0)
	l.nowFn = func() time.Time { return now }

	for i := 0; i < 2; i++ {
		if !l.Allow("a") {
			t.Fatalf("attempt %d: expect allowed within limit", i+1)
		}
	}
	if l.Allow("a") {
		t.Fatal("expect blocked after limit")
	}
	// 其他键不受影响。
	if !l.Allow("b") {
		t.Fatal("different key must not be affected")
	}

	// 窗口过期后计数重置。
	now = now.Add(time.Minute + time.Second)
	if !l.Allow("a") {
		t.Fatal("expect allowed after window rollover")
	}
}

// guestSessionAndToken 取一个匿名会话及其 CSRF 令牌。
func guestSessionAndToken(e *securityTestEnv) (string, string) {
	w := e.do(http.MethodGet, "/login", "", nil, "")
	m := csrfTokenRe.FindStringSubmatch(w.Body.String())
	if m == nil {
		return e.sessionCookie(w), ""
	}
	return e.sessionCookie(w), m[1]
}

// postJSONFrom 与 postJSON 相同，但可指定客户端 RemoteAddr 以模拟不同来源 IP。
func postJSONFrom(e *securityTestEnv, method, path, cookie, csrfToken, ip string, body interface{}) *httptest.ResponseRecorder {
	var buf bytes.Buffer
	_ = json.NewEncoder(&buf).Encode(body)
	req := httptest.NewRequest(method, path, &buf)
	req.Header.Set("Content-Type", "application/json")
	if csrfToken != "" {
		req.Header.Set("X-CSRF-Token", csrfToken)
	}
	if cookie != "" {
		req.Header.Set("Cookie", cookie)
	}
	req.RemoteAddr = ip + ":4321"
	w := httptest.NewRecorder()
	e.router.ServeHTTP(w, req)
	return w
}

// TestRegisterRateLimited 覆盖 SECURITY_TODO #27：同 IP 连续注册超过
// 阈值（10 次/小时）后返回 429/register_locked；其他 IP 不受影响。
func TestRegisterRateLimited(t *testing.T) {
	e := newSecurityTestEnv(t)
	if err := e.db.Model(&models.SiteSetting{}).Where("id = ?", 1).Update("allow_registration", true).Error; err != nil {
		t.Fatalf("enable registration: %v", err)
	}

	register := func(username, ip string) *httptest.ResponseRecorder {
		cookie, token := guestSessionAndToken(e)
		if token == "" {
			t.Fatal("login page did not render a CSRF token")
		}
		return postJSONFrom(e, http.MethodPost, "/api/auth/register", cookie, token, ip, gin.H{
			"username":         username,
			"password":         "secret1",
			"confirm_password": "secret1",
			"email":            username + "@example.com",
		})
	}

	const ipA = "198.51.100.10"
	for i := 0; i < registerLimitPerHour; i++ {
		w := register(fmt.Sprintf("reg%d", i), ipA)
		if w.Code != http.StatusOK || !respOK(w) {
			t.Fatalf("attempt %d: status = %d, body %s", i+1, w.Code, w.Body.String())
		}
	}

	// 下一次尝试（即使输入合法）被限流。
	w := register("reg-over", ipA)
	if w.Code != http.StatusTooManyRequests || respCode(w) != "register_locked" {
		t.Fatalf("rate-limited register: status = %d, code = %q, want 429/register_locked",
			w.Code, respCode(w))
	}

	// 其他 IP 不受影响。
	w = register("reg-other", "198.51.100.11")
	if w.Code != http.StatusOK || !respOK(w) {
		t.Fatalf("register from other IP: status = %d, body %s", w.Code, w.Body.String())
	}
}

// TestCommentRateLimited 覆盖 SECURITY_TODO #28：同 IP 高频提交评论超过
// 阈值（5 条/分钟）后返回 429/comments_locked；其他 IP 不受影响。
func TestCommentRateLimited(t *testing.T) {
	e := newSecurityTestEnv(t)

	comment := func(ip string) *httptest.ResponseRecorder {
		cookie, token := guestSessionAndToken(e)
		if token == "" {
			t.Fatal("login page did not render a CSRF token")
		}
		return postJSONFrom(e, http.MethodPost, "/api/article/alice-post/comments", cookie, token, ip, gin.H{
			"name":    "guest",
			"email":   "guest@example.com",
			"content": "nice post",
		})
	}

	const ipA = "198.51.100.20"
	for i := 0; i < commentLimitPerMin; i++ {
		w := comment(ipA)
		if w.Code != http.StatusOK || !respOK(w) {
			t.Fatalf("comment %d: status = %d, body %s", i+1, w.Code, w.Body.String())
		}
	}

	// 第六次提交被限流。
	w := comment(ipA)
	if w.Code != http.StatusTooManyRequests || respCode(w) != "comments_locked" {
		t.Fatalf("rate-limited comment: status = %d, code = %q, want 429/comments_locked",
			w.Code, respCode(w))
	}

	// 其他 IP 不受影响。
	w = comment("198.51.100.21")
	if w.Code != http.StatusOK || !respOK(w) {
		t.Fatalf("comment from other IP: status = %d, body %s", w.Code, w.Body.String())
	}
}