steve-chavez
608f7ca45a
This loadtests the jwt decoding logic. For this it adds an optional `-k`(kind) parameter to `postgrest-loadtest` and `postgrest-loadtest-against`. Old kind (default): ``` postgrest-loadtest -k mixed postgrest-loadtest-against -k mixed ``` New kind: ``` postgrest-loadtest -k jwt postgrest-loadtest-against -k jwt ``` Internally it uses a dynamically generated targets file using python which looks like: ``` GET http://postgrest/authors_only Authorization: Bearer <jwt> GET http://postgrest/authors_only Authorization: Bearer <another-jwt> ... ``` Then this is used to run vegeta with the `-lazy` option.
75 lines
2.1 KiB
Python
Executable File
75 lines
2.1 KiB
Python
Executable File
# generates a file to be used by the vegeta load testing tool
|
||
import time
|
||
import hmac
|
||
import hashlib
|
||
import base64
|
||
import json
|
||
import argparse
|
||
import sys
|
||
|
||
SECRET = b"reallyreallyreallyreallyverysafe"
|
||
URL = "http://postgrest"
|
||
JWT_DURATION = 60
|
||
TOTAL_TARGETS = 40000 # tuned by hand to reduce result variance
|
||
|
||
|
||
def base64url_encode(data: bytes) -> str:
|
||
"""URL-safe Base64 encode without padding."""
|
||
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
||
|
||
|
||
def generate_jwt() -> str:
|
||
"""Generate an HS256 JWT"""
|
||
# Header & payload
|
||
header = {"alg": "HS256", "typ": "JWT"}
|
||
now = int(time.time())
|
||
payload = {
|
||
"iat": now,
|
||
"exp": now + JWT_DURATION,
|
||
"role": "postgrest_test_author",
|
||
}
|
||
|
||
# Encode to JSON and then to Base64URL
|
||
header_b = json.dumps(header, separators=(",", ":")).encode()
|
||
payload_b = json.dumps(payload, separators=(",", ":")).encode()
|
||
header_b64 = base64url_encode(header_b)
|
||
payload_b64 = base64url_encode(payload_b)
|
||
|
||
# Sign (HMAC‑SHA256) the "<header>.<payload>" string
|
||
signing_input = f"{header_b64}.{payload_b64}".encode()
|
||
signature = hmac.new(SECRET, signing_input, hashlib.sha256).digest()
|
||
signature_b64 = base64url_encode(signature)
|
||
|
||
return f"{header_b64}.{payload_b64}.{signature_b64}"
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Generate Vegeta targets with unique JWTs"
|
||
)
|
||
parser.add_argument(
|
||
"output",
|
||
help="Path to write the generated targets file",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
lines = []
|
||
for _ in range(TOTAL_TARGETS):
|
||
token = generate_jwt()
|
||
lines.append(f"GET {URL}/authors_only")
|
||
lines.append(f"Authorization: Bearer {token}")
|
||
lines.append("") # blank line to separate requests
|
||
|
||
try:
|
||
with open(args.output, "w") as f:
|
||
f.write("\n".join(lines))
|
||
except IOError as e:
|
||
print(f"Error writing to {args.output}: {e}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
print(f"Generated {TOTAL_TARGETS} targets in {args.output}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|