new-api/relay/upstream_request_snapshot.go

package relay

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

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
)

const maxUpstreamRequestSnapshotBytes = 256 * 1024

func buildUpstreamRequestSnapshot(raw []byte) *model.TaskUpstreamRequestSnapshot {
	if len(raw) == 0 {
		return nil
	}
	snapshot := &model.TaskUpstreamRequestSnapshot{
		RawBytes: len(raw),
		SHA256:   sha256Hex(raw),
	}

	var body any
	if err := common.Unmarshal(raw, &body); err == nil {
		sanitized, redacted := sanitizeUpstreamRequestValue(body)
		snapshot.Body = sanitized
		snapshot.Redacted = redacted
	} else {
		snapshot.Body = map[string]any{
			"raw_preview": truncateStringToBytes(string(raw), maxUpstreamRequestSnapshotBytes/2),
			"format":      "non_json",
		}
		snapshot.Truncated = len(raw) > maxUpstreamRequestSnapshotBytes/2
	}

	finalizeUpstreamRequestSnapshot(snapshot, raw)
	return snapshot
}

func BuildUpstreamRequestSnapshotForTask(raw []byte) *model.TaskUpstreamRequestSnapshot {
	return buildUpstreamRequestSnapshot(raw)
}

func sanitizeUpstreamRequestValue(value any) (any, bool) {
	switch v := value.(type) {
	case map[string]any:
		out := make(map[string]any, len(v))
		redacted := false
		for key, item := range v {
			sanitized, itemRedacted := sanitizeUpstreamRequestValue(item)
			out[key] = sanitized
			redacted = redacted || itemRedacted
		}
		return out, redacted
	case []any:
		out := make([]any, len(v))
		redacted := false
		for i, item := range v {
			sanitized, itemRedacted := sanitizeUpstreamRequestValue(item)
			out[i] = sanitized
			redacted = redacted || itemRedacted
		}
		return out, redacted
	case string:
		if mediaType, ok := parseBase64DataURI(v); ok {
			return map[string]any{
				"redacted":   true,
				"kind":       "data_uri",
				"media_type": mediaType,
				"bytes":      len(v),
				"sha256":     sha256Hex([]byte(v)),
			}, true
		}
		return v, false
	default:
		return value, false
	}
}

func parseBase64DataURI(value string) (string, bool) {
	lower := strings.ToLower(value)
	if !strings.HasPrefix(lower, "data:") {
		return "", false
	}
	marker := ";base64,"
	idx := strings.Index(lower, marker)
	if idx <= len("data:") {
		return "", false
	}
	mediaType := value[len("data:"):idx]
	if !strings.HasPrefix(strings.ToLower(mediaType), "image/") &&
		!strings.HasPrefix(strings.ToLower(mediaType), "audio/") &&
		!strings.HasPrefix(strings.ToLower(mediaType), "video/") {
		return "", false
	}
	return mediaType, true
}

func finalizeUpstreamRequestSnapshot(snapshot *model.TaskUpstreamRequestSnapshot, raw []byte) {
	updateStoredBytes(snapshot)
	if snapshot.StoredBytes <= maxUpstreamRequestSnapshotBytes {
		return
	}

	bodyBytes, err := common.Marshal(snapshot.Body)
	previewSource := raw
	if err == nil {
		previewSource = bodyBytes
	}
	snapshot.Body = map[string]any{
		"truncated_preview": truncateStringToBytes(string(previewSource), maxUpstreamRequestSnapshotBytes/2),
	}
	snapshot.Truncated = true
	updateStoredBytes(snapshot)
	for snapshot.StoredBytes > maxUpstreamRequestSnapshotBytes {
		body, _ := snapshot.Body.(map[string]any)
		preview, _ := body["truncated_preview"].(string)
		if len(preview) == 0 {
			break
		}
		body["truncated_preview"] = truncateStringToBytes(preview, len(preview)/2)
		snapshot.Body = body
		updateStoredBytes(snapshot)
	}
}

func updateStoredBytes(snapshot *model.TaskUpstreamRequestSnapshot) {
	for i := 0; i < 4; i++ {
		data, err := common.Marshal(snapshot)
		if err != nil {
			return
		}
		if snapshot.StoredBytes == len(data) {
			return
		}
		snapshot.StoredBytes = len(data)
	}
}

func truncateStringToBytes(value string, maxBytes int) string {
	if maxBytes <= 0 || len(value) <= maxBytes {
		return value
	}
	return value[:maxBytes]
}

func sha256Hex(data []byte) string {
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:])
}