Rebase & Squash (#9070)

Signed-off-by: Yodo <pierre@azmed.co>
This commit is contained in:
Pierre 2024-01-02 21:05:36 +01:00 committed by GitHub
parent c7a8eef977
commit bbd68f3a50
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 57 additions and 8 deletions

1
go.mod
View File

@ -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

2
go.sum
View File

@ -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=

View File

@ -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