Merge pull request #102 from Snawoot/srv_select

Automatic server selection feature

Author: Snawoot

README.md

@@ -109,6 +109,10 @@ eu3.sec-tunnel.com,77.111.244.22,443
 | proxy | String | sets base proxy to use for all dial-outs. Format: `<http\|https\|socks5\|socks5h>://[login:password@]host[:port]` Examples: `http://user:[email protected]:3128`, `socks5://10.0.0.1:1080` |
 | refresh | Duration | login refresh interval (default 4h0m0s) |
 | refresh-retry | Duration | login refresh retry interval (default 5s) |
+| server-selection | Enum | server selection policy (first/random/fastest) (default fastest) |
+| server-selection-dl-limit | Number | restrict amount of downloaded data per connection by fastest server selection |
+| server-selection-test-url | String | URL used for download benchmark by fastest server selection policy (default `https://ajax.googleapis.com/ajax/libs/angularjs/1.8.2/angular.min.js`) |
+| server-selection-timeout | Duration | timeout given for server selection function to produce result (default 30s) |
 | timeout | Duration | timeout for network operations (default 10s) |
 | verbosity | Number | logging verbosity (10 - debug, 20 - info, 30 - warning, 40 - error, 50 - critical) (default 20) |
 | version | - | show program version and exit |

dialer/selection.go

@@ -0,0 +1,131 @@
+package dialer
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"math/rand/v2"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/hashicorp/go-multierror"
+)
+
+type ServerSelection int
+
+const (
+	_ = iota
+	ServerSelectionFirst
+	ServerSelectionRandom
+	ServerSelectionFastest
+)
+
+func (ss ServerSelection) String() string {
+	switch ss {
+	case ServerSelectionFirst:
+		return "first"
+	case ServerSelectionRandom:
+		return "random"
+	case ServerSelectionFastest:
+		return "fastest"
+	default:
+		return fmt.Sprintf("ServerSelection(%d)", int(ss))
+	}
+}
+
+func ParseServerSelection(s string) (ServerSelection, error) {
+	switch strings.ToLower(s) {
+	case "first":
+		return ServerSelectionFirst, nil
+	case "random":
+		return ServerSelectionRandom, nil
+	case "fastest":
+		return ServerSelectionFastest, nil
+	}
+	return 0, errors.New("unknown server selection strategy")
+}
+
+type SelectionFunc = func(ctx context.Context, dialers []ContextDialer) (ContextDialer, error)
+
+func SelectFirst(_ context.Context, dialers []ContextDialer) (ContextDialer, error) {
+	if len(dialers) == 0 {
+		return nil, errors.New("empty dialers list")
+	}
+	return dialers[0], nil
+}
+
+func SelectRandom(_ context.Context, dialers []ContextDialer) (ContextDialer, error) {
+	if len(dialers) == 0 {
+		return nil, errors.New("empty dialers list")
+	}
+	return dialers[rand.IntN(len(dialers))], nil
+}
+
+func probeDialer(ctx context.Context, dialer ContextDialer, url string, dlLimit int64) error {
+	httpClient := http.Client{
+		Transport: &http.Transport{
+			MaxIdleConns:          100,
+			IdleConnTimeout:       90 * time.Second,
+			TLSHandshakeTimeout:   10 * time.Second,
+			ExpectContinueTimeout: 1 * time.Second,
+			DialContext:           dialer.DialContext,
+		},
+	}
+	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
+	if err != nil {
+		return err
+	}
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("bad status code %d for URL %q", resp.StatusCode, url)
+	}
+	var rd io.Reader = resp.Body
+	if dlLimit > 0 {
+		rd = io.LimitReader(rd, dlLimit)
+	}
+	_, err = io.Copy(io.Discard, rd)
+	return err
+}
+
+func NewFastestServerSelectionFunc(url string, dlLimit int64) SelectionFunc {
+	return func(ctx context.Context, dialers []ContextDialer) (ContextDialer, error) {
+		var resErr error
+		masterNotInterested := make(chan struct{})
+		defer close(masterNotInterested)
+		errors := make(chan error)
+		success := make(chan ContextDialer)
+		for _, dialer := range dialers {
+			go func(dialer ContextDialer) {
+				err := probeDialer(ctx, dialer, url, dlLimit)
+				if err == nil {
+					select {
+					case success <- dialer:
+					case <-masterNotInterested:
+					}
+				} else {
+					select {
+					case errors <- err:
+					case <-masterNotInterested:
+					}
+				}
+			}(dialer)
+		}
+		for _ = range dialers {
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			case d := <-success:
+				return d, nil
+			case err := <-errors:
+				resErr = multierror.Append(resErr, err)
+			}
+		}
+		return nil, resErr
+	}
+}

dialer/upstream.go

@@ -221,6 +221,10 @@ func (d *ProxyDialer) Dial(network, address string) (net.Conn, error) {
 	return d.DialContext(context.Background(), network, address)
 }
 
+func (d *ProxyDialer) Address() (string, error) {
+	return d.address()
+}
+
 func readResponse(r io.Reader, req *http.Request) (*http.Response, error) {
 	endOfResponse := []byte("\r\n\r\n")
 	buf := &bytes.Buffer{}