Merge pull request #55 from kaitranntt/kai/fix/usage-detail-sanitization

fix: sanitize usage detail identifiers

Author: kaitranntt

internal/usage/logger_plugin.go

@@ -214,8 +214,8 @@ func (s *RequestStatistics) Record(ctx context.Context, record coreusage.Record)
 	s.updateAPIStats(stats, modelName, RequestDetail{
 		Timestamp: timestamp,
 		LatencyMs: normaliseLatency(record.Latency),
-		Source:    record.Source,
-		AuthIndex: record.AuthIndex,
+		Source:    sanitizeUsageDetailSource(record.Source),
+		AuthIndex: sanitizeUsageDetailAuthIndex(record.AuthIndex),
 		Tokens:    detail,
 		Failed:    failed,
 	})
@@ -308,6 +308,7 @@ func (s *RequestStatistics) Snapshot() StatisticsSnapshot {
 
 	result.APIs = make(map[string]APISnapshot, len(s.apis))
 	for apiName, stats := range s.apis {
+		apiName = sanitizeUsageAPIIdentifier(apiName)
 		apiSnapshot := APISnapshot{
 			TotalRequests: stats.TotalRequests,
 			TotalTokens:   stats.TotalTokens,
@@ -317,6 +318,9 @@ func (s *RequestStatistics) Snapshot() StatisticsSnapshot {
 		for modelName, modelStatsValue := range stats.Models {
 			requestDetails := make([]RequestDetail, len(modelStatsValue.Details))
 			copy(requestDetails, modelStatsValue.Details)
+			for i := range requestDetails {
+				requestDetails[i] = sanitizeRequestDetail(requestDetails[i])
+			}
 			apiSnapshot.Models[modelName] = ModelSnapshot{
 				TotalRequests: modelStatsValue.TotalRequests,
 				TotalTokens:   modelStatsValue.TotalTokens,
@@ -386,6 +390,7 @@ func (s *RequestStatistics) RestoreSnapshot(snapshot StatisticsSnapshot) MergeRe
 
 func (s *RequestStatistics) loadSnapshot(snapshot StatisticsSnapshot) {
 	for apiName, apiSnapshot := range snapshot.APIs {
+		apiName = sanitizeUsageAPIIdentifier(apiName)
 		stats := s.apiStatsForKey(apiName)
 		var apiRequests int64
 		var apiTokens int64
@@ -449,6 +454,7 @@ func normaliseRequestDetails(details []RequestDetail) ([]RequestDetail, detailAg
 		if detail.Timestamp.IsZero() {
 			detail.Timestamp = now
 		}
+		detail = sanitizeRequestDetail(detail)
 		out = append(out, detail)
 		totals.requests++
 		totals.tokens += detail.Tokens.TotalTokens
@@ -557,14 +563,127 @@ func normaliseUsageIdentifier(value string) string {
 	return trimRunes(value, maxUsageIdentifierRunes)
 }
 
+func sanitizeRequestDetail(detail RequestDetail) RequestDetail {
+	detail.Source = sanitizeUsageDetailSource(detail.Source)
+	detail.AuthIndex = sanitizeUsageDetailAuthIndex(detail.AuthIndex)
+	return detail
+}
+
+func sanitizeUsageAPIIdentifier(value string) string {
+	raw := strings.TrimSpace(value)
+	normalized := normaliseUsageIdentifier(raw)
+	if isSafeUsageIdentifier(normalized) {
+		return normalized
+	}
+	return "api-key:" + shortUsageHash(raw)
+}
+
+func sanitizeUsageDetailSource(value string) string {
+	raw := strings.TrimSpace(value)
+	normalized := normaliseUsageIdentifier(raw)
+	if normalized == unknownUsageBucket || isSafeUsageIdentifier(normalized) {
+		return normalized
+	}
+	return "source:" + shortUsageHash(raw)
+}
+
+func sanitizeUsageDetailAuthIndex(value string) string {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return ""
+	}
+	normalized := trimRunes(value, maxUsageIdentifierRunes)
+	lower := strings.ToLower(normalized)
+	if strings.HasPrefix(lower, "auth:") {
+		hash := strings.TrimPrefix(lower, "auth:")
+		if isHexIdentifier(hash, 12) || isHexIdentifier(hash, 16) || isHexIdentifier(hash, 64) {
+			return normalized
+		}
+		return "auth:" + shortUsageHash(value)
+	}
+	if isHexIdentifier(lower, 16) || isHexIdentifier(lower, 64) {
+		return normalized
+	}
+	return "auth:" + shortUsageHash(value)
+}
+
+func isSafeUsageIdentifier(value string) bool {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return true
+	}
+	lower := strings.ToLower(value)
+	if lower == unknownUsageBucket || lower == overflowUsageBucket {
+		return true
+	}
+	if strings.HasPrefix(lower, "api-key:") {
+		hash := strings.TrimPrefix(lower, "api-key:")
+		return isHexIdentifier(hash, 8) || isHexIdentifier(hash, 12) || isHexIdentifier(hash, 64)
+	}
+	if strings.HasPrefix(lower, "source:") {
+		hash := strings.TrimPrefix(lower, "source:")
+		return isHexIdentifier(hash, 12) || isHexIdentifier(hash, 64)
+	}
+	if isHTTPRoute(value) || isSafePathIdentifier(value) {
+		return true
+	}
+	switch lower {
+	case "gemini", "gemini-cli", "aistudio", "vertex", "claude", "codex", "openai",
+		"openai-compatibility", "openai-compatible", "antigravity", "github-copilot",
+		"gitlab", "cursor", "kiro", "kilo", "kimi", "iflow", "codebuddy", "local":
+		return true
+	default:
+		return false
+	}
+}
+
+func isSafePathIdentifier(value string) bool {
+	if !strings.HasPrefix(value, "/") {
+		return false
+	}
+	lower := strings.ToLower(value)
+	return !strings.Contains(value, "?") &&
+		!strings.Contains(lower, "key") &&
+		!strings.Contains(lower, "token") &&
+		!strings.Contains(lower, "auth")
+}
+
+func isHTTPRoute(value string) bool {
+	parts := strings.Fields(value)
+	if len(parts) != 2 || !strings.HasPrefix(parts[1], "/") {
+		return false
+	}
+	switch strings.ToUpper(parts[0]) {
+	case "GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS":
+		return true
+	default:
+		return false
+	}
+}
+
+func isHexIdentifier(value string, length int) bool {
+	if len(value) != length {
+		return false
+	}
+	for _, r := range value {
+		if (r < '0' || r > '9') && (r < 'a' || r > 'f') {
+			return false
+		}
+	}
+	return true
+}
+
 func secretUsageBucket(value string) string {
 	value = strings.TrimSpace(value)
 	if value == "" {
 		return ""
 	}
+	return fmt.Sprintf("api-key:%s", shortUsageHash(value)[:8])
+}
+
+func shortUsageHash(value string) string {
 	sum := sha256.Sum256([]byte(value))
-	hash := hex.EncodeToString(sum[:])[:8]
-	return fmt.Sprintf("api-key:%s", hash)
+	return hex.EncodeToString(sum[:])[:12]
 }
 
 func trimRunes(value string, maxRunes int) string {

internal/usage/logger_plugin_test.go

@@ -48,6 +48,77 @@ func TestRequestStatisticsRecordIncludesLatency(t *testing.T) {
 	}
 }
 
+func TestRequestStatisticsSanitizesDetailIdentifiers(t *testing.T) {
+	prevStatsEnabled := StatisticsEnabled()
+	SetStatisticsEnabled(true)
+	t.Cleanup(func() {
+		SetStatisticsEnabled(prevStatsEnabled)
+	})
+
+	const rawAPIKey = "sk-test-client-secret"
+	const rawSource = "person@example.com"
+	const rawAuthIndex = "raw-auth-index-secret"
+
+	stats := NewRequestStatistics()
+	stats.Record(context.Background(), coreusage.Record{
+		APIKey:      rawAPIKey,
+		Model:       "gpt-5.4",
+		Source:      rawSource,
+		AuthIndex:   rawAuthIndex,
+		RequestedAt: time.Date(2026, 3, 20, 12, 0, 0, 0, time.UTC),
+	})
+
+	snapshot := stats.Snapshot()
+	apiName := onlyAPIName(t, snapshot)
+	if strings.Contains(apiName, rawAPIKey) {
+		t.Fatalf("api bucket leaked raw key: %q", apiName)
+	}
+	details := onlyAPISnapshot(t, snapshot).Models["gpt-5.4"].Details
+	if len(details) != 1 {
+		t.Fatalf("details len = %d, want 1", len(details))
+	}
+	assertNoRawUsageIdentifier(t, details[0].Source, rawSource)
+	assertNoRawUsageIdentifier(t, details[0].AuthIndex, rawAuthIndex)
+}
+
+func TestRequestStatisticsRestoreSanitizesImportedDetails(t *testing.T) {
+	prevStatsEnabled := StatisticsEnabled()
+	SetStatisticsEnabled(true)
+	t.Cleanup(func() {
+		SetStatisticsEnabled(prevStatsEnabled)
+	})
+
+	const rawSource = "source:imported-person@example.com"
+	const rawAuthIndex = "auth:imported-auth-secret"
+	const rawAPIKey = "api-key:sk-imported-api-secret"
+
+	stats := NewRequestStatistics()
+	stats.RestoreSnapshot(StatisticsSnapshot{
+		APIs: map[string]APISnapshot{
+			rawAPIKey: {
+				Models: map[string]ModelSnapshot{
+					"gpt-5.4": {
+						Details: []RequestDetail{{
+							Timestamp: time.Date(2026, 3, 20, 12, 0, 0, 0, time.UTC),
+							Source:    rawSource,
+							AuthIndex: rawAuthIndex,
+						}},
+					},
+				},
+			},
+		},
+	})
+
+	snapshot := stats.Snapshot()
+	assertNoRawUsageIdentifier(t, onlyAPIName(t, snapshot), rawAPIKey)
+	details := onlyAPISnapshot(t, snapshot).Models["gpt-5.4"].Details
+	if len(details) != 1 {
+		t.Fatalf("details len = %d, want 1", len(details))
+	}
+	assertNoRawUsageIdentifier(t, details[0].Source, rawSource)
+	assertNoRawUsageIdentifier(t, details[0].AuthIndex, rawAuthIndex)
+}
+
 func TestRequestStatisticsUsesStableLoggingContext(t *testing.T) {
 	prevStatsEnabled := StatisticsEnabled()
 	SetStatisticsEnabled(true)
@@ -288,3 +359,16 @@ func requireSnapshotTotals(t *testing.T, got, want StatisticsSnapshot) {
 		t.Fatalf("tokens by hour[12] = %d, want %d", got.TokensByHour["12"], want.TokensByHour["12"])
 	}
 }
+
+func assertNoRawUsageIdentifier(t *testing.T, got, raw string) {
+	t.Helper()
+	if got == "" {
+		t.Fatalf("sanitized identifier is empty")
+	}
+	if got == raw || strings.Contains(got, raw) {
+		t.Fatalf("identifier leaked raw value: got %q raw %q", got, raw)
+	}
+	if !strings.HasPrefix(got, "source:") && !strings.HasPrefix(got, "auth:") && !strings.HasPrefix(got, "api-key:") {
+		t.Fatalf("identifier = %q, want stable redacted prefix", got)
+	}
+}

test/usage_logging_test.go

@@ -42,7 +42,10 @@ func TestGeminiExecutorRecordsSuccessfulZeroUsageInStatistics(t *testing.T) {
 	})
 
 	model, source := executeGeminiZeroUsage(t, "stats")
-	detail := waitForStatisticsDetail(t, "gemini", model, source)
+	detail := waitForStatisticsDetail(t, "gemini", model)
+	if detail.Source == source {
+		t.Fatalf("detail source leaked raw account identifier")
+	}
 	if detail.Failed {
 		t.Fatalf("detail failed = true, want false")
 	}
@@ -121,7 +124,7 @@ func waitForQueuedUsageModelTotalTokens(t *testing.T, wantProvider, wantModel st
 	t.Fatalf("timed out waiting for queued usage payload for provider=%q model=%q", wantProvider, wantModel)
 }
 
-func waitForStatisticsDetail(t *testing.T, apiName, model, source string) internalusage.RequestDetail {
+func waitForStatisticsDetail(t *testing.T, apiName, model string) internalusage.RequestDetail {
 	t.Helper()
 
 	deadline := time.Now().Add(2 * time.Second)
@@ -138,14 +141,14 @@ func waitForStatisticsDetail(t *testing.T, apiName, model, source string) intern
 			continue
 		}
 		for _, detail := range modelSnapshot.Details {
-			if detail.Source == source {
+			if detail.Source != "" {
 				return detail
 			}
 		}
 		time.Sleep(10 * time.Millisecond)
 	}
 
-	t.Fatalf("timed out waiting for statistics detail for api=%q model=%q source=%q", apiName, model, source)
+	t.Fatalf("timed out waiting for statistics detail for api=%q model=%q", apiName, model)
 	return internalusage.RequestDetail{}
 }