From 58237be608928aacc9add090645b2b2c6b7fadd0 Mon Sep 17 00:00:00 2001 From: steve-chavez Date: Sun, 20 Apr 2025 14:25:41 -0500 Subject: [PATCH] test: add loadtest for async purge of JWT cache --- .github/workflows/test.yaml | 5 +- .gitignore | 1 + nix/tools/generate_targets.py | 94 +++++++++++++++++++++++++++++++++++ nix/tools/loadtest.nix | 39 +++++++++++---- test/load/fixtures.sql | 13 ++++- 5 files changed, 141 insertions(+), 11 deletions(-) create mode 100644 nix/tools/generate_targets.py diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index ace8ce26d..cf2273fad 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -111,6 +111,9 @@ jobs: loadtest: + strategy: + matrix: + kind: ['mixed', 'jwt'] name: Loadtest runs-on: ubuntu-24.04 steps: @@ -128,7 +131,7 @@ jobs: prefix: v - name: Run loadtest run: | - postgrest-loadtest-against main ${{ steps.get-latest-tag.outputs.tag }} + postgrest-loadtest-against -k ${{ matrix.kind }} main ${{ steps.get-latest-tag.outputs.tag }} postgrest-loadtest-report >> "$GITHUB_STEP_SUMMARY" flake: diff --git a/.gitignore b/.gitignore index 986da1505..a4f750785 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,4 @@ coverage loadtest .history .docs-build +gen_targets.http diff --git a/nix/tools/generate_targets.py b/nix/tools/generate_targets.py new file mode 100644 index 000000000..347626ab5 --- /dev/null +++ b/nix/tools/generate_targets.py @@ -0,0 +1,94 @@ +# generates a file to be used by the vegeta load testing tool +import time +import hmac +import hashlib +import base64 +import json +import argparse +import sys +import random + +SECRET = b"reallyreallyreallyreallyverysafe" +URL = "http://postgrest" +JWT_DURATION = 120 +TOTAL_TARGETS = 50000 # tuned by hand to reduce result variance + + +def base64url_encode(data: bytes) -> str: + """URL-safe Base64 encode without padding.""" + return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii") + + +def generate_jwt(exp_inc: int) -> str: + """Generate an HS256 JWT""" + # Header & payload + header = {"alg": "HS256", "typ": "JWT"} + now = int(time.time()) + payload = { + "sub": f"user_{random.getrandbits(32)}", + "iat": now, + "exp": now + exp_inc, + "role": "postgrest_test_author", + } + + # Encode to JSON and then to Base64URL + header_b = json.dumps(header, separators=(",", ":")).encode() + payload_b = json.dumps(payload, separators=(",", ":")).encode() + header_b64 = base64url_encode(header_b) + payload_b64 = base64url_encode(payload_b) + + # Sign (HMAC‑SHA256) the "
." string + signing_input = f"{header_b64}.{payload_b64}".encode() + signature = hmac.new(SECRET, signing_input, hashlib.sha256).digest() + signature_b64 = base64url_encode(signature) + + return f"{header_b64}.{payload_b64}.{signature_b64}" + + +# We want to ensure 401 Unauthorized responses don't happen during +# JWT validation, this can happen when the jwt `exp` is too short. +# At the same time, we want to ensure the `exp` is not too big, +# so expires will occur and postgREST will have to clean cached expired JWTs. +def estimate_adequate_jwt_exp_increase(iteration: int) -> int: + # estimated time takes to build and run postgrest itself + build_run_postgrest_time = 2 + # estimated time it takes to generate the targets file + file_generation_time = TOTAL_TARGETS // (10**-5) + # estimated exp time so some JWTs will expire + dynamic_exp_inc = iteration // 1000 + + return build_run_postgrest_time + file_generation_time + dynamic_exp_inc + + +def main(): + parser = argparse.ArgumentParser( + description="Generate Vegeta targets with unique JWTs" + ) + parser.add_argument( + "output", + help="Path to write the generated targets file", + ) + args = parser.parse_args() + + lines = [] + start_time = time.time() + + for i in range(TOTAL_TARGETS): + token = generate_jwt(estimate_adequate_jwt_exp_increase(i)) + lines.append(f"OPTIONS {URL}/authors_only") + lines.append(f"Authorization: Bearer {token}") + lines.append("") # blank line to separate requests + + try: + with open(args.output, "w") as f: + 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 in {args.output} ({elapsed:.2f}s)") + + +if __name__ == "__main__": + main() diff --git a/nix/tools/loadtest.nix b/nix/tools/loadtest.nix index 79347bacb..7ac9f1cf8 100644 --- a/nix/tools/loadtest.nix +++ b/nix/tools/loadtest.nix @@ -41,6 +41,8 @@ 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_LEFTOVERS([additional vegeta arguments])" ]; workingDir = "/"; @@ -61,13 +63,30 @@ let mkdir -p "$(dirname "$_arg_output")" abs_output="$(realpath "$_arg_output")" - # shellcheck disable=SC2145 - ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ - ${withTools.withSlowPg} \ - ${withTools.withPgrst} \ - ${withTools.withSlowPgrst} \ - sh -c "cd \"$_arg_testdir\" && ${runner} -targets targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" - ${vegeta}/bin/vegeta report -type=text "$_arg_output" + case "$_arg_kind" in + jwt) + + ${genTargets} "$_arg_testdir"/gen_targets.http + + # 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" + ;; + + *) + + # shellcheck disable=SC2145 + ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \ + ${withTools.withSlowPg} \ + ${withTools.withPgrst} \ + ${withTools.withSlowPgrst} \ + sh -c "cd \"$_arg_testdir\" && ${runner} -targets targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\"" + ${vegeta}/bin/vegeta report -type=text "$_arg_output" + ;; + + esac ''; loadtestAgainst = @@ -85,6 +104,7 @@ let ''; args = [ "ARG_POSITIONAL_INF([target], [Commit-ish reference to compare with], 1)" + "ARG_OPTIONAL_SINGLE([kind], [k], [Kind of loadtest], [mixed])" ]; positionalCompletion = '' @@ -108,7 +128,7 @@ let # Save the results in the current working tree, too, # otherwise they'd be lost in the temporary working tree # created by withTools.withGit. - ${withTools.withGit} "$tgt" ${loadtest} --output "$PWD/loadtest/$tgt.bin" --testdir "$PWD/test/load" + ${withTools.withGit} "$tgt" ${loadtest} -k "$_arg_kind" --output "$PWD/loadtest/$tgt.bin" --testdir "$PWD/test/load" cat << EOF @@ -124,7 +144,7 @@ let EOF - ${loadtest} --output "$PWD/loadtest/head.bin" --testdir "$PWD/test/load" + ${loadtest} -k "$_arg_kind" --output "$PWD/loadtest/head.bin" --testdir "$PWD/test/load" cat << EOF @@ -180,6 +200,7 @@ let | ${toMarkdown} ''; + genTargets = writers.writePython3 "postgrest-gen-loadtest-targets" { } (builtins.readFile ./generate_targets.py); in buildToolbox { name = "postgrest-loadtest"; diff --git a/test/load/fixtures.sql b/test/load/fixtures.sql index eee7dfab5..3622dcf78 100644 --- a/test/load/fixtures.sql +++ b/test/load/fixtures.sql @@ -1,5 +1,7 @@ CREATE ROLE postgrest_test_anonymous; +CREATE ROLE postgrest_test_author; GRANT postgrest_test_anonymous TO :PGUSER; +GRANT postgrest_test_author TO :PGUSER; CREATE SCHEMA test; -- PUT+PATCH target needs one record and column to modify @@ -31,10 +33,19 @@ CREATE TABLE test.roles ( character TEXT ); + +CREATE TABLE test.authors_only (); + CREATE FUNCTION test.call_me (name TEXT) RETURNS TEXT STABLE LANGUAGE SQL AS $$ SELECT 'Hello ' || name || ', how are you?'; $$; -GRANT USAGE ON SCHEMA test TO postgrest_test_anonymous; +GRANT USAGE ON SCHEMA test TO postgrest_test_anonymous, postgrest_test_author; GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA test TO postgrest_test_anonymous; + +REVOKE ALL PRIVILEGES ON TABLE + authors_only +FROM postgrest_test_anonymous; + +GRANT ALL ON TABLE authors_only TO postgrest_test_author;