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" {