dns-server/gfw/manager.go

package gfw

import (
	"encoding/base64"
	"fmt"
	"os"
	"regexp"
	"strings"
	"sync"

	"dns-server/config"
	"dns-server/logger"
)

// regexRule 正则规则结构，包含编译后的表达式和原始字符串
type regexRule struct {
	pattern  *regexp.Regexp
	original string
}

// GFWListManager GFWList管理器
type GFWListManager struct {
	config      *config.GFWListConfig
	domainRules map[string]bool
	regexRules  []regexRule
	rulesMutex  sync.RWMutex
}

// NewGFWListManager 创建GFWList管理器实例
func NewGFWListManager(config *config.GFWListConfig) *GFWListManager {
	return &GFWListManager{
		config:      config,
		domainRules: make(map[string]bool),
		regexRules:  []regexRule{},
	}
}

// LoadRules 加载GFWList规则
func (m *GFWListManager) LoadRules() error {
	// 如果GFWList功能未启用，不加载规则
	if !m.config.Enabled {
		return nil
	}

	m.rulesMutex.Lock()
	defer m.rulesMutex.Unlock()

	// 清空现有规则
	m.domainRules = make(map[string]bool)
	m.regexRules = []regexRule{}

	if m.config.Content == "" {
		return nil // 没有GFWList内容，直接返回
	}

	// 从文件路径读取GFWList内容
	fileContent, err := os.ReadFile(m.config.Content)
	if err != nil {
		return fmt.Errorf("读取GFWList文件失败: %w", err)
	}

	rawContent := string(fileContent)
	var ruleContent string

	// 过滤注释行，收集可能的Base64内容
	var base64Content strings.Builder
	lines := strings.Split(rawContent, "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "!") || strings.HasPrefix(line, "[") {
			// 跳过注释行和头信息行
			continue
		}
		base64Content.WriteString(line)
	}

	// 尝试Base64解码
	decoded, err := base64.StdEncoding.DecodeString(base64Content.String())
	if err == nil {
		// 解码成功，使用解码后的内容
		ruleContent = string(decoded)
		logger.Info("GFWList文件为Base64编码，已成功解码")
	} else {
		// 解码失败，使用原始内容（可能是纯文本格式）
		ruleContent = rawContent
		logger.Info("GFWList文件为纯文本格式，直接解析")
	}

	// 按行解析规则内容
	ruleLines := strings.Split(ruleContent, "\n")
	for _, line := range ruleLines {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "!") || strings.HasPrefix(line, "[") {
			// 跳过空行、注释行和头信息行
			continue
		}
		m.parseRule(line)
	}

	logger.Info(fmt.Sprintf("GFWList规则加载完成，域名规则: %d, 正则规则: %d",
		len(m.domainRules), len(m.regexRules)))
	return nil
}

// parseRule 解析规则行
func (m *GFWListManager) parseRule(line string) {
	// 保存原始规则用于后续使用
	originalLine := line

	// 处理注释
	if strings.HasPrefix(line, "!") || strings.HasPrefix(line, "#") || line == "" {
		return
	}

	// 移除规则选项部分（暂时不处理规则选项）
	if strings.Contains(line, "$") {
		parts := strings.SplitN(line, "$", 2)
		line = parts[0]
		// 规则选项暂时不处理
	}

	// 处理不同类型的规则
	switch {
	case strings.HasPrefix(line, "||") && strings.HasSuffix(line, "^"):
		// AdGuardHome域名规则格式: ||example.com^
		domain := strings.TrimSuffix(strings.TrimPrefix(line, "||"), "^")
		m.addDomainRule(domain, originalLine)

	case strings.HasPrefix(line, "||"):
		// 域名片段匹配规则: ||google 匹配任何包含google的域名
		domain := strings.TrimPrefix(line, "||")
		// 添加精确域名匹配
		m.addDomainRule(domain, originalLine)
		// 同时添加正则表达式规则，匹配任何包含该域名片段的域名
		if re, err := regexp.Compile("(?i).*" + regexp.QuoteMeta(domain) + ".*"); err == nil {
			m.addRegexRule(re, originalLine)
		}

	case strings.HasPrefix(line, "*"):
		// 通配符规则，转换为正则表达式
		pattern := strings.ReplaceAll(line, "*", ".*")
		pattern = "^" + pattern + "$"
		if re, err := regexp.Compile(pattern); err == nil {
			// 保存原始规则字符串
			m.addRegexRule(re, originalLine)
		}

	case strings.HasPrefix(line, "/") && strings.HasSuffix(line, "/"):
		// 正则表达式匹配规则：/regex/ 格式，不区分大小写
		pattern := strings.TrimPrefix(strings.TrimSuffix(line, "/"), "/")
		// 编译为不区分大小写的正则表达式，确保能匹配域名中任意位置
		// 对于像 /domain/ 这样的规则，应该匹配包含 domain 字符串的任何域名
		if re, err := regexp.Compile("(?i).*" + regexp.QuoteMeta(pattern) + ".*"); err == nil {
			// 保存原始规则字符串
			m.addRegexRule(re, originalLine)
		}

	case strings.HasPrefix(line, "|") && strings.HasSuffix(line, "|"):
		// 完整URL匹配规则
		urlPattern := strings.TrimPrefix(strings.TrimSuffix(line, "|"), "|")
		// 将URL模式转换为正则表达式
		pattern := "^" + regexp.QuoteMeta(urlPattern) + "$"
		if re, err := regexp.Compile(pattern); err == nil {
			m.addRegexRule(re, originalLine)
		}

	case strings.HasPrefix(line, "|"):
		// URL开头匹配规则
		urlPattern := strings.TrimPrefix(line, "|")
		pattern := "^" + regexp.QuoteMeta(urlPattern)
		if re, err := regexp.Compile(pattern); err == nil {
			m.addRegexRule(re, originalLine)
		}

	case strings.HasSuffix(line, "|"):
		// URL结尾匹配规则
		urlPattern := strings.TrimSuffix(line, "|")
		pattern := regexp.QuoteMeta(urlPattern) + "$"
		if re, err := regexp.Compile(pattern); err == nil {
			m.addRegexRule(re, originalLine)
		}

	default:
		// 默认作为普通域名规则
		m.addDomainRule(line, originalLine)
	}
}

// addDomainRule 添加域名规则
func (m *GFWListManager) addDomainRule(domain string, original string) {
	m.domainRules[domain] = true
}

// addRegexRule 添加正则表达式规则
func (m *GFWListManager) addRegexRule(re *regexp.Regexp, original string) {
	rule := regexRule{
		pattern:  re,
		original: original,
	}
	m.regexRules = append(m.regexRules, rule)
}

// IsMatch 检查域名是否匹配GFWList规则
func (m *GFWListManager) IsMatch(domain string) bool {
	m.rulesMutex.RLock()
	defer m.rulesMutex.RUnlock()

	// 预处理域名，去除可能的端口号
	if strings.Contains(domain, ":") {
		parts := strings.Split(domain, ":")
		domain = parts[0]
	}

	// 检查精确域名匹配
	if m.domainRules[domain] {
		return true
	}

	// 检查子域名匹配
	parts := strings.Split(domain, ".")
	for i := 0; i < len(parts)-1; i++ {
		subdomain := strings.Join(parts[i:], ".")
		if m.domainRules[subdomain] {
			return true
		}
	}

	// 检查正则表达式匹配
	for _, re := range m.regexRules {
		if re.pattern.MatchString(domain) {
			return true
		}
	}

	return false
}

// GetGFWListIP 获取GFWList的目标IP地址
func (m *GFWListManager) GetGFWListIP() string {
	return m.config.IP
}