Skip to content

Commit 4a75a73

Browse files
kaitrannttkaitranntt
authored
fix(kiro): accept raw IDE token imports (#118)
1 parent 643ae0b commit 4a75a73

8 files changed

Lines changed: 581 additions & 42 deletions

File tree

‎internal/auth/kiro/aws.go‎

Lines changed: 34 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -495,17 +495,12 @@ func ExtractIDCIdentifier(startURL string) string {
495495
// when email is unavailable to prevent filename collisions.
496496
// Format: kiro-{authMethod}-{identifier}[-{seq}].json
497497
func GenerateTokenFileName(tokenData *KiroTokenData) string {
498-
authMethod := tokenData.AuthMethod
499-
if authMethod == "" {
500-
authMethod = "unknown"
501-
}
498+
authMethod := sanitizeTokenFileComponent(tokenData.AuthMethod, "unknown")
502499

503500
// Priority 1: Use email if available (no sequence needed, email is unique)
504501
if tokenData.Email != "" {
505502
// Sanitize email for filename (replace @ and . with -)
506-
sanitizedEmail := tokenData.Email
507-
sanitizedEmail = strings.ReplaceAll(sanitizedEmail, "@", "-")
508-
sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-")
503+
sanitizedEmail := sanitizeTokenFileComponent(tokenData.Email, "account")
509504
return fmt.Sprintf("kiro-%s-%s.json", authMethod, sanitizedEmail)
510505
}
511506

@@ -514,7 +509,7 @@ func GenerateTokenFileName(tokenData *KiroTokenData) string {
514509

515510
// Priority 2: For IDC, use startUrl identifier with sequence
516511
if authMethod == "idc" && tokenData.StartURL != "" {
517-
identifier := ExtractIDCIdentifier(tokenData.StartURL)
512+
identifier := sanitizeTokenFileComponent(ExtractIDCIdentifier(tokenData.StartURL), "")
518513
if identifier != "" {
519514
return fmt.Sprintf("kiro-%s-%s-%05d.json", authMethod, identifier, seq)
520515
}
@@ -524,6 +519,37 @@ func GenerateTokenFileName(tokenData *KiroTokenData) string {
524519
return fmt.Sprintf("kiro-%s-%05d.json", authMethod, seq)
525520
}
526521

522+
func sanitizeTokenFileComponent(value, fallback string) string {
523+
value = strings.TrimSpace(value)
524+
if value == "" {
525+
return fallback
526+
}
527+
528+
var builder strings.Builder
529+
lastDash := false
530+
for _, r := range value {
531+
keep := (r >= 'a' && r <= 'z') ||
532+
(r >= 'A' && r <= 'Z') ||
533+
(r >= '0' && r <= '9') ||
534+
r == '_' || r == '-' || r == '+'
535+
if keep {
536+
builder.WriteRune(r)
537+
lastDash = false
538+
continue
539+
}
540+
if !lastDash {
541+
builder.WriteByte('-')
542+
lastDash = true
543+
}
544+
}
545+
546+
safe := strings.Trim(builder.String(), "-_")
547+
if safe == "" {
548+
return fallback
549+
}
550+
return safe
551+
}
552+
527553
// DefaultKiroRegion is the fallback region when none is specified.
528554
const DefaultKiroRegion = "us-east-1"
529555

‎internal/auth/kiro/aws_auth.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,7 @@ func (k *KiroAuth) CreateTokenStorage(tokenData *KiroTokenData) *KiroTokenStorag
274274
LastRefresh: time.Now().Format(time.RFC3339),
275275
ClientID: tokenData.ClientID,
276276
ClientSecret: tokenData.ClientSecret,
277+
ClientIDHash: tokenData.ClientIDHash,
277278
Region: tokenData.Region,
278279
StartURL: tokenData.StartURL,
279280
Email: tokenData.Email,
@@ -314,6 +315,9 @@ func (k *KiroAuth) UpdateTokenStorage(storage *KiroTokenStorage, tokenData *Kiro
314315
if tokenData.ClientSecret != "" {
315316
storage.ClientSecret = tokenData.ClientSecret
316317
}
318+
if tokenData.ClientIDHash != "" {
319+
storage.ClientIDHash = tokenData.ClientIDHash
320+
}
317321
if tokenData.Region != "" {
318322
storage.Region = tokenData.Region
319323
}

‎internal/auth/kiro/aws_test.go‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,47 @@ func TestGenerateTokenFileName(t *testing.T) {
318318
}
319319
}
320320

321+
func TestGenerateTokenFileNameSanitizesUserControlledComponents(t *testing.T) {
322+
tests := []struct {
323+
name string
324+
tokenData *KiroTokenData
325+
}{
326+
{
327+
name: "auth method traversal",
328+
tokenData: &KiroTokenData{
329+
AuthMethod: "../evil",
330+
Email: "user@example.com",
331+
},
332+
},
333+
{
334+
name: "email traversal",
335+
tokenData: &KiroTokenData{
336+
AuthMethod: "social",
337+
Email: "../../owned@example.com",
338+
},
339+
},
340+
{
341+
name: "idc identifier traversal",
342+
tokenData: &KiroTokenData{
343+
AuthMethod: "idc",
344+
StartURL: "https://../../etc.awsapps.com/start",
345+
},
346+
},
347+
}
348+
349+
for _, tt := range tests {
350+
t.Run(tt.name, func(t *testing.T) {
351+
name := GenerateTokenFileName(tt.tokenData)
352+
if strings.Contains(name, "/") || strings.Contains(name, "\\") || strings.Contains(name, "..") {
353+
t.Fatalf("GenerateTokenFileName() = %q, contains unsafe path component", name)
354+
}
355+
if !strings.HasPrefix(name, "kiro-") || !strings.HasSuffix(name, ".json") {
356+
t.Fatalf("GenerateTokenFileName() = %q, want kiro-*.json", name)
357+
}
358+
})
359+
}
360+
}
361+
321362
func TestParseProfileARN(t *testing.T) {
322363
tests := []struct {
323364
name string

‎internal/auth/kiro/oauth_web.go‎

Lines changed: 84 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,16 @@ func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSess
396396
StartURL: session.startURL,
397397
}
398398

399+
if _, errSave := h.saveTokenToFile(tokenData); errSave != nil {
400+
log.Errorf("OAuth Web: failed to save token to file: %v", errSave)
401+
h.mu.Lock()
402+
session.status = statusFailed
403+
session.error = "Failed to save authentication tokens"
404+
session.completedAt = time.Now()
405+
h.mu.Unlock()
406+
return
407+
}
408+
399409
h.mu.Lock()
400410
session.status = statusSuccess
401411
session.completedAt = time.Now()
@@ -407,47 +417,51 @@ func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSess
407417
h.onTokenObtained(tokenData)
408418
}
409419

410-
// Save token to file
411-
h.saveTokenToFile(tokenData)
412-
413420
log.Infof("OAuth Web: authentication successful for %s", email)
414421
return
415422
}
416423
}
417424
}
418425

419-
// saveTokenToFile saves the token data to the auth directory
420-
func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) {
426+
// saveTokenToFile saves the token data to the auth directory.
427+
func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) (string, error) {
421428
// Get auth directory from config or use default
422429
authDir := ""
423430
if h.cfg != nil && h.cfg.AuthDir != "" {
424431
var err error
425432
authDir, err = util.ResolveAuthDir(h.cfg.AuthDir)
426433
if err != nil {
427-
log.Errorf("OAuth Web: failed to resolve auth directory: %v", err)
434+
return "", fmt.Errorf("failed to resolve auth directory: %w", err)
428435
}
429436
}
430437

431438
// Fall back to default location
432439
if authDir == "" {
433440
home, err := os.UserHomeDir()
434441
if err != nil {
435-
log.Errorf("OAuth Web: failed to get home directory: %v", err)
436-
return
442+
return "", fmt.Errorf("failed to get home directory: %w", err)
437443
}
438444
authDir = filepath.Join(home, ".cli-proxy-api")
439445
}
440446

441447
// Create directory if not exists
442448
if err := os.MkdirAll(authDir, 0700); err != nil {
443-
log.Errorf("OAuth Web: failed to create auth directory: %v", err)
444-
return
449+
return "", fmt.Errorf("failed to create auth directory: %w", err)
445450
}
446451

447452
// Generate filename using the unified function
448453
fileName := GenerateTokenFileName(tokenData)
449454

450455
authFilePath := filepath.Join(authDir, fileName)
456+
absAuthDir, errAuthDir := filepath.Abs(authDir)
457+
absAuthFilePath, errAuthFilePath := filepath.Abs(authFilePath)
458+
if errAuthDir != nil || errAuthFilePath != nil {
459+
return "", fmt.Errorf("failed to resolve token output path")
460+
}
461+
authDirPrefix := absAuthDir + string(os.PathSeparator)
462+
if absAuthFilePath != absAuthDir && !strings.HasPrefix(absAuthFilePath, authDirPrefix) {
463+
return "", fmt.Errorf("invalid token output path")
464+
}
451465

452466
// Convert to storage format and save
453467
storage := &KiroTokenStorage{
@@ -461,17 +475,18 @@ func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) {
461475
LastRefresh: time.Now().Format(time.RFC3339),
462476
ClientID: tokenData.ClientID,
463477
ClientSecret: tokenData.ClientSecret,
478+
ClientIDHash: tokenData.ClientIDHash,
464479
Region: tokenData.Region,
465480
StartURL: tokenData.StartURL,
466481
Email: tokenData.Email,
467482
}
468483

469484
if err := storage.SaveTokenToFile(authFilePath); err != nil {
470-
log.Errorf("OAuth Web: failed to save token to file: %v", err)
471-
return
485+
return "", fmt.Errorf("failed to save token to file: %w", err)
472486
}
473487

474488
log.Infof("OAuth Web: token saved to %s", authFilePath)
489+
return fileName, nil
475490
}
476491

477492
func (h *OAuthWebHandler) ssoClient(session *webAuthSession) *SSOOIDCClient {
@@ -591,6 +606,17 @@ func (h *OAuthWebHandler) handleSocialCallback(c *gin.Context) {
591606
Region: "us-east-1",
592607
}
593608

609+
if _, errSave := h.saveTokenToFile(tokenData); errSave != nil {
610+
log.Errorf("OAuth Web: failed to save social token: %v", errSave)
611+
h.mu.Lock()
612+
session.status = statusFailed
613+
session.error = "Failed to save authentication tokens"
614+
session.completedAt = time.Now()
615+
h.mu.Unlock()
616+
h.renderError(c, session.error)
617+
return
618+
}
619+
594620
h.mu.Lock()
595621
session.status = statusSuccess
596622
session.completedAt = time.Now()
@@ -606,9 +632,6 @@ func (h *OAuthWebHandler) handleSocialCallback(c *gin.Context) {
606632
h.onTokenObtained(tokenData)
607633
}
608634

609-
// Save token to file
610-
h.saveTokenToFile(tokenData)
611-
612635
log.Infof("OAuth Web: social authentication successful for %s via %s", email, provider)
613636
h.renderSuccess(c, session)
614637
}
@@ -751,18 +774,50 @@ type ImportTokenRequest struct {
751774
RefreshToken string `json:"refreshToken"`
752775
}
753776

754-
// handleImportToken handles manual refresh token import from Kiro IDE
777+
// handleImportToken handles manual token imports from Kiro IDE.
755778
func (h *OAuthWebHandler) handleImportToken(c *gin.Context) {
756-
var req ImportTokenRequest
757-
if err := c.ShouldBindJSON(&req); err != nil {
779+
data, err := c.GetRawData()
780+
if err != nil {
758781
c.JSON(http.StatusBadRequest, gin.H{
759782
"success": false,
760783
"error": "Invalid request body",
761784
})
762785
return
763786
}
764787

765-
refreshToken := strings.TrimSpace(req.RefreshToken)
788+
parsed, err := parseImportTokenPayload(data)
789+
if err != nil {
790+
c.JSON(http.StatusBadRequest, gin.H{
791+
"success": false,
792+
"error": err.Error(),
793+
})
794+
return
795+
}
796+
if parsed.rawKiroIDEJSON {
797+
tokenData := parsed.tokenData
798+
fileName, errSave := h.saveTokenToFile(tokenData)
799+
if errSave != nil {
800+
log.Errorf("OAuth Web: failed to save imported token: %v", errSave)
801+
c.JSON(http.StatusInternalServerError, gin.H{
802+
"success": false,
803+
"error": "Failed to save token",
804+
})
805+
return
806+
}
807+
if h.onTokenObtained != nil {
808+
h.onTokenObtained(tokenData)
809+
}
810+
811+
log.Infof("OAuth Web: raw Kiro IDE token imported successfully")
812+
c.JSON(http.StatusOK, gin.H{
813+
"success": true,
814+
"message": "Token imported successfully",
815+
"fileName": fileName,
816+
})
817+
return
818+
}
819+
820+
refreshToken := strings.TrimSpace(parsed.tokenData.RefreshToken)
766821
if refreshToken == "" {
767822
c.JSON(http.StatusBadRequest, gin.H{
768823
"success": false,
@@ -801,17 +856,20 @@ func (h *OAuthWebHandler) handleImportToken(c *gin.Context) {
801856
tokenData.AuthMethod = "social"
802857
tokenData.Provider = "imported"
803858

804-
// Notify callback if set
859+
// Save token to file
860+
fileName, errSave := h.saveTokenToFile(tokenData)
861+
if errSave != nil {
862+
log.Errorf("OAuth Web: failed to save imported token: %v", errSave)
863+
c.JSON(http.StatusInternalServerError, gin.H{
864+
"success": false,
865+
"error": "Failed to save token",
866+
})
867+
return
868+
}
805869
if h.onTokenObtained != nil {
806870
h.onTokenObtained(tokenData)
807871
}
808872

809-
// Save token to file
810-
h.saveTokenToFile(tokenData)
811-
812-
// Generate filename for response using the unified function
813-
fileName := GenerateTokenFileName(tokenData)
814-
815873
log.Infof("OAuth Web: token imported successfully")
816874
c.JSON(http.StatusOK, gin.H{
817875
"success": true,

0 commit comments

Comments
 (0)