chore: 添加限流回归测试脚本

2026-04-06 12:23:48 +08:00
parent 063430b644
commit 346f197ba1
1 changed files with 83 additions and 0 deletions
@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""测试 Envoy 限流是否生效"""
+
+import http.client
+import json
+import re
+
+BASE = "127.0.0.1"
+PORT = 18080
+
+ENDPOINTS = [
+    (
+        "/api/v1/auth/login",
+        10,
+        {"username": "test", "password": "test", "phone": "0", "remember": True},
+    ),
+    (
+        "/api/v1/auth/register",
+        5,
+        {"username": "test", "password": "test", "phone": "0"},
+    ),
+    ("/api/v1/auth/forgot-password/send", 3, {"email": "test@test.com"}),
+    ("/api/v1/email/verification-code/send", 3, {"email": "test@test.com"}),
+]
+
+
+def get_cookies():
+    conn = http.client.HTTPConnection(BASE, PORT)
+    conn.request("GET", "/healthz")
+    resp = conn.getresponse()
+    resp.read()
+    cookies = {}
+    for h, v in resp.getheaders():
+        if h.lower() == "set-cookie":
+            m = re.match(r"([^=]+)=([^;]+)", v)
+            if m:
+                cookies[m.group(1)] = m.group(2)
+    conn.close()
+    return cookies
+
+
+def post(cookies, path, body_dict):
+    conn = http.client.HTTPConnection(BASE, PORT)
+    token = cookies.get("__Host-XSRF-TOKEN", "")
+    guard = cookies.get("__Host-XSRF-GUARD", "")
+    headers = {
+        "Content-Type": "application/json",
+        "Cookie": f"__Host-XSRF-TOKEN={token}; __Host-XSRF-GUARD={guard}",
+        "xsrf-token": token,
+    }
+    conn.request("POST", path, body=json.dumps(body_dict), headers=headers)
+    resp = conn.getresponse()
+    resp.read()
+    rl_remain = resp.getheader("x-ratelimit-remaining", "-")
+    ratelimited = resp.getheader("x-envoy-ratelimited", "")
+    conn.close()
+    return resp.status, rl_remain, bool(ratelimited)
+
+
+def test_endpoint(cookies, path, limit, body_dict):
+    count = limit + 3
+    print(f"\n=== {path} (限额 {limit}/min, 发 {count} 次) ===")
+    first_429 = None
+    for i in range(1, count + 1):
+        status, remain, limited = post(cookies, path, body_dict)
+        tag = " ← 429!" if limited else ""
+        print(f"  #{i:02d}  status={status}  remaining={remain}{tag}")
+        if limited and first_429 is None:
+            first_429 = i
+    if first_429 == limit + 1:
+        print(f"  ✓ 第 {first_429} 次触发限流，符合预期")
+    elif first_429:
+        print(f"  ✗ 第 {first_429} 次触发限流，预期第 {limit + 1} 次")
+    else:
+        print(f"  ✗ 未触发限流！")
+
+
+print("获取 CSRF Cookie...")
+cookies = get_cookies()
+print(f"  token = {cookies.get('__Host-XSRF-TOKEN', '?')[:40]}...")
+
+for path, limit, body in ENDPOINTS:
+    test_endpoint(cookies, path, limit, body)