实现缓存功能

dns/cache.go

@@ -0,0 +1,127 @@
+package dns
+
+import (
+	"sync"
+	"time"
+
+	"github.com/miekg/dns"
+)
+
+// DNSCacheItem 表示缓存中的DNS响应项
+type DNSCacheItem struct {
+	Response *dns.Msg  // DNS响应消息
+	Expiry   time.Time // 过期时间
+}
+
+// DNSCache DNS缓存结构
+type DNSCache struct {
+	cache      map[string]*DNSCacheItem // 缓存映射表
+	mutex      sync.RWMutex             // 读写锁，保护缓存
+	defaultTTL time.Duration            // 默认缓存TTL
+}
+
+// NewDNSCache 创建新的DNS缓存实例
+func NewDNSCache(defaultTTL time.Duration) *DNSCache {
+	cache := &DNSCache{
+		cache:      make(map[string]*DNSCacheItem),
+		defaultTTL: defaultTTL,
+	}
+
+	// 启动缓存清理协程
+	go cache.startCleanupLoop()
+
+	return cache
+}
+
+// cacheKey 生成缓存键
+func cacheKey(qName string, qType uint16) string {
+	return qName + "|" + dns.TypeToString[qType]
+}
+
+// Set 设置缓存项
+func (c *DNSCache) Set(qName string, qType uint16, response *dns.Msg, ttl time.Duration) {
+	if ttl <= 0 {
+		ttl = c.defaultTTL
+	}
+
+	key := cacheKey(qName, qType)
+	item := &DNSCacheItem{
+		Response: response.Copy(), // 复制响应以避免外部修改
+		Expiry:   time.Now().Add(ttl),
+	}
+
+	c.mutex.Lock()
+	c.cache[key] = item
+	c.mutex.Unlock()
+}
+
+// Get 获取缓存项
+func (c *DNSCache) Get(qName string, qType uint16) (*dns.Msg, bool) {
+	key := cacheKey(qName, qType)
+
+	c.mutex.RLock()
+	item, found := c.cache[key]
+	if !found {
+		c.mutex.RUnlock()
+		return nil, false
+	}
+
+	// 检查是否过期
+	if time.Now().After(item.Expiry) {
+		c.mutex.RUnlock()
+		// 过期了，删除缓存项（在写锁中）
+		c.delete(key)
+		return nil, false
+	}
+
+	// 返回缓存的响应副本
+	response := item.Response.Copy()
+	c.mutex.RUnlock()
+
+	return response, true
+}
+
+// delete 删除缓存项
+func (c *DNSCache) delete(key string) {
+	c.mutex.Lock()
+	delete(c.cache, key)
+	c.mutex.Unlock()
+}
+
+// Clear 清空缓存
+func (c *DNSCache) Clear() {
+	c.mutex.Lock()
+	c.cache = make(map[string]*DNSCacheItem)
+	c.mutex.Unlock()
+}
+
+// Size 获取缓存大小
+func (c *DNSCache) Size() int {
+	c.mutex.RLock()
+	defer c.mutex.RUnlock()
+	return len(c.cache)
+}
+
+// startCleanupLoop 启动定期清理过期缓存的协程
+func (c *DNSCache) startCleanupLoop() {
+	ticker := time.NewTicker(time.Minute * 5) // 每5分钟清理一次
+	defer ticker.Stop()
+
+	for range ticker.C {
+		c.cleanupExpired()
+	}
+}
+
+// cleanupExpired 清理过期的缓存项
+func (c *DNSCache) cleanupExpired() {
+	now := time.Now()
+
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	for key, item := range c.cache {
+		if now.After(item.Expiry) {
+			delete(c.cache, key)
+		}
+	}
+}