Wolfgang Walther
e83a25a698
Merge the same code into genTargets, which removes the need to save the private key to disk and read it again.
131 lines
3.7 KiB
Python
131 lines
3.7 KiB
Python
# generates a file to be used by the vegeta load testing tool
|
|
|
|
# It includes a worst case scenario for the JWT cache:
|
|
# - all requests will have a unique JWT so no cache hits
|
|
# - all jwts have an expiration that will be long enough to be
|
|
# valid at time of request but short enough that already
|
|
# validated jwts will expire later during the loadtest run
|
|
# - the above guarantees JWT cache purging will happen
|
|
# - we want this to track resource consumption in the worst case
|
|
|
|
# And a more normal scenario where non-expiring JWTs are picked
|
|
# from an array
|
|
import time
|
|
import argparse
|
|
import random
|
|
import jwcrypto.jwt as jwt
|
|
from typing import Optional
|
|
from pathlib import Path
|
|
|
|
URL = "http://postgrest"
|
|
|
|
secret_key = "reallyreallyreallyreallyverysafe"
|
|
|
|
|
|
def generate_target(
|
|
now: int,
|
|
exp_inc: Optional[int],
|
|
key: jwt.JWK,
|
|
) -> list[str]:
|
|
"""Generate a target using an HS256 or RS256 JWT"""
|
|
headers = {
|
|
"sub": f"user_{random.getrandbits(32)}",
|
|
"iat": now,
|
|
}
|
|
|
|
if exp_inc is not None:
|
|
headers["exp"] = now + exp_inc
|
|
|
|
claims = {
|
|
"role": "postgrest_test_author",
|
|
}
|
|
|
|
headers["alg"] = "RS256" if key.get("kty") == "RSA" else "HS256"
|
|
|
|
token = jwt.JWT(headers, claims)
|
|
token.make_signed_token(key)
|
|
|
|
return [
|
|
f"OPTIONS {URL}/authors_only?{headers["alg"]}",
|
|
f"Authorization: Bearer {token.serialize()}",
|
|
"", # blank line to separate requests
|
|
]
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate Vegeta targets with unique JWTs"
|
|
)
|
|
parser.add_argument(
|
|
"generated_path",
|
|
metavar="GENERATED_PATH",
|
|
help="Path to write the generated files",
|
|
type=Path,
|
|
)
|
|
parser.add_argument(
|
|
"--worst",
|
|
dest="worst",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="Generate worst case targets for a JWT cache",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
targets_path = args.generated_path / "gen_targets.http"
|
|
|
|
hs = jwt.JWK.from_password(secret_key)
|
|
rsa = jwt.JWK.generate(kty="RSA", size=4096)
|
|
|
|
jwks = jwt.JWKSet()
|
|
jwks.add(hs)
|
|
jwks.add(rsa)
|
|
|
|
jwks_path = args.generated_path / "gen_jwks.json"
|
|
|
|
# Technically, this exports the private keys, because HS does not have the concept
|
|
# of a public key. This is not a problem for tests, though, PostgREST can verify
|
|
# tokens with the private key just as well.
|
|
jwks_path.write_text(jwks.export())
|
|
print(f"Created JWKSet on {jwks_path}")
|
|
|
|
nsamples = 500 # per algorithm
|
|
ntargets = 100000
|
|
|
|
print(f"Generating {ntargets} targets...")
|
|
|
|
now = int(time.time())
|
|
|
|
lines = []
|
|
|
|
# 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 needs to
|
|
# clean cached expired JWTs
|
|
if args.worst:
|
|
# estimated time it takes to run postgrest itself
|
|
run_postgrest_time = 2
|
|
|
|
for i in range(ntargets):
|
|
target = generate_target(
|
|
now, run_postgrest_time + i // 1000, rsa if i % 2 == 0 else hs
|
|
)
|
|
lines.extend(target)
|
|
|
|
else:
|
|
hs_targets = [generate_target(now, None, hs) for _ in range(nsamples)]
|
|
rsa_targets = [generate_target(now, None, rsa) for _ in range(nsamples)]
|
|
for i in range(ntargets):
|
|
target = random.choice(hs_targets if i % 2 == 0 else rsa_targets)
|
|
lines.extend(target)
|
|
|
|
with open(targets_path, "w") as f:
|
|
f.write("\n".join(lines))
|
|
|
|
print(f"Created {ntargets} targets", end=" ")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|