Skip to content

Commit

Permalink
[Release-1.27] Added support for env *_PROXY variables for agent load…
Browse files Browse the repository at this point in the history
…balancer (#9117)

Signed-off-by: Yodo <[email protected]>
Signed-off-by: Derek Nola <[email protected]>
Co-authored-by: Pierre <[email protected]>
  • Loading branch information
dereknola and pierre-az authored Jan 7, 2024
1 parent cf5fd5f commit 5ecf175
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 8 deletions.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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.11.0
github.com/onsi/gomega v1.27.10
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -900,6 +900,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=
Expand Down
62 changes: 54 additions & 8 deletions pkg/agent/loadbalancer/servers.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 5ecf175

Please sign in to comment.