|
|
|
@ -3,14 +3,47 @@ package loadbalancer
|
|
|
|
|
import ( |
|
|
|
|
"context" |
|
|
|
|
"errors" |
|
|
|
|
"fmt" |
|
|
|
|
"math/rand" |
|
|
|
|
"net" |
|
|
|
|
"net/http" |
|
|
|
|
"net/url" |
|
|
|
|
"os" |
|
|
|
|
"strconv" |
|
|
|
|
|
|
|
|
|
"github.com/k3s-io/k3s/pkg/version" |
|
|
|
|
http_dialer "github.com/mwitkow/go-http-dialer" |
|
|
|
|
"golang.org/x/net/proxy" |
|
|
|
|
|
|
|
|
|
"github.com/sirupsen/logrus" |
|
|
|
|
"k8s.io/apimachinery/pkg/util/sets" |
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
var defaultDialer = &net.Dialer{} |
|
|
|
|
var defaultDialer proxy.Dialer = &net.Dialer{} |
|
|
|
|
|
|
|
|
|
func init() { |
|
|
|
|
// Check if env variable for proxy is set
|
|
|
|
|
address := os.Getenv(version.ProgramUpper + "_URL") |
|
|
|
|
|
|
|
|
|
if useProxy, _ := strconv.ParseBool(os.Getenv(version.ProgramUpper + "_AGENT_HTTP_PROXY_ALLOWED")); !useProxy { |
|
|
|
|
return |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
req, err := http.NewRequest("GET", "https://"+address, nil) |
|
|
|
|
if err != nil { |
|
|
|
|
logrus.Errorf("Error creating request for address %s: %v", address, err) |
|
|
|
|
} |
|
|
|
|
proxyURL, err := http.ProxyFromEnvironment(req) |
|
|
|
|
if err != nil { |
|
|
|
|
logrus.Errorf("Error getting the proxy for address %s: %v", address, err) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
if dialer, err := proxyDialer(proxyURL); err != nil { |
|
|
|
|
logrus.Errorf("Error creating the proxyDialer for %s: %v", address, err) |
|
|
|
|
} else { |
|
|
|
|
defaultDialer = dialer |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (lb *LoadBalancer) setServers(serverAddresses []string) bool { |
|
|
|
|
serverAddresses, hasOriginalServer := sortServers(serverAddresses, lb.defaultServerAddress) |
|
|
|
@ -84,20 +117,33 @@ func (lb *LoadBalancer) nextServer(failedServer string) (string, error) {
|
|
|
|
|
return lb.currentServerAddress, nil |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// dialContext dials a new connection, and adds its wrapped connection to the map
|
|
|
|
|
// dialContext dials a new connection using the environment's proxy settings, and adds its wrapped connection to the map
|
|
|
|
|
func (s *server) dialContext(ctx context.Context, network, address string) (net.Conn, error) { |
|
|
|
|
conn, err := defaultDialer.DialContext(ctx, network, address) |
|
|
|
|
conn, err := defaultDialer.Dial(network, address) |
|
|
|
|
if err != nil { |
|
|
|
|
return nil, err |
|
|
|
|
} |
|
|
|
|
// don't lock until adding the connection to the map, otherwise we may block
|
|
|
|
|
// while waiting for the dial to time out
|
|
|
|
|
|
|
|
|
|
// Wrap the connection and add it to the server's connection map
|
|
|
|
|
s.mutex.Lock() |
|
|
|
|
defer s.mutex.Unlock() |
|
|
|
|
|
|
|
|
|
conn = &serverConn{server: s, Conn: conn} |
|
|
|
|
s.connections[conn] = struct{}{} |
|
|
|
|
return conn, nil |
|
|
|
|
wrappedConn := &serverConn{server: s, Conn: conn} |
|
|
|
|
s.connections[wrappedConn] = struct{}{} |
|
|
|
|
return wrappedConn, nil |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// proxyDialer creates a new proxy.Dialer that routes connections through the specified proxy.
|
|
|
|
|
func proxyDialer(proxyURL *url.URL) (proxy.Dialer, error) { |
|
|
|
|
if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" { |
|
|
|
|
// Create a new HTTP proxy dialer
|
|
|
|
|
httpProxyDialer := http_dialer.New(proxyURL) |
|
|
|
|
return httpProxyDialer, nil |
|
|
|
|
} else if proxyURL.Scheme == "socks5" { |
|
|
|
|
// For SOCKS5 proxies, use the proxy package's FromURL
|
|
|
|
|
return proxy.FromURL(proxyURL, proxy.Direct) |
|
|
|
|
} |
|
|
|
|
return nil, fmt.Errorf("unsupported proxy scheme: %s", proxyURL.Scheme) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// closeAll closes all connections to the server, and removes their entries from the map
|
|
|
|
|