优化修复

dns/cache.go

@@ -1,6 +1,9 @@
 package dns
 
 import (
+	"encoding/json"
+	"io/ioutil"
+	"os"
 	"sync"
 	"time"
 
@@ -12,6 +15,46 @@ type DNSCacheItem struct {
 	Response  *dns.Msg  // DNS响应消息
 	Expiry    time.Time // 过期时间
 	HasDNSSEC bool      // 是否包含DNSSEC记录
+	Size      int       // 缓存项大小（字节）
+}
+
+// SerializableDNSCacheItem 用于JSON序列化的缓存项结构
+type SerializableDNSCacheItem struct {
+	ResponseBytes []byte `json:"responseBytes"` // 二进制DNS响应
+	Expiry        int64  `json:"expiry"`        // 过期时间（纳秒）
+	HasDNSSEC     bool   `json:"hasDNSSEC"`     // 是否包含DNSSEC记录
+	Size          int    `json:"size"`          // 缓存项大小（字节）
+}
+
+// SerializableDNSCache 可序列化的缓存结构
+type SerializableDNSCache struct {
+	Items         map[string]*SerializableDNSCacheItem `json:"items"`         // 缓存项
+	TTL           int64                                `json:"ttl"`           // 默认TTL（纳秒）
+	MaxSize       int                                  `json:"maxSize"`       // 最大缓存大小
+	CacheMode     string                               `json:"cacheMode"`     // 缓存模式
+	CacheFilePath string                               `json:"cacheFilePath"` // 缓存文件路径
+}
+
+// DNSCache DNS缓存结构
+type DNSCache struct {
+	cache         map[string]*LRUNode // 缓存映射表，直接存储链表节点
+	mutex         sync.RWMutex        // 读写锁，保护缓存
+	ttl           time.Duration       // 默认缓存TTL
+	maxSize       int                 // 最大缓存条目数
+	cacheSize     int64               // 当前缓存大小（字节）
+	maxCacheSize  int64               // 最大缓存大小（字节）
+	cacheMode     string              // 缓存模式
+	cacheFilePath string              // 缓存文件路径
+	saveInterval  time.Duration       // 保存间隔
+	saveMutex     sync.Mutex          // 保存互斥锁
+	maxCacheTTL   time.Duration       // 最大缓存TTL
+	minCacheTTL   time.Duration       // 最小缓存TTL
+	saveStopCh    chan struct{}       // 保存循环停止通道
+	saveRunning   bool                // 保存循环是否运行
+	saveLoopMutex sync.Mutex          // 保护保存循环状态的互斥锁
+	// 双向链表头和尾指针，用于LRU淘汰
+	head *LRUNode // 头指针，指向最久未使用的节点
+	tail *LRUNode // 尾指针，指向最近使用的节点
 }
 
 // LRUNode 双向链表节点，用于LRU缓存
@@ -22,30 +65,41 @@ type LRUNode struct {
 	next  *LRUNode
 }
 
-// DNSCache DNS缓存结构
-type DNSCache struct {
-	cache      map[string]*LRUNode // 缓存映射表，直接存储链表节点
-	mutex      sync.RWMutex        // 读写锁，保护缓存
-	defaultTTL time.Duration       // 默认缓存TTL
-	maxSize    int                 // 最大缓存条目数
-	// 双向链表头和尾指针，用于LRU淘汰
-	head *LRUNode // 头指针，指向最久未使用的节点
-	tail *LRUNode // 尾指针，指向最近使用的节点
-}
-
 // NewDNSCache 创建新的DNS缓存实例
-func NewDNSCache(defaultTTL time.Duration) *DNSCache {
+func NewDNSCache(defaultTTL time.Duration, cacheMode string, cacheSizeMB int, cacheFilePath string, saveInterval time.Duration, maxCacheTTL, minCacheTTL time.Duration) *DNSCache {
+	// 计算最大缓存大小（字节）
+	maxCacheSize := int64(cacheSizeMB) * 1024 * 1024
+
 	cache := &DNSCache{
-		cache:      make(map[string]*LRUNode),
-		defaultTTL: defaultTTL,
-		maxSize:    10000, // 默认最大缓存10000条记录
-		head:       nil,
-		tail:       nil,
+		cache:         make(map[string]*LRUNode),
+		ttl:           defaultTTL,
+		maxSize:       10000, // 默认最大缓存10000条记录
+		cacheSize:     0,
+		maxCacheSize:  maxCacheSize,
+		cacheMode:     cacheMode,
+		cacheFilePath: cacheFilePath,
+		saveInterval:  saveInterval,
+		maxCacheTTL:   maxCacheTTL,
+		minCacheTTL:   minCacheTTL,
+		saveStopCh:    make(chan struct{}),
+		saveRunning:   false,
+		head:          nil,
+		tail:          nil,
+	}
+
+	// 加载现有缓存（如果存在）
+	if cacheMode == "file" {
+		cache.LoadFromFile()
 	}
 
 	// 启动缓存清理协程
 	go cache.startCleanupLoop()
 
+	// 启动定期保存协程（如果是文件缓存）
+	if cacheMode == "file" {
+		go cache.startSaveLoop()
+	}
+
 	return cache
 }
 
@@ -102,6 +156,16 @@ func cacheKey(qName string, qType uint16) string {
 	return qName + "|" + dns.TypeToString[qType]
 }
 
+// calculateItemSize 计算缓存项大小
+func calculateItemSize(item *DNSCacheItem) int {
+	// 序列化响应以计算大小
+	data, err := json.Marshal(item)
+	if err != nil {
+		return 0
+	}
+	return len(data)
+}
+
 // hasDNSSECRecords 检查响应是否包含DNSSEC记录
 func hasDNSSECRecords(response *dns.Msg) bool {
 	// 定义检查单个RR是否为DNSSEC记录的辅助函数
@@ -135,8 +199,17 @@ func hasDNSSECRecords(response *dns.Msg) bool {
 
 // Set 设置缓存项
 func (c *DNSCache) Set(qName string, qType uint16, response *dns.Msg, ttl time.Duration) {
+	// 设置默认TTL
 	if ttl <= 0 {
-		ttl = c.defaultTTL
+		ttl = c.ttl
+	}
+
+	// 应用maxCacheTTL和minCacheTTL约束
+	if c.maxCacheTTL > 0 && ttl > c.maxCacheTTL {
+		ttl = c.maxCacheTTL
+	}
+	if c.minCacheTTL > 0 && ttl < c.minCacheTTL {
+		ttl = c.minCacheTTL
 	}
 
 	key := cacheKey(qName, qType)
@@ -146,11 +219,15 @@ func (c *DNSCache) Set(qName string, qType uint16, response *dns.Msg, ttl time.D
 		HasDNSSEC: hasDNSSECRecords(response), // 检查并设置DNSSEC标志
 	}
 
+	// 计算缓存项大小
+	item.Size = calculateItemSize(item)
+
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
-	// 如果条目已存在，先从链表和缓存中移除
+	// 如果条目已存在，先从链表和缓存中移除，并更新缓存大小
 	if existingNode, found := c.cache[key]; found {
+		c.cacheSize -= int64(existingNode.value.Size)
 		c.removeNode(existingNode)
 		delete(c.cache, key)
 	}
@@ -162,17 +239,27 @@ func (c *DNSCache) Set(qName string, qType uint16, response *dns.Msg, ttl time.D
 	}
 	c.addNodeToTail(newNode)
 	c.cache[key] = newNode
+	c.cacheSize += int64(item.Size)
 
-	// 检查是否超过最大大小限制，如果超过则移除最久未使用的条目
+	// 检查是否超过最大条目数限制，如果超过则移除最久未使用的条目
 	if len(c.cache) > c.maxSize {
 		// 最久未使用的条目是链表的头节点
 		if c.head != nil {
+			c.cacheSize -= int64(c.head.value.Size)
 			oldestKey := c.head.key
 			// 从缓存和链表中移除头节点
 			delete(c.cache, oldestKey)
 			c.removeNode(c.head)
 		}
 	}
+
+	// 检查是否超过最大缓存大小，如果超过则继续移除最久未使用的条目
+	for c.cacheSize > c.maxCacheSize && c.head != nil {
+		c.cacheSize -= int64(c.head.value.Size)
+		oldestKey := c.head.key
+		delete(c.cache, oldestKey)
+		c.removeNode(c.head)
+	}
 }
 
 // Get 获取缓存项
@@ -255,6 +342,230 @@ func (c *DNSCache) startCleanupLoop() {
 	}
 }
 
+// startSaveLoop 启动定期保存缓存的协程
+func (c *DNSCache) startSaveLoop() {
+	c.saveLoopMutex.Lock()
+	// 如果已经在运行，直接返回
+	if c.saveRunning {
+		c.saveLoopMutex.Unlock()
+		return
+	}
+	// 重置停止通道
+	c.saveStopCh = make(chan struct{})
+	c.saveRunning = true
+	c.saveLoopMutex.Unlock()
+
+	go func() {
+		ticker := time.NewTicker(c.saveInterval) // 根据配置的间隔保存
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				// 检查缓存模式，如果不是file模式则不保存
+				c.mutex.RLock()
+				mode := c.cacheMode
+				c.mutex.RUnlock()
+				if mode == "file" {
+					c.SaveToFile()
+				}
+			case <-c.saveStopCh:
+				// 停止保存循环
+				c.saveLoopMutex.Lock()
+				c.saveRunning = false
+				c.saveLoopMutex.Unlock()
+				return
+			}
+		}
+	}()
+}
+
+// saveCacheToFile 保存缓存到文件的底层实现，不检查缓存模式
+func (c *DNSCache) saveCacheToFile() {
+	c.saveMutex.Lock()
+	defer c.saveMutex.Unlock()
+
+	c.mutex.RLock()
+	defer c.mutex.RUnlock()
+
+	// 收集有效的缓存项
+	validItems := make(map[string]*SerializableDNSCacheItem)
+	now := time.Now()
+
+	for key, node := range c.cache {
+		// 只保存未过期的缓存项
+		if now.Before(node.value.Expiry) {
+			// 序列化DNS响应为二进制
+			responseBytes, err := node.value.Response.Pack()
+			if err != nil {
+				continue // 跳过无法序列化的响应
+			}
+
+			// 创建可序列化的缓存项
+			validItems[key] = &SerializableDNSCacheItem{
+				ResponseBytes: responseBytes,
+				Expiry:        node.value.Expiry.UnixNano(),
+				HasDNSSEC:     node.value.HasDNSSEC,
+				Size:          node.value.Size,
+			}
+		}
+	}
+
+	// 创建可序列化的缓存结构
+	serializableCache := &SerializableDNSCache{
+		Items:         validItems,
+		TTL:           int64(c.ttl),
+		MaxSize:       c.maxSize,
+		CacheMode:     c.cacheMode,
+		CacheFilePath: c.cacheFilePath,
+	}
+
+	// 序列化到JSON
+	data, err := json.MarshalIndent(serializableCache, "", "  ")
+	if err != nil {
+		return
+	}
+
+	// 确保目录存在
+	os.MkdirAll(cacheDir(), 0755)
+
+	// 保存到文件
+	err = ioutil.WriteFile(c.cacheFilePath, data, 0644)
+	if err != nil {
+		return
+	}
+}
+
+// SaveToFile 保存缓存到文件
+func (c *DNSCache) SaveToFile() {
+	// 检查缓存模式，如果不是file模式，直接返回
+	c.mutex.RLock()
+	mode := c.cacheMode
+	c.mutex.RUnlock()
+	if mode != "file" {
+		return
+	}
+
+	// 调用底层保存逻辑
+	c.saveCacheToFile()
+}
+
+// LoadFromFile 从文件加载缓存
+func (c *DNSCache) LoadFromFile() {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	// 检查文件是否存在
+	if _, err := os.Stat(c.cacheFilePath); os.IsNotExist(err) {
+		return // 文件不存在，跳过加载
+	}
+
+	// 读取文件内容
+	data, err := ioutil.ReadFile(c.cacheFilePath)
+	if err != nil {
+		return
+	}
+
+	// 反序列化JSON
+	var serializableCache SerializableDNSCache
+	err = json.Unmarshal(data, &serializableCache)
+	if err != nil {
+		return
+	}
+
+	// 加载缓存项
+	now := time.Now()
+	for key, serializableItem := range serializableCache.Items {
+		// 转换过期时间
+		expiry := time.Unix(0, serializableItem.Expiry)
+		// 只加载未过期的缓存项
+		if now.Before(expiry) {
+			// 反序列化二进制DNS响应
+			response := &dns.Msg{}
+			err := response.Unpack(serializableItem.ResponseBytes)
+			if err != nil {
+				continue // 跳过无法反序列化的响应
+			}
+
+			// 创建缓存项
+			item := &DNSCacheItem{
+				Response:  response,
+				Expiry:    expiry,
+				HasDNSSEC: serializableItem.HasDNSSEC,
+				Size:      serializableItem.Size,
+			}
+
+			// 创建新的链表节点并添加到尾部
+			newNode := &LRUNode{
+				key:   key,
+				value: item,
+			}
+			c.addNodeToTail(newNode)
+			c.cache[key] = newNode
+			c.cacheSize += int64(item.Size)
+		}
+	}
+}
+
+// cacheDir 返回缓存目录
+func cacheDir() string {
+	return "data"
+}
+
+// SetMaxCacheTTL 设置最大缓存TTL
+func (c *DNSCache) SetMaxCacheTTL(ttl time.Duration) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	c.maxCacheTTL = ttl
+}
+
+// SetMinCacheTTL 设置最小缓存TTL
+func (c *DNSCache) SetMinCacheTTL(ttl time.Duration) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	c.minCacheTTL = ttl
+}
+
+// SetCacheMode 设置缓存模式
+func (c *DNSCache) SetCacheMode(mode string) {
+	c.mutex.Lock()
+	oldMode := c.cacheMode
+	c.mutex.Unlock()
+
+	// 根据模式变化决定是否启动或停止保存循环
+	if oldMode != mode {
+		if oldMode == "file" {
+			// 从file模式切换到其他模式，先保存当前缓存到文件
+			// 直接调用底层保存逻辑，不检查缓存模式
+			c.saveCacheToFile()
+		}
+
+		c.mutex.Lock()
+		c.cacheMode = mode
+		c.mutex.Unlock()
+
+		if mode == "file" {
+			// 切换到file模式，启动保存循环
+			c.startSaveLoop()
+		} else {
+			// 切换到非file模式，停止保存循环
+			c.saveLoopMutex.Lock()
+			if c.saveRunning {
+				close(c.saveStopCh)
+				c.saveRunning = false
+			}
+			c.saveLoopMutex.Unlock()
+		}
+	}
+}
+
+// SetMaxCacheSize 设置最大缓存大小
+func (c *DNSCache) SetMaxCacheSize(size int64) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	c.maxCacheSize = size
+}
+
 // cleanupExpired 清理过期的缓存项
 func (c *DNSCache) cleanupExpired() {
 	now := time.Now()

dns/server.go

@@ -176,8 +176,13 @@ type Stats struct {
 func NewServer(config *config.DNSConfig, shieldConfig *config.ShieldConfig, shieldManager *shield.ShieldManager, gfwConfig *config.GFWListConfig, gfwManager *gfw.GFWListManager) *Server {
 	ctx, cancel := context.WithCancel(context.Background())
 
-	// 从配置中读取DNS缓存TTL值（分钟）
-	cacheTTL := time.Duration(config.CacheTTL) * time.Minute
+	// 从配置中读取DNS缓存TTL值（秒）
+	cacheTTL := time.Duration(config.CacheTTL) * time.Second
+	// 保存间隔（秒）
+	saveInterval := time.Duration(config.SaveInterval) * time.Second
+	// 最大和最小缓存TTL（秒）
+	maxCacheTTL := time.Duration(config.MaxCacheTTL) * time.Second
+	minCacheTTL := time.Duration(config.MinCacheTTL) * time.Second
 
 	server := &Server{
 		config:        config,
@@ -220,7 +225,7 @@ func NewServer(config *config.DNSConfig, shieldConfig *config.ShieldConfig, shie
 		stopped:         false, // 初始化为未停止状态
 
 		// DNS查询缓存初始化
-		DnsCache: NewDNSCache(cacheTTL),
+		DnsCache: NewDNSCache(cacheTTL, config.CacheMode, config.CacheSize, config.CacheFilePath, saveInterval, maxCacheTTL, minCacheTTL),
 		// 初始化域名DNSSEC状态映射表
 		domainDNSSECStatus: make(map[string]bool),
 		// 初始化服务器状态跟踪
@@ -2645,10 +2650,16 @@ func (s *Server) GetTopBlockedDomains(limit int) []BlockedDomain {
 	s.blockedDomainsMutex.RLock()
 	defer s.blockedDomainsMutex.RUnlock()
 
-	// 转换为切片
+	// 计算30天前的时间戳
+	thirtyDaysAgo := time.Now().Add(-30 * 24 * time.Hour).Unix()
+
+	// 转换为切片并过滤最近30天的数据
 	domains := make([]BlockedDomain, 0, len(s.blockedDomains))
 	for _, entry := range s.blockedDomains {
-		domains = append(domains, *entry)
+		// 只包含最近30天的数据
+		if entry.LastSeen >= thirtyDaysAgo {
+			domains = append(domains, *entry)
+		}
 	}
 
 	// 按计数排序
@@ -2668,10 +2679,16 @@ func (s *Server) GetTopResolvedDomains(limit int) []BlockedDomain {
 	s.resolvedDomainsMutex.RLock()
 	defer s.resolvedDomainsMutex.RUnlock()
 
-	// 转换为切片
+	// 计算30天前的时间戳
+	thirtyDaysAgo := time.Now().Add(-30 * 24 * time.Hour).Unix()
+
+	// 转换为切片并过滤最近30天的数据
 	domains := make([]BlockedDomain, 0, len(s.resolvedDomains))
 	for _, entry := range s.resolvedDomains {
-		domains = append(domains, *entry)
+		// 只包含最近30天的数据
+		if entry.LastSeen >= thirtyDaysAgo {
+			domains = append(domains, *entry)
+		}
 	}
 
 	// 按数量排序