feat: add AWS Identity Center (IDC) authentication support #57
Conversation
Add persistRefreshedAuth function to write refreshed tokens back to the auth file after inline token refresh. This prevents repeated token refreshes on every request when the token expires. Changes: - Add persistRefreshedAuth() to kiro_executor.go - Call persist after all token refresh paths (401, 403, pre-request) - Remove unused log import from sdk/auth/kiro.go
Summary of ChangesHello @PancakeZik, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the authentication capabilities by integrating AWS Identity Center (IDC) support. It provides a robust alternative to Builder ID, catering to organizational SSO requirements. The changes encompass the entire authentication lifecycle, from initial login and client registration to region-aware token refreshing and persistent storage of credentials. This ensures a more flexible and secure authentication experience for users, aligning with modern AWS identity management practices. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Closes #56 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for AWS Identity Center (IDC) authentication, providing an alternative to Builder ID. The changes are extensive, touching authentication flows, request execution, and token management. My review focuses on improving maintainability by reducing code duplication, enhancing error handling robustness, and addressing hardcoded values. I've identified several areas where logic is repeated across different files and functions, such as token refresh and persistence, which could be centralized. Additionally, I've suggested improvements to make error handling less brittle and to move magic strings into constants.
internal/auth/kiro/sso_oidc.go
Outdated
| input, _ := reader.ReadString('\n') | ||
| input = strings.TrimSpace(input) | ||
| if input == "" { | ||
| return defaultValue | ||
| } | ||
| return input |
There was a problem hiding this comment.
The error from reader.ReadString is ignored. While this might be acceptable for a simple CLI tool, it's better practice to handle potential I/O errors from os.Stdin. For example, if the input stream is closed or encounters an error, this could lead to unexpected behavior. Handling the error improves the robustness of the CLI.
| input, _ := reader.ReadString('\n') | |
| input = strings.TrimSpace(input) | |
| if input == "" { | |
| return defaultValue | |
| } | |
| return input | |
| input, err := reader.ReadString('\n') | |
| if err != nil { | |
| log.Warnf("Error reading input: %v", err) | |
| return defaultValue | |
| } | |
| input = strings.TrimSpace(input) | |
| if input == "" { | |
| return defaultValue | |
| } | |
| return input |
internal/auth/kiro/sso_oidc.go
Outdated
| if _, err := fmt.Sscanf(input, "%d", &selection); err != nil || selection < 1 || selection > len(options) { | ||
| return 0 // Default to first option | ||
| } |
There was a problem hiding this comment.
On invalid input (non-numeric, out of range), the function silently defaults to the first option. This could be confusing for the user. Consider providing feedback on invalid input and re-prompting for a valid selection in a loop to improve the user experience.
| if errResp.Error == "authorization_pending" { | ||
| return nil, fmt.Errorf("authorization_pending") | ||
| } | ||
| if errResp.Error == "slow_down" { | ||
| return nil, fmt.Errorf("slow_down") | ||
| } |
There was a problem hiding this comment.
Using fmt.Errorf to create errors that are later checked with strings.Contains is brittle. If the error message changes, this logic will break. It's better to use sentinel errors. You can define package-level error variables (e.g., var ErrAuthorizationPending = errors.New("authorization_pending")) and return them here. Then, the calling function can use errors.Is for robust error checking.
| if errResp.Error == "authorization_pending" { | |
| return nil, fmt.Errorf("authorization_pending") | |
| } | |
| if errResp.Error == "slow_down" { | |
| return nil, fmt.Errorf("slow_down") | |
| } | |
| if errResp.Error == "authorization_pending" { | |
| return nil, ErrAuthorizationPending | |
| } | |
| if errResp.Error == "slow_down" { | |
| return nil, ErrSlowDown | |
| } |
internal/auth/kiro/sso_oidc.go
Outdated
| req.Header.Set("Content-Type", "application/json") | ||
| req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region)) | ||
| req.Header.Set("Connection", "keep-alive") | ||
| req.Header.Set("x-amz-user-agent", "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE") |
There was a problem hiding this comment.
This x-amz-user-agent header, along with others in this function, contains hardcoded values like SDK versions (3.738.0). This makes them difficult to find and update. To improve maintainability, consider defining these magic strings as constants at the package level.
| auth = refreshedAuth | ||
| // Persist the refreshed auth to file so subsequent requests use it | ||
| if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { | ||
| log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) | ||
| } | ||
| accessToken, profileArn = kiroCredentials(auth) | ||
| log.Infof("kiro: token refreshed successfully before request") | ||
| } |
There was a problem hiding this comment.
The logic for persisting a refreshed token and updating credentials is duplicated in multiple places in this file. This makes the code harder to maintain. Consider extracting this logic into a helper function. For example:
func (e *KiroExecutor) applyRefreshedAuth(auth *cliproxyauth.Auth) (*cliproxyauth.Auth, string, string) {
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
}
accessToken, profileArn := kiroCredentials(auth)
log.Infof("kiro: token refreshed successfully")
return auth, accessToken, profileArn
}This would simplify the call sites significantly.
- Handle errors in promptInput instead of ignoring them - Improve promptSelect to provide feedback on invalid input and re-prompt - Use sentinel errors (ErrAuthorizationPending, ErrSlowDown) instead of string-based error checking with strings.Contains - Move hardcoded x-amz-user-agent header to idcAmzUserAgent constant Addresses code review feedback from Gemini Code Assist.
2 participants
Summary
x-amzn-kiro-agent-mode: spec) for IDC requestsChanges
internal/auth/kiro/sso_oidc.go: Added IDC login flow, region-specific token refresh, and method selection promptinternal/runtime/executor/kiro_executor.go: Added IDC header detection and Kiro IDE user-agent for IDC authsdk/auth/kiro.go: Added IDC auth method handling and token persistence after refreshinternal/auth/kiro/aws.go: Added IDC-specific fields to KiroTokenData structTest plan
--method idcflag