nix(loadtest): remove pyjwt dependency
Let's use the same dependency when creating key material and when creating tokens, no need to carry multiple different dependencies and converting keys between them.
This commit is contained in:
@@ -14,41 +14,45 @@ import time
|
|||||||
import argparse
|
import argparse
|
||||||
import sys
|
import sys
|
||||||
import random
|
import random
|
||||||
import jwt
|
import jwcrypto.jwt as jwt
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
URL = "http://postgrest"
|
URL = "http://postgrest"
|
||||||
|
|
||||||
secret_key = b"reallyreallyreallyreallyverysafe"
|
secret_key = "reallyreallyreallyreallyverysafe"
|
||||||
|
|
||||||
|
|
||||||
def generate_target(
|
def generate_target(
|
||||||
now: int,
|
now: int,
|
||||||
exp_inc: Optional[int],
|
exp_inc: Optional[int],
|
||||||
rsa_private_key: Optional[jwt.algorithms.RSAAlgorithm],
|
rsa_private_key: Optional[jwt.JWK],
|
||||||
) -> list[str]:
|
) -> list[str]:
|
||||||
"""Generate a target using an HS256 or RS256 JWT"""
|
"""Generate a target using an HS256 or RS256 JWT"""
|
||||||
payload = {
|
headers = {
|
||||||
"sub": f"user_{random.getrandbits(32)}",
|
"sub": f"user_{random.getrandbits(32)}",
|
||||||
"iat": now,
|
"iat": now,
|
||||||
"role": "postgrest_test_author",
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if exp_inc is not None:
|
if exp_inc is not None:
|
||||||
payload["exp"] = now + exp_inc
|
headers["exp"] = now + exp_inc
|
||||||
|
|
||||||
|
claims = {
|
||||||
|
"role": "postgrest_test_author",
|
||||||
|
}
|
||||||
|
|
||||||
if rsa_private_key is None:
|
if rsa_private_key is None:
|
||||||
key = secret_key
|
key = jwt.JWK.from_password(secret_key)
|
||||||
alg = "HS256"
|
headers["alg"] = "HS256"
|
||||||
else:
|
else:
|
||||||
key = rsa_private_key
|
key = rsa_private_key
|
||||||
alg = "RS256"
|
headers["alg"] = "RS256"
|
||||||
token = jwt.encode(payload, key, alg)
|
token = jwt.JWT(headers, claims)
|
||||||
|
token.make_signed_token(key)
|
||||||
|
|
||||||
return [
|
return [
|
||||||
f"OPTIONS {URL}/authors_only?{alg}",
|
f"OPTIONS {URL}/authors_only?{headers["alg"]}",
|
||||||
f"Authorization: Bearer {token}",
|
f"Authorization: Bearer {token.serialize()}",
|
||||||
"", # blank line to separate requests
|
"", # blank line to separate requests
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -82,7 +86,7 @@ def main():
|
|||||||
|
|
||||||
targets_path = args.generated_path / "gen_targets.http"
|
targets_path = args.generated_path / "gen_targets.http"
|
||||||
|
|
||||||
rsa_private_key: Optional[jwt.algorithms.RSAAlgorithm] = None
|
rsa_private_key: Optional[jwt.JWK] = None
|
||||||
|
|
||||||
nsamples = 500 # per algorithm
|
nsamples = 500 # per algorithm
|
||||||
ntargets = 100000
|
ntargets = 100000
|
||||||
@@ -98,7 +102,7 @@ def main():
|
|||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
rsa_private_key = jwt.algorithms.RSAAlgorithm.from_jwk(private_key_data)
|
rsa_private_key = jwt.JWK.from_json(private_key_data)
|
||||||
except Exception as exc: # broad exception to capture parsing errors
|
except Exception as exc: # broad exception to capture parsing errors
|
||||||
err = f"Error loading RSA private key from {args.private_key_path}: " f"{exc}"
|
err = f"Error loading RSA private key from {args.private_key_path}: " f"{exc}"
|
||||||
print(err, file=sys.stderr)
|
print(err, file=sys.stderr)
|
||||||
|
|||||||
@@ -312,7 +312,7 @@ let
|
|||||||
genTargets =
|
genTargets =
|
||||||
writers.writePython3 "postgrest-gen-loadtest-targets"
|
writers.writePython3 "postgrest-gen-loadtest-targets"
|
||||||
{
|
{
|
||||||
libraries = [ python3Packages.pyjwt python3Packages.jwcrypto ];
|
libraries = [ python3Packages.jwcrypto ];
|
||||||
doCheck = false; # postgrest-style conflicts with this
|
doCheck = false; # postgrest-style conflicts with this
|
||||||
}
|
}
|
||||||
(builtins.readFile ./generate_targets.py);
|
(builtins.readFile ./generate_targets.py);
|
||||||
|
|||||||
Reference in New Issue
Block a user