jimmy/postgrest
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

nix: jwt loadtest chain commands and split rsa gen

Browse Source 
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
2025-12-16 21:15:28 -05:00
committed by Steve Chavez
parent 4cdc4c4861
commit 8f34afd66e
4 changed files with 202 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -26,3 +26,4 @@ loadtest
.docs-build
gen_targets.http
gen_jwk.json
gen_private.json
nix/tools/gen_rsa_materials.py
+52
View File
@@ -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
+80 -22
View File
@@ -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__":
nix/tools/loadtest.nix
+69 -27
View File
@@ -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"
{