nix: jwt loadtest chain commands and split rsa gen
Chaining the target generation command after the build is done ensures that if a build takes too long, the dynamic target generation for vegeta won't cause 401 Unauthorized errors due to already expired JWTs. Also split the rsa materials writing to another python program for easier maintenance.
This commit is contained in:
steve-chavez
committed by
Steve Chavez
Steve Chavez
parent
4cdc4c4861
commit
8f34afd66e
@@ -12,10 +12,10 @@
|
||||
# from an array
|
||||
import time
|
||||
import argparse
|
||||
import subprocess
|
||||
import sys
|
||||
import random
|
||||
import jwt
|
||||
import jwcrypto.jwk as jwk
|
||||
from typing import Optional
|
||||
from pathlib import Path
|
||||
from enum import Enum
|
||||
@@ -24,12 +24,12 @@ URL = "http://postgrest"
|
||||
|
||||
secret_key = b"reallyreallyreallyreallyverysafe"
|
||||
|
||||
key = jwk.JWK.generate(kty="RSA", size=4096)
|
||||
private_key = jwt.algorithms.RSAAlgorithm.from_jwk(key.export_private())
|
||||
public_key = key.export_public()
|
||||
|
||||
|
||||
def generate_jwt(now: int, exp_inc: Optional[int], is_hs: bool) -> str:
|
||||
def generate_jwt(
|
||||
now: int,
|
||||
exp_inc: Optional[int],
|
||||
rsa_private_key: Optional[jwt.algorithms.RSAAlgorithm],
|
||||
) -> str:
|
||||
"""Generate an HS256 or RS256 JWT"""
|
||||
payload = {
|
||||
"sub": f"user_{random.getrandbits(32)}",
|
||||
@@ -40,9 +40,13 @@ def generate_jwt(now: int, exp_inc: Optional[int], is_hs: bool) -> str:
|
||||
if exp_inc is not None:
|
||||
payload["exp"] = now + exp_inc
|
||||
|
||||
k = secret_key if is_hs else private_key
|
||||
alg = "HS256" if is_hs else "RS256"
|
||||
return jwt.encode(payload, k, alg)
|
||||
if rsa_private_key is None:
|
||||
key = secret_key
|
||||
alg = "HS256"
|
||||
else:
|
||||
key = rsa_private_key
|
||||
alg = "RS256"
|
||||
return jwt.encode(payload, key, alg)
|
||||
|
||||
|
||||
HTTP_METHODS = (
|
||||
@@ -64,14 +68,44 @@ def append_targets(lines: list[str], token: str, http_method: HttpMethod):
|
||||
lines.append("") # blank line to separate requests
|
||||
|
||||
|
||||
# we use this to chain commands on loadtest.nix
|
||||
def run_command(command: list[str]):
|
||||
if not command:
|
||||
return
|
||||
|
||||
if command[0] == "--":
|
||||
command = command[1:]
|
||||
|
||||
if not command:
|
||||
return
|
||||
|
||||
try:
|
||||
subprocess.run(command, check=True)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
print(
|
||||
f"Error executing command {' '.join(command)}: {exc}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(exc.returncode)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate Vegeta targets with unique JWTs"
|
||||
)
|
||||
parser.add_argument(
|
||||
"output",
|
||||
"targets_path",
|
||||
metavar="TARGETS_PATH",
|
||||
help="Path to write the generated targets file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--private-key",
|
||||
dest="private_key_path",
|
||||
metavar="PRIVATE_KEY_PATH",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Path to the RSA private key file (required when --rsa is used)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--worst",
|
||||
dest="worst",
|
||||
@@ -85,23 +119,31 @@ def main():
|
||||
metavar="JWK_PATH",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Path for generating a RSA JWK file to sign tokens with",
|
||||
help="Path to an existing RSA JWK file used for signing tokens",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
dest="http_method",
|
||||
choices=list(HTTP_METHODS),
|
||||
default=None,
|
||||
required=True,
|
||||
help="HTTP method for the vegeta targets",
|
||||
)
|
||||
parser.add_argument(
|
||||
"command",
|
||||
nargs=argparse.REMAINDER,
|
||||
help="Command (and arguments) to run after generating the targets",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
rsa_private_key: Optional[jwt.algorithms.RSAAlgorithm] = None
|
||||
|
||||
is_hs = args.jwk_path is None
|
||||
|
||||
http_method = HttpMethod(args.http_method)
|
||||
|
||||
nsamples = 1000
|
||||
|
||||
if is_hs:
|
||||
ntargets = 200000
|
||||
else:
|
||||
@@ -109,12 +151,25 @@ def main():
|
||||
ntargets = 50000
|
||||
|
||||
if not is_hs:
|
||||
if args.private_key_path is None:
|
||||
parser.error("--rsa requires the --private-key option")
|
||||
try:
|
||||
with open(args.jwk_path, "w") as jwk:
|
||||
jwk.write(public_key)
|
||||
print(f"Created {args.jwk_path} file containing the RSA JWK")
|
||||
except IOError as e:
|
||||
print(f"Error writing to {args.jwk_path}: {e}", file=sys.stderr)
|
||||
private_key_data = args.private_key_path.read_text()
|
||||
except OSError as e:
|
||||
err = (
|
||||
f"Error reading RSA private key from {args.private_key_path}: "
|
||||
f"{e}. Generate RSA materials first with gen_rsa_materials.py."
|
||||
)
|
||||
print(err, file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
rsa_private_key = jwt.algorithms.RSAAlgorithm.from_jwk(private_key_data)
|
||||
except Exception as exc: # broad exception to capture parsing errors
|
||||
err = (
|
||||
f"Error loading RSA private key from {args.private_key_path}: " f"{exc}"
|
||||
)
|
||||
print(err, file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Generating {ntargets} targets...")
|
||||
@@ -133,6 +188,7 @@ def main():
|
||||
if args.worst:
|
||||
# estimated time takes to build and run postgrest itself
|
||||
build_run_postgrest_time = 2
|
||||
|
||||
# estimated time it takes to generate the targets file
|
||||
# the division numbers are tuned by hand
|
||||
if is_hs: # hs generation is much faster
|
||||
@@ -144,25 +200,27 @@ def main():
|
||||
inc = build_run_postgrest_time + gen_time
|
||||
|
||||
for i in range(ntargets):
|
||||
token = generate_jwt(now, inc + i // 1000, is_hs)
|
||||
token = generate_jwt(now, inc + i // 1000, rsa_private_key)
|
||||
append_targets(lines, token, http_method)
|
||||
|
||||
else:
|
||||
tokens = [generate_jwt(now, None, is_hs) for _ in range(nsamples)]
|
||||
tokens = [generate_jwt(now, None, rsa_private_key) for _ in range(nsamples)]
|
||||
for i in range(ntargets):
|
||||
token = random.choice(tokens)
|
||||
append_targets(lines, token, http_method)
|
||||
|
||||
try:
|
||||
with open(args.output, "w") as f:
|
||||
with open(args.targets_path, "w") as f:
|
||||
f.write("\n".join(lines))
|
||||
except IOError as e:
|
||||
print(f"Error writing to {args.output}: {e}", file=sys.stderr)
|
||||
print(f"Error writing to {args.targets_path}: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"Created {ntargets} targets", end=" ")
|
||||
print(f"in {args.output} ({elapsed:.2f}s)")
|
||||
print(f"in {args.targets_path} ({elapsed:.2f}s)")
|
||||
|
||||
run_command(args.command)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user