Support power-of-<arbitrary number>

DNSCrypt · Mar 20, 2020 · 34d83f0 · 34d83f0
1 parent b57cc19
commit 34d83f0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git 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
@@ -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 {