nix: add loadtest with unique JWTs

This loadtests the jwt decoding logic. For this it adds an optional `-k`(kind) parameter to `postgrest-loadtest` and `postgrest-loadtest-against`. Old kind (default): ``` postgrest-loadtest -k mixed postgrest-loadtest-against -k mixed ``` New kind: ``` postgrest-loadtest -k jwt postgrest-loadtest-against -k jwt ``` Internally it uses a dynamically generated targets file using python which looks like: ``` GET http://postgrest/authors_only Authorization: Bearer <jwt> GET http://postgrest/authors_only Authorization: Bearer <another-jwt> ... ``` Then this is used to run vegeta with the `-lazy` option.
2025-04-17 23:11:33 -05:00
parent a37ec1e1a5
commit 608f7ca45a
5 changed files with 122 additions and 14 deletions
@@ -0,0 +1,74 @@
+# 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
+
+SECRET = b"reallyreallyreallyreallyverysafe"
+URL = "http://postgrest"
+JWT_DURATION = 60
+TOTAL_TARGETS = 40000  # 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() -> str:
+    """Generate an HS256 JWT"""
+    # Header & payload
+    header = {"alg": "HS256", "typ": "JWT"}
+    now = int(time.time())
+    payload = {
+        "iat": now,
+        "exp": now + JWT_DURATION,
+        "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 "<header>.<payload>" 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}"
+
+
+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 = []
+    for _ in range(TOTAL_TARGETS):
+        token = generate_jwt()
+        lines.append(f"GET {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)
+
+    print(f"Generated {TOTAL_TARGETS} targets in {args.output}")
+
+
+if __name__ == "__main__":
+    main()
@@ -29,7 +29,6 @@ let
                                     -max-workers 1 \
                                     -workers 1 \
                                     -rate 0 \
-                                     -duration 60s \
                                     "''${_arg_leftovers[@]}"
      '';

@@ -41,6 +40,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 +62,29 @@ 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} -duration 60s -targets targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\""
+            ${vegeta}/bin/vegeta report -type=text "$_arg_output"
+            ;;
+
+        esac
      '';

  loadtestAgainst =
@@ -85,6 +102,7 @@ let
          '';
        args = [
          "ARG_POSITIONAL_INF([target], [Commit-ish reference to compare with], 1)"
+          "ARG_OPTIONAL_SINGLE([kind], [k], [Kind of loadtest], [mixed])"
        ];
        positionalCompletion =
          ''
@@ -99,7 +117,7 @@ let

        cat << EOF

-        Running loadtest on "$tgt"...
+        Running "$_arg_kind" loadtest on "$tgt"...

        EOF

@@ -108,7 +126,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

@@ -120,11 +138,11 @@ let

        cat << EOF

-        Running loadtest on HEAD...
+        Running $_arg_kind" loadtest on HEAD...

        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 +198,7 @@ let
          | ${toMarkdown}
      '';

+  genTargets = writers.writePython3 "postgrest-gen-loadtest-targets" { } (builtins.readFile ./generate_targets.py);
 in
 buildToolbox {
  name = "postgrest-loadtest";