sub2api/backend/internal/service/gateway_attribution.go

package service

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"

	"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.88.
// Source: src/constants/system.ts + src/utils/fingerprint.ts
const (
	// fingerprintSalt must match the hardcoded salt in the real CLI.
	// Source: extracted/src/utils/fingerprint.ts:8
	fingerprintSalt = "59cf53e54c78"
)

// 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) string {
	version := cliVersion + "." + fingerprint
	// 注意：cch 字段由 Bun 的 NATIVE_CLIENT_ATTESTATION 编译时 feature 控制。
	// npm 安装版本（非原生二进制）此 feature 为 false，所以不包含 cch 字段。
	// 只有原生二进制安装（Bun 打包）才会有 cch，且其值会被 Bun 的 Zig 层替换为真实 hash。
	// 我们模拟 npm 安装版本的行为：不包含 cch。
	return fmt.Sprintf("x-anthropic-billing-header: cc_version=%s; cc_entrypoint=cli;", version)
}

// 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) []byte {
	// Compute fingerprint from the first user message
	firstMsgText := extractFirstUserMessageText(body)
	fingerprint := computeAttributionFingerprint(firstMsgText, cliVersion)
	attribution := buildAttributionBlock(cliVersion, fingerprint)

	// 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
	}
}

// generateSessionIDForAccount generates a deterministic per-account session UUID
// that remains stable within a process-like timeframe.
// Uses instanceSalt + accountID to ensure uniqueness across sub2api instances.
func generateSessionIDForAccount(instanceSalt string, accountID int64) string {
	// Use a per-account stable UUID (like real CLI's per-process UUID).
	// We use accountID as the base — each account gets a different "session".
	seed := fmt.Sprintf("session:%s:%d", instanceSalt, accountID)
	hash := sha256.Sum256([]byte(seed))
	sessionUUID, err := uuid.FromBytes(hash[:16])
	if err != nil {
		return uuid.New().String()
	}
	// Set UUID v4 variant
	sessionUUID[6] = (sessionUUID[6] & 0x0f) | 0x40
	sessionUUID[8] = (sessionUUID[8] & 0x3f) | 0x80
	return sessionUUID.String()
}