-
Notifications
You must be signed in to change notification settings - Fork 0
/
cache.go
115 lines (95 loc) · 1.91 KB
/
cache.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
package resolver
import (
"fmt"
"sync"
"time"
"github.com/dolfly/core/logger"
"github.com/miekg/dns"
)
const (
defaultTTL = 60 * time.Second
)
type CacheKey string
// NewCacheKey generates resolver cache key from question of dns query.
func NewCacheKey(q *dns.Question) CacheKey {
if q == nil {
return ""
}
key := fmt.Sprintf("%s%s.%s", q.Name, dns.Class(q.Qclass).String(), dns.Type(q.Qtype).String())
return CacheKey(key)
}
type cacheItem struct {
msg *dns.Msg
ts time.Time
ttl time.Duration
}
type Cache struct {
m sync.Map
logger logger.Logger
}
func NewCache() *Cache {
return &Cache{}
}
func (c *Cache) WithLogger(logger logger.Logger) *Cache {
c.logger = logger
return c
}
func (c *Cache) Load(key CacheKey) (msg *dns.Msg, ttl time.Duration) {
v, ok := c.m.Load(key)
if !ok {
return
}
item, ok := v.(*cacheItem)
if !ok {
return
}
msg = item.msg.Copy()
for i := range msg.Answer {
d := uint32(time.Since(item.ts).Seconds())
if msg.Answer[i].Header().Ttl > d {
msg.Answer[i].Header().Ttl -= d
} else {
msg.Answer[i].Header().Ttl = 1
}
}
ttl = item.ttl - time.Since(item.ts)
c.logger.Debugf("resolver cache hit: %s, ttl: %v", key, ttl)
return
}
func (c *Cache) Store(key CacheKey, mr *dns.Msg, ttl time.Duration) {
if key == "" || mr == nil || ttl < 0 {
return
}
if ttl == 0 {
for _, answer := range mr.Answer {
v := time.Duration(answer.Header().Ttl) * time.Second
if ttl == 0 || ttl > v {
ttl = v
}
}
if ttl == 0 {
ttl = defaultTTL
}
} else {
for i := range mr.Answer {
mr.Answer[i].Header().Ttl = uint32(ttl.Seconds())
}
}
c.m.Store(key, &cacheItem{
msg: mr.Copy(),
ts: time.Now(),
ttl: ttl,
})
c.logger.Debugf("resolver cache store: %s, ttl: %v", key, ttl)
}
func (c *Cache) RefreshTTL(key CacheKey) {
v, ok := c.m.Load(key)
if !ok {
return
}
item, ok := v.(*cacheItem)
if !ok {
return
}
item.ts = time.Now()
}