dns-server/dns/cache.go

package dns

import (
	"sync"
	"time"

	"github.com/miekg/dns"
)

// DNSCacheItem 表示缓存中的DNS响应项
type DNSCacheItem struct {
	Response *dns.Msg  // DNS响应消息
	Expiry   time.Time // 过期时间
	HasDNSSEC bool     // 是否包含DNSSEC记录
}

// 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]
}

// hasDNSSECRecords 检查响应是否包含DNSSEC记录
func hasDNSSECRecords(response *dns.Msg) bool {
	// 检查响应中是否包含DNSSEC相关记录（DNSKEY、RRSIG、DS、NSEC、NSEC3等）
	for _, rr := range response.Answer {
		if _, ok := rr.(*dns.DNSKEY); ok {
			return true
		}
		if _, ok := rr.(*dns.RRSIG); ok {
			return true
		}
		if _, ok := rr.(*dns.DS); ok {
			return true
		}
		if _, ok := rr.(*dns.NSEC); ok {
			return true
		}
		if _, ok := rr.(*dns.NSEC3); ok {
			return true
		}
	}
	for _, rr := range response.Ns {
		if _, ok := rr.(*dns.DNSKEY); ok {
			return true
		}
		if _, ok := rr.(*dns.RRSIG); ok {
			return true
		}
		if _, ok := rr.(*dns.DS); ok {
			return true
		}
		if _, ok := rr.(*dns.NSEC); ok {
			return true
		}
		if _, ok := rr.(*dns.NSEC3); ok {
			return true
		}
	}
	for _, rr := range response.Extra {
		if _, ok := rr.(*dns.DNSKEY); ok {
			return true
		}
		if _, ok := rr.(*dns.RRSIG); ok {
			return true
		}
		if _, ok := rr.(*dns.DS); ok {
			return true
		}
		if _, ok := rr.(*dns.NSEC); ok {
			return true
		}
		if _, ok := rr.(*dns.NSEC3); ok {
			return true
		}
	}
	return false
}

// 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),
		HasDNSSEC: hasDNSSECRecords(response), // 检查并设置DNSSEC标志
	}

	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)
		}
	}
}