zetaloop
84 lines
2.5 KiB
Python
Executable File
84 lines
2.5 KiB
Python
Executable File
#!/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)
|