feat: dns caching

dns/cache/dns_cache.go

@@ -0,0 +1,60 @@
+package cache
+
+import (
+	"sync"
+	"time"
+
+	"github.com/xvzc/SpoofDPI/util"
+)
+
+type DNSCacheEntry struct {
+	ip           string
+	expiry_timer *time.Timer
+}
+
+type DNSCache struct {
+	cacheMap  map[string]DNSCacheEntry
+	cacheLock sync.RWMutex
+}
+
+var dnsCache *DNSCache
+var once sync.Once
+
+func GetCache() *DNSCache {
+	once.Do(func() {
+		dnsCache = &DNSCache{
+			cacheMap:  make(map[string]DNSCacheEntry),
+			cacheLock: sync.RWMutex{},
+		}
+	})
+	return dnsCache
+}
+
+// interface function for the inner map basically.
+func (d *DNSCache) Set(key string, value string) {
+	d.cacheLock.Lock()
+	defer d.cacheLock.Unlock()
+	if _, ok := d.cacheMap[key]; ok {
+		return
+	}
+
+	d.cacheMap[key] = DNSCacheEntry{
+		ip: value,
+		expiry_timer: time.AfterFunc(time.Duration(*&util.GetConfig().DnsCacheTTL)*time.Second, func() {
+			d.cacheLock.Lock()
+			defer d.cacheLock.Unlock()
+			delete(d.cacheMap, key)
+		}),
+	}
+}
+
+func (d *DNSCache) Get(key string) (string, bool) {
+	d.cacheLock.RLock()
+	defer d.cacheLock.RUnlock()
+
+	if value, ok := d.cacheMap[key]; ok {
+		return value.ip, true
+	}
+
+	return "", false
+}

dns/dns.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/miekg/dns"
+	"github.com/xvzc/SpoofDPI/dns/cache"
 	"github.com/xvzc/SpoofDPI/dns/resolver"
 	"github.com/xvzc/SpoofDPI/util"
 	"github.com/xvzc/SpoofDPI/util/log"
@@ -23,6 +24,7 @@ type Resolver interface {
 type Dns struct {
 	host          string
 	port          string
+	dnsCache      *cache.DNSCache
 	systemClient  Resolver
 	generalClient Resolver
 	dohClient     Resolver
@@ -41,6 +43,7 @@ func NewDns(config *util.Config) *Dns {
 	return &Dns{
 		host:          config.DnsAddr,
 		port:          port,
+		dnsCache:      cache.GetCache(),
 		systemClient:  resolver.NewSystemResolver(),
 		generalClient: resolver.NewGeneralResolver(net.JoinHostPort(addr, port)),
 		dohClient:     resolver.NewDOHResolver(addr),
@@ -56,6 +59,11 @@ func (d *Dns) ResolveHost(ctx context.Context, host string, enableDoh bool, useS
 		return ip.String(), nil
 	}
 
+	if r, ok := d.dnsCache.Get(host); ok {
+		logger.Debug().Msgf("found %s in cache", host)
+		return r, nil
+	}
+
 	clt := d.clientFactory(enableDoh, useSystemDns)
 	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 	defer cancel()
@@ -71,8 +79,11 @@ func (d *Dns) ResolveHost(ctx context.Context, host string, enableDoh bool, useS
 	}
 
 	if len(addrs) > 0 {
-		d := time.Since(t).Milliseconds()
-		logger.Debug().Msgf("resolved %s from %s in %d ms", addrs[0].String(), host, d)
+		elapsed := time.Since(t).Milliseconds()
+		logger.Debug().Msgf("resolved %s from %s in %d ms", addrs[0].String(), host, elapsed)
+
+		d.dnsCache.Set(host, addrs[0].String())
+
 		return addrs[0].String(), nil
 	}
 

util/args.go

@@ -14,6 +14,7 @@ type Args struct {
 	DnsAddr        string
 	DnsPort        uint16
 	DnsIPv4Only    bool
+	DnsCacheTTL    uint16
 	EnableDoh      bool
 	Debug          bool
 	Silent         bool
@@ -42,6 +43,7 @@ func ParseArgs() *Args {
 	uintNVar(&args.Port, "port", 8080, "port")
 	flag.StringVar(&args.DnsAddr, "dns-addr", "8.8.8.8", "dns address")
 	uintNVar(&args.DnsPort, "dns-port", 53, "port number for dns")
+	uintNVar(&args.DnsCacheTTL, "dns-cache-ttl", 300, "time-to-live in seconds for dns cache")
 	flag.BoolVar(&args.EnableDoh, "enable-doh", false, "enable 'dns-over-https'")
 	flag.BoolVar(&args.Debug, "debug", false, "enable debug output")
 	flag.BoolVar(&args.Silent, "silent", false, "do not show the banner and server information at start up")

util/config.go

@@ -14,6 +14,7 @@ type Config struct {
 	DnsAddr         string
 	DnsPort         int
 	DnsIPv4Only     bool
+	DnsCacheTTL     int
 	EnableDoh       bool
 	Debug           bool
 	Silent          bool
@@ -38,6 +39,7 @@ func (c *Config) Load(args *Args) {
 	c.DnsAddr = args.DnsAddr
 	c.DnsPort = int(args.DnsPort)
 	c.DnsIPv4Only = args.DnsIPv4Only
+	c.DnsCacheTTL = int(args.DnsCacheTTL)
 	c.Debug = args.Debug
 	c.EnableDoh = args.EnableDoh
 	c.Silent = args.Silent