Merge pull request moby#1062 from sanimej/fixes

Limit number of concurrent DNS queries
shethchintan7 · Mar 30, 2016 · fa69adb · fa69adb
2 parents 45fe6f0 + f9b23d3
commit fa69adb
Showing 1 changed file with 41 additions and 1 deletion.
diff --git a/resolver.go b/resolver.go
@@ -47,6 +47,8 @@ const (
 	maxExtDNS       = 3 //max number of external servers to try
 	extIOTimeout    = 3 * time.Second
 	defaultRespSize = 512
+	maxConcurrent   = 50
+	logInterval     = 2 * time.Second
 )
 
 type extDNSEntry struct {
@@ -64,6 +66,9 @@ type resolver struct {
 	tcpServer  *dns.Server
 	tcpListen  *net.TCPListener
 	err        error
+	count      int32
+	tStamp     time.Time
+	queryLock  sync.Mutex
 }
 
 func init() {
@@ -162,6 +167,9 @@ func (r *resolver) Stop() {
 	r.conn = nil
 	r.tcpServer = nil
 	r.err = fmt.Errorf("setup not done yet")
+	r.tStamp = time.Time{}
+	r.count = 0
+	r.queryLock = sync.Mutex{}
 }
 
 func (r *resolver) SetExtServers(dns []string) {
@@ -328,7 +336,8 @@ func (r *resolver) ServeDNS(w dns.ResponseWriter, query *dns.Msg) {
 			if extDNS.ipStr == "" {
 				break
 			}
-			log.Debugf("Querying ext dns %s:%s for %s[%d]", proto, extDNS.ipStr, name, query.Question[0].Qtype)
+			log.Debugf("Query %s[%d] from %s, forwarding to %s:%s", name, query.Question[0].Qtype,
+				w.LocalAddr().String(), proto, extDNS.ipStr)
 
 			extConnect := func() {
 				addr := fmt.Sprintf("%s:%d", extDNS.ipStr, 53)
@@ -366,18 +375,29 @@ func (r *resolver) ServeDNS(w dns.ResponseWriter, query *dns.Msg) {
 			extConn.SetDeadline(time.Now().Add(extIOTimeout))
 			co := &dns.Conn{Conn: extConn}
 
+			if r.concurrentQueryInc() == false {
+				old := r.tStamp
+				r.tStamp = time.Now()
+				if r.tStamp.Sub(old) > logInterval {
+					log.Errorf("More than %v concurrent queries from %s", maxConcurrent, w.LocalAddr().String())
+				}
+				continue
+			}
+
 			defer func() {
 				if proto == "tcp" {
 					co.Close()
 				}
 			}()
 			err = co.WriteMsg(query)
 			if err != nil {
+				r.concurrentQueryDec()
 				log.Debugf("Send to DNS server failed, %s", err)
 				continue
 			}
 
 			resp, err = co.ReadMsg()
+			r.concurrentQueryDec()
 			if err != nil {
 				log.Debugf("Read from DNS server failed, %s", err)
 				continue
@@ -397,3 +417,23 @@ func (r *resolver) ServeDNS(w dns.ResponseWriter, query *dns.Msg) {
 		log.Errorf("error writing resolver resp, %s", err)
 	}
 }
+
+func (r *resolver) concurrentQueryInc() bool {
+	r.queryLock.Lock()
+	defer r.queryLock.Unlock()
+	if r.count == maxConcurrent {
+		return false
+	}
+	r.count++
+	return true
+}
+
+func (r *resolver) concurrentQueryDec() bool {
+	r.queryLock.Lock()
+	defer r.queryLock.Unlock()
+	if r.count == 0 {
+		return false
+	}
+	r.count--
+	return true
+}