Move http/socks proxy stuff to separate file

Signed-off-by: Brad Davidson <[email protected]>

pkg/agent/loadbalancer/httpproxy.go

@@ -0,0 +1,70 @@
+package loadbalancer
+
+import (
+	"fmt"
+	"net"
+	"net/url"
+	"os"
+	"strconv"
+	"time"
+
+	"github.com/k3s-io/k3s/pkg/version"
+	http_dialer "github.com/mwitkow/go-http-dialer"
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/net/http/httpproxy"
+	"golang.org/x/net/proxy"
+)
+
+var defaultDialer proxy.Dialer = &net.Dialer{
+	Timeout:   10 * time.Second,
+	KeepAlive: 30 * time.Second,
+}
+
+// SetHTTPProxy configures a proxy-enabled dialer to be used for all loadbalancer connections,
+// if the agent has been configured to allow use of a HTTP proxy, and the environment has been configured
+// to indicate use of a HTTP proxy for the server URL.
+func SetHTTPProxy(address string) error {
+	// Check if env variable for proxy is set
+	if useProxy, _ := strconv.ParseBool(os.Getenv(version.ProgramUpper + "_AGENT_HTTP_PROXY_ALLOWED")); !useProxy || address == "" {
+		return nil
+	}
+
+	serverURL, err := url.Parse(address)
+	if err != nil {
+		return errors.Wrapf(err, "failed to parse address %s", address)
+	}
+
+	// Call this directly instead of using the cached environment used by http.ProxyFromEnvironment to allow for testing
+	proxyFromEnvironment := httpproxy.FromEnvironment().ProxyFunc()
+	proxyURL, err := proxyFromEnvironment(serverURL)
+	if err != nil {
+		return errors.Wrapf(err, "failed to get proxy for address %s", address)
+	}
+	if proxyURL == nil {
+		logrus.Debug(version.ProgramUpper + "_AGENT_HTTP_PROXY_ALLOWED is true but no proxy is configured for URL " + serverURL.String())
+		return nil
+	}
+
+	dialer, err := proxyDialer(proxyURL, defaultDialer)
+	if err != nil {
+		return errors.Wrapf(err, "failed to create proxy dialer for %s", proxyURL)
+	}
+
+	defaultDialer = dialer
+	logrus.Debugf("Using proxy %s for agent connection to %s", proxyURL, serverURL)
+	return nil
+}
+
+// proxyDialer creates a new proxy.Dialer that routes connections through the specified proxy.
+func proxyDialer(proxyURL *url.URL, forward proxy.Dialer) (proxy.Dialer, error) {
+	if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" {
+		// Create a new HTTP proxy dialer
+		httpProxyDialer := http_dialer.New(proxyURL, http_dialer.WithDialer(forward.(*net.Dialer)))
+		return httpProxyDialer, nil
+	} else if proxyURL.Scheme == "socks5" {
+		// For SOCKS5 proxies, use the proxy package's FromURL
+		return proxy.FromURL(proxyURL, forward)
+	}
+	return nil, fmt.Errorf("unsupported proxy scheme: %s", proxyURL.Scheme)
+}

pkg/agent/loadbalancer/servers.go

@@ -2,66 +2,18 @@ package loadbalancer
 
 import (
 	"context"
-	"fmt"
 	"math/rand"
 	"net"
-	"net/url"
-	"os"
 	"slices"
-	"strconv"
 	"time"
 
-	"github.com/k3s-io/k3s/pkg/version"
-	http_dialer "github.com/mwitkow/go-http-dialer"
 	"github.com/pkg/errors"
-	"golang.org/x/net/http/httpproxy"
-	"golang.org/x/net/proxy"
 
 	"github.com/sirupsen/logrus"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/apimachinery/pkg/util/wait"
 )
 
-var defaultDialer proxy.Dialer = &net.Dialer{
-	Timeout:   10 * time.Second,
-	KeepAlive: 30 * time.Second,
-}
-
-// SetHTTPProxy configures a proxy-enabled dialer to be used for all loadbalancer connections,
-// if the agent has been configured to allow use of a HTTP proxy, and the environment has been configured
-// to indicate use of a HTTP proxy for the server URL.
-func SetHTTPProxy(address string) error {
-	// Check if env variable for proxy is set
-	if useProxy, _ := strconv.ParseBool(os.Getenv(version.ProgramUpper + "_AGENT_HTTP_PROXY_ALLOWED")); !useProxy || address == "" {
-		return nil
-	}
-
-	serverURL, err := url.Parse(address)
-	if err != nil {
-		return errors.Wrapf(err, "failed to parse address %s", address)
-	}
-
-	// Call this directly instead of using the cached environment used by http.ProxyFromEnvironment to allow for testing
-	proxyFromEnvironment := httpproxy.FromEnvironment().ProxyFunc()
-	proxyURL, err := proxyFromEnvironment(serverURL)
-	if err != nil {
-		return errors.Wrapf(err, "failed to get proxy for address %s", address)
-	}
-	if proxyURL == nil {
-		logrus.Debug(version.ProgramUpper + "_AGENT_HTTP_PROXY_ALLOWED is true but no proxy is configured for URL " + serverURL.String())
-		return nil
-	}
-
-	dialer, err := proxyDialer(proxyURL, defaultDialer)
-	if err != nil {
-		return errors.Wrapf(err, "failed to create proxy dialer for %s", proxyURL)
-	}
-
-	defaultDialer = dialer
-	logrus.Debugf("Using proxy %s for agent connection to %s", proxyURL, serverURL)
-	return nil
-}
-
 func (lb *LoadBalancer) setServers(serverAddresses []string) bool {
 	serverAddresses, hasDefaultServer := sortServers(serverAddresses, lb.defaultServerAddress)
 	if len(serverAddresses) == 0 {
@@ -174,19 +126,6 @@ func (s *server) dialContext(ctx context.Context, network, address string) (net.
 	return wrappedConn, nil
 }
 
-// proxyDialer creates a new proxy.Dialer that routes connections through the specified proxy.
-func proxyDialer(proxyURL *url.URL, forward proxy.Dialer) (proxy.Dialer, error) {
-	if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" {
-		// Create a new HTTP proxy dialer
-		httpProxyDialer := http_dialer.New(proxyURL, http_dialer.WithDialer(forward.(*net.Dialer)))
-		return httpProxyDialer, nil
-	} else if proxyURL.Scheme == "socks5" {
-		// For SOCKS5 proxies, use the proxy package's FromURL
-		return proxy.FromURL(proxyURL, forward)
-	}
-	return nil, fmt.Errorf("unsupported proxy scheme: %s", proxyURL.Scheme)
-}
-
 // closeAll closes all connections to the server, and removes their entries from the map
 func (s *server) closeAll() {
 	s.mutex.Lock()