diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 30a39a280..537579cbc 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -113,7 +113,7 @@ jobs: loadtest: strategy: matrix: - kind: ['mixed', 'jwt'] + kind: ['mixed', 'jwt-hs-50k', 'jwt-rsa-1k'] name: Loadtest runs-on: ubuntu-24.04 steps: diff --git a/.gitignore b/.gitignore index a4f750785..7794fdd3c 100644 --- a/.gitignore +++ b/.gitignore @@ -25,3 +25,4 @@ loadtest .history .docs-build gen_targets.http +gen_jwk.http diff --git a/nix/tools/generate_targets_rsa.py b/nix/tools/generate_targets_rsa.py new file mode 100644 index 000000000..3de7b209b --- /dev/null +++ b/nix/tools/generate_targets_rsa.py @@ -0,0 +1,68 @@ +# generates a file to be used by the vegeta load testing tool +# It runs TOTAL_TARGETS amount of requests with +# TOTAL_JWTS amount of JWTs randomly spread among them +import time +import argparse +import sys +import random +import jwt +import jwcrypto.jwk as jwk + +URL = "http://postgrest" +TOTAL_JWTS = 1000 +TOTAL_TARGETS = 50000 # tuned by hand to reduce result variance + +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() -> str: + """Generate an RS256 JWT""" + now = int(time.time()) + payload = { + "sub": f"user_{random.getrandbits(32)}", + "iat": now, + "role": "postgrest_test_author", + } + + return jwt.encode(payload, private_key, "RS256") + + +def main(): + parser = argparse.ArgumentParser( + description="Generate Vegeta targets with unique JWTs" + ) + parser.add_argument( + "output", + help="Path to write the generated targets file", + ) + parser.add_argument("jwk", help="Path to write the generated JWK") + args = parser.parse_args() + + tokens = [generate_jwt() for _ in range(TOTAL_JWTS)] + lines = [] + start_time = time.time() + + for i in range(TOTAL_TARGETS): + token = random.choice(tokens) + lines.append(f"OPTIONS {URL}/authors_only") + lines.append(f"Authorization: Bearer {token}") + lines.append("") # blank line to separate requests + + try: + with open(args.jwk, "w") as jwk, open(args.output, "w") as f: + jwk.write(public_key) + f.write("\n".join(lines)) + except IOError as e: + print(f"Error writing to {args.output}: {e}", file=sys.stderr) + sys.exit(1) + + elapsed = time.time() - start_time + print(f"Created {TOTAL_TARGETS} targets with unique JWTs", end=" ") + print(f"in {args.output} ({elapsed:.2f}s)") + + +if __name__ == "__main__": + main() diff --git a/nix/tools/loadtest.nix b/nix/tools/loadtest.nix index 7ac9f1cf8..02fee267e 100644 --- a/nix/tools/loadtest.nix +++ b/nix/tools/loadtest.nix @@ -41,8 +41,10 @@ let args = [ "ARG_OPTIONAL_SINGLE([output], [o], [Filename to dump json output to], [./loadtest/result.bin])" "ARG_OPTIONAL_SINGLE([testdir], [t], [Directory to load tests and fixtures from], [./test/load])" - "ARG_OPTIONAL_SINGLE([kind], [k], [Kind of loadtest (mixed: repeat mixed requests, jwt: run once over many requests with unique jwts)], [mixed])" - "ARG_TYPE_GROUP_SET([KIND], [KIND], [kind], [mixed,jwt])" + "ARG_OPTIONAL_SINGLE([kind], [k], [Kind of loadtest], [mixed])" + "ARG_OPTIONAL_SINGLE([jwtcache], [], [JWT Cache on/off], [on])" + "ARG_TYPE_GROUP_SET([KIND], [KIND], [kind], [mixed,jwt-hs-50k,jwt-rsa-1k])" + "ARG_TYPE_GROUP_SET([JWTCACHE], [JWTCACHE], [jwtcache], [on,off])" "ARG_LEFTOVERS([additional vegeta arguments])" ]; workingDir = "/"; @@ -64,9 +66,30 @@ let abs_output="$(realpath "$_arg_output")" case "$_arg_kind" in - jwt) + jwt-hs-50k) - ${genTargets} "$_arg_testdir"/gen_targets.http + ${genTargets ./generate_targets.py} "$_arg_testdir"/gen_targets.http + + if [ "$_arg_jwtcache" = "off" ]; then + export PGRST_JWT_CACHE_MAX_LIFETIME="0" + fi + + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withPgrst} \ + sh -c "cd \"$_arg_testdir\" && ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" + ${vegeta}/bin/vegeta report -type=text "$_arg_output" + ;; + + jwt-rsa-1k) + + ${genTargets ./generate_targets_rsa.py} "$_arg_testdir"/gen_targets.http "$_arg_testdir"/gen_jwk.http + + export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.http" + + if [ "$_arg_jwtcache" = "off" ]; then + export PGRST_JWT_CACHE_MAX_LIFETIME="0" + fi # shellcheck disable=SC2145 ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ @@ -200,7 +223,12 @@ let | ${toMarkdown} ''; - genTargets = writers.writePython3 "postgrest-gen-loadtest-targets" { } (builtins.readFile ./generate_targets.py); + genTargets = genTargetsScript: + writers.writePython3 "postgrest-gen-loadtest-targets" + { + libraries = [ python3Packages.pyjwt python3Packages.jwcrypto ]; + } + (builtins.readFile genTargetsScript); in buildToolbox { name = "postgrest-loadtest";