onesvm-browser-server/server/internal/safetyscan/safetyscan.go

// Package safetyscan 响应侧合规扫描：关键词分类、PII 打码、提示词注入检测、
// 注入包裹 delimiter（design-arch §5.4 响应侧行 / Contract A1-safetyscan）。
//
// 词表来源：本包 wordlist.yaml（样例，部署前由合规方扩充）；NFKC 归一 +
// 前 64KB 限制；title/description/markdown 分别扫（fit 结果与元数据不混扫）。
package safetyscan

import (
	_ "embed"
	"fmt"
	"strings"

	"gopkg.in/yaml.v3"
)

//go:embed wordlist.yaml
var wordlistRaw []byte

// Wordlist 词表结构（对应 wordlist.yaml）。
type Wordlist struct {
	Version           string   `yaml:"version"`
	HighRisk          []string `yaml:"high_risk"`
	InjectionPatterns []string `yaml:"injection_patterns"`
	PII               struct {
		IDCard   string `yaml:"id_card"`
		Mobile   string `yaml:"mobile"`
		BankCard string `yaml:"bank_card"`
	} `yaml:"pii"`
}

// compiled 编译后的扫描器（构造后只读，Scanner 并发安全）。
type compiled struct {
	wordlist  *Wordlist
	highRisk  []term       // 高危词（NFKC 归一后子串匹配）
	injection []*regexWrap // 注入正则（归一文本上跑）
	piiIDCard *regexWrap   // 身份证（原文上替换）
	piiMobile *regexWrap
	piiBank   *regexWrap
}

// term 高危词条目。
type term struct {
	word string // 归一后的匹配词
	raw  string // 原词（审计展示）
}

// ScanResult 扫描结论（Warnings/Hits 恒非 nil，信封纪律 [] 非 null）。
type ScanResult struct {
	Blocked     bool
	Warnings    []string
	Redacted    bool
	RedactCount int
	Hits        []string
	WordlistVer string
}

// Scanner 线程安全扫描器。
type Scanner struct{ c *compiled }

// Version 返回词表版本。
func (s *Scanner) Version() string { return s.c.wordlist.Version }

// Load 用内嵌词表构造扫描器。
func Load() (*Scanner, error) { return load(wordlistRaw) }

// load 解析并编译词表。
func load(raw []byte) (*Scanner, error) {
	var wl Wordlist
	if err := yaml.Unmarshal(raw, &wl); err != nil {
		return nil, fmt.Errorf("safetyscan: 词表解析失败: %w", err)
	}
	if wl.Version == "" {
		return nil, fmt.Errorf("safetyscan: 词表缺少 version 字段")
	}
	c := &compiled{wordlist: &wl}
	for _, w := range wl.HighRisk {
		c.highRisk = append(c.highRisk, term{word: normalize(w), raw: w})
	}
	for _, p := range wl.InjectionPatterns {
		re, err := compile(p)
		if err != nil {
			return nil, fmt.Errorf("safetyscan: 注入正则 %q 编译失败: %w", p, err)
		}
		c.injection = append(c.injection, re)
	}
	piiSpecs := []struct {
		raw string
		dst **regexWrap
		msg string
	}{
		{wl.PII.IDCard, &c.piiIDCard, "身份证"},
		{wl.PII.Mobile, &c.piiMobile, "手机号"},
		{wl.PII.BankCard, &c.piiBank, "银行卡"},
	}
	for _, sp := range piiSpecs {
		if sp.raw == "" {
			continue
		}
		re, err := compile(sp.raw)
		if err != nil {
			return nil, fmt.Errorf("safetyscan: %s正则编译失败: %w", sp.msg, err)
		}
		*sp.dst = re
	}
	return &Scanner{c: c}, nil
}

// scanLimit 前 64KB 扫描（design §5.4：前 64KB 解码扫描）。
const scanLimit = 64 * 1024

// redactPlaceholder PII 打码占位（redact 不丢弃）。
var redactPlaceholder = map[string]string{
	"id_card":   "[身份证已脱敏]",
	"mobile":    "[手机号已脱敏]",
	"bank_card": "[银行卡号已脱敏]",
}

// Scan 扫描 title/description/markdown 三字段（分别扫，不混扫）。
// 返回结论与打码后的字段副本；Blocked 时调用方置 error.code=blocked。
func (s *Scanner) Scan(title, description, markdown string) (res ScanResult, outTitle, outDesc, outMD string) {
	res = ScanResult{WordlistVer: s.c.wordlist.Version, Warnings: []string{}, Hits: []string{}}
	scanOne := func(field, text string) (string, bool) {
		if text == "" {
			return text, false
		}
		head := text
		if len(head) > scanLimit {
			head = head[:scanLimit]
		}
		norm := normalize(head)
		dirty := false
		redacted := text
		// PII redact：正则在原文上替换（样例词表为 ASCII 数字正则，归一不影响定位）。
		for _, p := range []struct {
			name string
			re   *regexWrap
			key  string
		}{
			{"pii.id_card", s.c.piiIDCard, "id_card"},
			{"pii.mobile", s.c.piiMobile, "mobile"},
			{"pii.bank_card", s.c.piiBank, "bank_card"},
		} {
			if p.re == nil {
				continue
			}
			if cnt := p.re.count(redacted); cnt > 0 {
				if p.key == "mobile" {
					// mobile 正则带边界捕获组：只替换组内号码，保留前后字符。
					redacted = p.re.replaceGroup(redacted, "$1", redactPlaceholder[p.key])
				} else {
					redacted = p.re.replaceAll(redacted, redactPlaceholder[p.key])
				}
				res.RedactCount += cnt
				res.Redacted = true
				dirty = true
				res.Hits = append(res.Hits, p.name)
			}
		}
		// 高危词：归一后子串匹配
		for _, t := range s.c.highRisk {
			if strings.Contains(norm, t.word) {
				res.Blocked = true
				res.Hits = append(res.Hits, "high_risk:"+t.raw)
			}
		}
		// 注入检测：命中计入 warnings
		for _, r := range s.c.injection {
			if r.matchString(norm) {
				res.Warnings = append(res.Warnings, fmt.Sprintf("prompt_injection_detected in %s", field))
				res.Hits = append(res.Hits, "injection:"+r.pattern)
			}
		}
		return redacted, dirty
	}
	outTitle, _ = scanOne("title", title)
	outDesc, _ = scanOne("description", description)
	outMD, _ = scanOne("markdown", markdown)
	return res, outTitle, outDesc, outMD
}