Allow multiple bootstrap servers

This commit is contained in:
Star Brilliant
2017-11-25 15:46:29 +08:00
parent 2577b720c7
commit a0bd8eb8a1
3 changed files with 22 additions and 11 deletions

View File

@@ -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
},
}

View File

@@ -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()
}