feat(auth): ✨ implement PKCE security and endpoint fallback strategies #80
Conversation
- Integrate PKCE (Proof Key for Code Exchange) into the OAuth flow in `GoogleOAuthBase`, including state parameter validation for enhanced security. - Implement robust endpoint fallback chains (Production ↔ Daily Sandbox) in `AntigravityAuthBase` for project discovery and user onboarding. - Update `User-Agent` and `X-Goog-Api-Client` headers to align with recent Node.js client versions. - Refactor project ID extraction and onboarding polling logic to handle variable API response formats and errors more gracefully.
Updates authentication and provider headers to strictly mimic the official Antigravity Electron client and VSCode Cloud Shell editor. This resolves issues with onboarding and specific model availability. - Update `User-Agent` and `X-Goog-Api-Client` headers in `AntigravityAuthBase` and `AntigravityProvider` to match known working client signatures. - Change `CALLBACK_PATH` from `/oauthcallback` to `/oauth-callback` to match the native flow. - Set `ideType` to `IDE_UNSPECIFIED` in client metadata payloads for compatibility. - Remove extraneous headers from the OAuth token exchange request in `GoogleOAuthBase` to align with the reference TypeScript implementation.
|
I'm starting my review of the PKCE security implementation and endpoint fallback strategies. I'll take a close look at the authentication flow and the new polling logic in |
![mirrobot-agent[bot]](https://avatars.githubusercontent.com/in/2140342?s=60&v=4){kind=small}
There was a problem hiding this comment.
Overall Assessment
This PR introduces significant improvements to the authentication security and robustness of the Antigravity provider. The implementation of PKCE (RFC 7636) is a great security addition, and the endpoint fallback strategy for project discovery and onboarding makes the provider much more resilient to environment-specific issues.
The code is well-structured and includes helpful comments explaining the importance of specific headers and endpoint behaviors.
Architectural Feedback
- PKCE Implementation: The use of
S256code challenge method and the generation of thecode_verifierusingsecretsare correct and follow best practices. - Endpoint Fallback: Moving constants to
gemini_shared_utilsand implementing a consistent retry/fallback logic acrossloadCodeAssistandonboardUseris a solid architectural choice.
Key Suggestions
- State Validation: I recommend aborting the flow if the returned OAuth
statedoesn't match the expected verifier, rather than falling back to the original verifier. This ensures the CSRF protection remains effective. - Onboarding Timeout: Please verify if 60 seconds is sufficient for all onboarding scenarios, as some server-side operations can be slower depending on the region or project state.
Questions for the Author
- Was the reduction of the onboarding timeout from 5 minutes to 60 seconds driven by any specific performance metrics?
This review was generated by an AI assistant.
| # Poll for onboarding completion (up to 5 minutes) | ||
| for i in range(150): # 150 × 2s = 5 minutes | ||
| # Poll for onboarding completion (up to 60 seconds) | ||
| for i in range(30): # 30 × 2s = 60 seconds |
There was a problem hiding this comment.
The onboarding timeout has been reduced from 5 minutes (150 attempts) to 60 seconds (30 attempts). While 60 seconds is often enough, some Long Running Operations (LROs) for onboarding can occasionally take longer. Was this reduction based on specific telemetry or observed behavior?
There was a problem hiding this comment.
More than enough to verify
| if decoded_verifier != code_verifier: | ||
| lib_logger.warning( | ||
| "OAuth state verifier mismatch - possible CSRF attempt. " | ||
| "Using original verifier." | ||
| ) |
There was a problem hiding this comment.
If the received_state decodes to a verifier that doesn’t match the original code_verifier, this is a strong indicator of a CSRF attempt or a session mismatch. Instead of falling back to the original verifier and proceeding, it would be safer to raise an exception and abort the authentication flow.
| def _encode_oauth_state(self, code_verifier: str) -> str: | ||
| """ | ||
| Encode OAuth state parameter containing PKCE verifier. | ||
| The state parameter provides CSRF protection and carries the PKCE verifier | ||
| so it can be recovered after the OAuth callback. | ||
| Args: | ||
| code_verifier: The PKCE code verifier to encode | ||
| Returns: | ||
| Base64url-encoded state string | ||
| """ | ||
| import base64 | ||
|
|
||
| state_data = {"v": code_verifier} # Minimal - just verifier | ||
| json_bytes = json.dumps(state_data, separators=(",", ":")).encode("utf-8") | ||
| return base64.urlsafe_b64encode(json_bytes).decode("ascii").rstrip("=") |
There was a problem hiding this comment.
Encoding the code_verifier directly into the state parameter is an interesting approach to avoid local state management. However, since the state is not signed or encrypted, it doesn’t provide much beyond what PKCE already provides. Usually, state should be a random nonce validated against a local store. If local storage is not an option, consider adding a signature to the state to ensure it wasn’t tampered with, although PKCE itself mitigates authorization code injection.
GoogleOAuthBase, including state parameter validation for enhanced security.AntigravityAuthBasefor project discovery and user onboarding.User-AgentandX-Goog-Api-Clientheaders to align with recent Node.js client versions.Important
Implement PKCE security in OAuth flow, add endpoint fallback strategies, update headers, and refactor project ID extraction in
AntigravityAuthBaseandGoogleOAuthBase.GoogleOAuthBasefor OAuth flow, including state parameter validation._generate_pkce(),_encode_oauth_state(), and_decode_oauth_state()for PKCE support.AntigravityAuthBaseforloadCodeAssistandonboardUser.ANTIGRAVITY_LOAD_ENDPOINT_ORDERandANTIGRAVITY_ENDPOINT_FALLBACKSfor fallback logic.User-AgentandX-Goog-Api-Clientheaders inantigravity_auth_base.pyandgoogle_oauth_base.py.AntigravityAuthBaseto handle variable API response formats.This description was created by
for fa4d94b. You can customize this summary. It will automatically update as commits are pushed.