128 lines
2.5 KiB
Go
128 lines
2.5 KiB
Go
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)
|
|
}
|
|
}
|
|
}
|