add dns cache

dns_cache.go

@@ -0,0 +1,75 @@
+package main
+
+import (
+	"errors"
+	"log"
+	"net"
+	"sync"
+	"time"
+)
+
+var (
+	ErrNoRecords = errors.New("no records")
+)
+
+type DNSCache struct {
+	records map[string]DNSRecord
+	mutex   sync.Mutex
+}
+
+type DNSRecord struct {
+	addresses []net.IP
+	time      time.Time
+}
+
+func (dns *DNSCache) Put(domain string, addresses []net.IP) {
+	if len(addresses) > 0 {
+		dns.mutex.Lock()
+		defer dns.mutex.Unlock()
+		_, contains := dns.records[domain]
+		if contains {
+			return
+		}
+		dns.records[domain] = DNSRecord{
+			addresses: addresses,
+			time:      time.Now(),
+		}
+	}
+}
+
+func (dns *DNSCache) Get(domain string) ([]net.IP, error) {
+	dns.mutex.Lock()
+	defer dns.mutex.Unlock()
+	_, contains := dns.records[domain]
+	if contains {
+		return dns.records[domain].addresses, nil
+	}
+	return nil, ErrNoRecords
+}
+
+func (dns *DNSCache) CleanUp() {
+	log.Printf("DNSCache: Starting DNS Cache CleanUp")
+	records := 0
+	for range time.Tick(time.Minute) {
+		dns.mutex.Lock()
+		for k, v := range dns.records {
+			if time.Since(v.time).Minutes() > 10 {
+				delete(dns.records, k)
+				records++
+			}
+		}
+		dns.mutex.Unlock()
+		if records > 0 {
+			log.Printf("DNSCache: %d dns records has been removed", records)
+		}
+		records = 0
+	}
+}
+
+func NewDNSCache() *DNSCache {
+	cache := &DNSCache{
+		records: make(map[string]DNSRecord),
+	}
+	go cache.CleanUp()
+	return cache
+}

proxy_dialer.go

@@ -14,10 +14,31 @@ import (
 )
 
 type ProxyDialer struct {
+	dnsCache *DNSCache
 	proxy    C.Proxy
 	resolver *net.Resolver
 }
 
+func (p *ProxyDialer) resolve(address string) (net.IP, error) {
+	ips, err := p.dnsCache.Get(address)
+	if err != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		ips, err := p.resolver.LookupIP(ctx, "ip", address)
+		if err != nil {
+			return nil, err
+		}
+		if len(ips) == 0 {
+			return nil, fmt.Errorf("no address associated with this domain %s", address)
+		}
+		p.dnsCache.Put(address, ips)
+		return ips[0], nil
+	}
+	ip := ips[0]
+	log.Printf("Cache: %s is available in cache : %s\n", address, ip.String())
+	return ip, nil
+}
+
 func (p *ProxyDialer) Dial(network, addr string) (c net.Conn, err error) {
 	log.Printf("Dial: network: %s, addr: %s\n", network, addr)
 
@@ -32,16 +53,9 @@ func (p *ProxyDialer) Dial(network, addr string) (c net.Conn, err error) {
 	ip := net.ParseIP(address)
 
 	if ip == nil {
-		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		defer cancel()
-		ips, err := p.resolver.LookupIP(ctx, "ip", address)
-		if err != nil {
+		if ip, err = p.resolve(address); err != nil {
 			return nil, err
 		}
-		if len(ips) == 0 {
-			return nil, fmt.Errorf("no address associated with this domain %s", addr)
-		}
-		ip = ips[0]
 	}
 
 	metadata := &C.Metadata{
@@ -83,16 +97,9 @@ func (p *ProxyDialer) DialUDP(network, addr string) (pc net.PacketConn, writeTo
 	ip := net.ParseIP(address)
 
 	if ip == nil {
-		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		defer cancel()
-		ips, err := p.resolver.LookupIP(ctx, "ip", address)
-		if err != nil {
+		if ip, err = p.resolve(address); err != nil {
 			return nil, nil, err
 		}
-		if len(ips) == 0 {
-			return nil, nil, fmt.Errorf("no address associated with this domain %s", addr)
-		}
-		ip = ips[0]
 	}
 
 	metadata := &C.Metadata{
@@ -130,7 +137,7 @@ func (p *ProxyDialer) Addr() string {
 	return ""
 }
 
-func NewProxyDialer(p C.Proxy, Dns string) (*ProxyDialer, error) {
+func NewProxyDialer(p C.Proxy, Dns string, dnsCache *DNSCache) (*ProxyDialer, error) {
 	c := strings.LastIndex(Dns, ":")
 	if c == -1 {
 		return nil, errors.New("bad dns address")
@@ -149,7 +156,8 @@ func NewProxyDialer(p C.Proxy, Dns string) (*ProxyDialer, error) {
 	}
 
 	pd := &ProxyDialer{
-		proxy: p,
+		proxy:    p,
+		dnsCache: dnsCache,
 	}
 
 	pd.resolver = &net.Resolver{

ssr_client.go

@@ -12,6 +12,7 @@ import (
 type SSRClient struct {
 	httpServer   proxy.Server
 	socks5server proxy.Server
+	dnsCache     *DNSCache
 }
 
 func (s *SSRClient) ListenAndServe() {
@@ -33,7 +34,11 @@ func NewSSRClient(
 	protocolParam string,
 	dns string,
 ) (*SSRClient, error) {
-	client := &SSRClient{}
+	dnsCache := NewDNSCache()
+
+	client := &SSRClient{
+		dnsCache: dnsCache,
+	}
 
 	ssrNode := map[string]interface{}{
 		"name":           "ssr",
@@ -51,7 +56,7 @@ func NewSSRClient(
 
 	p, _ := outbound.ParseProxy(ssrNode)
 
-	pr, err := NewProxyDialer(p, dns)
+	pr, err := NewProxyDialer(p, dns, dnsCache)
 
 	if err != nil {
 		return nil, err