mirror of
https://github.com/githubhjs/CLIProxyAPIPlus.git
synced 2026-07-12 01:25:13 +00:00
chore(build): 更新 .gitignore 文件
- 添加 *.bak 文件扩展名到忽略列表
This commit is contained in:
@@ -50,3 +50,4 @@ _bmad-output/*
|
|||||||
# macOS
|
# macOS
|
||||||
.DS_Store
|
.DS_Store
|
||||||
._*
|
._*
|
||||||
|
*.bak
|
||||||
|
|||||||
@@ -1,305 +0,0 @@
|
|||||||
// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API.
|
|
||||||
// It includes interfaces and implementations for token storage and authentication methods.
|
|
||||||
package kiro
|
|
||||||
|
|
||||||
import (
|
|
||||||
"encoding/base64"
|
|
||||||
"encoding/json"
|
|
||||||
"fmt"
|
|
||||||
"os"
|
|
||||||
"path/filepath"
|
|
||||||
"strings"
|
|
||||||
)
|
|
||||||
|
|
||||||
// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow
|
|
||||||
type PKCECodes struct {
|
|
||||||
// CodeVerifier is the cryptographically random string used to correlate
|
|
||||||
// the authorization request to the token request
|
|
||||||
CodeVerifier string `json:"code_verifier"`
|
|
||||||
// CodeChallenge is the SHA256 hash of the code verifier, base64url-encoded
|
|
||||||
CodeChallenge string `json:"code_challenge"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// KiroTokenData holds OAuth token information from AWS CodeWhisperer (Kiro)
|
|
||||||
type KiroTokenData struct {
|
|
||||||
// AccessToken is the OAuth2 access token for API access
|
|
||||||
AccessToken string `json:"accessToken"`
|
|
||||||
// RefreshToken is used to obtain new access tokens
|
|
||||||
RefreshToken string `json:"refreshToken"`
|
|
||||||
// ProfileArn is the AWS CodeWhisperer profile ARN
|
|
||||||
ProfileArn string `json:"profileArn"`
|
|
||||||
// ExpiresAt is the timestamp when the token expires
|
|
||||||
ExpiresAt string `json:"expiresAt"`
|
|
||||||
// AuthMethod indicates the authentication method used (e.g., "builder-id", "social")
|
|
||||||
AuthMethod string `json:"authMethod"`
|
|
||||||
// Provider indicates the OAuth provider (e.g., "AWS", "Google")
|
|
||||||
Provider string `json:"provider"`
|
|
||||||
// ClientID is the OIDC client ID (needed for token refresh)
|
|
||||||
ClientID string `json:"clientId,omitempty"`
|
|
||||||
// ClientSecret is the OIDC client secret (needed for token refresh)
|
|
||||||
ClientSecret string `json:"clientSecret,omitempty"`
|
|
||||||
// Email is the user's email address (used for file naming)
|
|
||||||
Email string `json:"email,omitempty"`
|
|
||||||
// StartURL is the IDC/Identity Center start URL (only for IDC auth method)
|
|
||||||
StartURL string `json:"startUrl,omitempty"`
|
|
||||||
// Region is the AWS region for IDC authentication (only for IDC auth method)
|
|
||||||
Region string `json:"region,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// KiroAuthBundle aggregates authentication data after OAuth flow completion
|
|
||||||
type KiroAuthBundle struct {
|
|
||||||
// TokenData contains the OAuth tokens from the authentication flow
|
|
||||||
TokenData KiroTokenData `json:"token_data"`
|
|
||||||
// LastRefresh is the timestamp of the last token refresh
|
|
||||||
LastRefresh string `json:"last_refresh"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// KiroUsageInfo represents usage information from CodeWhisperer API
|
|
||||||
type KiroUsageInfo struct {
|
|
||||||
// SubscriptionTitle is the subscription plan name (e.g., "KIRO FREE")
|
|
||||||
SubscriptionTitle string `json:"subscription_title"`
|
|
||||||
// CurrentUsage is the current credit usage
|
|
||||||
CurrentUsage float64 `json:"current_usage"`
|
|
||||||
// UsageLimit is the maximum credit limit
|
|
||||||
UsageLimit float64 `json:"usage_limit"`
|
|
||||||
// NextReset is the timestamp of the next usage reset
|
|
||||||
NextReset string `json:"next_reset"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// KiroModel represents a model available through the CodeWhisperer API
|
|
||||||
type KiroModel struct {
|
|
||||||
// ModelID is the unique identifier for the model
|
|
||||||
ModelID string `json:"modelId"`
|
|
||||||
// ModelName is the human-readable name
|
|
||||||
ModelName string `json:"modelName"`
|
|
||||||
// Description is the model description
|
|
||||||
Description string `json:"description"`
|
|
||||||
// RateMultiplier is the credit multiplier for this model
|
|
||||||
RateMultiplier float64 `json:"rateMultiplier"`
|
|
||||||
// RateUnit is the unit for rate calculation (e.g., "credit")
|
|
||||||
RateUnit string `json:"rateUnit"`
|
|
||||||
// MaxInputTokens is the maximum input token limit
|
|
||||||
MaxInputTokens int `json:"maxInputTokens,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// KiroIDETokenFile is the default path to Kiro IDE's token file
|
|
||||||
const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json"
|
|
||||||
|
|
||||||
// LoadKiroIDEToken loads token data from Kiro IDE's token file.
|
|
||||||
func LoadKiroIDEToken() (*KiroTokenData, error) {
|
|
||||||
homeDir, err := os.UserHomeDir()
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to get home directory: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
tokenPath := filepath.Join(homeDir, KiroIDETokenFile)
|
|
||||||
data, err := os.ReadFile(tokenPath)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to read Kiro IDE token file (%s): %w", tokenPath, err)
|
|
||||||
}
|
|
||||||
|
|
||||||
var token KiroTokenData
|
|
||||||
if err := json.Unmarshal(data, &token); err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to parse Kiro IDE token: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if token.AccessToken == "" {
|
|
||||||
return nil, fmt.Errorf("access token is empty in Kiro IDE token file")
|
|
||||||
}
|
|
||||||
|
|
||||||
return &token, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// LoadKiroTokenFromPath loads token data from a custom path.
|
|
||||||
// This supports multiple accounts by allowing different token files.
|
|
||||||
func LoadKiroTokenFromPath(tokenPath string) (*KiroTokenData, error) {
|
|
||||||
// Expand ~ to home directory
|
|
||||||
if len(tokenPath) > 0 && tokenPath[0] == '~' {
|
|
||||||
homeDir, err := os.UserHomeDir()
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to get home directory: %w", err)
|
|
||||||
}
|
|
||||||
tokenPath = filepath.Join(homeDir, tokenPath[1:])
|
|
||||||
}
|
|
||||||
|
|
||||||
data, err := os.ReadFile(tokenPath)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to read token file (%s): %w", tokenPath, err)
|
|
||||||
}
|
|
||||||
|
|
||||||
var token KiroTokenData
|
|
||||||
if err := json.Unmarshal(data, &token); err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to parse token file: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if token.AccessToken == "" {
|
|
||||||
return nil, fmt.Errorf("access token is empty in token file")
|
|
||||||
}
|
|
||||||
|
|
||||||
return &token, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// ListKiroTokenFiles lists all Kiro token files in the cache directory.
|
|
||||||
// This supports multiple accounts by finding all token files.
|
|
||||||
func ListKiroTokenFiles() ([]string, error) {
|
|
||||||
homeDir, err := os.UserHomeDir()
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to get home directory: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache")
|
|
||||||
|
|
||||||
// Check if directory exists
|
|
||||||
if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
|
|
||||||
return nil, nil // No token files
|
|
||||||
}
|
|
||||||
|
|
||||||
entries, err := os.ReadDir(cacheDir)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to read cache directory: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
var tokenFiles []string
|
|
||||||
for _, entry := range entries {
|
|
||||||
if entry.IsDir() {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
name := entry.Name()
|
|
||||||
// Look for kiro token files only (avoid matching unrelated AWS SSO cache files)
|
|
||||||
if strings.HasSuffix(name, ".json") && strings.HasPrefix(name, "kiro") {
|
|
||||||
tokenFiles = append(tokenFiles, filepath.Join(cacheDir, name))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return tokenFiles, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// LoadAllKiroTokens loads all Kiro tokens from the cache directory.
|
|
||||||
// This supports multiple accounts.
|
|
||||||
func LoadAllKiroTokens() ([]*KiroTokenData, error) {
|
|
||||||
files, err := ListKiroTokenFiles()
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
var tokens []*KiroTokenData
|
|
||||||
for _, file := range files {
|
|
||||||
token, err := LoadKiroTokenFromPath(file)
|
|
||||||
if err != nil {
|
|
||||||
// Skip invalid token files
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
tokens = append(tokens, token)
|
|
||||||
}
|
|
||||||
|
|
||||||
return tokens, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// JWTClaims represents the claims we care about from a JWT token.
|
|
||||||
// JWT tokens from Kiro/AWS contain user information in the payload.
|
|
||||||
type JWTClaims struct {
|
|
||||||
Email string `json:"email,omitempty"`
|
|
||||||
Sub string `json:"sub,omitempty"`
|
|
||||||
PreferredUser string `json:"preferred_username,omitempty"`
|
|
||||||
Name string `json:"name,omitempty"`
|
|
||||||
Iss string `json:"iss,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// ExtractEmailFromJWT extracts the user's email from a JWT access token.
|
|
||||||
// JWT tokens typically have format: header.payload.signature
|
|
||||||
// The payload is base64url-encoded JSON containing user claims.
|
|
||||||
func ExtractEmailFromJWT(accessToken string) string {
|
|
||||||
if accessToken == "" {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
|
|
||||||
// JWT format: header.payload.signature
|
|
||||||
parts := strings.Split(accessToken, ".")
|
|
||||||
if len(parts) != 3 {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
|
|
||||||
// Decode the payload (second part)
|
|
||||||
payload := parts[1]
|
|
||||||
|
|
||||||
// Add padding if needed (base64url requires padding)
|
|
||||||
switch len(payload) % 4 {
|
|
||||||
case 2:
|
|
||||||
payload += "=="
|
|
||||||
case 3:
|
|
||||||
payload += "="
|
|
||||||
}
|
|
||||||
|
|
||||||
decoded, err := base64.URLEncoding.DecodeString(payload)
|
|
||||||
if err != nil {
|
|
||||||
// Try RawURLEncoding (no padding)
|
|
||||||
decoded, err = base64.RawURLEncoding.DecodeString(parts[1])
|
|
||||||
if err != nil {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
var claims JWTClaims
|
|
||||||
if err := json.Unmarshal(decoded, &claims); err != nil {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
|
|
||||||
// Return email if available
|
|
||||||
if claims.Email != "" {
|
|
||||||
return claims.Email
|
|
||||||
}
|
|
||||||
|
|
||||||
// Fallback to preferred_username (some providers use this)
|
|
||||||
if claims.PreferredUser != "" && strings.Contains(claims.PreferredUser, "@") {
|
|
||||||
return claims.PreferredUser
|
|
||||||
}
|
|
||||||
|
|
||||||
// Fallback to sub if it looks like an email
|
|
||||||
if claims.Sub != "" && strings.Contains(claims.Sub, "@") {
|
|
||||||
return claims.Sub
|
|
||||||
}
|
|
||||||
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
|
|
||||||
// SanitizeEmailForFilename sanitizes an email address for use in a filename.
|
|
||||||
// Replaces special characters with underscores and prevents path traversal attacks.
|
|
||||||
// Also handles URL-encoded characters to prevent encoded path traversal attempts.
|
|
||||||
func SanitizeEmailForFilename(email string) string {
|
|
||||||
if email == "" {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
|
|
||||||
result := email
|
|
||||||
|
|
||||||
// First, handle URL-encoded path traversal attempts (%2F, %2E, %5C, etc.)
|
|
||||||
// This prevents encoded characters from bypassing the sanitization.
|
|
||||||
// Note: We replace % last to catch any remaining encodings including double-encoding (%252F)
|
|
||||||
result = strings.ReplaceAll(result, "%2F", "_") // /
|
|
||||||
result = strings.ReplaceAll(result, "%2f", "_")
|
|
||||||
result = strings.ReplaceAll(result, "%5C", "_") // \
|
|
||||||
result = strings.ReplaceAll(result, "%5c", "_")
|
|
||||||
result = strings.ReplaceAll(result, "%2E", "_") // .
|
|
||||||
result = strings.ReplaceAll(result, "%2e", "_")
|
|
||||||
result = strings.ReplaceAll(result, "%00", "_") // null byte
|
|
||||||
result = strings.ReplaceAll(result, "%", "_") // Catch remaining % to prevent double-encoding attacks
|
|
||||||
|
|
||||||
// Replace characters that are problematic in filenames
|
|
||||||
// Keep @ and . in middle but replace other special characters
|
|
||||||
for _, char := range []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " ", "\x00"} {
|
|
||||||
result = strings.ReplaceAll(result, char, "_")
|
|
||||||
}
|
|
||||||
|
|
||||||
// Prevent path traversal: replace leading dots in each path component
|
|
||||||
// This handles cases like "../../../etc/passwd" → "_.._.._.._etc_passwd"
|
|
||||||
parts := strings.Split(result, "_")
|
|
||||||
for i, part := range parts {
|
|
||||||
for strings.HasPrefix(part, ".") {
|
|
||||||
part = "_" + part[1:]
|
|
||||||
}
|
|
||||||
parts[i] = part
|
|
||||||
}
|
|
||||||
result = strings.Join(parts, "_")
|
|
||||||
|
|
||||||
return result
|
|
||||||
}
|
|
||||||
@@ -1,385 +0,0 @@
|
|||||||
// Package kiro provides OAuth Web authentication for Kiro.
|
|
||||||
package kiro
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"crypto/rand"
|
|
||||||
"encoding/base64"
|
|
||||||
"fmt"
|
|
||||||
"html/template"
|
|
||||||
"net/http"
|
|
||||||
"sync"
|
|
||||||
"time"
|
|
||||||
|
|
||||||
"github.com/gin-gonic/gin"
|
|
||||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
|
||||||
log "github.com/sirupsen/logrus"
|
|
||||||
)
|
|
||||||
|
|
||||||
const (
|
|
||||||
defaultSessionExpiry = 10 * time.Minute
|
|
||||||
pollIntervalSeconds = 5
|
|
||||||
)
|
|
||||||
|
|
||||||
type authSessionStatus string
|
|
||||||
|
|
||||||
const (
|
|
||||||
statusPending authSessionStatus = "pending"
|
|
||||||
statusSuccess authSessionStatus = "success"
|
|
||||||
statusFailed authSessionStatus = "failed"
|
|
||||||
)
|
|
||||||
|
|
||||||
type webAuthSession struct {
|
|
||||||
stateID string
|
|
||||||
deviceCode string
|
|
||||||
userCode string
|
|
||||||
authURL string
|
|
||||||
verificationURI string
|
|
||||||
expiresIn int
|
|
||||||
interval int
|
|
||||||
status authSessionStatus
|
|
||||||
startedAt time.Time
|
|
||||||
completedAt time.Time
|
|
||||||
expiresAt time.Time
|
|
||||||
error string
|
|
||||||
tokenData *KiroTokenData
|
|
||||||
ssoClient *SSOOIDCClient
|
|
||||||
clientID string
|
|
||||||
clientSecret string
|
|
||||||
region string
|
|
||||||
cancelFunc context.CancelFunc
|
|
||||||
}
|
|
||||||
|
|
||||||
type OAuthWebHandler struct {
|
|
||||||
cfg *config.Config
|
|
||||||
sessions map[string]*webAuthSession
|
|
||||||
mu sync.RWMutex
|
|
||||||
onTokenObtained func(*KiroTokenData)
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewOAuthWebHandler(cfg *config.Config) *OAuthWebHandler {
|
|
||||||
return &OAuthWebHandler{
|
|
||||||
cfg: cfg,
|
|
||||||
sessions: make(map[string]*webAuthSession),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) SetTokenCallback(callback func(*KiroTokenData)) {
|
|
||||||
h.onTokenObtained = callback
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) RegisterRoutes(router gin.IRouter) {
|
|
||||||
oauth := router.Group("/v0/oauth/kiro")
|
|
||||||
{
|
|
||||||
oauth.GET("/start", h.handleStart)
|
|
||||||
oauth.GET("/callback", h.handleCallback)
|
|
||||||
oauth.GET("/status", h.handleStatus)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func generateStateID() (string, error) {
|
|
||||||
b := make([]byte, 16)
|
|
||||||
if _, err := rand.Read(b); err != nil {
|
|
||||||
return "", err
|
|
||||||
}
|
|
||||||
return base64.RawURLEncoding.EncodeToString(b), nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) handleStart(c *gin.Context) {
|
|
||||||
stateID, err := generateStateID()
|
|
||||||
if err != nil {
|
|
||||||
h.renderError(c, "Failed to generate state parameter")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
region := defaultIDCRegion
|
|
||||||
startURL := builderIDStartURL
|
|
||||||
|
|
||||||
ssoClient := NewSSOOIDCClient(h.cfg)
|
|
||||||
|
|
||||||
regResp, err := ssoClient.RegisterClientWithRegion(c.Request.Context(), region)
|
|
||||||
if err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to register client: %v", err)
|
|
||||||
h.renderError(c, fmt.Sprintf("Failed to register client: %v", err))
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
authResp, err := ssoClient.StartDeviceAuthorizationWithIDC(
|
|
||||||
c.Request.Context(),
|
|
||||||
regResp.ClientID,
|
|
||||||
regResp.ClientSecret,
|
|
||||||
startURL,
|
|
||||||
region,
|
|
||||||
)
|
|
||||||
if err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to start device authorization: %v", err)
|
|
||||||
h.renderError(c, fmt.Sprintf("Failed to start device authorization: %v", err))
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(authResp.ExpiresIn)*time.Second)
|
|
||||||
|
|
||||||
session := &webAuthSession{
|
|
||||||
stateID: stateID,
|
|
||||||
deviceCode: authResp.DeviceCode,
|
|
||||||
userCode: authResp.UserCode,
|
|
||||||
authURL: authResp.VerificationURIComplete,
|
|
||||||
verificationURI: authResp.VerificationURI,
|
|
||||||
expiresIn: authResp.ExpiresIn,
|
|
||||||
interval: authResp.Interval,
|
|
||||||
status: statusPending,
|
|
||||||
startedAt: time.Now(),
|
|
||||||
ssoClient: ssoClient,
|
|
||||||
clientID: regResp.ClientID,
|
|
||||||
clientSecret: regResp.ClientSecret,
|
|
||||||
region: region,
|
|
||||||
cancelFunc: cancel,
|
|
||||||
}
|
|
||||||
|
|
||||||
h.mu.Lock()
|
|
||||||
h.sessions[stateID] = session
|
|
||||||
h.mu.Unlock()
|
|
||||||
|
|
||||||
go h.pollForToken(ctx, session)
|
|
||||||
|
|
||||||
h.renderStartPage(c, session)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSession) {
|
|
||||||
defer session.cancelFunc()
|
|
||||||
|
|
||||||
interval := time.Duration(session.interval) * time.Second
|
|
||||||
if interval < time.Duration(pollIntervalSeconds)*time.Second {
|
|
||||||
interval = time.Duration(pollIntervalSeconds) * time.Second
|
|
||||||
}
|
|
||||||
|
|
||||||
ticker := time.NewTicker(interval)
|
|
||||||
defer ticker.Stop()
|
|
||||||
|
|
||||||
for {
|
|
||||||
select {
|
|
||||||
case <-ctx.Done():
|
|
||||||
h.mu.Lock()
|
|
||||||
if session.status == statusPending {
|
|
||||||
session.status = statusFailed
|
|
||||||
session.error = "Authentication timed out"
|
|
||||||
}
|
|
||||||
h.mu.Unlock()
|
|
||||||
return
|
|
||||||
case <-ticker.C:
|
|
||||||
tokenResp, err := h.ssoClient(session).CreateTokenWithRegion(
|
|
||||||
ctx,
|
|
||||||
session.clientID,
|
|
||||||
session.clientSecret,
|
|
||||||
session.deviceCode,
|
|
||||||
session.region,
|
|
||||||
)
|
|
||||||
|
|
||||||
if err != nil {
|
|
||||||
errStr := err.Error()
|
|
||||||
if errStr == ErrAuthorizationPending.Error() {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
if errStr == ErrSlowDown.Error() {
|
|
||||||
interval += 5 * time.Second
|
|
||||||
ticker.Reset(interval)
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
h.mu.Lock()
|
|
||||||
session.status = statusFailed
|
|
||||||
session.error = errStr
|
|
||||||
session.completedAt = time.Now()
|
|
||||||
h.mu.Unlock()
|
|
||||||
|
|
||||||
log.Errorf("OAuth Web: token polling failed: %v", err)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
|
|
||||||
profileArn := session.ssoClient.fetchProfileArn(ctx, tokenResp.AccessToken)
|
|
||||||
email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken)
|
|
||||||
|
|
||||||
tokenData := &KiroTokenData{
|
|
||||||
AccessToken: tokenResp.AccessToken,
|
|
||||||
RefreshToken: tokenResp.RefreshToken,
|
|
||||||
ProfileArn: profileArn,
|
|
||||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
|
||||||
AuthMethod: "builder-id",
|
|
||||||
Provider: "AWS",
|
|
||||||
ClientID: session.clientID,
|
|
||||||
ClientSecret: session.clientSecret,
|
|
||||||
Email: email,
|
|
||||||
}
|
|
||||||
|
|
||||||
h.mu.Lock()
|
|
||||||
session.status = statusSuccess
|
|
||||||
session.completedAt = time.Now()
|
|
||||||
session.expiresAt = expiresAt
|
|
||||||
session.tokenData = tokenData
|
|
||||||
h.mu.Unlock()
|
|
||||||
|
|
||||||
if h.onTokenObtained != nil {
|
|
||||||
h.onTokenObtained(tokenData)
|
|
||||||
}
|
|
||||||
|
|
||||||
log.Infof("OAuth Web: authentication successful for %s", email)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) ssoClient(session *webAuthSession) *SSOOIDCClient {
|
|
||||||
return session.ssoClient
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) handleCallback(c *gin.Context) {
|
|
||||||
stateID := c.Query("state")
|
|
||||||
errParam := c.Query("error")
|
|
||||||
|
|
||||||
if errParam != "" {
|
|
||||||
h.renderError(c, errParam)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
if stateID == "" {
|
|
||||||
h.renderError(c, "Missing state parameter")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
h.mu.RLock()
|
|
||||||
session, exists := h.sessions[stateID]
|
|
||||||
h.mu.RUnlock()
|
|
||||||
|
|
||||||
if !exists {
|
|
||||||
h.renderError(c, "Invalid or expired session")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
if session.status == statusSuccess {
|
|
||||||
h.renderSuccess(c, session)
|
|
||||||
} else if session.status == statusFailed {
|
|
||||||
h.renderError(c, session.error)
|
|
||||||
} else {
|
|
||||||
c.Redirect(http.StatusFound, "/v0/oauth/kiro/start")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) handleStatus(c *gin.Context) {
|
|
||||||
stateID := c.Query("state")
|
|
||||||
if stateID == "" {
|
|
||||||
c.JSON(http.StatusBadRequest, gin.H{"error": "missing state parameter"})
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
h.mu.RLock()
|
|
||||||
session, exists := h.sessions[stateID]
|
|
||||||
h.mu.RUnlock()
|
|
||||||
|
|
||||||
if !exists {
|
|
||||||
c.JSON(http.StatusNotFound, gin.H{"error": "session not found"})
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
response := gin.H{
|
|
||||||
"status": string(session.status),
|
|
||||||
}
|
|
||||||
|
|
||||||
switch session.status {
|
|
||||||
case statusPending:
|
|
||||||
elapsed := time.Since(session.startedAt).Seconds()
|
|
||||||
remaining := float64(session.expiresIn) - elapsed
|
|
||||||
if remaining < 0 {
|
|
||||||
remaining = 0
|
|
||||||
}
|
|
||||||
response["remaining_seconds"] = int(remaining)
|
|
||||||
case statusSuccess:
|
|
||||||
response["completed_at"] = session.completedAt.Format(time.RFC3339)
|
|
||||||
response["expires_at"] = session.expiresAt.Format(time.RFC3339)
|
|
||||||
case statusFailed:
|
|
||||||
response["error"] = session.error
|
|
||||||
response["failed_at"] = session.completedAt.Format(time.RFC3339)
|
|
||||||
}
|
|
||||||
|
|
||||||
c.JSON(http.StatusOK, response)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) renderStartPage(c *gin.Context, session *webAuthSession) {
|
|
||||||
tmpl, err := template.New("start").Parse(oauthWebStartPageHTML)
|
|
||||||
if err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to parse template: %v", err)
|
|
||||||
c.String(http.StatusInternalServerError, "Template error")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
data := map[string]interface{}{
|
|
||||||
"AuthURL": session.authURL,
|
|
||||||
"UserCode": session.userCode,
|
|
||||||
"ExpiresIn": session.expiresIn,
|
|
||||||
"StateID": session.stateID,
|
|
||||||
}
|
|
||||||
|
|
||||||
c.Header("Content-Type", "text/html; charset=utf-8")
|
|
||||||
if err := tmpl.Execute(c.Writer, data); err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to render template: %v", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) renderError(c *gin.Context, errMsg string) {
|
|
||||||
tmpl, err := template.New("error").Parse(oauthWebErrorPageHTML)
|
|
||||||
if err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to parse error template: %v", err)
|
|
||||||
c.String(http.StatusInternalServerError, "Template error")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
data := map[string]interface{}{
|
|
||||||
"Error": errMsg,
|
|
||||||
}
|
|
||||||
|
|
||||||
c.Header("Content-Type", "text/html; charset=utf-8")
|
|
||||||
c.Status(http.StatusBadRequest)
|
|
||||||
if err := tmpl.Execute(c.Writer, data); err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to render error template: %v", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) renderSuccess(c *gin.Context, session *webAuthSession) {
|
|
||||||
tmpl, err := template.New("success").Parse(oauthWebSuccessPageHTML)
|
|
||||||
if err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to parse success template: %v", err)
|
|
||||||
c.String(http.StatusInternalServerError, "Template error")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
data := map[string]interface{}{
|
|
||||||
"ExpiresAt": session.expiresAt.Format(time.RFC3339),
|
|
||||||
}
|
|
||||||
|
|
||||||
c.Header("Content-Type", "text/html; charset=utf-8")
|
|
||||||
if err := tmpl.Execute(c.Writer, data); err != nil {
|
|
||||||
log.Errorf("OAuth Web: failed to render success template: %v", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) CleanupExpiredSessions() {
|
|
||||||
h.mu.Lock()
|
|
||||||
defer h.mu.Unlock()
|
|
||||||
|
|
||||||
now := time.Now()
|
|
||||||
for id, session := range h.sessions {
|
|
||||||
if session.status != statusPending && now.Sub(session.completedAt) > 30*time.Minute {
|
|
||||||
delete(h.sessions, id)
|
|
||||||
} else if session.status == statusPending && now.Sub(session.startedAt) > defaultSessionExpiry {
|
|
||||||
session.cancelFunc()
|
|
||||||
delete(h.sessions, id)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *OAuthWebHandler) GetSession(stateID string) (*webAuthSession, bool) {
|
|
||||||
h.mu.RLock()
|
|
||||||
defer h.mu.RUnlock()
|
|
||||||
session, exists := h.sessions[stateID]
|
|
||||||
return session, exists
|
|
||||||
}
|
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -1,97 +0,0 @@
|
|||||||
package common
|
|
||||||
|
|
||||||
import (
|
|
||||||
"unicode/utf8"
|
|
||||||
)
|
|
||||||
|
|
||||||
type UTF8StreamParser struct {
|
|
||||||
buffer []byte
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewUTF8StreamParser() *UTF8StreamParser {
|
|
||||||
return &UTF8StreamParser{
|
|
||||||
buffer: make([]byte, 0, 64),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *UTF8StreamParser) Write(data []byte) {
|
|
||||||
p.buffer = append(p.buffer, data...)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *UTF8StreamParser) Read() (string, bool) {
|
|
||||||
if len(p.buffer) == 0 {
|
|
||||||
return "", false
|
|
||||||
}
|
|
||||||
|
|
||||||
validLen := p.findValidUTF8End(p.buffer)
|
|
||||||
if validLen == 0 {
|
|
||||||
return "", false
|
|
||||||
}
|
|
||||||
|
|
||||||
result := string(p.buffer[:validLen])
|
|
||||||
p.buffer = p.buffer[validLen:]
|
|
||||||
|
|
||||||
return result, true
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *UTF8StreamParser) Flush() string {
|
|
||||||
if len(p.buffer) == 0 {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
result := string(p.buffer)
|
|
||||||
p.buffer = p.buffer[:0]
|
|
||||||
return result
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *UTF8StreamParser) Reset() {
|
|
||||||
p.buffer = p.buffer[:0]
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *UTF8StreamParser) findValidUTF8End(data []byte) int {
|
|
||||||
if len(data) == 0 {
|
|
||||||
return 0
|
|
||||||
}
|
|
||||||
|
|
||||||
end := len(data)
|
|
||||||
for i := 1; i <= 3 && i <= len(data); i++ {
|
|
||||||
b := data[len(data)-i]
|
|
||||||
if b&0x80 == 0 {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
if b&0xC0 == 0xC0 {
|
|
||||||
size := p.utf8CharSize(b)
|
|
||||||
available := i
|
|
||||||
if size > available {
|
|
||||||
end = len(data) - i
|
|
||||||
}
|
|
||||||
break
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if end > 0 && !utf8.Valid(data[:end]) {
|
|
||||||
for i := end - 1; i >= 0; i-- {
|
|
||||||
if utf8.Valid(data[:i+1]) {
|
|
||||||
return i + 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return 0
|
|
||||||
}
|
|
||||||
|
|
||||||
return end
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *UTF8StreamParser) utf8CharSize(b byte) int {
|
|
||||||
if b&0x80 == 0 {
|
|
||||||
return 1
|
|
||||||
}
|
|
||||||
if b&0xE0 == 0xC0 {
|
|
||||||
return 2
|
|
||||||
}
|
|
||||||
if b&0xF0 == 0xE0 {
|
|
||||||
return 3
|
|
||||||
}
|
|
||||||
if b&0xF8 == 0xF0 {
|
|
||||||
return 4
|
|
||||||
}
|
|
||||||
return 1
|
|
||||||
}
|
|
||||||
@@ -1,402 +0,0 @@
|
|||||||
package common
|
|
||||||
|
|
||||||
import (
|
|
||||||
"strings"
|
|
||||||
"sync"
|
|
||||||
"testing"
|
|
||||||
"unicode/utf8"
|
|
||||||
)
|
|
||||||
|
|
||||||
func TestNewUTF8StreamParser(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
if p == nil {
|
|
||||||
t.Fatal("expected non-nil UTF8StreamParser")
|
|
||||||
}
|
|
||||||
if p.buffer == nil {
|
|
||||||
t.Error("expected non-nil buffer")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestWrite(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("hello"))
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok to be true")
|
|
||||||
}
|
|
||||||
if result != "hello" {
|
|
||||||
t.Errorf("expected 'hello', got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestWrite_MultipleWrites(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("hel"))
|
|
||||||
p.Write([]byte("lo"))
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok to be true")
|
|
||||||
}
|
|
||||||
if result != "hello" {
|
|
||||||
t.Errorf("expected 'hello', got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestRead_EmptyBuffer(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
result, ok := p.Read()
|
|
||||||
if ok {
|
|
||||||
t.Error("expected ok to be false for empty buffer")
|
|
||||||
}
|
|
||||||
if result != "" {
|
|
||||||
t.Errorf("expected empty string, got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestRead_IncompleteUTF8(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
|
|
||||||
// Write incomplete multi-byte UTF-8 character
|
|
||||||
// 中 (U+4E2D) = E4 B8 AD
|
|
||||||
p.Write([]byte{0xE4, 0xB8})
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if ok {
|
|
||||||
t.Error("expected ok to be false for incomplete UTF-8")
|
|
||||||
}
|
|
||||||
if result != "" {
|
|
||||||
t.Errorf("expected empty string, got '%s'", result)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Complete the character
|
|
||||||
p.Write([]byte{0xAD})
|
|
||||||
result, ok = p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok to be true after completing UTF-8")
|
|
||||||
}
|
|
||||||
if result != "中" {
|
|
||||||
t.Errorf("expected '中', got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestRead_MixedASCIIAndUTF8(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("Hello 世界"))
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok to be true")
|
|
||||||
}
|
|
||||||
if result != "Hello 世界" {
|
|
||||||
t.Errorf("expected 'Hello 世界', got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestRead_PartialMultibyteAtEnd(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
// "Hello" + partial "世" (E4 B8 96)
|
|
||||||
p.Write([]byte("Hello"))
|
|
||||||
p.Write([]byte{0xE4, 0xB8})
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok to be true for valid portion")
|
|
||||||
}
|
|
||||||
if result != "Hello" {
|
|
||||||
t.Errorf("expected 'Hello', got '%s'", result)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Complete the character
|
|
||||||
p.Write([]byte{0x96})
|
|
||||||
result, ok = p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok to be true after completing")
|
|
||||||
}
|
|
||||||
if result != "世" {
|
|
||||||
t.Errorf("expected '世', got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestFlush(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("hello"))
|
|
||||||
|
|
||||||
result := p.Flush()
|
|
||||||
if result != "hello" {
|
|
||||||
t.Errorf("expected 'hello', got '%s'", result)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Verify buffer is cleared
|
|
||||||
result2, ok := p.Read()
|
|
||||||
if ok {
|
|
||||||
t.Error("expected ok to be false after flush")
|
|
||||||
}
|
|
||||||
if result2 != "" {
|
|
||||||
t.Errorf("expected empty string after flush, got '%s'", result2)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestFlush_EmptyBuffer(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
result := p.Flush()
|
|
||||||
if result != "" {
|
|
||||||
t.Errorf("expected empty string, got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestFlush_IncompleteUTF8(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte{0xE4, 0xB8})
|
|
||||||
|
|
||||||
result := p.Flush()
|
|
||||||
// Flush returns everything including incomplete bytes
|
|
||||||
if len(result) != 2 {
|
|
||||||
t.Errorf("expected 2 bytes flushed, got %d", len(result))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestReset(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("hello"))
|
|
||||||
p.Reset()
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if ok {
|
|
||||||
t.Error("expected ok to be false after reset")
|
|
||||||
}
|
|
||||||
if result != "" {
|
|
||||||
t.Errorf("expected empty string after reset, got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestUtf8CharSize(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
|
|
||||||
testCases := []struct {
|
|
||||||
b byte
|
|
||||||
expected int
|
|
||||||
}{
|
|
||||||
{0x00, 1}, // ASCII
|
|
||||||
{0x7F, 1}, // ASCII max
|
|
||||||
{0xC0, 2}, // 2-byte start
|
|
||||||
{0xDF, 2}, // 2-byte start
|
|
||||||
{0xE0, 3}, // 3-byte start
|
|
||||||
{0xEF, 3}, // 3-byte start
|
|
||||||
{0xF0, 4}, // 4-byte start
|
|
||||||
{0xF7, 4}, // 4-byte start
|
|
||||||
{0x80, 1}, // Continuation byte (fallback)
|
|
||||||
}
|
|
||||||
|
|
||||||
for _, tc := range testCases {
|
|
||||||
size := p.utf8CharSize(tc.b)
|
|
||||||
if size != tc.expected {
|
|
||||||
t.Errorf("utf8CharSize(0x%X) = %d, expected %d", tc.b, size, tc.expected)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestStreamingScenario(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
|
|
||||||
// Simulate streaming: "Hello, 世界! 🌍"
|
|
||||||
chunks := [][]byte{
|
|
||||||
[]byte("Hello, "),
|
|
||||||
{0xE4, 0xB8}, // partial 世
|
|
||||||
{0x96, 0xE7}, // complete 世, partial 界
|
|
||||||
{0x95, 0x8C}, // complete 界
|
|
||||||
[]byte("! "),
|
|
||||||
{0xF0, 0x9F}, // partial 🌍
|
|
||||||
{0x8C, 0x8D}, // complete 🌍
|
|
||||||
}
|
|
||||||
|
|
||||||
var results []string
|
|
||||||
for _, chunk := range chunks {
|
|
||||||
p.Write(chunk)
|
|
||||||
if result, ok := p.Read(); ok {
|
|
||||||
results = append(results, result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
combined := strings.Join(results, "")
|
|
||||||
if combined != "Hello, 世界! 🌍" {
|
|
||||||
t.Errorf("expected 'Hello, 世界! 🌍', got '%s'", combined)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestValidUTF8Output(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
|
|
||||||
testStrings := []string{
|
|
||||||
"Hello World",
|
|
||||||
"你好世界",
|
|
||||||
"こんにちは",
|
|
||||||
"🎉🎊🎁",
|
|
||||||
"Mixed 混合 Текст ტექსტი",
|
|
||||||
}
|
|
||||||
|
|
||||||
for _, s := range testStrings {
|
|
||||||
p.Reset()
|
|
||||||
p.Write([]byte(s))
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Errorf("expected ok for '%s'", s)
|
|
||||||
}
|
|
||||||
if !utf8.ValidString(result) {
|
|
||||||
t.Errorf("invalid UTF-8 output for input '%s'", s)
|
|
||||||
}
|
|
||||||
if result != s {
|
|
||||||
t.Errorf("expected '%s', got '%s'", s, result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestLargeData(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
|
|
||||||
// Generate large UTF-8 string
|
|
||||||
var builder strings.Builder
|
|
||||||
for i := 0; i < 1000; i++ {
|
|
||||||
builder.WriteString("Hello 世界! ")
|
|
||||||
}
|
|
||||||
largeString := builder.String()
|
|
||||||
|
|
||||||
p.Write([]byte(largeString))
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok for large data")
|
|
||||||
}
|
|
||||||
if result != largeString {
|
|
||||||
t.Error("large data mismatch")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestByteByByteWriting(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
input := "Hello 世界"
|
|
||||||
inputBytes := []byte(input)
|
|
||||||
|
|
||||||
var results []string
|
|
||||||
for _, b := range inputBytes {
|
|
||||||
p.Write([]byte{b})
|
|
||||||
if result, ok := p.Read(); ok {
|
|
||||||
results = append(results, result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
combined := strings.Join(results, "")
|
|
||||||
if combined != input {
|
|
||||||
t.Errorf("expected '%s', got '%s'", input, combined)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestEmoji4ByteUTF8(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
|
|
||||||
// 🎉 = F0 9F 8E 89
|
|
||||||
emoji := "🎉"
|
|
||||||
emojiBytes := []byte(emoji)
|
|
||||||
|
|
||||||
for i := 0; i < len(emojiBytes)-1; i++ {
|
|
||||||
p.Write(emojiBytes[i : i+1])
|
|
||||||
result, ok := p.Read()
|
|
||||||
if ok && result != "" {
|
|
||||||
t.Errorf("unexpected output before emoji complete: '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
p.Write(emojiBytes[len(emojiBytes)-1:])
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok {
|
|
||||||
t.Error("expected ok after completing emoji")
|
|
||||||
}
|
|
||||||
if result != emoji {
|
|
||||||
t.Errorf("expected '%s', got '%s'", emoji, result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestContinuationBytesOnly(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
|
|
||||||
// Write only continuation bytes (invalid UTF-8)
|
|
||||||
p.Write([]byte{0x80, 0x80, 0x80})
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
// Should handle gracefully - either return nothing or return the bytes
|
|
||||||
_ = result
|
|
||||||
_ = ok
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestUTF8StreamParser_ConcurrentSafety(t *testing.T) {
|
|
||||||
// Note: UTF8StreamParser doesn't have built-in locks,
|
|
||||||
// so this test verifies it works with external synchronization
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
var mu sync.Mutex
|
|
||||||
const numGoroutines = 10
|
|
||||||
const numOperations = 100
|
|
||||||
|
|
||||||
var wg sync.WaitGroup
|
|
||||||
wg.Add(numGoroutines)
|
|
||||||
|
|
||||||
for i := 0; i < numGoroutines; i++ {
|
|
||||||
go func() {
|
|
||||||
defer wg.Done()
|
|
||||||
for j := 0; j < numOperations; j++ {
|
|
||||||
mu.Lock()
|
|
||||||
switch j % 4 {
|
|
||||||
case 0:
|
|
||||||
p.Write([]byte("test"))
|
|
||||||
case 1:
|
|
||||||
p.Read()
|
|
||||||
case 2:
|
|
||||||
p.Flush()
|
|
||||||
case 3:
|
|
||||||
p.Reset()
|
|
||||||
}
|
|
||||||
mu.Unlock()
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
|
|
||||||
wg.Wait()
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestConsecutiveReads(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("hello"))
|
|
||||||
|
|
||||||
result1, ok1 := p.Read()
|
|
||||||
if !ok1 || result1 != "hello" {
|
|
||||||
t.Error("first read failed")
|
|
||||||
}
|
|
||||||
|
|
||||||
result2, ok2 := p.Read()
|
|
||||||
if ok2 || result2 != "" {
|
|
||||||
t.Error("second read should return empty")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestFlushThenWrite(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("first"))
|
|
||||||
p.Flush()
|
|
||||||
p.Write([]byte("second"))
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok || result != "second" {
|
|
||||||
t.Errorf("expected 'second', got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestResetThenWrite(t *testing.T) {
|
|
||||||
p := NewUTF8StreamParser()
|
|
||||||
p.Write([]byte("first"))
|
|
||||||
p.Reset()
|
|
||||||
p.Write([]byte("second"))
|
|
||||||
|
|
||||||
result, ok := p.Read()
|
|
||||||
if !ok || result != "second" {
|
|
||||||
t.Errorf("expected 'second', got '%s'", result)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,470 +0,0 @@
|
|||||||
package auth
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"fmt"
|
|
||||||
"strings"
|
|
||||||
"time"
|
|
||||||
|
|
||||||
kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
|
|
||||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
|
||||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
|
||||||
)
|
|
||||||
|
|
||||||
// extractKiroIdentifier extracts a meaningful identifier for file naming.
|
|
||||||
// Returns account name if provided, otherwise profile ARN ID.
|
|
||||||
// All extracted values are sanitized to prevent path injection attacks.
|
|
||||||
func extractKiroIdentifier(accountName, profileArn string) string {
|
|
||||||
// Priority 1: Use account name if provided
|
|
||||||
if accountName != "" {
|
|
||||||
return kiroauth.SanitizeEmailForFilename(accountName)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Priority 2: Use profile ARN ID part (sanitized to prevent path injection)
|
|
||||||
if profileArn != "" {
|
|
||||||
parts := strings.Split(profileArn, "/")
|
|
||||||
if len(parts) >= 2 {
|
|
||||||
// Sanitize the ARN component to prevent path traversal
|
|
||||||
return kiroauth.SanitizeEmailForFilename(parts[len(parts)-1])
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Fallback: timestamp
|
|
||||||
return fmt.Sprintf("%d", time.Now().UnixNano()%100000)
|
|
||||||
}
|
|
||||||
|
|
||||||
// KiroAuthenticator implements OAuth authentication for Kiro with Google login.
|
|
||||||
type KiroAuthenticator struct{}
|
|
||||||
|
|
||||||
// NewKiroAuthenticator constructs a Kiro authenticator.
|
|
||||||
func NewKiroAuthenticator() *KiroAuthenticator {
|
|
||||||
return &KiroAuthenticator{}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Provider returns the provider key for the authenticator.
|
|
||||||
func (a *KiroAuthenticator) Provider() string {
|
|
||||||
return "kiro"
|
|
||||||
}
|
|
||||||
|
|
||||||
// RefreshLead indicates how soon before expiry a refresh should be attempted.
|
|
||||||
// Set to 5 minutes to match Antigravity and avoid frequent refresh checks while still ensuring timely token refresh.
|
|
||||||
func (a *KiroAuthenticator) RefreshLead() *time.Duration {
|
|
||||||
d := 5 * time.Minute
|
|
||||||
return &d
|
|
||||||
}
|
|
||||||
|
|
||||||
// createAuthRecord creates an auth record from token data.
|
|
||||||
func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, source string) (*coreauth.Auth, error) {
|
|
||||||
// Parse expires_at
|
|
||||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
|
||||||
if err != nil {
|
|
||||||
expiresAt = time.Now().Add(1 * time.Hour)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Extract identifier for file naming
|
|
||||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
|
||||||
|
|
||||||
// Determine label based on auth method
|
|
||||||
label := fmt.Sprintf("kiro-%s", source)
|
|
||||||
if tokenData.AuthMethod == "idc" {
|
|
||||||
label = "kiro-idc"
|
|
||||||
}
|
|
||||||
|
|
||||||
now := time.Now()
|
|
||||||
fileName := fmt.Sprintf("%s-%s.json", label, idPart)
|
|
||||||
|
|
||||||
metadata := map[string]any{
|
|
||||||
"type": "kiro",
|
|
||||||
"access_token": tokenData.AccessToken,
|
|
||||||
"refresh_token": tokenData.RefreshToken,
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"expires_at": tokenData.ExpiresAt,
|
|
||||||
"auth_method": tokenData.AuthMethod,
|
|
||||||
"provider": tokenData.Provider,
|
|
||||||
"client_id": tokenData.ClientID,
|
|
||||||
"client_secret": tokenData.ClientSecret,
|
|
||||||
"email": tokenData.Email,
|
|
||||||
}
|
|
||||||
|
|
||||||
// Add IDC-specific fields if present
|
|
||||||
if tokenData.StartURL != "" {
|
|
||||||
metadata["start_url"] = tokenData.StartURL
|
|
||||||
}
|
|
||||||
if tokenData.Region != "" {
|
|
||||||
metadata["region"] = tokenData.Region
|
|
||||||
}
|
|
||||||
|
|
||||||
attributes := map[string]string{
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"source": source,
|
|
||||||
"email": tokenData.Email,
|
|
||||||
}
|
|
||||||
|
|
||||||
// Add IDC-specific attributes if present
|
|
||||||
if tokenData.AuthMethod == "idc" {
|
|
||||||
attributes["source"] = "aws-idc"
|
|
||||||
if tokenData.StartURL != "" {
|
|
||||||
attributes["start_url"] = tokenData.StartURL
|
|
||||||
}
|
|
||||||
if tokenData.Region != "" {
|
|
||||||
attributes["region"] = tokenData.Region
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
record := &coreauth.Auth{
|
|
||||||
ID: fileName,
|
|
||||||
Provider: "kiro",
|
|
||||||
FileName: fileName,
|
|
||||||
Label: label,
|
|
||||||
Status: coreauth.StatusActive,
|
|
||||||
CreatedAt: now,
|
|
||||||
UpdatedAt: now,
|
|
||||||
Metadata: metadata,
|
|
||||||
Attributes: attributes,
|
|
||||||
// NextRefreshAfter is aligned with RefreshLead (5min)
|
|
||||||
NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
|
|
||||||
}
|
|
||||||
|
|
||||||
if tokenData.Email != "" {
|
|
||||||
fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email)
|
|
||||||
} else {
|
|
||||||
fmt.Println("\n✓ Kiro authentication completed successfully!")
|
|
||||||
}
|
|
||||||
|
|
||||||
return record, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// Login performs OAuth login for Kiro with AWS (Builder ID or IDC).
|
|
||||||
// This shows a method selection prompt and handles both flows.
|
|
||||||
func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
|
||||||
if cfg == nil {
|
|
||||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
|
||||||
}
|
|
||||||
|
|
||||||
// Use the unified method selection flow (Builder ID or IDC)
|
|
||||||
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
|
|
||||||
tokenData, err := ssoClient.LoginWithMethodSelection(ctx)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("login failed: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
return a.createAuthRecord(tokenData, "aws")
|
|
||||||
}
|
|
||||||
|
|
||||||
// LoginWithAuthCode performs OAuth login for Kiro with AWS Builder ID using authorization code flow.
|
|
||||||
// This provides a better UX than device code flow as it uses automatic browser callback.
|
|
||||||
func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
|
||||||
if cfg == nil {
|
|
||||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
|
||||||
}
|
|
||||||
|
|
||||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
|
||||||
|
|
||||||
// Use AWS Builder ID authorization code flow
|
|
||||||
tokenData, err := oauth.LoginWithBuilderIDAuthCode(ctx)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("login failed: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Parse expires_at
|
|
||||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
|
||||||
if err != nil {
|
|
||||||
expiresAt = time.Now().Add(1 * time.Hour)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Extract identifier for file naming
|
|
||||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
|
||||||
|
|
||||||
now := time.Now()
|
|
||||||
fileName := fmt.Sprintf("kiro-aws-%s.json", idPart)
|
|
||||||
|
|
||||||
record := &coreauth.Auth{
|
|
||||||
ID: fileName,
|
|
||||||
Provider: "kiro",
|
|
||||||
FileName: fileName,
|
|
||||||
Label: "kiro-aws",
|
|
||||||
Status: coreauth.StatusActive,
|
|
||||||
CreatedAt: now,
|
|
||||||
UpdatedAt: now,
|
|
||||||
Metadata: map[string]any{
|
|
||||||
"type": "kiro",
|
|
||||||
"access_token": tokenData.AccessToken,
|
|
||||||
"refresh_token": tokenData.RefreshToken,
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"expires_at": tokenData.ExpiresAt,
|
|
||||||
"auth_method": tokenData.AuthMethod,
|
|
||||||
"provider": tokenData.Provider,
|
|
||||||
"client_id": tokenData.ClientID,
|
|
||||||
"client_secret": tokenData.ClientSecret,
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
Attributes: map[string]string{
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"source": "aws-builder-id-authcode",
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
// NextRefreshAfter is aligned with RefreshLead (5min)
|
|
||||||
NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
|
|
||||||
}
|
|
||||||
|
|
||||||
if tokenData.Email != "" {
|
|
||||||
fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email)
|
|
||||||
} else {
|
|
||||||
fmt.Println("\n✓ Kiro authentication completed successfully!")
|
|
||||||
}
|
|
||||||
|
|
||||||
return record, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// LoginWithGoogle performs OAuth login for Kiro with Google.
|
|
||||||
// This uses a custom protocol handler (kiro://) to receive the callback.
|
|
||||||
func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
|
||||||
if cfg == nil {
|
|
||||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
|
||||||
}
|
|
||||||
|
|
||||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
|
||||||
|
|
||||||
// Use Google OAuth flow with protocol handler
|
|
||||||
tokenData, err := oauth.LoginWithGoogle(ctx)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("google login failed: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Parse expires_at
|
|
||||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
|
||||||
if err != nil {
|
|
||||||
expiresAt = time.Now().Add(1 * time.Hour)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Extract identifier for file naming
|
|
||||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
|
||||||
|
|
||||||
now := time.Now()
|
|
||||||
fileName := fmt.Sprintf("kiro-google-%s.json", idPart)
|
|
||||||
|
|
||||||
record := &coreauth.Auth{
|
|
||||||
ID: fileName,
|
|
||||||
Provider: "kiro",
|
|
||||||
FileName: fileName,
|
|
||||||
Label: "kiro-google",
|
|
||||||
Status: coreauth.StatusActive,
|
|
||||||
CreatedAt: now,
|
|
||||||
UpdatedAt: now,
|
|
||||||
Metadata: map[string]any{
|
|
||||||
"type": "kiro",
|
|
||||||
"access_token": tokenData.AccessToken,
|
|
||||||
"refresh_token": tokenData.RefreshToken,
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"expires_at": tokenData.ExpiresAt,
|
|
||||||
"auth_method": tokenData.AuthMethod,
|
|
||||||
"provider": tokenData.Provider,
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
Attributes: map[string]string{
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"source": "google-oauth",
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
// NextRefreshAfter is aligned with RefreshLead (5min)
|
|
||||||
NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
|
|
||||||
}
|
|
||||||
|
|
||||||
if tokenData.Email != "" {
|
|
||||||
fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email)
|
|
||||||
} else {
|
|
||||||
fmt.Println("\n✓ Kiro Google authentication completed successfully!")
|
|
||||||
}
|
|
||||||
|
|
||||||
return record, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// LoginWithGitHub performs OAuth login for Kiro with GitHub.
|
|
||||||
// This uses a custom protocol handler (kiro://) to receive the callback.
|
|
||||||
func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
|
||||||
if cfg == nil {
|
|
||||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
|
||||||
}
|
|
||||||
|
|
||||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
|
||||||
|
|
||||||
// Use GitHub OAuth flow with protocol handler
|
|
||||||
tokenData, err := oauth.LoginWithGitHub(ctx)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("github login failed: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Parse expires_at
|
|
||||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
|
||||||
if err != nil {
|
|
||||||
expiresAt = time.Now().Add(1 * time.Hour)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Extract identifier for file naming
|
|
||||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
|
||||||
|
|
||||||
now := time.Now()
|
|
||||||
fileName := fmt.Sprintf("kiro-github-%s.json", idPart)
|
|
||||||
|
|
||||||
record := &coreauth.Auth{
|
|
||||||
ID: fileName,
|
|
||||||
Provider: "kiro",
|
|
||||||
FileName: fileName,
|
|
||||||
Label: "kiro-github",
|
|
||||||
Status: coreauth.StatusActive,
|
|
||||||
CreatedAt: now,
|
|
||||||
UpdatedAt: now,
|
|
||||||
Metadata: map[string]any{
|
|
||||||
"type": "kiro",
|
|
||||||
"access_token": tokenData.AccessToken,
|
|
||||||
"refresh_token": tokenData.RefreshToken,
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"expires_at": tokenData.ExpiresAt,
|
|
||||||
"auth_method": tokenData.AuthMethod,
|
|
||||||
"provider": tokenData.Provider,
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
Attributes: map[string]string{
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"source": "github-oauth",
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
// NextRefreshAfter is aligned with RefreshLead (5min)
|
|
||||||
NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
|
|
||||||
}
|
|
||||||
|
|
||||||
if tokenData.Email != "" {
|
|
||||||
fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email)
|
|
||||||
} else {
|
|
||||||
fmt.Println("\n✓ Kiro GitHub authentication completed successfully!")
|
|
||||||
}
|
|
||||||
|
|
||||||
return record, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// ImportFromKiroIDE imports token from Kiro IDE's token file.
|
|
||||||
func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.Config) (*coreauth.Auth, error) {
|
|
||||||
tokenData, err := kiroauth.LoadKiroIDEToken()
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to load Kiro IDE token: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Parse expires_at
|
|
||||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
|
||||||
if err != nil {
|
|
||||||
expiresAt = time.Now().Add(1 * time.Hour)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Extract email from JWT if not already set (for imported tokens)
|
|
||||||
if tokenData.Email == "" {
|
|
||||||
tokenData.Email = kiroauth.ExtractEmailFromJWT(tokenData.AccessToken)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Extract identifier for file naming
|
|
||||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
|
||||||
// Sanitize provider to prevent path traversal (defense-in-depth)
|
|
||||||
provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider)))
|
|
||||||
if provider == "" {
|
|
||||||
provider = "imported" // Fallback for legacy tokens without provider
|
|
||||||
}
|
|
||||||
|
|
||||||
now := time.Now()
|
|
||||||
fileName := fmt.Sprintf("kiro-%s-%s.json", provider, idPart)
|
|
||||||
|
|
||||||
record := &coreauth.Auth{
|
|
||||||
ID: fileName,
|
|
||||||
Provider: "kiro",
|
|
||||||
FileName: fileName,
|
|
||||||
Label: fmt.Sprintf("kiro-%s", provider),
|
|
||||||
Status: coreauth.StatusActive,
|
|
||||||
CreatedAt: now,
|
|
||||||
UpdatedAt: now,
|
|
||||||
Metadata: map[string]any{
|
|
||||||
"type": "kiro",
|
|
||||||
"access_token": tokenData.AccessToken,
|
|
||||||
"refresh_token": tokenData.RefreshToken,
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"expires_at": tokenData.ExpiresAt,
|
|
||||||
"auth_method": tokenData.AuthMethod,
|
|
||||||
"provider": tokenData.Provider,
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
Attributes: map[string]string{
|
|
||||||
"profile_arn": tokenData.ProfileArn,
|
|
||||||
"source": "kiro-ide-import",
|
|
||||||
"email": tokenData.Email,
|
|
||||||
},
|
|
||||||
// NextRefreshAfter is aligned with RefreshLead (5min)
|
|
||||||
NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
|
|
||||||
}
|
|
||||||
|
|
||||||
// Display the email if extracted
|
|
||||||
if tokenData.Email != "" {
|
|
||||||
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s, Account: %s)\n", tokenData.Provider, tokenData.Email)
|
|
||||||
} else {
|
|
||||||
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s)\n", tokenData.Provider)
|
|
||||||
}
|
|
||||||
|
|
||||||
return record, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// Refresh refreshes an expired Kiro token using AWS SSO OIDC.
|
|
||||||
func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
|
||||||
if auth == nil || auth.Metadata == nil {
|
|
||||||
return nil, fmt.Errorf("invalid auth record")
|
|
||||||
}
|
|
||||||
|
|
||||||
refreshToken, ok := auth.Metadata["refresh_token"].(string)
|
|
||||||
if !ok || refreshToken == "" {
|
|
||||||
return nil, fmt.Errorf("refresh token not found")
|
|
||||||
}
|
|
||||||
|
|
||||||
clientID, _ := auth.Metadata["client_id"].(string)
|
|
||||||
clientSecret, _ := auth.Metadata["client_secret"].(string)
|
|
||||||
authMethod, _ := auth.Metadata["auth_method"].(string)
|
|
||||||
startURL, _ := auth.Metadata["start_url"].(string)
|
|
||||||
region, _ := auth.Metadata["region"].(string)
|
|
||||||
|
|
||||||
var tokenData *kiroauth.KiroTokenData
|
|
||||||
var err error
|
|
||||||
|
|
||||||
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
|
|
||||||
|
|
||||||
// Use SSO OIDC refresh for AWS Builder ID or IDC, otherwise use Kiro's OAuth refresh endpoint
|
|
||||||
switch {
|
|
||||||
case clientID != "" && clientSecret != "" && authMethod == "idc" && region != "":
|
|
||||||
// IDC refresh with region-specific endpoint
|
|
||||||
tokenData, err = ssoClient.RefreshTokenWithRegion(ctx, clientID, clientSecret, refreshToken, region, startURL)
|
|
||||||
case clientID != "" && clientSecret != "" && authMethod == "builder-id":
|
|
||||||
// Builder ID refresh with default endpoint
|
|
||||||
tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
|
|
||||||
default:
|
|
||||||
// Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub)
|
|
||||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
|
||||||
tokenData, err = oauth.RefreshToken(ctx, refreshToken)
|
|
||||||
}
|
|
||||||
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("token refresh failed: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Parse expires_at
|
|
||||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
|
||||||
if err != nil {
|
|
||||||
expiresAt = time.Now().Add(1 * time.Hour)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Clone auth to avoid mutating the input parameter
|
|
||||||
updated := auth.Clone()
|
|
||||||
now := time.Now()
|
|
||||||
updated.UpdatedAt = now
|
|
||||||
updated.LastRefreshedAt = now
|
|
||||||
updated.Metadata["access_token"] = tokenData.AccessToken
|
|
||||||
updated.Metadata["refresh_token"] = tokenData.RefreshToken
|
|
||||||
updated.Metadata["expires_at"] = tokenData.ExpiresAt
|
|
||||||
updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization
|
|
||||||
// NextRefreshAfter is aligned with RefreshLead (5min)
|
|
||||||
updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute)
|
|
||||||
|
|
||||||
return updated, nil
|
|
||||||
}
|
|
||||||
Reference in New Issue
Block a user