From 04ce0ac0a91027c457e3a7aaa70360d55d1e8cba Mon Sep 17 00:00:00 2001 From: Pierre <129078893+pierre-az@users.noreply.github.com> Date: Tue, 2 Jan 2024 21:05:36 +0100 Subject: [PATCH] Rebase & Squash (#9070) Signed-off-by: Yodo Signed-off-by: Derek Nola --- go.mod | 1 + go.sum | 2 + pkg/agent/loadbalancer/servers.go | 62 +++++++++++++++++++++++++++---- 3 files changed, 57 insertions(+), 8 deletions(-) diff --git a/go.mod b/go.mod index 87d02507d8..535b7d1fc7 100644 --- a/go.mod +++ b/go.mod @@ -127,6 +127,7 @@ require ( github.com/libp2p/go-libp2p v0.30.0 github.com/mattn/go-sqlite3 v1.14.17 github.com/minio/minio-go/v7 v7.0.33 + github.com/mwitkow/go-http-dialer v0.0.0-20161116154839-378f744fb2b8 github.com/natefinch/lumberjack v2.0.0+incompatible github.com/onsi/ginkgo/v2 v2.11.0 github.com/onsi/gomega v1.27.10 diff --git a/go.sum b/go.sum index d94339cabe..da2aff63c4 100644 --- a/go.sum +++ b/go.sum @@ -1260,6 +1260,8 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= +github.com/mwitkow/go-http-dialer v0.0.0-20161116154839-378f744fb2b8 h1:BhQQWYKJwXPtAhm12d4gQU4LKS9Yov22yOrDc2QA7ho= +github.com/mwitkow/go-http-dialer v0.0.0-20161116154839-378f744fb2b8/go.mod h1:ntWhh7pzdiiRKBMxUB5iG+Q2gmZBxGxpX1KyK6N8kX8= github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f h1:y5//uYreIhSUg3J1GEMiLbxo1LJaP8RfCpH6pymGZus= github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f/go.mod h1:ZdcZmHo+o7JKHSa8/e818NopupXU1YMK5fe1lsApnBw= github.com/natefinch/lumberjack v2.0.0+incompatible h1:4QJd3OLAMgj7ph+yZTuX13Ld4UpgHp07nNdFX7mqFfM= diff --git a/pkg/agent/loadbalancer/servers.go b/pkg/agent/loadbalancer/servers.go index e8e9f83a15..8dfb648bab 100644 --- a/pkg/agent/loadbalancer/servers.go +++ b/pkg/agent/loadbalancer/servers.go @@ -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