|
|
|
|
@ -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
|
|
|
|
|
|
Pierre
11 months ago
committed by