Add dns cache

Author: net-byte

README.md

@@ -5,14 +5,14 @@ DNS proxy server over socks5(udp)
 # Usage
 ```
 Usage of ./bin/dns2socks:
+  -c cache dns type a (default true)
   -d string
     	remote dns server address (default "8.8.8.8:53")
   -l string
     	local dns server address (default "127.0.0.1:53")
   -s string
     	socks5(udp) proxy address (default "127.0.0.1:1080")
 
-
 ```
 # Run on docker  
 ```

cmd/dns2socks.go

@@ -4,13 +4,17 @@ import (
 	"io"
 	"log"
 	"net"
+	"time"
 
 	"github.com/miekg/dns"
+	"github.com/patrickmn/go-cache"
 	"github.com/thinkgos/go-socks5/ccsocks5"
 )
 
+var c = cache.New(30*time.Minute, 10*time.Minute)
+
 //StartServer acts start a dns proxy server
-func StartServer(localAddr *string, socksAddr *string, dnsAddr *string) {
+func StartServer(localAddr *string, socksAddr *string, dnsAddr *string, cached *bool) {
 	addr, err := net.ResolveUDPAddr("udp", *localAddr)
 	if nil != err {
 		log.Fatalln("Unable to get UDP socket:", err)
@@ -21,6 +25,7 @@ func StartServer(localAddr *string, socksAddr *string, dnsAddr *string) {
 	}
 	defer conn.Close()
 	log.Printf("dns2socks started on %v", *localAddr)
+
 	client := ccsocks5.NewClient(*socksAddr)
 	defer client.Close()
 	buf := make([]byte, 4096)
@@ -30,7 +35,17 @@ func StartServer(localAddr *string, socksAddr *string, dnsAddr *string) {
 			continue
 		}
 		b := buf[:n]
-		printQuestion(b)
+		if *cached {
+			answer := getAnswerFromCache(b)
+			if answer != nil {
+				respMsg := new(dns.Msg)
+				respMsg.Unpack(b)
+				respMsg.Answer = append(respMsg.Answer, *answer)
+				respData, _ := respMsg.Pack()
+				conn.WriteToUDP(respData, fromAddr)
+				continue
+			}
+		}
 		proxyConn, err := client.Dial("udp", *dnsAddr)
 		if err != nil {
 			log.Println(err)
@@ -45,21 +60,35 @@ func StartServer(localAddr *string, socksAddr *string, dnsAddr *string) {
 					break
 				}
 				b := buf[:n]
-				printAnswer(b)
+				if *cached {
+					setAnswerCache(b)
+				}
 				conn.WriteToUDP(b, fromAddr)
 			}
 		}()
 	}
 }
 
-func printQuestion(data []byte) {
+func getAnswerFromCache(data []byte) *dns.RR {
 	msg := new(dns.Msg)
 	msg.Unpack(data)
-	log.Printf("dns question:%v", msg.Question)
+	q := msg.Question[0]
+	if q.Qtype == dns.TypeA {
+		if v, found := c.Get(q.Name); found {
+			log.Printf("query dns from cache:%v", msg.Question)
+			ret := v.(*dns.RR)
+			return ret
+		}
+	}
+	return nil
 }
 
-func printAnswer(data []byte) {
+func setAnswerCache(data []byte) {
 	msg := new(dns.Msg)
 	msg.Unpack(data)
-	log.Printf("dns answer:%v", msg.Answer)
+	q := msg.Question[0]
+	if q.Qtype == dns.TypeA && len(msg.Answer) > 0 {
+		c.Set(q.Name, &msg.Answer[len(msg.Answer)-1], cache.DefaultExpiration)
+		log.Printf("set dns cache:%v", q.Name)
+	}
 }

go.mod

@@ -4,6 +4,7 @@ go 1.16
 
 require (
 	github.com/miekg/dns v1.1.40
+	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/thinkgos/go-socks5 v0.2.2 // indirect
 	golang.org/x/net v0.0.0-20210224082022-3d97a244fca7
 )

go.sum

@@ -1,6 +1,8 @@
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/miekg/dns v1.1.40 h1:pyyPFfGMnciYUk/mXpKkVmeMQjfXqt3FAJ2hy7tPiLA=
 github.com/miekg/dns v1.1.40/go.mod h1:KNUDUusw/aVsxyTYZM1oqvCicbwhgbNgztCETuNZ7xM=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=

main.go

@@ -10,6 +10,7 @@ func main() {
 	localAddr := flag.String("l", "127.0.0.1:53", "local dns server address")
 	socksAddr := flag.String("s", "127.0.0.1:1080", "socks5(udp) proxy address")
 	dnsAddr := flag.String("d", "8.8.8.8:53", "remote dns server address")
+	cached := flag.Bool("c", true, "cache dns type a")
 	flag.Parse()
-	cmd.StartServer(localAddr, socksAddr, dnsAddr)
+	cmd.StartServer(localAddr, socksAddr, dnsAddr, cached)
 }