Token verification with previous public token #83
Conversation
tmikula-dev
requested review from
ABMC831,
Zejnilovic,
oto-macenauer-absa and
petr-pokorny-absa
as code owners
WalkthroughReplaces single-key JWT handling with a HandlerToken-based flow: config key renamed to Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant EventGate
participant LoginService
participant HandlerToken as "HandlerToken"
rect rgb(230, 245, 255)
Note over EventGate,LoginService: Key fetch & store
EventGate->>LoginService: GET TOKEN_PUBLIC_KEYS_URL (/token/public-keys)
LoginService-->>EventGate: { keys: [ {key: K1}, {key: K0} ] } or { key: K0 }
EventGate->>HandlerToken: load_public_keys() -> parse & store RSA keys
end
Client->>EventGate: POST /event (Authorization: Bearer <token>)
EventGate->>HandlerToken: extract_token(headers)
EventGate->>HandlerToken: decode_jwt(token)
rect rgb(255, 250, 230)
Note over HandlerToken: Try keys sequentially
HandlerToken->>HandlerToken: verify with key[0] (may fail)
HandlerToken->>HandlerToken: verify with key[1] (may succeed)
end
alt verification succeeds
HandlerToken-->>EventGate: claims
EventGate-->>Client: 200 OK
else all keys fail
HandlerToken-->>EventGate: verification error
EventGate-->>Client: 401 Unauthorized
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 2 inconclusive)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Tip 📝 Customizable high-level summaries are now available in beta!You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.
Example instruction:
Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
README.md(2 hunks)conf/config.json(1 hunks)src/event_gate_lambda.py(4 hunks)tests/test_conf_validation.py(1 hunks)tests/test_event_gate_lambda.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tests/test_event_gate_lambda.py (1)
src/event_gate_lambda.py (2)
lambda_handler(274-302)decode_jwt_all(225-236)
🪛 GitHub Actions: Python Check
src/event_gate_lambda.py
[error] 103-103: mypy: List comprehension has incompatible type List[DHPublicKey | DSAPublicKey | RSAPublicKey | EllipticCurvePublicKey | Ed25519PublicKey | Ed448PublicKey | X25519PublicKey | X448PublicKey]; expected List[RSAPublicKey]
🪛 Ruff (0.14.4)
tests/test_event_gate_lambda.py
384-384: Unused function argument: token
(ARG001)
384-384: Unused function argument: algorithms
(ARG001)
386-386: Avoid specifying long messages outside the exception class
(TRY003)
399-399: Unused function argument: token
(ARG001)
399-399: Unused function argument: key
(ARG001)
399-399: Unused function argument: algorithms
(ARG001)
400-400: Avoid specifying long messages outside the exception class
(TRY003)
src/event_gate_lambda.py
89-89: Probable use of requests call with verify=False disabling SSL certificate checks
(S501)
100-100: Abstract raise to an inner function
(TRY301)
100-100: Avoid specifying long messages outside the exception class
(TRY003)
108-108: Avoid specifying long messages outside the exception class
(TRY003)
236-236: Avoid specifying long messages outside the exception class
(TRY003)
|
Trivy has completed a full security repository scan ✅ You can find the analysis results for this PR branch on this overview.
|
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
src/event_gate_lambda.py (2)
102-104: Runtime type validation is still missing.Using
cast()satisfies the type checker but doesn't prevent non-RSA keys from enteringTOKEN_PUBLIC_KEYS, which will cause runtime failures injwt.decode. The previous review's concern remains valid: if the identity service serves an EC/DSA key during rotation, EventGate will accept it during initialization but fail unpredictably during token verification.Replace the list comprehension with an explicit loop that validates each key's type at runtime, as suggested in the previous review.
106-108: Add TypeError to exception handling.The exception handler should catch
TypeErrorto handle the runtime type validation suggested in the previous comment, ensuring consistent error reporting for all key initialization failures.Apply this diff:
-except (requests.RequestException, ValueError, KeyError, UnsupportedAlgorithm) as exc: +except (requests.RequestException, ValueError, KeyError, UnsupportedAlgorithm, TypeError) as exc: logger.exception("Failed to fetch or deserialize token public key from %s", TOKEN_PUBLIC_KEYS_URL) raise RuntimeError("Token public key initialization failed") from exc
🧹 Nitpick comments (1)
src/event_gate_lambda.py (1)
225-237: Consider adding debug logging for key rotation observability.The multi-key verification logic is correct and implements the expected behavior. However, for operational observability during key rotations, consider logging when verification fails with one key and falls back to another.
Apply this diff to add debug logging:
def decode_jwt_all(token_encoded: str) -> Dict[str, Any]: """Decode JWT using any of the loaded public keys. Args: token_encoded: Encoded bearer JWT token string. """ - for public_key in TOKEN_PUBLIC_KEYS: + for idx, public_key in enumerate(TOKEN_PUBLIC_KEYS): try: return jwt.decode(token_encoded, public_key, algorithms=["RS256"]) except jwt.PyJWTError: + logger.debug("Token verification failed with public key %d, trying next", idx) continue raise jwt.PyJWTError("Verification failed for all public keys")This would help operators confirm that key rotation is working as expected when monitoring logs.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/event_gate_lambda.py(4 hunks)
🧰 Additional context used
🪛 Ruff (0.14.4)
src/event_gate_lambda.py
89-89: Probable use of requests call with verify=False disabling SSL certificate checks
(S501)
100-100: Abstract raise to an inner function
(TRY301)
100-100: Avoid specifying long messages outside the exception class
(TRY003)
108-108: Avoid specifying long messages outside the exception class
(TRY003)
236-236: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (3)
src/event_gate_lambda.py (3)
23-23: Imports support multi-key verification.The new imports (
cast,RSAPublicKey) enable the multi-key JWT verification feature, aligning with the PR objectives.Also applies to: 31-31
84-86: Backward-compatible configuration loading.The fallback from
token_public_keys_urltotoken_public_key_urlensures compatibility during migration, which aligns with the PR's objective to support key rotation without breaking existing deployments.
183-183: Multi-key verification correctly integrated.The switch to
decode_jwt_all()enables verification against multiple public keys, fulfilling the PR's core objective of supporting key rotation without service interruption.
jakipatryk
left a comment
jakipatryk
left a comment
There was a problem hiding this comment.
I'd add a scheduled update of the key set as well, as there is also a risk from the other direction (EventGate not yet having the new key if single Lambda instance exists for too long).
Login Service does the rotation, attaches the new key to the key set, and then for 30min doesn't yet use it for signing JWTs. So updating the key set on EventGate side with interval <30 min should be the optimal solution.
oto-macenauer-absa
left a comment
oto-macenauer-absa
left a comment
There was a problem hiding this comment.
looks good, not sure if we need to keep the legacy behavior for the keys on the other hand it's better to be on the safe side
src/event_gate_lambda.py
Outdated
| TOKEN_PROVIDER_URL = CONFIG["token_provider_url"] | ||
| # Add timeout to avoid hanging requests; wrap in robust error handling so failures are explicit | ||
| # Initialize token public keys | ||
| TOKEN_PROVIDER_URL = CONFIG.get("token_provider_url") |
There was a problem hiding this comment.
put the token handling to separate module pls
There was a problem hiding this comment.
This was change has a bigger impact, please see it in commit: d23819c. The project is now better readable. Need a thorough revision please.
|
Trivy has completed a full security repository scan ✅ You can find the analysis results for this PR branch on this overview.
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/handlers/test_handler_token.py (1)
85-91: Consider adding a test for theAuthorization: Bearer <token>flow.The current tests cover the direct
Bearerheader and empty headers, but the standardAuthorizationheader with Bearer scheme (the more common pattern) isn't explicitly tested here.def test_extract_token_authorization_header(): token = HandlerToken.extract_token({"Authorization": "Bearer mytoken123"}) assert token == "mytoken123" def test_extract_token_authorization_header_case_insensitive(): token = HandlerToken.extract_token({"authorization": "bearer MYTOKEN"}) assert token == "MYTOKEN"src/handlers/handler_token.py (2)
44-47: Add validation for missingpublic_keys_url.If neither
TOKEN_PUBLIC_KEYS_URLnorTOKEN_PUBLIC_KEY_URLis present in the config,public_keys_urlwill beNone. This will causeload_public_keysto fail with an unclear error whenrequests.get(None, ...)is called.Consider validating during initialization:
def __init__(self, config): self.provider_url: str = config.get(TOKEN_PROVIDER_URL, "") self.public_keys_url: str = config.get(TOKEN_PUBLIC_KEYS_URL) or config.get(TOKEN_PUBLIC_KEY_URL) + if not self.public_keys_url: + raise ValueError("Missing required config: token_public_keys_url or token_public_key_url") self.public_keys: list[RSAPublicKey] = []
94-100: Consider differentiatingExpiredSignatureErrorfor faster failure.Currently, all
PyJWTErrorexceptions (includingExpiredSignatureError) cause the loop to try the next key. However, if a token is expired, it will be expired regardless of which key is used. You could short-circuit onExpiredSignatureErrorto avoid unnecessary iterations.for public_key in self.public_keys: try: return jwt.decode(token_encoded, public_key, algorithms=["RS256"]) + except jwt.ExpiredSignatureError: + raise # Expired tokens won't validate with any key except jwt.PyJWTError: continue raise jwt.PyJWTError("Verification failed for all public keys")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
src/event_gate_lambda.py(5 hunks)src/handlers/__init__.py(1 hunks)src/handlers/handler_token.py(1 hunks)src/utils/constants.py(1 hunks)tests/conftest.py(1 hunks)tests/handlers/__init__.py(1 hunks)tests/handlers/test_handler_token.py(1 hunks)tests/test_event_gate_lambda.py(1 hunks)tests/test_event_gate_lambda_local_access.py(1 hunks)tests/utils/test_extract_token.py(0 hunks)
💤 Files with no reviewable changes (1)
- tests/utils/test_extract_token.py
✅ Files skipped from review due to trivial changes (2)
- src/handlers/init.py
- tests/handlers/init.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_event_gate_lambda.py
🧰 Additional context used
🧬 Code graph analysis (3)
tests/handlers/test_handler_token.py (2)
src/handlers/handler_token.py (3)
HandlerToken(39-142)decode_jwt(84-100)extract_token(110-142)src/event_gate_lambda.py (1)
lambda_handler(192-220)
tests/conftest.py (2)
src/writers/writer_kafka.py (1)
KafkaException(35-36)tests/test_event_gate_lambda_local_access.py (1)
Bucket(46-47)
src/event_gate_lambda.py (1)
src/handlers/handler_token.py (5)
HandlerToken(39-142)load_public_keys(49-82)decode_jwt(84-100)get_token(102-107)extract_token(110-142)
🪛 Ruff (0.14.5)
tests/handlers/test_handler_token.py
61-61: Unused function argument: token
(ARG001)
61-61: Unused function argument: algorithms
(ARG001)
61-61: Unused noqa directive (non-enabled: D401)
Remove unused noqa directive
(RUF100)
63-63: Avoid specifying long messages outside the exception class
(TRY003)
76-76: Unused function argument: token
(ARG001)
76-76: Unused function argument: key
(ARG001)
76-76: Unused function argument: algorithms
(ARG001)
76-76: Unused noqa directive (non-enabled: D401)
Remove unused noqa directive
(RUF100)
77-77: Avoid specifying long messages outside the exception class
(TRY003)
91-91: Possible hardcoded password assigned to: "token"
(S105)
src/handlers/handler_token.py
60-60: Probable use of requests call with verify=False disabling SSL certificate checks
(S501)
72-72: Abstract raise to an inner function
(TRY301)
72-72: Avoid specifying long messages outside the exception class
(TRY003)
79-79: Consider moving this statement to an else block
(TRY300)
82-82: Avoid specifying long messages outside the exception class
(TRY003)
100-100: Avoid specifying long messages outside the exception class
(TRY003)
tests/conftest.py
61-61: Unused noqa directive (non-enabled: D401)
Remove unused noqa directive
(RUF100)
97-97: Unused noqa directive (non-enabled: D401)
Remove unused noqa directive
(RUF100)
101-101: Unused noqa directive (non-enabled: D401)
Remove unused noqa directive
(RUF100)
src/utils/constants.py
22-22: Possible hardcoded password assigned to: "TOKEN_PROVIDER_URL"
(S105)
23-23: Possible hardcoded password assigned to: "TOKEN_PUBLIC_KEY_URL"
(S105)
24-24: Possible hardcoded password assigned to: "TOKEN_PUBLIC_KEYS_URL"
(S105)
🔇 Additional comments (14)
src/utils/constants.py (1)
21-24: LGTM!The constants are well-defined config key names. The static analysis warnings (S105) are false positives—these are dictionary keys for configuration lookup, not hardcoded secrets.
tests/test_event_gate_lambda_local_access.py (1)
57-57: LGTM on the attribute rename.The change from
CONFIGtoconfigaligns with the module refactoring.However, note that
local_configat line 15 still uses"token_public_key_url"instead of the new"token_public_keys_url". This works due to the fallback inHandlerToken.__init__, but consider updating it for consistency with the renamed config key.tests/handlers/test_handler_token.py (2)
25-51: Well-structured tests for core token handling.The tests cover the key scenarios: redirect, expired token, and multi-key verification. The mocking approach using
event_gate_modulefixture is clean.
54-82: Good coverage of multi-key JWT verification logic.These tests validate the core feature: falling back to the second key when the first fails, and raising an aggregate error when all keys fail.
tests/conftest.py (2)
33-121: Well-designed centralized fixture for module testing.The
event_gate_modulefixture effectively isolates external dependencies and provides a clean testing environment. The approach of conditionally creating dummy modules only when the real ones are unavailable is pragmatic.One minor observation: the fixture uses both
started_patcheslist andExitStackfor cleanup. This works correctly, but theExitStackcould handle all cleanup uniformly if the patches were entered viaexit_stack.enter_context(p)instead ofp.start().
124-149: Clean helper fixtures.
make_eventandvalid_payloadprovide good reusable test utilities.src/handlers/handler_token.py (1)
109-142: Well-implemented token extraction with robust header handling.The case-insensitive header normalization and support for both direct
Bearerheader and standardAuthorization: Bearerscheme is well thought out.src/event_gate_lambda.py (7)
30-30: LGTM!Clean import of the new
HandlerTokenmodule, which consolidates token-related functionality as requested in past reviews.
64-64: Consistent naming convention applied.The rename from
CONFIGtoconfigaligns with Python naming conventions (lowercase for module-level variables that aren't true constants). Ensure any external references to this module's config are updated if applicable.
70-77: LGTM!Config access updated consistently to use the renamed
configvariable.
80-86: Clean initialization with appropriate fail-fast behavior.The chained
HandlerToken(config).load_public_keys()is clean and ensures the Lambda fails at cold start if public keys cannot be loaded—appropriate for a security-critical path. The handler'sload_public_keys()already logs and raises a descriptiveRuntimeErroron failure.Note: Keys are loaded once at module initialization. If keys rotate while the Lambda instance is warm, it won't pick up new keys until the instance is recycled. This is typically acceptable for Lambda's lifecycle, but worth documenting if not already covered.
150-152: Core change for key rotation support looks correct.The
decode_jwt()method inHandlerTokeniterates through all loaded public keys, enabling seamless verification during key rotation. The exception handling correctly catchesjwt.PyJWTErrorwhen all keys fail verification.
202-202: LGTM!Delegation to
handler_token.get_token()is clean and maintains the same 303 redirect behavior.
213-213: LGTM!Delegation to the static
extract_token()method is appropriate. The method handles both theBearerheader andAuthorization: Bearer <token>formats with proper case-insensitive lookup.
| try: | ||
| response_json = requests.get(self.public_keys_url, verify=False, timeout=5).json() |
There was a problem hiding this comment.
Security concern: SSL certificate verification is disabled.
Using verify=False disables SSL certificate verification, making this request vulnerable to man-in-the-middle attacks. This is particularly concerning when fetching cryptographic public keys, as an attacker could inject malicious keys.
Consider making SSL verification configurable or defaulting to verify=True:
- response_json = requests.get(self.public_keys_url, verify=False, timeout=5).json()
+ response_json = requests.get(self.public_keys_url, timeout=5).json()If there's a legitimate need to disable verification in certain environments (e.g., development with self-signed certs), consider making it configurable via an environment variable or config parameter rather than unconditionally disabling it.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| response_json = requests.get(self.public_keys_url, verify=False, timeout=5).json() | |
| try: | |
| response_json = requests.get(self.public_keys_url, timeout=5).json() |
🧰 Tools
🪛 Ruff (0.14.5)
60-60: Probable use of requests call with verify=False disabling SSL certificate checks
(S501)
🤖 Prompt for AI Agents
In src/handlers/handler_token.py around lines 59-60, the requests.get call
disables SSL certificate verification (verify=False) which is insecure; change
it to use a configurable verification flag (e.g., read from an environment
variable or config with default True) and pass that flag to requests.get
(verify=<config_flag>) instead of False, ensure the default is True, add a log
warning if verification is explicitly disabled, and do not leave verify=False
hard-coded in the repository.
|
@jakipatryk Thank you for the comment! I did create an issue describing your comment, which is now in our ToDo section with a high priority #87. I will send you the solution for revision, when it is ready to. |
| """ | ||
|
|
||
| # Token related constants | ||
| TOKEN_PROVIDER_URL = "token_provider_url" |
There was a problem hiding this comment.
Maybe it could be more explicitly said that these are constants for config keys? On the first glance, it looks like the constant holds actual URL.
| continue | ||
| raise jwt.PyJWTError("Verification failed for all public keys") | ||
|
|
||
| def get_token(self) -> Dict[str, Any]: |
There was a problem hiding this comment.
get_token_provider_info maybe, or something similar? get_token suggests actual token of some kind is returned, at least to me.
Overview
Adding a logic for verification to solve a new pair of public token to fetch current as well as previous one.
Release Notes:
Related
Closes #65
Summary by CodeRabbit
New Features
Documentation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.