From a0bd8eb8a1880cb4980192ba4bf88f2dfbf48ebc Mon Sep 17 00:00:00 2001
From: Star Brilliant <coder@poorlab.com>
Date: Sat, 25 Nov 2017 15:46:29 +0800
Subject: [PATCH] Allow multiple bootstrap servers

---
 doh-client/client.go       | 24 +++++++++++++++---------
 doh-client/main.go         |  7 ++++++-
 systemd/doh-client.service |  2 +-
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/doh-client/client.go b/doh-client/client.go
index d942d1e..4d76cc9 100644
--- a/doh-client/client.go
+++ b/doh-client/client.go
@@ -22,6 +22,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"math/rand"
 	"io/ioutil"
 	"log"
 	"net"
@@ -37,7 +38,7 @@ import (
 type Client struct {
 	addr			string
 	upstream		string
-	bootstrap		string
+	bootstraps		[]string
 	timeout			uint
 	noECS			bool
 	verbose			bool
@@ -46,10 +47,11 @@ type Client struct {
 	httpClient		*http.Client
 }
 
-func NewClient(addr, upstream, bootstrap string, timeout uint, noECS, verbose bool) (c *Client, err error) {
+func NewClient(addr, upstream string, bootstraps []string, timeout uint, noECS, verbose bool) (c *Client, err error) {
 	c = &Client {
 		addr: addr,
 		upstream: upstream,
+		bootstraps: bootstraps,
 		timeout: timeout,
 		noECS: noECS,
 		verbose: verbose,
@@ -66,18 +68,22 @@ func NewClient(addr, upstream, bootstrap string, timeout uint, noECS, verbose bo
 		Handler: dns.HandlerFunc(c.tcpHandlerFunc),
 	}
 	bootResolver := net.DefaultResolver
-	if bootstrap != "" {
-		bootstrapAddr, err := net.ResolveUDPAddr("udp", bootstrap)
-		if err != nil {
-			bootstrapAddr, err = net.ResolveUDPAddr("udp", "[" + bootstrap + "]:53")
+	if len(c.bootstraps) != 0 {
+		for i, bootstrap := range c.bootstraps {
+			bootstrapAddr, err := net.ResolveUDPAddr("udp", bootstrap)
+			if err != nil {
+				bootstrapAddr, err = net.ResolveUDPAddr("udp", "[" + bootstrap + "]:53")
+			}
+			if err != nil { return nil, err }
+			c.bootstraps[i] = bootstrapAddr.String()
 		}
-		if err != nil { return nil, err }
-		c.bootstrap = bootstrapAddr.String()
 		bootResolver = &net.Resolver {
 			PreferGo: true,
 			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
 				var d net.Dialer
-				conn, err := d.DialContext(ctx, network, c.bootstrap)
+				num_servers := len(c.bootstraps)
+				bootstrap := c.bootstraps[rand.Intn(num_servers)]
+				conn, err := d.DialContext(ctx, network, bootstrap)
 				return conn, err
 			},
 		}
diff --git a/doh-client/main.go b/doh-client/main.go
index d37299b..967b961 100644
--- a/doh-client/main.go
+++ b/doh-client/main.go
@@ -21,6 +21,7 @@ package main
 import (
 	"flag"
 	"log"
+	"strings"
 )
 
 func main() {
@@ -32,7 +33,11 @@ func main() {
 	verbose := flag.Bool("verbose", false, "Enable logging")
 	flag.Parse()
 
-	client, err := NewClient(*addr, *upstream, *bootstrap, *timeout, *noECS, *verbose)
+	bootstraps := []string {}
+	if *bootstrap != "" {
+		bootstraps = strings.Split(*bootstrap, ",")
+	}
+	client, err := NewClient(*addr, *upstream, bootstraps, *timeout, *noECS, *verbose)
 	if err != nil { log.Fatalln(err) }
 	_ = client.Start()
 }
diff --git a/systemd/doh-client.service b/systemd/doh-client.service
index fba0626..eb3cf2f 100644
--- a/systemd/doh-client.service
+++ b/systemd/doh-client.service
@@ -4,7 +4,7 @@ After=network.target
 
 [Service]
 AmbientCapabilities=CAP_NET_BIND_SERVICE
-ExecStart=/usr/local/bin/doh-client -addr 127.0.0.1:53 -upstream https://dns.google.com/resolve -bootstrap 8.8.8.8:53
+ExecStart=/usr/local/bin/doh-client -addr 127.0.0.1:53 -upstream https://dns.google.com/resolve -bootstrap 8.8.8.8:53,8.8.4.4:53
 LimitNOFILE=1048576
 Restart=always
 RestartSec=3