You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
85 lines
1.3 KiB
Go
85 lines
1.3 KiB
Go
package main
|
|
|
|
import (
|
|
"crypto/md5"
|
|
"encoding/binary"
|
|
"encoding/hex"
|
|
"github.com/miekg/dns"
|
|
"log"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type cache struct {
|
|
m map[string]*elem
|
|
lock *sync.RWMutex
|
|
ttl int64
|
|
max int
|
|
}
|
|
|
|
type elem struct {
|
|
m *dns.Msg
|
|
t int64
|
|
}
|
|
|
|
func newCache(max int, ttl int64) *cache {
|
|
return &cache{
|
|
max: max,
|
|
ttl: ttl,
|
|
m: map[string]*elem{},
|
|
lock: new(sync.RWMutex),
|
|
}
|
|
}
|
|
|
|
func key(m *dns.Msg) string {
|
|
d := m.Question[0].Name
|
|
b := []byte(d)
|
|
b1 := make([]byte, 4)
|
|
binary.BigEndian.PutUint16(b1[0:], m.Question[0].Qclass)
|
|
binary.BigEndian.PutUint16(b1[2:], m.Question[0].Qtype)
|
|
b = append(b, b1...)
|
|
h := md5.New()
|
|
h.Write(b)
|
|
s1 := hex.EncodeToString(h.Sum(nil))
|
|
return s1
|
|
}
|
|
|
|
func (c *cache) get(m *dns.Msg) *dns.Msg {
|
|
c.lock.RLock()
|
|
defer c.lock.RUnlock()
|
|
k := key(m)
|
|
if m1, ok := c.m[k]; ok {
|
|
t := time.Now().Unix()
|
|
if t < m1.t {
|
|
return m1.m
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *cache) set(m *dns.Msg) {
|
|
c.lock.Lock()
|
|
defer c.lock.Unlock()
|
|
|
|
if len(c.m) >= c.max {
|
|
log.Printf("clean the old cache")
|
|
c.cleanOld()
|
|
}
|
|
|
|
k := key(m)
|
|
c.m[k] = &elem{
|
|
m.Copy(),
|
|
time.Now().Unix() + c.ttl,
|
|
}
|
|
}
|
|
|
|
// must hold the write lock
|
|
func (c *cache) cleanOld() {
|
|
t1 := time.Now().Unix()
|
|
for k, v := range c.m {
|
|
if v.t >= t1 {
|
|
delete(c.m, k)
|
|
}
|
|
}
|
|
}
|