easyai-ai-gateway/apps/api/internal/store/json_storage_guard.go

package store

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
)

const (
	storageBinaryPrefixChars      = 16
	storageGenericBase64MinLength = 4096
	storageInvalidBinaryMinLength = 512
	storageJSONSanitizerMaxDepth  = 64
	storagePlaceholderPrefix      = "[GatewayBinary:v1;"
	storageBufferObjectType       = "buffer"
	storageDataURLPrefix          = "data:"
	storageDataURLBase64Marker    = ";base64"
	storageDataURLMaxContentType  = 64
)

var (
	ErrTaskRequestBinaryNotMaterialized = errors.New("task request binary payload was not materialized")
	ErrTaskResultBinaryNotMaterialized  = errors.New("task result binary payload was not materialized")
)

type taskPayloadBinaryError struct {
	target error
	code   string
	count  int
}

func (e *taskPayloadBinaryError) Error() string {
	return fmt.Sprintf("%s: detected %d inline binary value(s)", e.code, e.count)
}

func (e *taskPayloadBinaryError) ErrorCode() string {
	return e.code
}

func (e *taskPayloadBinaryError) Is(target error) bool {
	return target == e.target
}

type storageSanitizeReport struct {
	Value       any
	BinaryCount int
}

// sanitizeJSONForStorage is the final task-domain persistence guard. It always
// returns a detached JSON-compatible value and replaces inline binary payloads
// with a bounded, deterministic placeholder.
func sanitizeJSONForStorage(value any) any {
	return sanitizeJSONForStorageWithReport(value).Value
}

func sanitizeJSONForStorageWithReport(value any) storageSanitizeReport {
	next, count := sanitizeJSONStorageValue(value, nil, 0)
	return storageSanitizeReport{Value: next, BinaryCount: count}
}

func sanitizeJSONStorageValue(value any, path []string, depth int) (any, int) {
	if depth >= storageJSONSanitizerMaxDepth {
		return "[JSON,max-depth]", 0
	}
	switch typed := value.(type) {
	case map[string]any:
		if payload, contentType, ok := storageBufferObjectBytes(typed); ok {
			return storageBinaryPlaceholder(payload, contentType, "buffer"), 1
		}
		next := make(map[string]any, len(typed))
		count := 0
		for key, child := range typed {
			sanitized, childCount := sanitizeJSONStorageValue(child, appendStoragePath(path, key), depth+1)
			next[key] = sanitized
			count += childCount
		}
		return next, count
	case []any:
		if storagePathIsBinary(path) {
			if payload, ok := storageNumberArrayBytes(typed); ok {
				return storageBinaryPlaceholder(payload, "", "buffer"), 1
			}
		}
		next := make([]any, len(typed))
		count := 0
		for index, child := range typed {
			sanitized, childCount := sanitizeJSONStorageValue(child, path, depth+1)
			next[index] = sanitized
			count += childCount
		}
		return next, count
	case []byte:
		return storageBinaryPlaceholder(typed, "", "buffer"), 1
	case string:
		if placeholder, ok := storageStringPlaceholder(typed, path); ok {
			return placeholder, 1
		}
		return typed, 0
	default:
		return value, 0
	}
}

func appendStoragePath(path []string, key string) []string {
	next := make([]string, len(path)+1)
	copy(next, path)
	next[len(path)] = key
	return next
}

func storageStringPlaceholder(value string, path []string) (string, bool) {
	raw := strings.TrimSpace(value)
	if raw == "" || strings.HasPrefix(raw, storagePlaceholderPrefix) {
		return "", false
	}
	encoded, contentType, dataURL := storageBase64StringParts(raw)
	strict := dataURL || storagePathIsBinary(path)
	if !strict && len(encoded) < storageGenericBase64MinLength {
		return "", false
	}
	payload, ok := storageDecodeBase64(encoded)
	if ok {
		encoding := "raw"
		if dataURL {
			encoding = "data-uri"
		}
		return storageBinaryPlaceholder(payload, contentType, encoding), true
	}
	if strict && len(raw) >= storageInvalidBinaryMinLength {
		return storageBinaryPlaceholder([]byte(raw), contentType, "raw"), true
	}
	return "", false
}

func storageBase64StringParts(value string) (encoded string, contentType string, dataURL bool) {
	if !strings.HasPrefix(strings.ToLower(value), storageDataURLPrefix) {
		return value, "", false
	}
	prefix, payload, ok := strings.Cut(value, ",")
	if !ok || !strings.Contains(strings.ToLower(prefix), storageDataURLBase64Marker) {
		return value, "", false
	}
	mediaType := strings.TrimSpace(prefix[len(storageDataURLPrefix):])
	if before, _, found := strings.Cut(mediaType, ";"); found {
		mediaType = before
	}
	if len(mediaType) > storageDataURLMaxContentType || !storageSafeContentType(mediaType) {
		mediaType = ""
	}
	return payload, strings.ToLower(mediaType), true
}

func storageSafeContentType(value string) bool {
	if value == "" {
		return true
	}
	for _, char := range value {
		switch {
		case char >= 'a' && char <= 'z':
		case char >= 'A' && char <= 'Z':
		case char >= '0' && char <= '9':
		case char == '/', char == '.', char == '+', char == '-':
		default:
			return false
		}
	}
	return true
}

func storageDecodeBase64(value string) ([]byte, bool) {
	normalized := removeStorageASCIIWhitespace(value)
	if normalized == "" {
		return nil, false
	}
	for _, encoding := range []*base64.Encoding{
		base64.StdEncoding,
		base64.RawStdEncoding,
		base64.URLEncoding,
		base64.RawURLEncoding,
	} {
		payload, err := encoding.DecodeString(normalized)
		if err != nil || len(payload) == 0 {
			continue
		}
		canonical := encoding.EncodeToString(payload)
		if strings.TrimRight(normalized, "=") == strings.TrimRight(canonical, "=") {
			return payload, true
		}
	}
	return nil, false
}

func removeStorageASCIIWhitespace(value string) string {
	return strings.Map(func(char rune) rune {
		switch char {
		case ' ', '\n', '\r', '\t':
			return -1
		default:
			return char
		}
	}, value)
}

func storagePathIsBinary(path []string) bool {
	if len(path) == 0 {
		return false
	}
	key := normalizeStorageBinaryKey(path[len(path)-1])
	if storageBinaryKey(key) {
		return true
	}
	if len(path) < 2 {
		return false
	}
	parent := normalizeStorageBinaryKey(path[len(path)-2])
	return (parent == "inlinedata" || parent == "binary" || parent == "media") &&
		(key == "data" || key == "content")
}

func normalizeStorageBinaryKey(value string) string {
	return strings.Map(func(char rune) rune {
		switch {
		case char >= 'a' && char <= 'z':
			return char
		case char >= 'A' && char <= 'Z':
			return char + ('a' - 'A')
		case char >= '0' && char <= '9':
			return char
		default:
			return -1
		}
	}, value)
}

func storageBinaryKey(key string) bool {
	return key == "b64" ||
		key == "b64json" ||
		key == "base64" ||
		key == "buffer" ||
		key == "bytes" ||
		strings.Contains(key, "base64") ||
		strings.Contains(key, "buffer") ||
		strings.Contains(key, "binary") ||
		strings.HasSuffix(key, "b64") ||
		strings.HasSuffix(key, "bytes")
}

func storageBufferObjectBytes(value map[string]any) ([]byte, string, bool) {
	kind, _ := value["type"].(string)
	if normalizeStorageBinaryKey(kind) != storageBufferObjectType {
		return nil, "", false
	}
	contentType := firstNonEmpty(
		stringFromAny(value["mime_type"]),
		stringFromAny(value["mimeType"]),
		stringFromAny(value["contentType"]),
	)
	switch data := value["data"].(type) {
	case []byte:
		if len(data) == 0 {
			return nil, "", false
		}
		return append([]byte(nil), data...), contentType, true
	case []any:
		payload, ok := storageNumberArrayBytes(data)
		return payload, contentType, ok
	default:
		return nil, "", false
	}
}

func storageNumberArrayBytes(values []any) ([]byte, bool) {
	if len(values) == 0 {
		return nil, false
	}
	payload := make([]byte, len(values))
	for index, value := range values {
		next, ok := storageByteFromAny(value)
		if !ok {
			return nil, false
		}
		payload[index] = next
	}
	return payload, true
}

func storageByteFromAny(value any) (byte, bool) {
	switch typed := value.(type) {
	case byte:
		return typed, true
	case int:
		if typed >= 0 && typed <= 255 {
			return byte(typed), true
		}
	case int32:
		if typed >= 0 && typed <= 255 {
			return byte(typed), true
		}
	case int64:
		if typed >= 0 && typed <= 255 {
			return byte(typed), true
		}
	case float64:
		asInt := int(typed)
		if typed == float64(asInt) && asInt >= 0 && asInt <= 255 {
			return byte(asInt), true
		}
	case json.Number:
		asInt, err := strconv.ParseInt(string(typed), 10, 16)
		if err == nil && asInt >= 0 && asInt <= 255 {
			return byte(asInt), true
		}
	}
	return 0, false
}

func storageBinaryPlaceholder(payload []byte, contentType string, encoding string) string {
	digest := sha256.Sum256(payload)
	prefix := base64.StdEncoding.EncodeToString(payload)
	if len(prefix) > storageBinaryPrefixChars {
		prefix = prefix[:storageBinaryPrefixChars]
	}
	contentType = strings.ToLower(strings.TrimSpace(contentType))
	if contentType == "" || len(contentType) > 32 || !storageSafeContentType(contentType) {
		contentType = "application/octet-stream"
	}
	switch encoding {
	case "data-uri", "buffer":
	default:
		encoding = "raw"
	}
	return fmt.Sprintf(
		"[GatewayBinary:v1;prefix=%s;sha256=%x;bytes=%d;mime=%s;encoding=%s]",
		prefix,
		digest,
		len(payload),
		contentType,
		encoding,
	)
}