refactor(auth): extract token filename generation into unified function

Add ExtractIDCIdentifier and GenerateTokenFileName functions to centralize
token filename generation logic. This improves code maintainability by:

- Extracting IDC identifier from startUrl for unique token file naming
- Supporting priority-based filename generation (email > startUrl > authMethod)
- Removing duplicate filename generation code from oauth_web.go
- Adding comprehensive unit tests for the new functions
This commit is contained in:
yuechenglong.5
2026-01-26 13:54:32 +08:00
parent fe8c7a62aa
commit f74a688fb9
3 changed files with 223 additions and 27 deletions

View File

@@ -360,7 +360,7 @@ func SanitizeEmailForFilename(email string) string {
}
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)
@@ -378,7 +378,7 @@ func SanitizeEmailForFilename(email string) string {
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, "_")
@@ -389,6 +389,65 @@ func SanitizeEmailForFilename(email string) string {
parts[i] = part
}
result = strings.Join(parts, "_")
return result
}
// ExtractIDCIdentifier extracts a unique identifier from IDC startUrl.
// Examples:
// - "https://d-1234567890.awsapps.com/start" -> "d-1234567890"
// - "https://my-company.awsapps.com/start" -> "my-company"
// - "https://acme-corp.awsapps.com/start" -> "acme-corp"
func ExtractIDCIdentifier(startURL string) string {
if startURL == "" {
return ""
}
// Remove protocol prefix
url := strings.TrimPrefix(startURL, "https://")
url = strings.TrimPrefix(url, "http://")
// Extract subdomain (first part before the first dot)
// Format: {identifier}.awsapps.com/start
parts := strings.Split(url, ".")
if len(parts) > 0 && parts[0] != "" {
identifier := parts[0]
// Sanitize for filename safety
identifier = strings.ReplaceAll(identifier, "/", "_")
identifier = strings.ReplaceAll(identifier, "\\", "_")
identifier = strings.ReplaceAll(identifier, ":", "_")
return identifier
}
return ""
}
// GenerateTokenFileName generates a unique filename for token storage.
// Priority: email > startUrl identifier (for IDC) > authMethod only
// Format: kiro-{authMethod}-{identifier}.json
func GenerateTokenFileName(tokenData *KiroTokenData) string {
authMethod := tokenData.AuthMethod
if authMethod == "" {
authMethod = "unknown"
}
// Priority 1: Use email if available
if tokenData.Email != "" {
// Sanitize email for filename (replace @ and . with -)
sanitizedEmail := tokenData.Email
sanitizedEmail = strings.ReplaceAll(sanitizedEmail, "@", "-")
sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-")
return fmt.Sprintf("kiro-%s-%s.json", authMethod, sanitizedEmail)
}
// Priority 2: For IDC, use startUrl identifier
if authMethod == "idc" && tokenData.StartURL != "" {
identifier := ExtractIDCIdentifier(tokenData.StartURL)
if identifier != "" {
return fmt.Sprintf("kiro-%s-%s.json", authMethod, identifier)
}
}
// Priority 3: Fallback to authMethod only
return fmt.Sprintf("kiro-%s.json", authMethod)
}

View File

@@ -151,11 +151,161 @@ func TestSanitizeEmailForFilename(t *testing.T) {
// createTestJWT creates a test JWT token with the given claims
func createTestJWT(claims map[string]any) string {
header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
payloadBytes, _ := json.Marshal(claims)
payload := base64.RawURLEncoding.EncodeToString(payloadBytes)
signature := base64.RawURLEncoding.EncodeToString([]byte("fake-signature"))
return header + "." + payload + "." + signature
}
func TestExtractIDCIdentifier(t *testing.T) {
tests := []struct {
name string
startURL string
expected string
}{
{
name: "Empty URL",
startURL: "",
expected: "",
},
{
name: "Standard IDC URL with d- prefix",
startURL: "https://d-1234567890.awsapps.com/start",
expected: "d-1234567890",
},
{
name: "IDC URL with company name",
startURL: "https://my-company.awsapps.com/start",
expected: "my-company",
},
{
name: "IDC URL with simple name",
startURL: "https://acme-corp.awsapps.com/start",
expected: "acme-corp",
},
{
name: "IDC URL without https",
startURL: "http://d-9876543210.awsapps.com/start",
expected: "d-9876543210",
},
{
name: "IDC URL with subdomain only",
startURL: "https://test.awsapps.com/start",
expected: "test",
},
{
name: "Builder ID URL",
startURL: "https://view.awsapps.com/start",
expected: "view",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ExtractIDCIdentifier(tt.startURL)
if result != tt.expected {
t.Errorf("ExtractIDCIdentifier() = %q, want %q", result, tt.expected)
}
})
}
}
func TestGenerateTokenFileName(t *testing.T) {
tests := []struct {
name string
tokenData *KiroTokenData
expected string
}{
{
name: "IDC with email",
tokenData: &KiroTokenData{
AuthMethod: "idc",
Email: "user@example.com",
StartURL: "https://d-1234567890.awsapps.com/start",
},
expected: "kiro-idc-user-example-com.json",
},
{
name: "IDC without email but with startUrl",
tokenData: &KiroTokenData{
AuthMethod: "idc",
Email: "",
StartURL: "https://d-1234567890.awsapps.com/start",
},
expected: "kiro-idc-d-1234567890.json",
},
{
name: "IDC with company name in startUrl",
tokenData: &KiroTokenData{
AuthMethod: "idc",
Email: "",
StartURL: "https://my-company.awsapps.com/start",
},
expected: "kiro-idc-my-company.json",
},
{
name: "IDC without email and without startUrl",
tokenData: &KiroTokenData{
AuthMethod: "idc",
Email: "",
StartURL: "",
},
expected: "kiro-idc.json",
},
{
name: "Builder ID with email",
tokenData: &KiroTokenData{
AuthMethod: "builder-id",
Email: "user@gmail.com",
StartURL: "https://view.awsapps.com/start",
},
expected: "kiro-builder-id-user-gmail-com.json",
},
{
name: "Builder ID without email",
tokenData: &KiroTokenData{
AuthMethod: "builder-id",
Email: "",
StartURL: "https://view.awsapps.com/start",
},
expected: "kiro-builder-id.json",
},
{
name: "Social auth with email",
tokenData: &KiroTokenData{
AuthMethod: "google",
Email: "user@gmail.com",
},
expected: "kiro-google-user-gmail-com.json",
},
{
name: "Empty auth method",
tokenData: &KiroTokenData{
AuthMethod: "",
Email: "",
},
expected: "kiro-unknown.json",
},
{
name: "Email with special characters",
tokenData: &KiroTokenData{
AuthMethod: "idc",
Email: "user.name+tag@sub.example.com",
StartURL: "https://d-1234567890.awsapps.com/start",
},
expected: "kiro-idc-user-name+tag-sub-example-com.json",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := GenerateTokenFileName(tt.tokenData)
if result != tt.expected {
t.Errorf("GenerateTokenFileName() = %q, want %q", result, tt.expected)
}
})
}
}

View File

@@ -421,7 +421,7 @@ func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) {
log.Errorf("OAuth Web: failed to resolve auth directory: %v", err)
}
}
// Fall back to default location
if authDir == "" {
home, err := os.UserHomeDir()
@@ -431,24 +431,16 @@ func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) {
}
authDir = filepath.Join(home, ".cli-proxy-api")
}
// Create directory if not exists
if err := os.MkdirAll(authDir, 0700); err != nil {
log.Errorf("OAuth Web: failed to create auth directory: %v", err)
return
}
// Generate filename based on auth method
// Format: kiro-{authMethod}.json or kiro-{authMethod}-{email}.json
fileName := fmt.Sprintf("kiro-%s.json", tokenData.AuthMethod)
if tokenData.Email != "" {
// Sanitize email for filename (replace @ and . with -)
sanitizedEmail := tokenData.Email
sanitizedEmail = strings.ReplaceAll(sanitizedEmail, "@", "-")
sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-")
fileName = fmt.Sprintf("kiro-%s-%s.json", tokenData.AuthMethod, sanitizedEmail)
}
// Generate filename using the unified function
fileName := GenerateTokenFileName(tokenData)
authFilePath := filepath.Join(authDir, fileName)
// Convert to storage format and save
@@ -811,13 +803,8 @@ func (h *OAuthWebHandler) handleImportToken(c *gin.Context) {
// Save token to file
h.saveTokenToFile(tokenData)
// Generate filename for response
fileName := fmt.Sprintf("kiro-%s.json", tokenData.AuthMethod)
if tokenData.Email != "" {
sanitizedEmail := strings.ReplaceAll(tokenData.Email, "@", "-")
sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-")
fileName = fmt.Sprintf("kiro-%s-%s.json", tokenData.AuthMethod, sanitizedEmail)
}
// Generate filename for response using the unified function
fileName := GenerateTokenFileName(tokenData)
log.Infof("OAuth Web: token imported successfully")
c.JSON(http.StatusOK, gin.H{