add dns cache for dns proxy

cmd/sower/main.go

@@ -80,7 +80,8 @@ func init() {
 			Msg("Load config")
 	}
 
-	conf.Router.Direct.Rules = append(conf.Router.Direct.Rules, conf.Remote.Addr)
+	conf.Router.Direct.Rules = append(conf.Router.Direct.Rules,
+		conf.Remote.Addr, "**.in-addr.arpa", "**.ip6.arpa")
 	log.Info().
 		Str("version", version).
 		Str("date", date).

pkg/mem/mem.go

@@ -2,7 +2,6 @@ package mem
 
 import (
 	"errors"
-	"fmt"
 	"reflect"
 	"sync"
 	"time"
@@ -12,7 +11,7 @@ import (
 
 // Data define the type which can speed up by mem cache
 type Data interface {
-	Fulfill(key interface{}) error
+	Fulfill(key string) (err error)
 }
 
 // Cache is the definition of cache, be careful of the memory usage
@@ -28,12 +27,12 @@ type Cache struct {
 var DefaultCache = New(time.Minute)
 
 // Remember is a surge, it provides a quite simple way to use cache
-func Remember(dst Data, key interface{}) error {
+func Remember(dst Data, key string) error {
 	return DefaultCache.Remember(dst, key)
 }
 
 // Delete is a surge, it delete a specified data in DefaultCache
-func Delete(dst Data, key interface{}) {
+func Delete(dst Data, key string) {
 	DefaultCache.Delete(dst, key)
 }
 
@@ -49,7 +48,7 @@ func New(rotateInterval time.Duration) *Cache {
 }
 
 // Remember automatically save and retrieve data from a cache entity
-func (c *Cache) Remember(dst Data, key interface{}) error {
+func (c *Cache) Remember(dst Data, key string) error {
 	rv := reflect.ValueOf(dst)
 	if rv.Kind() != reflect.Ptr {
 		panic("invalid not pointor type: " + reflect.TypeOf(dst).Name())
@@ -72,17 +71,16 @@ func (c *Cache) Remember(dst Data, key interface{}) error {
 	}
 
 	// First: load from cache
-	cacheKey := fmt.Sprintf("%T%v", dst, key)
-	if val, ok := c.now.Load(cacheKey); ok {
+	if val, ok := c.now.Load(key); ok {
 		return deepcopier.Copy(val).To(dst)
 	}
 
 	// Second: load from old cache, or waitting the sigle groutine getting data
 	ch := make(chan struct{})
-	if chVal, ok := c.barrier.LoadOrStore(cacheKey, ch); ok {
+	if chVal, ok := c.barrier.LoadOrStore(key, ch); ok {
 		close(ch) // the ch is not used
 
-		if val, ok := c.old.Load(cacheKey); ok {
+		if val, ok := c.old.Load(key); ok {
 			return deepcopier.Copy(val).To(dst)
 		}
 
@@ -90,12 +88,12 @@ func (c *Cache) Remember(dst Data, key interface{}) error {
 		// type error: already failed
 		if ch, ok = chVal.(chan struct{}); ok {
 			<-ch
-			if val, ok := c.now.Load(cacheKey); ok {
+			if val, ok := c.now.Load(key); ok {
 				return deepcopier.Copy(val).To(dst)
 			}
 		}
 
-		val, _ := c.barrier.Load(cacheKey)
+		val, _ := c.barrier.Load(key)
 		if err, ok := val.(error); ok {
 			return err
 		}
@@ -106,25 +104,24 @@ func (c *Cache) Remember(dst Data, key interface{}) error {
 	// Third: getting data from CacheType, maybe from db
 	err := dst.Fulfill(key)
 	if err != nil {
-		c.barrier.Store(cacheKey, err)
+		c.barrier.Store(key, err)
 		return err
 	}
 
-	c.now.Store(cacheKey, dst)
+	c.now.Store(key, dst)
 	close(ch) // broadcast, wakeup all waiting groutine
 
 	return nil
 }
 
 // Delete immediately specified the cached content to expire
-func (c *Cache) Delete(dst Data, key interface{}) {
+func (c *Cache) Delete(dst Data, key string) {
 	c.rwmutex.Lock()
 	defer c.rwmutex.Unlock()
 
-	cacheKey := fmt.Sprintf("%T%v", dst, key)
-	c.old.Delete(cacheKey)
-	c.now.Delete(cacheKey)
-	c.barrier.Store(cacheKey, errors.New(cacheKey+"is deleted"))
+	c.old.Delete(key)
+	c.now.Delete(key)
+	c.barrier.Store(key, errors.New(key+" is deleted"))
 }
 
 // Rotate force refresh cached data

router/dns.go

@@ -9,12 +9,6 @@ import (
 )
 
 func (r *Router) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
-	// *Msg r has an TSIG record and it was validated
-	if req.IsTsig() != nil && w.TsigStatus() == nil {
-		lastTsig := req.Extra[len(req.Extra)-1].(*dns.TSIG)
-		req.SetTsig(lastTsig.Hdr.Name, dns.HmacMD5, 300, time.Now().Unix())
-	}
-
 	// https://stackoverflow.com/questions/4082081/requesting-a-and-aaaa-records-in-single-dns-query/4083071#4083071
 	if len(req.Question) == 0 {
 		w.WriteMsg(r.dnsFail(req, dns.RcodeFormatError))
@@ -34,28 +28,15 @@ func (r *Router) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
 		return
 	}
 
-	conn := <-r.dns.connCh
-	resp, rtt, err := r.dns.ExchangeWithConn(req, conn)
-	deferlog.Std.DebugWarn(err).
-		Dur("rtt", rtt).
-		Str("domain", domain).
-		Msg("exchange dns record")
-
-	if err != nil {
-		conn.Close()
+	c := &dnsCache{Router: r, Req: req}
+	if err := r.dns.cache.Remember(c, req.Question[0].String()); err != nil {
 		w.WriteMsg(r.dnsFail(req, dns.RcodeServerFailure))
 		return
 	}
 
-	select {
-	case r.dns.connCh <- conn:
-	default:
-		conn.Close()
-	}
-
-	resp.Id = req.Id
-	resp.Compress = true
-	w.WriteMsg(resp)
+	c.Resp.SetReply(req)
+	c.Resp.Compress = true
+	w.WriteMsg(c.Resp)
 }
 
 func (r *Router) dnsFail(req *dns.Msg, rcode int) *dns.Msg {
@@ -82,3 +63,29 @@ func (r *Router) dnsProxyA(domain string, localIP net.IP, req *dns.Msg) *dns.Msg
 	}
 	return m
 }
+
+type dnsCache struct {
+	*Router
+	Req, Resp *dns.Msg
+}
+
+func (r *dnsCache) Fulfill(question string) (err error) {
+	conn := <-r.dns.connCh
+
+	var rtt time.Duration
+	r.Resp, rtt, err = r.dns.ExchangeWithConn(r.Req, conn)
+	if err != nil {
+		r.Resp, rtt, err = r.dns.ExchangeWithConn(r.Req, conn)
+	}
+	deferlog.Std.DebugWarn(err).
+		Dur("rtt", rtt).
+		Str("question", question).
+		Msg("exchange dns record")
+
+	select {
+	case r.dns.connCh <- conn:
+	default:
+		conn.Close()
+	}
+	return err
+}

router/ping.go

@@ -19,7 +19,7 @@ func (r *Router) isAccess(domain string, port uint16) bool {
 	}
 
 	p := &ping{}
-	r.cache.Remember(p, domain)
+	r.accessCache.Remember(p, domain)
 	return p.isAccess
 }
 
@@ -27,8 +27,8 @@ type ping struct {
 	isAccess bool
 }
 
-func (p *ping) Fulfill(key interface{}) error {
-	_, err := http.Head(net.JoinHostPort(key.(string), "80"))
+func (p *ping) Fulfill(key string) error {
+	_, err := http.Head(net.JoinHostPort(key, "80"))
 	p.isAccess = (err == nil)
 	return nil
 }

router/router.go

@@ -17,17 +17,18 @@ import (
 
 type ProxyDialFn func(network, host string, port uint16) (net.Conn, error)
 type Router struct {
-	blockRule  *util.Node
-	directRule *util.Node
-	proxyRule  *util.Node
-	ProxyDial  ProxyDialFn
-	cache      *mem.Cache
+	blockRule   *util.Node
+	directRule  *util.Node
+	proxyRule   *util.Node
+	ProxyDial   ProxyDialFn
+	accessCache *mem.Cache
 
 	dns struct {
 		fallbackDNS string
 		serveIP     net.IP
 		dns.Client
 		connCh chan *dns.Conn
+		cache  *mem.Cache
 	}
 
 	country struct {
@@ -38,13 +39,14 @@ type Router struct {
 
 func NewRouter(serveIP, fallbackDNS, mmdbFile string, proxyDial ProxyDialFn) *Router {
 	r := Router{
-		ProxyDial: proxyDial,
-		cache:     mem.New(time.Hour), // TODO: config
+		ProxyDial:   proxyDial,
+		accessCache: mem.New(time.Hour), // TODO: config
 	}
 
 	r.dns.serveIP = net.ParseIP(serveIP)
 	r.dns.fallbackDNS = fallbackDNS
 	r.dns.connCh = make(chan *dns.Conn, 1)
+	r.dns.cache = mem.New(5 * time.Minute) // Tll: 10 minutes
 	go r.dialDNSConn()
 
 	var err error