From 8f34afd66ead58fd25e0ff52906c07c3f55cf6c6 Mon Sep 17 00:00:00 2001 From: steve-chavez Date: Tue, 16 Dec 2025 13:52:58 -0500 Subject: [PATCH] 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. --- .gitignore | 1 + nix/tools/gen_rsa_materials.py | 52 +++++++++++++++++ nix/tools/generate_targets.py | 102 ++++++++++++++++++++++++++------- nix/tools/loadtest.nix | 96 ++++++++++++++++++++++--------- 4 files changed, 202 insertions(+), 49 deletions(-) create mode 100644 nix/tools/gen_rsa_materials.py diff --git a/.gitignore b/.gitignore index af81dde65..230716db7 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,4 @@ loadtest .docs-build gen_targets.http gen_jwk.json +gen_private.json diff --git a/nix/tools/gen_rsa_materials.py b/nix/tools/gen_rsa_materials.py new file mode 100644 index 000000000..d961d88b5 --- /dev/null +++ b/nix/tools/gen_rsa_materials.py @@ -0,0 +1,52 @@ +# Generate RSA JWK/public material for loadtests. + +import argparse +import sys +from pathlib import Path + +import jwcrypto.jwk as jwk + + +def main(): + parser = argparse.ArgumentParser( + description="Generate RSA JWK/private key pair for loadtests" + ) + parser.add_argument( + "--rsa", + dest="jwk_path", + metavar="JWK_PATH", + type=Path, + required=True, + help="Path to write the RSA JWK file", + ) + parser.add_argument( + "--private-key", + dest="private_key_path", + metavar="PRIVATE_KEY_PATH", + type=Path, + required=True, + help="Path to write the RSA private key file", + ) + + args = parser.parse_args() + + key = jwk.JWK.generate(kty="RSA", size=4096) + private_jwk, public_jwk = key.export_private(), key.export_public() + + try: + args.jwk_path.write_text(public_jwk) + print(f"Created RSA JWK on {args.jwk_path}") + except OSError as e: + print(f"Error writing to {args.jwk_path}:{e}", file=sys.stderr) + sys.exit(1) + + try: + args.private_key_path.write_text(private_jwk) + print(f"Created private key on {args.private_key_path}") + except OSError as e: + print(f"Error writing to {args.private_key_path}:{e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/nix/tools/generate_targets.py b/nix/tools/generate_targets.py index 6de5a6a1e..417a80267 100644 --- a/nix/tools/generate_targets.py +++ b/nix/tools/generate_targets.py @@ -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__": diff --git a/nix/tools/loadtest.nix b/nix/tools/loadtest.nix index 97f82971c..d66934478 100644 --- a/nix/tools/loadtest.nix +++ b/nix/tools/loadtest.nix @@ -18,6 +18,8 @@ let ]; } '' + echo "Starting vegeta loadtest..." + # ARG_USE_ENV only adds defaults or docs for environment variables # We manually implement a required check here # See also: https://github.com/matejak/argbash/issues/80 @@ -71,59 +73,90 @@ let case "$_arg_kind" in jwt-hs) - ${genTargets} --method "$_arg_method" "$_arg_testdir"/gen_targets.http export PGRST_JWT_CACHE_MAX_ENTRIES="0" export PGRST_JWT_CACHE_MAX_LIFETIME="0" + + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withPgrst} -m "$_arg_monitor" \ + ${withGenTargets} --method "$_arg_method" "$_arg_testdir"/gen_targets.http \ + sh -c "cd \"$_arg_testdir\" && \ + ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" ;; jwt-hs-cache) - ${genTargets} --method "$_arg_method" "$_arg_testdir"/gen_targets.http + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withPgrst} -m "$_arg_monitor" \ + ${withGenTargets} --method "$_arg_method" "$_arg_testdir"/gen_targets.http \ + sh -c "cd \"$_arg_testdir\" && \ + ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" ;; jwt-hs-cache-worst) - ${genTargets} --method "$_arg_method" --worst "$_arg_testdir"/gen_targets.http + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withPgrst} -m "$_arg_monitor" \ + ${withGenTargets} --method "$_arg_method" --worst "$_arg_testdir"/gen_targets.http \ + sh -c "cd \"$_arg_testdir\" && \ + ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" ;; jwt-rsa) - ${genTargets} --method "$_arg_method" --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http export PGRST_JWT_CACHE_MAX_ENTRIES="0" export PGRST_JWT_CACHE_MAX_LIFETIME="0" + + ${genRsaMaterials} --rsa="$_arg_testdir"/gen_jwk.json --private-key="$_arg_testdir"/gen_private.json export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json" + + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withPgrst} -m "$_arg_monitor" \ + ${withGenTargets} --method "$_arg_method" --rsa="$_arg_testdir"/gen_jwk.json --private-key="$_arg_testdir"/gen_private.json "$_arg_testdir"/gen_targets.http \ + sh -c "cd \"$_arg_testdir\" && \ + ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" ;; jwt-rsa-cache) - ${genTargets} --method "$_arg_method" --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http + ${genRsaMaterials} --rsa="$_arg_testdir"/gen_jwk.json --private-key="$_arg_testdir"/gen_private.json export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json" + + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withPgrst} -m "$_arg_monitor" \ + ${withGenTargets} --method "$_arg_method" --rsa="$_arg_testdir"/gen_jwk.json --private-key="$_arg_testdir"/gen_private.json "$_arg_testdir"/gen_targets.http \ + sh -c "cd \"$_arg_testdir\" && \ + ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" ;; jwt-rsa-cache-worst) - ${genTargets} --method "$_arg_method" --worst --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json" + + ${genRsaMaterials} --rsa="$_arg_testdir"/gen_jwk.json --private-key="$_arg_testdir"/gen_private.json + export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json" + + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withPgrst} -m "$_arg_monitor" \ + ${withGenTargets} --method "$_arg_method" --worst --rsa="$_arg_testdir"/gen_jwk.json --private-key="$_arg_testdir"/gen_private.json "$_arg_testdir"/gen_targets.http \ + sh -c "cd \"$_arg_testdir\" && \ + ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" ;; - *) + mixed) + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withSlowPg} \ + ${withTools.withPgrst} -m "$_arg_monitor" \ + ${withTools.withSlowPgrst} \ + sh -c "cd \"$_arg_testdir\" && \ + ${runner} -targets targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" ;; esac - if [ "$_arg_kind" == "mixed" ]; then - # shellcheck disable=SC2145 - ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ - ${withTools.withSlowPg} \ - ${withTools.withPgrst} -m "$_arg_monitor" \ - ${withTools.withSlowPgrst} \ - sh -c "cd \"$_arg_testdir\" && \ - ${runner} -targets targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" - - ${vegeta}/bin/vegeta report -type=text "$_arg_output" - else - # shellcheck disable=SC2145 - ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ - ${withTools.withPgrst} -m "$_arg_monitor" \ - sh -c "cd \"$_arg_testdir\" && \ - ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" - - ${vegeta}/bin/vegeta report -type=text "$_arg_output" + ${vegeta}/bin/vegeta report -type=text "$_arg_output" + if [ "$_arg_kind" != "mixed" ]; then # fail in case 401 happened on jwt loadtests unauthorized_count="$(${vegeta}/bin/vegeta report -type=json "$_arg_output" \ | ${jq}/bin/jq -r '.status_codes["401"] // 0')" @@ -280,13 +313,22 @@ let | ${mergeMonitorResults} ''; - genTargets = - writers.writePython3 "postgrest-gen-loadtest-targets" + withGenTargets = + writers.writePython3 "postgrest-with-gen-loadtest-targets" { libraries = [ python3Packages.pyjwt python3Packages.jwcrypto ]; + doCheck = false; # postgrest-style conflicts with this } (builtins.readFile ./generate_targets.py); + genRsaMaterials = + writers.writePython3 "postgrest-gen-rsa-materials" + { + libraries = [ python3Packages.jwcrypto ]; + doCheck = false; # postgrest-style conflicts with this + } + (builtins.readFile ./gen_rsa_materials.py); + mergeMonitorResults = writers.writePython3 "postgrest-merge-monitor-results" {