photowall/go/pkg/sensitive/sensitive.go

package sensitive

import (
	"photowall/internal/model"
	"strings"
	"sync"

	"gorm.io/gorm"
)

// DFA 敏感词过滤器
type Filter struct {
	root  *trieNode
	mu    sync.RWMutex
	ready bool
}

type trieNode struct {
	children map[rune]*trieNode
	isEnd    bool
	word     string
}

var globalFilter = &Filter{root: newNode()}

func newNode() *trieNode {
	return &trieNode{children: make(map[rune]*trieNode)}
}

// Init 从数据库加载敏感词
func Init(db *gorm.DB) {
	var words []model.SensitiveWord
	db.Where("enabled = ?", true).Find(&words)
	globalFilter.mu.Lock()
	defer globalFilter.mu.Unlock()
	globalFilter.root = newNode()
	for _, w := range words {
		globalFilter.addWord(w.Word)
	}
	globalFilter.ready = true
}

// Reload 重新加载敏感词
func Reload(db *gorm.DB) {
	Init(db)
}

func (f *Filter) addWord(word string) {
	node := f.root
	for _, ch := range []rune(word) {
		if _, ok := node.children[ch]; !ok {
			node.children[ch] = newNode()
		}
		node = node.children[ch]
	}
	node.isEnd = true
	node.word = word
}

// Check 检测文本中是否包含敏感词，返回第一个匹配的词
func Check(text string) (bool, string) {
	if !globalFilter.ready || text == "" {
		return false, ""
	}
	globalFilter.mu.RLock()
	defer globalFilter.mu.RUnlock()
	runes := []rune(text)
	for i := 0; i < len(runes); i++ {
		node := globalFilter.root
		for j := i; j < len(runes); j++ {
			next, ok := node.children[runes[j]]
			if !ok {
				break
			}
			node = next
			if node.isEnd {
				return true, node.word
			}
		}
	}
	return false, ""
}

// Replace 替换文本中的敏感词为 ***
func Replace(text string) string {
	if !globalFilter.ready || text == "" {
		return text
	}
	globalFilter.mu.RLock()
	defer globalFilter.mu.RUnlock()
	runes := []rune(text)
	var result strings.Builder
	i := 0
	for i < len(runes) {
		node := globalFilter.root
		matchLen := 0
		for j := i; j < len(runes); j++ {
			next, ok := node.children[runes[j]]
			if !ok {
				break
			}
			node = next
			if node.isEnd {
				matchLen = j - i + 1
			}
		}
		if matchLen > 0 {
			result.WriteString("***")
			i += matchLen
		} else {
			result.WriteRune(runes[i])
			i++
		}
	}
	return result.String()
}

// CheckAndReplace 检测并过滤，返回是否包含敏感词和过滤后的文本
func CheckAndReplace(text string) (bool, string) {
	found, _ := Check(text)
	if !found {
		return false, text
	}
	return true, Replace(text)
}