sub2api/backend/internal/service/gateway_attribution.go

package service

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	claude "github.com/Wei-Shaw/sub2api/internal/pkg/claude"
	"github.com/google/uuid"
	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"

	"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
)

// Attribution block constants matching real Claude Code 2.1.89.
// Source: src/constants/system.ts + src/utils/fingerprint.ts

type attributionBlockOptions struct {
	Entrypoint string
	Workload   string
	OmitCCH    bool
}

// computeAttributionFingerprint computes a 3-character hex fingerprint
// matching the algorithm in the real Claude Code CLI.
//
// Algorithm: SHA256(SALT + msg[4] + msg[7] + msg[20] + version)[:3]
// Source: extracted/src/utils/fingerprint.ts:50-63
func computeAttributionFingerprint(firstUserMessageText, cliVersion string) string {
	indices := [3]int{4, 7, 20}
	chars := make([]byte, 0, 3)
	for _, i := range indices {
		if i < len(firstUserMessageText) {
			chars = append(chars, firstUserMessageText[i])
		} else {
			chars = append(chars, '0')
		}
	}

	input := fmt.Sprintf("%s%s%s", fingerprintSalt, string(chars), cliVersion)
	hash := sha256.Sum256([]byte(input))
	return hex.EncodeToString(hash[:])[:3]
}

// extractFirstUserMessageText extracts text from the first user message in the body.
// Handles both string content and array content (text blocks).
func extractFirstUserMessageText(body []byte) string {
	messages := gjson.GetBytes(body, "messages")
	if !messages.Exists() || !messages.IsArray() {
		return ""
	}

	var firstText string
	messages.ForEach(func(_, msg gjson.Result) bool {
		if msg.Get("role").String() != "user" {
			return true // continue
		}
		content := msg.Get("content")
		if content.Type == gjson.String {
			firstText = content.String()
			return false // break
		}
		if content.IsArray() {
			content.ForEach(func(_, block gjson.Result) bool {
				if block.Get("type").String() == "text" {
					firstText = block.Get("text").String()
					return false
				}
				return true
			})
			return false
		}
		return true
	})
	return firstText
}

// buildAttributionBlock builds the x-anthropic-billing-header attribution string
// that real Claude Code injects as the first system text block.
//
// Format: x-anthropic-billing-header: cc_version=<VERSION>.<fingerprint>; cc_entrypoint=cli; cch=00000;
// Source: extracted/src/constants/system.ts:73-95
func buildAttributionBlock(cliVersion, fingerprint string, opts attributionBlockOptions) string {
	if claude.AttributionHeaderDisabled() {
		return ""
	}
	version := cliVersion + "." + fingerprint
	entrypoint := strings.TrimSpace(opts.Entrypoint)
	if entrypoint == "" {
		entrypoint = claude.CurrentEntrypoint()
	}
	workload := strings.TrimSpace(opts.Workload)
	if workload == "" {
		workload = claude.CurrentWorkload()
	}

	var b strings.Builder
	b.Grow(96)
	fmt.Fprintf(&b, "x-anthropic-billing-header: cc_version=%s; cc_entrypoint=%s;", version, entrypoint)
	if !opts.OmitCCH {
		// 2.1.89+ 的 Claude Code 在 1P / standard providers 下保留 cch=00000 占位符，
		// 由下游 attestation / signing 逻辑在需要时替换。
		b.WriteString(" cch=00000;")
	}
	if workload != "" {
		fmt.Fprintf(&b, " cc_workload=%s;", workload)
	}
	return b.String()
}

// injectAttributionBlock prepends the x-anthropic-billing-header attribution block
// as the very first system text block in the request body.
// This must come BEFORE the "You are Claude Code" block.
//
// The real CLI injects this as system[0] with cache_control: {type: "ephemeral"}.
func injectAttributionBlock(body []byte, cliVersion string, opts attributionBlockOptions) []byte {
	// Compute fingerprint from the first user message
	firstMsgText := extractFirstUserMessageText(body)
	fingerprint := computeAttributionFingerprint(firstMsgText, cliVersion)
	attribution := buildAttributionBlock(cliVersion, fingerprint, opts)
	if attribution == "" {
		return body
	}

	// Build the attribution text block as JSON
	attrBlock, err := marshalAnthropicSystemTextBlock(attribution, true)
	if err != nil {
		logger.LegacyPrintf("service.gateway", "Warning: failed to build attribution block: %v", err)
		return body
	}

	systemResult := gjson.GetBytes(body, "system")

	// Handle the different system formats
	switch {
	case !systemResult.Exists() || systemResult.Type == gjson.Null:
		// No system field — inject just the attribution block
		newBody, err := sjson.SetRawBytes(body, "system", buildJSONArrayRaw([][]byte{attrBlock}))
		if err != nil {
			return body
		}
		return newBody

	case systemResult.Type == gjson.String:
		// String system — convert to array: [attribution, original]
		origBlock, err := marshalAnthropicSystemTextBlock(systemResult.String(), false)
		if err != nil {
			return body
		}
		newBody, setErr := sjson.SetRawBytes(body, "system", buildJSONArrayRaw([][]byte{attrBlock, origBlock}))
		if setErr != nil {
			return body
		}
		return newBody

	case systemResult.IsArray():
		// Array system — check if attribution already exists, prepend if not
		var items [][]byte
		alreadyHasAttribution := false
		systemResult.ForEach(func(_, item gjson.Result) bool {
			if item.Get("type").String() == "text" {
				text := item.Get("text").String()
				if len(text) > 30 && text[:30] == "x-anthropic-billing-header: cc" {
					alreadyHasAttribution = true
				}
			}
			return true
		})
		if alreadyHasAttribution {
			return body
		}

		items = append(items, attrBlock)
		systemResult.ForEach(func(_, item gjson.Result) bool {
			items = append(items, []byte(item.Raw))
			return true
		})
		newBody, setErr := sjson.SetRawBytes(body, "system", buildJSONArrayRaw(items))
		if setErr != nil {
			return body
		}
		return newBody

	default:
		return body
	}
}

// cliSessionEntry holds a cached session UUID with an expiration time.
type cliSessionEntry struct {
	id        string
	expiresAt time.Time
}

// cliSessionCache stores per-account session UUIDs that rotate on a TTL.
// Real CLI creates a new random UUID per process invocation; we approximate
// this by rotating every 30-60 minutes (jittered per account).
var (
	cliSessionCache   = make(map[int64]cliSessionEntry)
	cliSessionCacheMu sync.Mutex
)

// sessionTTLBase is the base TTL for session ID rotation.
const sessionTTLBase = 30 * time.Minute

// generateSessionIDForAccount returns a per-account session UUID that rotates
// periodically. Each account gets a random TTL jitter (0-30 min on top of
// the 30 min base) so accounts don't all rotate simultaneously.
func generateSessionIDForAccount(instanceSalt string, accountID int64) string {
	cliSessionCacheMu.Lock()
	defer cliSessionCacheMu.Unlock()

	now := time.Now()
	if entry, ok := cliSessionCache[accountID]; ok && now.Before(entry.expiresAt) {
		return entry.id
	}

	// Compute per-account jitter from a hash so the same account always gets
	// the same jitter within a process (avoids re-rolling on every rotation).
	jitterSeed := fmt.Sprintf("jitter:%s:%d", instanceSalt, accountID)
	h := sha256.Sum256([]byte(jitterSeed))
	jitterMinutes := int(h[0]) % 31 // 0-30 minutes
	ttl := sessionTTLBase + time.Duration(jitterMinutes)*time.Minute

	newID := uuid.New().String()
	cliSessionCache[accountID] = cliSessionEntry{
		id:        newID,
		expiresAt: now.Add(ttl),
	}
	return newID
}

// reUserHome matches /Users/<username>/ or /home/<username>/ path segments.
// Captures the prefix (/Users/ or /home/) so we can preserve it while replacing the username.
var reUserHome = regexp.MustCompile(`(/(Users|home)/)[^/\s"']+/`)

// reEnvLine matches lines of the form "Key: value" for the environment block
// fields injected by Claude Code's CLAUDE.md / sysprompt machinery.
var reEnvLine = regexp.MustCompile(`(?m)^(Platform|Shell|OS Version|Working directory):.*$`)

// canonicalEnvValues maps environment block keys to their canonical replacements.
// Values mirror cc-gateway's prompt_env config and represent a stock macOS dev machine.
var canonicalEnvValues = map[string]string{
	"Platform":          "Platform: darwin",
	"Shell":             "Shell: zsh",
	"OS Version":        "OS Version: Darwin 24.4.0",
	"Working directory": "Working directory: /Users/user/project",
}

// NormalizeSystemPromptEnv rewrites environment-specific fields in a system
// prompt text block to canonical values, preventing real machine fingerprinting.
//
// Handles two classes of leakage (matching cc-gateway rewriter.ts:rewritePromptText):
//  1. "Platform: Windows / Linux / Darwin 25.x" → canonical darwin/zsh/Darwin 24.4.0
//  2. "/Users/alice/" or "/home/bob/" → "/Users/user/"
//
// Only called on system prompt text blocks, never on user message content.
func NormalizeSystemPromptEnv(text string) string {
	// Replace env-info lines with canonical values
	text = reEnvLine.ReplaceAllStringFunc(text, func(line string) string {
		for key, canonical := range canonicalEnvValues {
			if len(line) >= len(key) && line[:len(key)] == key {
				return canonical
			}
		}
		return line
	})

	// Redact real usernames in home directory paths
	// e.g. /Users/alice/project -> /Users/user/project
	text = reUserHome.ReplaceAllString(text, "${1}user/")

	return text
}