Add backend weight round robin select

.gitignore

@@ -12,3 +12,5 @@
 
 # Project-local glide cache, RE: https://github.com/Masterminds/glide/issues/736
 .glide/
+
+.idea/

doh-client/client.go

@@ -36,14 +36,16 @@ import (
 	"sync"
 	"time"
 
+	"github.com/m13253/dns-over-https/doh-client/config"
+	"github.com/m13253/dns-over-https/doh-client/selector"
 	"github.com/m13253/dns-over-https/json-dns"
 	"github.com/miekg/dns"
 	"golang.org/x/net/http2"
 	"golang.org/x/net/idna"
 )
 
 type Client struct {
-	conf                 *config
+	conf                 *config.Config
 	bootstrap            []string
 	passthrough          []string
 	udpClient            *dns.Client
@@ -56,6 +58,7 @@ type Client struct {
 	httpTransport        *http.Transport
 	httpClient           *http.Client
 	httpClientLastCreate time.Time
+	Selector             selector.Selector
 }
 
 type DNSRequest struct {
@@ -68,7 +71,7 @@ type DNSRequest struct {
 	err               error
 }
 
-func NewClient(conf *config) (c *Client, err error) {
+func NewClient(conf *config.Config) (c *Client, err error) {
 	c = &Client{
 		conf: conf,
 	}
@@ -78,11 +81,11 @@ func NewClient(conf *config) (c *Client, err error) {
 	c.udpClient = &dns.Client{
 		Net:     "udp",
 		UDPSize: dns.DefaultMsgSize,
-		Timeout: time.Duration(conf.Timeout) * time.Second,
+		Timeout: time.Duration(conf.Other.Timeout) * time.Second,
 	}
 	c.tcpClient = &dns.Client{
 		Net:     "tcp",
-		Timeout: time.Duration(conf.Timeout) * time.Second,
+		Timeout: time.Duration(conf.Other.Timeout) * time.Second,
 	}
 	for _, addr := range conf.Listen {
 		c.udpServers = append(c.udpServers, &dns.Server{
@@ -98,9 +101,9 @@ func NewClient(conf *config) (c *Client, err error) {
 		})
 	}
 	c.bootstrapResolver = net.DefaultResolver
-	if len(conf.Bootstrap) != 0 {
-		c.bootstrap = make([]string, len(conf.Bootstrap))
-		for i, bootstrap := range conf.Bootstrap {
+	if len(conf.Other.Bootstrap) != 0 {
+		c.bootstrap = make([]string, len(conf.Other.Bootstrap))
+		for i, bootstrap := range conf.Other.Bootstrap {
 			bootstrapAddr, err := net.ResolveUDPAddr("udp", bootstrap)
 			if err != nil {
 				bootstrapAddr, err = net.ResolveUDPAddr("udp", "["+bootstrap+"]:53")
@@ -120,9 +123,9 @@ func NewClient(conf *config) (c *Client, err error) {
 				return conn, err
 			},
 		}
-		if len(conf.Passthrough) != 0 {
-			c.passthrough = make([]string, len(conf.Passthrough))
-			for i, passthrough := range conf.Passthrough {
+		if len(conf.Other.Passthrough) != 0 {
+			c.passthrough = make([]string, len(conf.Other.Passthrough))
+			for i, passthrough := range conf.Other.Passthrough {
 				if punycode, err := idna.ToASCII(passthrough); err != nil {
 					passthrough = punycode
 				}
@@ -133,7 +136,7 @@ func NewClient(conf *config) (c *Client, err error) {
 	// Most CDNs require Cookie support to prevent DDoS attack.
 	// Disabling Cookie does not effectively prevent tracking,
 	// so I will leave it on to make anti-DDoS services happy.
-	if !c.conf.NoCookies {
+	if !c.conf.Other.NoCookies {
 		c.cookieJar, err = cookiejar.New(nil)
 		if err != nil {
 			return nil, err
@@ -147,23 +150,59 @@ func NewClient(conf *config) (c *Client, err error) {
 	if err != nil {
 		return nil, err
 	}
+
+	switch c.conf.Upstream.UpstreamSelector {
+	default:
+		// if selector is invalid or random, use random selector, or should we stop program and let user knows he is wrong?
+		s := selector.NewRandomSelector()
+		for _, u := range c.conf.Upstream.UpstreamGoogle {
+			if err := s.Add(u.Url, selector.Google); err != nil {
+				return nil, err
+			}
+		}
+
+		for _, u := range c.conf.Upstream.UpstreamIETF {
+			if err := s.Add(u.Url, selector.IETF); err != nil {
+				return nil, err
+			}
+		}
+
+		c.Selector = s
+
+	case config.WeightedRoundRobin:
+		s := selector.NewWeightRoundRobinSelector(time.Duration(c.conf.Other.Timeout) * time.Second)
+		for _, u := range c.conf.Upstream.UpstreamGoogle {
+			if err := s.Add(u.Url, selector.Google, u.Weight); err != nil {
+				return nil, err
+			}
+		}
+
+		for _, u := range c.conf.Upstream.UpstreamIETF {
+			if err := s.Add(u.Url, selector.IETF, u.Weight); err != nil {
+				return nil, err
+			}
+		}
+
+		c.Selector = s
+	}
+
 	return c, nil
 }
 
 func (c *Client) newHTTPClient() error {
 	c.httpClientMux.Lock()
 	defer c.httpClientMux.Unlock()
-	if !c.httpClientLastCreate.IsZero() && time.Since(c.httpClientLastCreate) < time.Duration(c.conf.Timeout)*time.Second {
+	if !c.httpClientLastCreate.IsZero() && time.Since(c.httpClientLastCreate) < time.Duration(c.conf.Other.Timeout)*time.Second {
 		return nil
 	}
 	if c.httpTransport != nil {
 		c.httpTransport.CloseIdleConnections()
 	}
 	dialer := &net.Dialer{
-		Timeout:   time.Duration(c.conf.Timeout) * time.Second,
+		Timeout:   time.Duration(c.conf.Other.Timeout) * time.Second,
 		KeepAlive: 30 * time.Second,
-		DualStack: true,
-		Resolver:  c.bootstrapResolver,
+		// DualStack: true,
+		Resolver: c.bootstrapResolver,
 	}
 	c.httpTransport = &http.Transport{
 		DialContext:           dialer.DialContext,
@@ -172,9 +211,9 @@ func (c *Client) newHTTPClient() error {
 		MaxIdleConns:          100,
 		MaxIdleConnsPerHost:   10,
 		Proxy:                 http.ProxyFromEnvironment,
-		TLSHandshakeTimeout:   time.Duration(c.conf.Timeout) * time.Second,
+		TLSHandshakeTimeout:   time.Duration(c.conf.Other.Timeout) * time.Second,
 	}
-	if c.conf.NoIPv6 {
+	if c.conf.Other.NoIPv6 {
 		c.httpTransport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
 			if strings.HasPrefix(network, "tcp") {
 				network = "tcp4"
@@ -213,11 +252,15 @@ func (c *Client) Start() error {
 		}
 	}
 	close(results)
+
+	// start evaluate goroutine
+	go c.Selector.Evaluate()
+
 	return nil
 }
 
 func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.conf.Timeout)*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.conf.Other.Timeout)*time.Second)
 	defer cancel()
 
 	if r.Response {
@@ -246,7 +289,7 @@ func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
 	} else {
 		questionType = strconv.FormatUint(uint64(question.Qtype), 10)
 	}
-	if c.conf.Verbose {
+	if c.conf.Other.Verbose {
 		fmt.Printf("%s - - [%s] \"%s %s %s\"\n", w.RemoteAddr(), time.Now().Format("02/Jan/2006:15:04:05 -0700"), questionName, questionClass, questionType)
 	}
 
@@ -284,33 +327,25 @@ func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
 		return
 	}
 
-	requestType := ""
-	if len(c.conf.UpstreamIETF) == 0 {
-		requestType = "application/dns-json"
-	} else if len(c.conf.UpstreamGoogle) == 0 {
-		requestType = "application/dns-message"
-	} else {
-		numServers := len(c.conf.UpstreamGoogle) + len(c.conf.UpstreamIETF)
-		random := rand.Intn(numServers)
-		if random < len(c.conf.UpstreamGoogle) {
-			requestType = "application/dns-json"
-		} else {
-			requestType = "application/dns-message"
-		}
+	upstream := c.Selector.Get()
+	requestType := upstream.RequestType
+
+	if c.conf.Other.Verbose {
+		log.Println("choose upstream:", upstream)
 	}
 
 	var req *DNSRequest
 	if requestType == "application/dns-json" {
-		req = c.generateRequestGoogle(ctx, w, r, isTCP)
+		req = c.generateRequestGoogle(ctx, w, r, isTCP, upstream)
 	} else if requestType == "application/dns-message" {
-		req = c.generateRequestIETF(ctx, w, r, isTCP)
+		req = c.generateRequestIETF(ctx, w, r, isTCP, upstream)
 	} else {
 		panic("Unknown request Content-Type")
 	}
 
 	if req.response != nil {
 		defer req.response.Body.Close()
-		for _, header := range c.conf.DebugHTTPHeaders {
+		for _, header := range c.conf.Other.DebugHTTPHeaders {
 			if value := req.response.Header.Get(header); value != "" {
 				log.Printf("%s: %s\n", header, value)
 			}
@@ -360,7 +395,7 @@ var (
 
 func (c *Client) findClientIP(w dns.ResponseWriter, r *dns.Msg) (ednsClientAddress net.IP, ednsClientNetmask uint8) {
 	ednsClientNetmask = 255
-	if c.conf.NoECS {
+	if c.conf.Other.NoECS {
 		return net.IPv4(0, 0, 0, 0), 0
 	}
 	if opt := r.IsEdns0(); opt != nil {

doh-client/config.go → doh-client/config/config.go

@@ -21,18 +21,31 @@
    DEALINGS IN THE SOFTWARE.
 */
 
-package main
+package config
 
 import (
 	"fmt"
 
 	"github.com/BurntSushi/toml"
 )
 
-type config struct {
-	Listen           []string `toml:"listen"`
-	UpstreamGoogle   []string `toml:"upstream_google"`
-	UpstreamIETF     []string `toml:"upstream_ietf"`
+const (
+	Random             = "random"
+	WeightedRoundRobin = "weighted_round_robin"
+)
+
+type upstreamDetail struct {
+	Url    string `toml:"url"`
+	Weight int    `toml:"weight"`
+}
+
+type upstream struct {
+	UpstreamGoogle   []upstreamDetail `toml:"upstream_google"`
+	UpstreamIETF     []upstreamDetail `toml:"upstream_ietf"`
+	UpstreamSelector string           `toml:"upstream_selector"` // usable: random or weighted_random
+}
+
+type others struct {
 	Bootstrap        []string `toml:"bootstrap"`
 	Passthrough      []string `toml:"passthrough"`
 	Timeout          uint     `toml:"timeout"`
@@ -43,8 +56,14 @@ type config struct {
 	DebugHTTPHeaders []string `toml:"debug_http_headers"`
 }
 
-func loadConfig(path string) (*config, error) {
-	conf := &config{}
+type Config struct {
+	Listen   []string `toml:"listen"`
+	Upstream upstream `toml:"upstream"`
+	Other    others   `toml:"others"`
+}
+
+func LoadConfig(path string) (*Config, error) {
+	conf := &Config{}
 	metaData, err := toml.DecodeFile(path, conf)
 	if err != nil {
 		return nil, err
@@ -56,11 +75,15 @@ func loadConfig(path string) (*config, error) {
 	if len(conf.Listen) == 0 {
 		conf.Listen = []string{"127.0.0.1:53", "[::1]:53"}
 	}
-	if len(conf.UpstreamGoogle) == 0 && len(conf.UpstreamIETF) == 0 {
-		conf.UpstreamGoogle = []string{"https://dns.google.com/resolve"}
+	if len(conf.Upstream.UpstreamGoogle) == 0 && len(conf.Upstream.UpstreamIETF) == 0 {
+		conf.Upstream.UpstreamGoogle = []upstreamDetail{{Url: "https://dns.google.com/resolve", Weight: 50}}
 	}
-	if conf.Timeout == 0 {
-		conf.Timeout = 10
+	if conf.Other.Timeout == 0 {
+		conf.Other.Timeout = 10
+	}
+
+	if conf.Upstream.UpstreamSelector == "" {
+		conf.Upstream.UpstreamSelector = Random
 	}
 
 	return conf, nil