From b57cc19d70e8347141962079409f0f1a7390f782 Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Fri, 20 Mar 2020 17:37:34 +0100
Subject: [PATCH 1/2] Use an interface for load-balancing strategies

---
 dnscrypt-proxy/config.go      | 11 ++++----
 dnscrypt-proxy/serversInfo.go | 48 ++++++++++++++++++++---------------
 2 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/dnscrypt-proxy/config.go b/dnscrypt-proxy/config.go
index 5c8abbea..d01c0aa0 100644
--- a/dnscrypt-proxy/config.go
+++ b/dnscrypt-proxy/config.go
@@ -358,20 +358,19 @@ func ConfigLoad(proxy *Proxy, flags *ConfigFlags) error {
 	if len(config.ListenAddresses) == 0 && len(config.LocalDoH.ListenAddresses) == 0 {
 		dlog.Debug("No local IP/port configured")
 	}
-
-	lbStrategy := DefaultLBStrategy
+	lbStrategy := LBStrategy(DefaultLBStrategy)
 	switch strings.ToLower(config.LBStrategy) {
 	case "":
 		// default
 	case "p2":
-		lbStrategy = LBStrategyP2
+		lbStrategy = LBStrategyP2{}
 	case "ph":
-		lbStrategy = LBStrategyPH
+		lbStrategy = LBStrategyPH{}
 	case "fastest":
 	case "first":
-		lbStrategy = LBStrategyFirst
+		lbStrategy = LBStrategyFirst{}
 	case "random":
-		lbStrategy = LBStrategyRandom
+		lbStrategy = LBStrategyRandom{}
 	default:
 		dlog.Warnf("Unknown load balancing strategy: [%s]", config.LBStrategy)
 	}
diff --git a/dnscrypt-proxy/serversInfo.go b/dnscrypt-proxy/serversInfo.go
index 6327f7c0..e6de3028 100644
--- a/dnscrypt-proxy/serversInfo.go
+++ b/dnscrypt-proxy/serversInfo.go
@@ -62,17 +62,35 @@ type ServerInfo struct {
 	DOHClientCreds     DOHClientCreds
 }
 
-type LBStrategy int
+type LBStrategy interface {
+	getCandidate(serversCount int) int
+}
 
-const (
-	LBStrategyNone = LBStrategy(iota)
-	LBStrategyP2
-	LBStrategyPH
-	LBStrategyFirst
-	LBStrategyRandom
-)
+type LBStrategyP2 struct{}
 
-const DefaultLBStrategy = LBStrategyP2
+func (LBStrategyP2) getCandidate(serversCount int) int {
+	return rand.Intn(Min(serversCount, 2))
+}
+
+type LBStrategyPH struct{}
+
+func (LBStrategyPH) getCandidate(serversCount int) int {
+	return rand.Intn(Max(Min(serversCount, 2), serversCount/2))
+}
+
+type LBStrategyFirst struct{}
+
+func (LBStrategyFirst) getCandidate(int) int {
+	return 0
+}
+
+type LBStrategyRandom struct{}
+
+func (LBStrategyRandom) getCandidate(serversCount int) int {
+	return rand.Intn(serversCount)
+}
+
+var DefaultLBStrategy = LBStrategyP2{}
 
 type ServersInfo struct {
 	sync.RWMutex
@@ -209,17 +227,7 @@ func (serversInfo *ServersInfo) getOne() *ServerInfo {
 	if serversInfo.lbEstimator {
 		serversInfo.estimatorUpdate()
 	}
-	var candidate int
-	switch serversInfo.lbStrategy {
-	case LBStrategyFirst:
-		candidate = 0
-	case LBStrategyPH:
-		candidate = rand.Intn(Max(Min(serversCount, 2), serversCount/2))
-	case LBStrategyRandom:
-		candidate = rand.Intn(serversCount)
-	default:
-		candidate = rand.Intn(Min(serversCount, 2))
-	}
+	candidate := serversInfo.lbStrategy.getCandidate(serversCount)
 	serverInfo := serversInfo.inner[candidate]
 	dlog.Debugf("Using candidate [%s] RTT: %d", (*serverInfo).Name, int((*serverInfo).rtt.Value()))
 	serversInfo.Unlock()

From 34d83f027f16754b6524899fb809695864123dbf Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Fri, 20 Mar 2020 17:48:54 +0100
Subject: [PATCH 2/2] Support power-of-<arbitrary number>

---
 dnscrypt-proxy/config.go      | 12 ++++++++++--
 dnscrypt-proxy/serversInfo.go |  6 ++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/dnscrypt-proxy/config.go b/dnscrypt-proxy/config.go
index d01c0aa0..bc959dd2 100644
--- a/dnscrypt-proxy/config.go
+++ b/dnscrypt-proxy/config.go
@@ -359,7 +359,7 @@ func ConfigLoad(proxy *Proxy, flags *ConfigFlags) error {
 		dlog.Debug("No local IP/port configured")
 	}
 	lbStrategy := LBStrategy(DefaultLBStrategy)
-	switch strings.ToLower(config.LBStrategy) {
+	switch lbStrategyStr := strings.ToLower(config.LBStrategy); lbStrategyStr {
 	case "":
 		// default
 	case "p2":
@@ -372,7 +372,15 @@ func ConfigLoad(proxy *Proxy, flags *ConfigFlags) error {
 	case "random":
 		lbStrategy = LBStrategyRandom{}
 	default:
-		dlog.Warnf("Unknown load balancing strategy: [%s]", config.LBStrategy)
+		if strings.HasPrefix(lbStrategyStr, "p") {
+			n, err := strconv.ParseInt(strings.TrimPrefix(lbStrategyStr, "p"), 10, 32)
+			if err != nil || n <= 0 {
+				dlog.Fatalf("Invalid load balancing strategy: [%s]", config.LBStrategy)
+			}
+			lbStrategy = LBStrategyPN{n: int(n)}
+		} else {
+			dlog.Warnf("Unknown load balancing strategy: [%s]", config.LBStrategy)
+		}
 	}
 	proxy.serversInfo.lbStrategy = lbStrategy
 	proxy.serversInfo.lbEstimator = config.LBEstimator
diff --git a/dnscrypt-proxy/serversInfo.go b/dnscrypt-proxy/serversInfo.go
index e6de3028..66fb48ff 100644
--- a/dnscrypt-proxy/serversInfo.go
+++ b/dnscrypt-proxy/serversInfo.go
@@ -72,6 +72,12 @@ func (LBStrategyP2) getCandidate(serversCount int) int {
 	return rand.Intn(Min(serversCount, 2))
 }
 
+type LBStrategyPN struct{ n int }
+
+func (s LBStrategyPN) getCandidate(serversCount int) int {
+	return rand.Intn(Min(serversCount, s.n))
+}
+
 type LBStrategyPH struct{}
 
 func (LBStrategyPH) getCandidate(serversCount int) int {