From ffe5f187d3672d5386b0e05cb658fa6765f45111 Mon Sep 17 00:00:00 2001 From: Tommy Xiao Date: Mon, 21 Oct 2024 21:43:39 +0800 Subject: [PATCH] update agent/cli --- pkg/agent/cri/cri.go | 2 +- pkg/agent/cri/cri_linux.go | 4 +- pkg/agent/cri/cri_windows.go | 4 +- pkg/agent/loadbalancer/servers.go | 4 +- pkg/agent/loadbalancer/servers_test.go | 156 +++++++++++++++++++++++++ 5 files changed, 163 insertions(+), 7 deletions(-) create mode 100644 pkg/agent/loadbalancer/servers_test.go diff --git a/pkg/agent/cri/cri.go b/pkg/agent/cri/cri.go index cd92f639..8a1579b4 100644 --- a/pkg/agent/cri/cri.go +++ b/pkg/agent/cri/cri.go @@ -33,4 +33,4 @@ func WaitForService(ctx context.Context, address string, service string) error { } logrus.Infof("%s is now running", service) return nil -} +} \ No newline at end of file diff --git a/pkg/agent/cri/cri_linux.go b/pkg/agent/cri/cri_linux.go index 7d4774aa..8b8d5c43 100644 --- a/pkg/agent/cri/cri_linux.go +++ b/pkg/agent/cri/cri_linux.go @@ -9,7 +9,7 @@ import ( "google.golang.org/grpc" runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1" - k8sutil "k8s.io/kubernetes/pkg/kubelet/util" + k8sutil "k8s.io/cri-client/pkg/util" ) const socketPrefix = "unix://" @@ -36,4 +36,4 @@ func Connection(ctx context.Context, address string) (*grpc.ClientConn, error) { } return conn, nil -} +} \ No newline at end of file diff --git a/pkg/agent/cri/cri_windows.go b/pkg/agent/cri/cri_windows.go index ab700b6b..b8d978f0 100644 --- a/pkg/agent/cri/cri_windows.go +++ b/pkg/agent/cri/cri_windows.go @@ -9,7 +9,7 @@ import ( "google.golang.org/grpc" runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1" - "k8s.io/kubernetes/pkg/kubelet/util" + "k8s.io/cri-client/pkg/util" ) // Connection connects to a CRI socket at the given path. @@ -34,4 +34,4 @@ func Connection(ctx context.Context, address string) (*grpc.ClientConn, error) { } return conn, nil -} +} \ No newline at end of file diff --git a/pkg/agent/loadbalancer/servers.go b/pkg/agent/loadbalancer/servers.go index 1cdc9e9b..0fa96f36 100644 --- a/pkg/agent/loadbalancer/servers.go +++ b/pkg/agent/loadbalancer/servers.go @@ -10,9 +10,9 @@ import ( "strconv" "time" + "github.com/xiaods/k8e/pkg/version" http_dialer "github.com/mwitkow/go-http-dialer" "github.com/pkg/errors" - "github.com/xiaods/k8e/pkg/version" "golang.org/x/net/http/httpproxy" "golang.org/x/net/proxy" @@ -255,4 +255,4 @@ func (lb *LoadBalancer) runHealthChecks(ctx context.Context) { } }, time.Second, ctx.Done()) logrus.Debugf("Stopped health checking for load balancer %s", lb.serviceName) -} +} \ No newline at end of file diff --git a/pkg/agent/loadbalancer/servers_test.go b/pkg/agent/loadbalancer/servers_test.go new file mode 100644 index 00000000..5c156f0c --- /dev/null +++ b/pkg/agent/loadbalancer/servers_test.go @@ -0,0 +1,156 @@ +package loadbalancer + +import ( + "fmt" + "net" + "os" + "strings" + "testing" + + "github.com/xiaods/k8e/pkg/version" + "github.com/sirupsen/logrus" +) + +var defaultEnv map[string]string +var proxyEnvs = []string{version.ProgramUpper + "_AGENT_HTTP_PROXY_ALLOWED", "HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "no_proxy"} + +func init() { + logrus.SetLevel(logrus.DebugLevel) +} + +func prepareEnv(env ...string) { + defaultDialer = &net.Dialer{} + defaultEnv = map[string]string{} + for _, e := range proxyEnvs { + if v, ok := os.LookupEnv(e); ok { + defaultEnv[e] = v + os.Unsetenv(e) + } + } + for _, e := range env { + k, v, _ := strings.Cut(e, "=") + os.Setenv(k, v) + } +} + +func restoreEnv() { + for _, e := range proxyEnvs { + if v, ok := defaultEnv[e]; ok { + os.Setenv(e, v) + } else { + os.Unsetenv(e) + } + } +} + +func Test_UnitSetHTTPProxy(t *testing.T) { + type args struct { + address string + } + tests := []struct { + name string + args args + setup func() error + teardown func() error + wantErr bool + wantDialer string + }{ + { + name: "Default Proxy", + args: args{address: "https://1.2.3.4:6443"}, + wantDialer: "*net.Dialer", + setup: func() error { + prepareEnv(version.ProgramUpper+"_AGENT_HTTP_PROXY_ALLOWED=", "HTTP_PROXY=", "HTTPS_PROXY=", "NO_PROXY=") + return nil + }, + teardown: func() error { + restoreEnv() + return nil + }, + }, + { + name: "Agent Proxy Enabled", + args: args{address: "https://1.2.3.4:6443"}, + wantDialer: "*http_dialer.HttpTunnel", + setup: func() error { + prepareEnv(version.ProgramUpper+"_AGENT_HTTP_PROXY_ALLOWED=true", "HTTP_PROXY=http://proxy:8080", "HTTPS_PROXY=http://proxy:8080", "NO_PROXY=") + return nil + }, + teardown: func() error { + restoreEnv() + return nil + }, + }, + { + name: "Agent Proxy Enabled with Bogus Proxy", + args: args{address: "https://1.2.3.4:6443"}, + wantDialer: "*net.Dialer", + wantErr: true, + setup: func() error { + prepareEnv(version.ProgramUpper+"_AGENT_HTTP_PROXY_ALLOWED=true", "HTTP_PROXY=proxy proxy", "HTTPS_PROXY=proxy proxy", "NO_PROXY=") + return nil + }, + teardown: func() error { + restoreEnv() + return nil + }, + }, + { + name: "Agent Proxy Enabled with Bogus Server", + args: args{address: "https://1.2.3.4:k8e"}, + wantDialer: "*net.Dialer", + wantErr: true, + setup: func() error { + prepareEnv(version.ProgramUpper+"_AGENT_HTTP_PROXY_ALLOWED=true", "HTTP_PROXY=http://proxy:8080", "HTTPS_PROXY=http://proxy:8080", "NO_PROXY=") + return nil + }, + teardown: func() error { + restoreEnv() + return nil + }, + }, + { + name: "Agent Proxy Enabled but IP Excluded", + args: args{address: "https://1.2.3.4:6443"}, + wantDialer: "*net.Dialer", + setup: func() error { + prepareEnv(version.ProgramUpper+"_AGENT_HTTP_PROXY_ALLOWED=true", "HTTP_PROXY=http://proxy:8080", "HTTPS_PROXY=http://proxy:8080", "NO_PROXY=1.2.0.0/16") + return nil + }, + teardown: func() error { + restoreEnv() + return nil + }, + }, + { + name: "Agent Proxy Enabled but Domain Excluded", + args: args{address: "https://server.example.com:6443"}, + wantDialer: "*net.Dialer", + setup: func() error { + prepareEnv(version.ProgramUpper+"_AGENT_HTTP_PROXY_ALLOWED=true", "HTTP_PROXY=http://proxy:8080", "HTTPS_PROXY=http://proxy:8080", "NO_PROXY=*.example.com") + return nil + }, + teardown: func() error { + restoreEnv() + return nil + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + defer tt.teardown() + if err := tt.setup(); err != nil { + t.Errorf("Setup for SetHTTPProxy() failed = %v", err) + return + } + err := SetHTTPProxy(tt.args.address) + t.Logf("SetHTTPProxy() error = %v", err) + if (err != nil) != tt.wantErr { + t.Errorf("SetHTTPProxy() error = %v, wantErr %v", err, tt.wantErr) + } + if dialerType := fmt.Sprintf("%T", defaultDialer); dialerType != tt.wantDialer { + t.Errorf("Got wrong dialer type %s, wanted %s", dialerType, tt.wantDialer) + } + }) + } +} \ No newline at end of file