foreward upstream, in memory cache, udp listener, handler, cli flags

1 files changed, 80 insertions, 0 deletions

diff --git a/internal/cache/cache.go b/internal/cache/cache.go
new file mode 100644
index 0000000..4d6b45c
--- /dev/null
+++ b/internal/cache/cache.go
@@ -0,0 +1,80 @@
+package cache
+
+import (
+	"sync"
+	"time"
+
+	"linum/internal/dns"
+)
+
+type Entry struct {
+	Msg *dns.Msg
+	Expires time.Time
+}
+
+func (e *Entry) expired() bool {
+	return time.Now().After(e.Expires)
+}
+
+type Cache struct {
+	mu sync.RWMutex
+	data map[string]*Entry
+}
+
+func New() *Cache {
+	return &Cache{
+		data: make(map[string]*Entry),
+	}
+}
+
+func makeKey(qname dns.Name, qtype, qclass uint16) string {
+	return qname.String() + ":" + itoa(qtype) + ":" + itoa(qclass)
+}
+
+func itoa(v uint16) string {
+	if v == 0 {
+		return "0"
+	}
+
+	var buf [5]byte
+	i := len(buf)
+	for v > 0 {
+		i--
+		buf[i] = byte(v%10) + '0'
+		v /= 10
+	}
+	return string(buf[i:])
+}
+
+func (c *Cache) Get(qname dns.Name, qtype, qclass uint16) (*dns.Msg, bool) {
+	key := makeKey(qname, qtype, qclass)
+	c.mu.RLock()
+	e, ok := c.data[key]
+	c.mu.RUnlock()
+	if !ok || e.expired() {
+		return nil, false
+	}
+	return e.Msg, true
+}
+
+func (c *Cache) Set(qname dns.Name, qtype, qclass uint16, msg *dns.Msg, ttl uint32) {
+	if ttl == 0 {
+		return
+	}
+	if ttl > 3600 {
+		ttl = 3600
+	}
+	key := makeKey(qname, qtype, qclass)
+	c.mu.Lock()
+	c.data[key] = &Entry{
+		Msg: msg,
+		Expires: time.Now().Add(time.Duration(ttl) * time.Second),
+	}
+	c.mu.Unlock()
+}
+
+func (c *Cache) Len() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return len(c.data)
+}