From bbd68f3a507903442ea093ae5226a7026b307964 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 <pierre@azmed.co>
---
 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 2ba6b0317b..93e672c7e6 100644
--- a/go.mod
+++ b/go.mod
@@ -117,6 +117,7 @@ require (
 	github.com/lib/pq v1.10.2
 	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.13.2
 	github.com/onsi/gomega v1.30.0
diff --git a/go.sum b/go.sum
index 83b9e6462c..fea419146c 100644
--- a/go.sum
+++ b/go.sum
@@ -1121,6 +1121,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