FranklinWH-Toolbox/service/cmd/mcp/tools/helper/pii_stream.go

package helper

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
)

// MaskJSONBytes 对已序列化的 JSON 字节流做流式脱敏，返回脱敏后的 JSON 字节。
//
// 实现要点（优化 B）：
//   - 使用 json.Decoder 的 token 流式解析，不构建完整的 interface{} 对象树，
//     内存占用为常数级（栈深度 = JSON 嵌套深度）。
//   - 单次遍历：读到 object 的 key 命中敏感规则时，对其对应的 value 做脱敏；
//     其余 token 原样透传。
//   - 若敏感 key 的 value 是对象或数组，则整个子树被完全脱敏（fail-safe）。
//   - 解析失败时返回 error（fail-closed，由上层拒绝输出，不泄露原始数据）。
func MaskJSONBytes(in []byte) ([]byte, error) {
	dec := json.NewDecoder(bytes.NewReader(in))
	dec.UseNumber()

	var out bytes.Buffer
	w := &jsonWriter{buf: &out}

	if err := maskValue(dec, w, nil); err != nil {
		return nil, fmt.Errorf("PII 脱敏失败: %w", err)
	}

	// 确认没有多余内容（例如尾部垃圾），保证输入是单个合法 JSON 值。
	if dec.More() {
		return nil, fmt.Errorf("PII 脱敏失败: JSON 存在多余的尾部内容")
	}

	return out.Bytes(), nil
}

// maskValue 读取下一个 JSON 值并写入 w。
// forceMask 非 nil 时，表示当前值处于某个敏感 key 之下，整个值（含子树）需完全脱敏。
func maskValue(dec *json.Decoder, w *jsonWriter, forceMask *MaskKind) error {
	tok, err := dec.Token()
	if err != nil {
		return err
	}

	switch t := tok.(type) {
	case json.Delim:
		switch t {
		case '{':
			return maskObject(dec, w, forceMask)
		case '[':
			return maskArray(dec, w, forceMask)
		default:
			return fmt.Errorf("非预期的分隔符: %v", t)
		}
	default:
		// 标量值（string / json.Number / bool / nil）
		return writeScalar(w, tok, forceMask)
	}
}

// maskObject 处理一个 JSON 对象。forceMask 非 nil 时整个对象子树被脱敏。
func maskObject(dec *json.Decoder, w *jsonWriter, forceMask *MaskKind) error {
	w.openObject()
	for dec.More() {
		// 读 key
		keyTok, err := dec.Token()
		if err != nil {
			return err
		}
		key, ok := keyTok.(string)
		if !ok {
			return fmt.Errorf("对象 key 非字符串: %T", keyTok)
		}
		w.writeKey(key)

		// 决定该字段 value 的脱敏方式：
		//   - 若已处于强制脱敏子树中，继续沿用父级 forceMask。
		//   - 否则查规则表。
		var childForce *MaskKind
		if forceMask != nil {
			childForce = forceMask
		} else if kind, hit := maskKindFor(key); hit {
			k := kind
			childForce = &k
		}

		if err := maskValue(dec, w, childForce); err != nil {
			return err
		}
	}
	// 消费结束的 '}'
	if _, err := dec.Token(); err != nil {
		return err
	}
	w.closeObject()
	return nil
}

// maskArray 处理一个 JSON 数组。forceMask 非 nil 时数组内每个元素都被脱敏。
func maskArray(dec *json.Decoder, w *jsonWriter, forceMask *MaskKind) error {
	w.openArray()
	for dec.More() {
		if err := maskValue(dec, w, forceMask); err != nil {
			return err
		}
	}
	// 消费结束的 ']'
	if _, err := dec.Token(); err != nil {
		return err
	}
	w.closeArray()
	return nil
}

// writeScalar 写入一个标量值，必要时脱敏。
func writeScalar(w *jsonWriter, tok json.Token, forceMask *MaskKind) error {
	if forceMask != nil {
		// 处于敏感字段之下：
		//   - 字符串按具体 MaskKind 处理
		//   - 其余标量（数字/bool，例如经纬度以 number 存储）统一替换为掩码字符串
		if s, ok := tok.(string); ok {
			return w.writeValueRaw(quote(applyMaskString(*forceMask, s)))
		}
		if tok == nil {
			// null 保持 null（无信息泄露）
			return w.writeValueRaw("null")
		}
		return w.writeValueRaw(quote(maskFull()))
	}

	// 非敏感字段：原样写出
	switch v := tok.(type) {
	case string:
		return w.writeValueRaw(quote(v))
	case json.Number:
		return w.writeValueRaw(v.String())
	case bool:
		return w.writeValueRaw(strconv.FormatBool(v))
	case nil:
		return w.writeValueRaw("null")
	default:
		return fmt.Errorf("非预期的标量类型: %T", tok)
	}
}

// quote 将字符串按 JSON 规则转义并加引号。
func quote(s string) string {
	b, _ := json.Marshal(s)
	return string(b)
}

// indentJSON 将紧凑 JSON 重新缩进为两空格格式，便于阅读并与优化 A 的输出保持一致。
// 缩进失败（理论上不会发生，因输入已是合法 JSON）时返回原始字节。
func indentJSON(b []byte) []byte {
	var out bytes.Buffer
	if err := json.Indent(&out, b, "", "  "); err != nil {
		return b
	}
	return out.Bytes()
}

// jsonWriter 负责按结构写出 JSON，管理逗号与 key/value 状态。
type jsonWriter struct {
	buf   *bytes.Buffer
	stack []*frame
}

type frame struct {
	isObject   bool
	needComma  bool // 是否需要在下一个元素/键前加逗号
	expectKey  bool // 对象中：下一个 writeKey 之后期待 value
	pendingKey bool // 已写入 key，等待 value（value 前不加逗号）
}

func (w *jsonWriter) top() *frame {
	if len(w.stack) == 0 {
		return nil
	}
	return w.stack[len(w.stack)-1]
}

// beforeValue 在写入一个值（标量或容器起始）之前处理逗号。
func (w *jsonWriter) beforeValue() {
	f := w.top()
	if f == nil {
		return
	}
	if f.isObject {
		// 对象中值紧跟 key，key 已处理逗号，这里不加
		if f.pendingKey {
			f.pendingKey = false
			return
		}
	}
	if f.needComma {
		w.buf.WriteByte(',')
	}
	f.needComma = true
}

func (w *jsonWriter) openObject() {
	w.beforeValue()
	w.buf.WriteByte('{')
	w.stack = append(w.stack, &frame{isObject: true})
}

func (w *jsonWriter) closeObject() {
	w.buf.WriteByte('}')
	w.stack = w.stack[:len(w.stack)-1]
}

func (w *jsonWriter) openArray() {
	w.beforeValue()
	w.buf.WriteByte('[')
	w.stack = append(w.stack, &frame{isObject: false})
}

func (w *jsonWriter) closeArray() {
	w.buf.WriteByte(']')
	w.stack = w.stack[:len(w.stack)-1]
}

func (w *jsonWriter) writeKey(key string) {
	f := w.top()
	if f != nil && f.needComma {
		w.buf.WriteByte(',')
	}
	w.buf.WriteString(quote(key))
	w.buf.WriteByte(':')
	if f != nil {
		f.needComma = true
		f.pendingKey = true
	}
}

func (w *jsonWriter) writeValueRaw(raw string) error {
	w.beforeValue()
	_, err := w.buf.WriteString(raw)
	return err
}