From 6f78bd1ff07b50157c6094b527a0519563a05a05 Mon Sep 17 00:00:00 2001
From: Star Brilliant <m13253@hotmail.com>
Date: Thu, 16 Nov 2017 15:07:23 +0800
Subject: [PATCH] Add bootstrap server option

---
 Readme.md                  |  6 -----
 doh-client/client.go       | 46 ++++++++++++++++++++++++++++++++------
 doh-client/main.go         |  6 ++++-
 systemd/doh-client.service |  2 +-
 4 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/Readme.md b/Readme.md
index f098e66..188e6f4 100644
--- a/Readme.md
+++ b/Readme.md
@@ -27,12 +27,6 @@ default settings, type:
     sudo cp /usr/lib/systemd/system/doh-client.service /etc/systemd/system/
     sudoedit /etc/systemd/system/doh-client.service
 
-Perform a DNS lookup to `dns.google.com`, and store the result in `/etc/hosts`.
-This is used for bootstrap.
-
-    dig +short dns.google.com
-    echo 'XXX.XXX.XXX.XXX dns.google.com' | sudo tee -a /etc/hosts
-
 To automatically start DNS-over-HTTPS client as a system service, type:
 
     sudo systemctl start doh-client.service
diff --git a/doh-client/client.go b/doh-client/client.go
index 1eb30a7..5f95a2c 100644
--- a/doh-client/client.go
+++ b/doh-client/client.go
@@ -19,6 +19,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -33,16 +34,20 @@ import (
 )
 
 type Client struct {
-	addr		string
-	upstream	string
-	udpServer	*dns.Server
-	tcpServer	*dns.Server
+	addr			string
+	upstream		string
+	bootstrap		string
+	timeout			uint
+	udpServer		*dns.Server
+	tcpServer		*dns.Server
+	httpClient		*http.Client
 }
 
-func NewClient(addr, upstream string) (c *Client) {
+func NewClient(addr, upstream, bootstrap string, timeout uint) (c *Client, err error) {
 	c = &Client {
 		addr: addr,
 		upstream: upstream,
+		timeout: timeout,
 	}
 	c.udpServer = &dns.Server {
 		Addr: addr,
@@ -55,7 +60,34 @@ func NewClient(addr, upstream string) (c *Client) {
 		Net: "tcp",
 		Handler: dns.HandlerFunc(c.tcpHandlerFunc),
 	}
-	return
+	bootResolver := net.DefaultResolver
+	if bootstrap != "" {
+		bootstrapAddr, err := net.ResolveUDPAddr("udp", bootstrap)
+		if err != nil {
+			bootstrapAddr, err = net.ResolveUDPAddr("udp", "[" + bootstrap + "]:53")
+		}
+		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)
+				return conn, err
+			},
+		}
+	}
+	httpTransport := *http.DefaultTransport.(*http.Transport)
+	httpTransport.DialContext = (&net.Dialer {
+		Timeout: time.Duration(c.timeout) * time.Second,
+		KeepAlive: 30 * time.Second,
+		DualStack: true,
+		Resolver: bootResolver,
+	}).DialContext
+	c.httpClient = &http.Client {
+		Transport: &httpTransport,
+	}
+	return c, nil
 }
 
 func (c *Client) Start() error {
@@ -131,7 +163,7 @@ func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
 		return
 	}
 	req.Header.Set("User-Agent", "DNS-over-HTTPS/1.0 (+https://github.com/m13253/dns-over-https)")
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := c.httpClient.Do(req)
 	if err != nil {
 		log.Println(err)
 		reply.Rcode = dns.RcodeServerFailure
diff --git a/doh-client/main.go b/doh-client/main.go
index 2271c1e..1242572 100644
--- a/doh-client/main.go
+++ b/doh-client/main.go
@@ -20,13 +20,17 @@ package main
 
 import (
 	"flag"
+	"log"
 )
 
 func main() {
 	addr := flag.String("addr", "127.0.0.1:53", "DNS listen port")
 	upstream := flag.String("upstream", "https://dns.google.com/resolve", "HTTP path for upstream resolver")
+	bootstrap := flag.String("bootstrap", "", "The bootstrap DNS server to resolve the address of the upstream resolver")
+	timeout := flag.Uint("timeout", 10, "Timeout for upstream request")
 	flag.Parse()
 
-	client := NewClient(*addr, *upstream)
+	client, err := NewClient(*addr, *upstream, *bootstrap, *timeout)
+	if err != nil { log.Fatalln(err) }
 	_ = client.Start()
 }
diff --git a/systemd/doh-client.service b/systemd/doh-client.service
index 1bb4f6a..fba0626 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
+ExecStart=/usr/local/bin/doh-client -addr 127.0.0.1:53 -upstream https://dns.google.com/resolve -bootstrap 8.8.8.8:53
 LimitNOFILE=1048576
 Restart=always
 RestartSec=3