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

@@ -26,3 +26,4 @@ loadtest
 .docs-build
 gen_targets.http
 gen_jwk.json
+gen_private.json

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()

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(
-def generate_jwt(now: int, exp_inc: Optional[int], is_hs: bool) -> str:
+    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
+    if rsa_private_key is None:
-    alg = "HS256" if is_hs else "RS256"
+        key = secret_key
-    return jwt.encode(payload, k, alg)
+        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:
+            private_key_data = args.private_key_path.read_text()
-                jwk.write(public_key)
+        except OSError as e:
-                print(f"Created {args.jwk_path} file containing the RSA JWK")
+            err = (
-        except IOError as e:
+                f"Error reading RSA private key from {args.private_key_path}: "
-            print(f"Error writing to {args.jwk_path}: {e}", file=sys.stderr)
+                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__":

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
+        ${vegeta}/bin/vegeta report -type=text "$_arg_output"
-          # 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"
 
+        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 =
+  withGenTargets =
-    writers.writePython3 "postgrest-gen-loadtest-targets"
+    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"
       {