add resolver

pkg/internal/util/resolver/cache.go

@@ -0,0 +1,86 @@
+package resolver
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/go-gost/gost/pkg/logger"
+	"github.com/miekg/dns"
+)
+
+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) *dns.Msg {
+	v, ok := c.m.Load(key)
+	if !ok {
+		return nil
+	}
+
+	item, ok := v.(*cacheItem)
+	if !ok {
+		return nil
+	}
+
+	elapsed := time.Since(item.ts)
+	if item.ttl > 0 {
+		if elapsed > item.ttl {
+			c.m.Delete(key)
+			return nil
+		}
+	} else {
+		for _, rr := range item.msg.Answer {
+			if elapsed > time.Duration(rr.Header().Ttl)*time.Second {
+				c.m.Delete(key)
+				return nil
+			}
+		}
+	}
+
+	c.logger.Debugf("resolver cache hit %s", key)
+
+	return item.msg.Copy()
+}
+
+func (c *Cache) Store(key CacheKey, mr *dns.Msg, ttl time.Duration) {
+	if key == "" || mr == nil || ttl < 0 {
+		return
+	}
+
+	c.m.Store(key, &cacheItem{
+		msg: mr.Copy(),
+		ts:  time.Now(),
+		ttl: ttl,
+	})
+
+	c.logger.Debugf("resolver cache store %s", key)
+}

pkg/internal/util/resolver/resolver.go

@@ -0,0 +1,30 @@
+package resolver
+
+import (
+	"net"
+
+	"github.com/miekg/dns"
+)
+
+func AddSubnetOpt(m *dns.Msg, ip net.IP) {
+	if m == nil || ip == nil {
+		return
+	}
+
+	opt := new(dns.OPT)
+	opt.Hdr.Name = "."
+	opt.Hdr.Rrtype = dns.TypeOPT
+	e := new(dns.EDNS0_SUBNET)
+	e.Code = dns.EDNS0SUBNET
+	if ip := ip.To4(); ip != nil {
+		e.Family = 1
+		e.SourceNetmask = 24
+		e.Address = ip
+	} else {
+		e.Family = 2
+		e.SourceNetmask = 128
+		e.Address = ip.To16()
+	}
+	opt.Option = append(opt.Option, e)
+	m.Extra = append(m.Extra, opt)
+}