From 732c451dd4ff13303296d7df86896714e0fc53dd Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Wed, 24 Jan 2018 16:51:26 +0100
Subject: [PATCH] Add max_clients to cap the maximum number of client queries

---
 dnscrypt-proxy/config.go           |  3 +++
 dnscrypt-proxy/dnscrypt-proxy.toml |  5 +++++
 dnscrypt-proxy/main.go             | 33 ++++++++++++++++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/dnscrypt-proxy/config.go b/dnscrypt-proxy/config.go
index fe89ce69..adbc9a42 100644
--- a/dnscrypt-proxy/config.go
+++ b/dnscrypt-proxy/config.go
@@ -40,6 +40,7 @@ type Config struct {
 	SourceRequireNoLog  bool                    `toml:"require_nolog"`
 	SourceIPv4          bool                    `toml:"ipv4_servers"`
 	SourceIPv6          bool                    `toml:"ipv6_servers"`
+	MaxClients          uint32                  `toml:"max_clients"`
 }
 
 func newConfig() Config {
@@ -57,6 +58,7 @@ func newConfig() Config {
 		SourceRequireNoLog:  true,
 		SourceIPv4:          true,
 		SourceIPv6:          false,
+		MaxClients:          100,
 	}
 }
 
@@ -128,6 +130,7 @@ func ConfigLoad(proxy *Proxy, svcFlag *string, config_file string) error {
 		dlog.UseLogFile(*config.LogFile)
 	}
 	proxy.timeout = time.Duration(config.Timeout) * time.Millisecond
+	proxy.maxClients = config.MaxClients
 	proxy.mainProto = "udp"
 	if config.ForceTCP {
 		proxy.mainProto = "tcp"
diff --git a/dnscrypt-proxy/dnscrypt-proxy.toml b/dnscrypt-proxy/dnscrypt-proxy.toml
index d764ca9d..f7c1a924 100644
--- a/dnscrypt-proxy/dnscrypt-proxy.toml
+++ b/dnscrypt-proxy/dnscrypt-proxy.toml
@@ -23,6 +23,11 @@
 listen_addresses = ['127.0.0.1:53', '[::1]:53']
 
 
+## Maximum number of simultaneous client connections to accept
+
+max_clients = 100
+
+
 ## Require servers (from static definitions and from remote sources) to satisfy specific properties
 
 # Use servers reachable over IPv4
diff --git a/dnscrypt-proxy/main.go b/dnscrypt-proxy/main.go
index d741839b..d4a4fcb1 100644
--- a/dnscrypt-proxy/main.go
+++ b/dnscrypt-proxy/main.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/jedisct1/dlog"
@@ -50,6 +51,8 @@ type Proxy struct {
 	forwardFile                  string
 	pluginsGlobals               PluginsGlobals
 	urlsToPrefetch               []URLToPrefetch
+	clientsCount                 uint32
+	maxClients                   uint32
 }
 
 type App struct {
@@ -218,6 +221,11 @@ func (proxy *Proxy) udpListener(clientPc *net.UDPConn) {
 		}
 		packet := buffer[:length]
 		go func() {
+			if !proxy.clientsCountInc() {
+				dlog.Warnf("Too many connections (max=%d)", proxy.maxClients)
+				return
+			}
+			defer proxy.clientsCountDec()
 			proxy.processIncomingQuery(proxy.serversInfo.getOne(), "udp", proxy.mainProto, packet, &clientAddr, clientPc)
 		}()
 	}
@@ -242,6 +250,11 @@ func (proxy *Proxy) tcpListener(acceptPc *net.TCPListener) {
 		}
 		go func() {
 			defer clientPc.Close()
+			if !proxy.clientsCountInc() {
+				dlog.Warnf("Too many connections (max=%d)", proxy.maxClients)
+				return
+			}
+			defer proxy.clientsCountDec()
 			clientPc.SetDeadline(time.Now().Add(proxy.timeout))
 			packet, err := ReadPrefixed(clientPc.(*net.TCPConn))
 			if err != nil || len(packet) < MinDNSPacketSize {
@@ -300,6 +313,26 @@ func (proxy *Proxy) exchangeWithTCPServer(serverInfo *ServerInfo, encryptedQuery
 	return proxy.Decrypt(serverInfo, encryptedResponse, clientNonce)
 }
 
+func (proxy *Proxy) clientsCountInc() bool {
+	for {
+		count := proxy.clientsCount
+		if count >= proxy.maxClients {
+			return false
+		}
+		if atomic.CompareAndSwapUint32(&proxy.clientsCount, count, count+1) {
+			return true
+		}
+	}
+}
+
+func (proxy *Proxy) clientsCountDec() {
+	for {
+		if count := proxy.clientsCount; count == 0 || atomic.CompareAndSwapUint32(&proxy.clientsCount, count, count-1) {
+			break
+		}
+	}
+}
+
 func (proxy *Proxy) processIncomingQuery(serverInfo *ServerInfo, clientProto string, serverProto string, query []byte, clientAddr *net.Addr, clientPc net.Conn) {
 	if len(query) < MinDNSPacketSize || serverInfo == nil {
 		return