From 2d1fe186e298e71b5fd576c65e2f12aaa820e98a Mon Sep 17 00:00:00 2001 From: Jonathan Melitski Date: Wed, 3 Sep 2025 13:22:01 -0400 Subject: [PATCH 1/4] Initial serial verification code, need to test --- Dockerfile | 2 +- src/auth.py | 44 ++++++++++++++++++++++++++------------------ src/config.py | 7 +++++-- 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/Dockerfile b/Dockerfile index fd9522e..f6c122d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM python:3.11-buster +FROM python:3.11-bullseye LABEL maintainer="Penn Labs" diff --git a/src/auth.py b/src/auth.py index 5a7d9cf..06035d0 100644 --- a/src/auth.py +++ b/src/auth.py @@ -9,19 +9,26 @@ JWKS_URL = settings.JWKS_URL -def get_jwk(): +def get_jwks(): if settings.JWKS_CACHE: - key = settings.JWKS_CACHE - return key + # Check to make sure we have a cached JWK for each key, otherwise refetch all. + missing = False + for key in settings.JWKS_URL.keys(): + if key not in settings.JWKS_CACHE: + missing = True + + if not missing: + return settings.JWKS_CACHE # Make a request to get the JWKS - try: - response = requests.get(JWKS_URL) - jwks = jwk.JWKSet.from_json(response.text) - settings.JWKS_CACHE = jwks - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) - + for key in settings.JWKS_URL: + try: + response = requests.get(JWKS_URL) + jwks = jwk.JWKSet.from_json(response.text) + settings.JWKS_CACHE[key] = jwks + except Exception as e: + del settings.JWKS_CACHE[key] + raise HTTPException(status_code=500, detail=str(e)) return settings.JWKS_CACHE @@ -41,11 +48,12 @@ def get_token_from_header(request: Request): def verify_jwt(token: str = Depends(get_token_from_header)): - try: - # Load the public key - public_key = get_jwk() - # Decode and verify the JWT - decoded_token = jwt.JWT(key=public_key, jwt=token) - return decoded_token.claims - except Exception as e: - raise HTTPException(status_code=401, detail=str(e)) + public_keys = get_jwks() + for key in public_keys: + try: + decoded_token = jwt.JWT(key=key, jwt=token) + return decoded_token.claims + except: + pass + + raise HTTPException(status_code=401, detail="Failed to verify JWT token against public keys array.") diff --git a/src/config.py b/src/config.py index f20c759..4412f60 100644 --- a/src/config.py +++ b/src/config.py @@ -11,8 +11,11 @@ class Config(BaseSettings): DATABASE_URL: PostgresDsn REDIS_URL: RedisDsn - JWKS_CACHE: JWKSet | None = None - JWKS_URL: str = "https://platform.pennlabs.org/identity/jwks/" + JWKS_CACHE: list[JWKSet] | None = None + JWKS_URL: dict[str, str] = { + "b2b":"https://platform.pennlabs.org/identity/jwks/", + "user":"https://platform.pennlabs.org/accounts/.well-known/jwks.json" + } SITE_DOMAIN: str = "analytics.pennlabs.org" From 5098a020dace79c28cfb3a39a3917322f7000a63 Mon Sep 17 00:00:00 2001 From: Jonathan Melitski Date: Wed, 3 Sep 2025 13:26:01 -0400 Subject: [PATCH 2/4] Lint --- src/auth.py | 7 ++----- src/config.py | 5 ++--- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/auth.py b/src/auth.py index 06035d0..7366581 100644 --- a/src/auth.py +++ b/src/auth.py @@ -16,10 +16,8 @@ def get_jwks(): for key in settings.JWKS_URL.keys(): if key not in settings.JWKS_CACHE: missing = True - if not missing: return settings.JWKS_CACHE - # Make a request to get the JWKS for key in settings.JWKS_URL: try: @@ -53,7 +51,6 @@ def verify_jwt(token: str = Depends(get_token_from_header)): try: decoded_token = jwt.JWT(key=key, jwt=token) return decoded_token.claims - except: + except Exception: pass - - raise HTTPException(status_code=401, detail="Failed to verify JWT token against public keys array.") + raise HTTPException(status_code=401, detail="Failed to verify JWT token") diff --git a/src/config.py b/src/config.py index 4412f60..1b0f056 100644 --- a/src/config.py +++ b/src/config.py @@ -13,10 +13,9 @@ class Config(BaseSettings): JWKS_CACHE: list[JWKSet] | None = None JWKS_URL: dict[str, str] = { - "b2b":"https://platform.pennlabs.org/identity/jwks/", - "user":"https://platform.pennlabs.org/accounts/.well-known/jwks.json" + "b2b": "https://platform.pennlabs.org/identity/jwks/", + "user": "https://platform.pennlabs.org/accounts/.well-known/jwks.json", } - SITE_DOMAIN: str = "analytics.pennlabs.org" ENVIRONMENT: Environment = Environment.PRODUCTION From c42a7f996702438acbfccf911fa80fdc78cbfdf7 Mon Sep 17 00:00:00 2001 From: Jonathan Melitski Date: Wed, 3 Sep 2025 14:05:26 -0400 Subject: [PATCH 3/4] fix bugs but performance is an issue --- src/auth.py | 11 ++++------- src/config.py | 2 +- tests/test_load.py | 5 +++-- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/auth.py b/src/auth.py index 7366581..bcca141 100644 --- a/src/auth.py +++ b/src/auth.py @@ -5,10 +5,6 @@ from src.config import settings -# The URL to the JWKS endpoint -JWKS_URL = settings.JWKS_URL - - def get_jwks(): if settings.JWKS_CACHE: # Check to make sure we have a cached JWK for each key, otherwise refetch all. @@ -17,15 +13,16 @@ def get_jwks(): if key not in settings.JWKS_CACHE: missing = True if not missing: + print("Can used cached keys") return settings.JWKS_CACHE # Make a request to get the JWKS for key in settings.JWKS_URL: try: - response = requests.get(JWKS_URL) + response = requests.get(settings.JWKS_URL[key]) jwks = jwk.JWKSet.from_json(response.text) settings.JWKS_CACHE[key] = jwks except Exception as e: - del settings.JWKS_CACHE[key] + print(str(e)) raise HTTPException(status_code=500, detail=str(e)) return settings.JWKS_CACHE @@ -49,7 +46,7 @@ def verify_jwt(token: str = Depends(get_token_from_header)): public_keys = get_jwks() for key in public_keys: try: - decoded_token = jwt.JWT(key=key, jwt=token) + decoded_token = jwt.JWT(key=public_keys[key], jwt=token) return decoded_token.claims except Exception: pass diff --git a/src/config.py b/src/config.py index 1b0f056..253ba21 100644 --- a/src/config.py +++ b/src/config.py @@ -11,7 +11,7 @@ class Config(BaseSettings): DATABASE_URL: PostgresDsn REDIS_URL: RedisDsn - JWKS_CACHE: list[JWKSet] | None = None + JWKS_CACHE: dict[str, JWKSet] = {} JWKS_URL: dict[str, str] = { "b2b": "https://platform.pennlabs.org/identity/jwks/", "user": "https://platform.pennlabs.org/accounts/.well-known/jwks.json", diff --git a/tests/test_load.py b/tests/test_load.py index 13eb980..5df4bc8 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -5,7 +5,8 @@ from datetime import datetime import requests -from test_token import get_tokens + +from tests.test_token import get_tokens # Runtime should be less that 3 seconds for most laptops @@ -19,7 +20,7 @@ def make_request(): access_token, _ = get_tokens() - url = "http://localhost:8000/analytics" + url = "http://localhost:80/analytics/" payload = json.dumps( { "product": random.randint(1, 10), From a7dae6f8ffb37219962806e5e3b38719cdbf8383 Mon Sep 17 00:00:00 2001 From: Jonathan Melitski Date: Wed, 3 Sep 2025 14:19:06 -0400 Subject: [PATCH 4/4] Improve testing --- src/auth.py | 1 - tests/test_load.py | 12 ++++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/src/auth.py b/src/auth.py index bcca141..7d43ce1 100644 --- a/src/auth.py +++ b/src/auth.py @@ -13,7 +13,6 @@ def get_jwks(): if key not in settings.JWKS_CACHE: missing = True if not missing: - print("Can used cached keys") return settings.JWKS_CACHE # Make a request to get the JWKS for key in settings.JWKS_URL: diff --git a/tests/test_load.py b/tests/test_load.py index 5df4bc8..11fac0a 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -17,10 +17,8 @@ THREADS = 16 -def make_request(): - access_token, _ = get_tokens() - - url = "http://localhost:80/analytics/" +def make_request(token: str): + url = "http://localhost:8000/analytics/" payload = json.dumps( { "product": random.randint(1, 10), @@ -42,7 +40,7 @@ def make_request(): ) headers = { "Content-Type": "application/json", - "Authorization": f"Bearer {access_token}", + "Authorization": f"Bearer {token}", } try: @@ -56,8 +54,10 @@ def make_request(): def run_threads(): with ThreadPoolExecutor(max_workers=THREADS) as executor: + # Fetch token once to improve performance + access_token, _ = get_tokens() for _ in range(NUMBER_OF_REQUESTS): - executor.submit(make_request) + executor.submit(make_request, access_token) def test_load():