implement the DestHostBackendManager

Author: charleszheng44

cmd/agent/main.go

@@ -70,6 +70,7 @@ type GrpcProxyAgentOptions struct {
 	adminServerPort  int
 
 	agentID       string
+	agentHost     string
 	syncInterval  time.Duration
 	probeInterval time.Duration
 
@@ -81,6 +82,7 @@ func (o *GrpcProxyAgentOptions) ClientSetConfig(dialOption grpc.DialOption) *age
 	return &agent.ClientSetConfig{
 		Address:                 fmt.Sprintf("%s:%d", o.proxyServerHost, o.proxyServerPort),
 		AgentID:                 o.agentID,
+		AgentHost:               o.agentHost,
 		SyncInterval:            o.syncInterval,
 		ProbeInterval:           o.probeInterval,
 		DialOption:              dialOption,
@@ -101,6 +103,7 @@ func (o *GrpcProxyAgentOptions) Flags() *pflag.FlagSet {
 	flags.DurationVar(&o.syncInterval, "sync-interval", o.syncInterval, "The initial interval by which the agent periodically checks if it has connections to all instances of the proxy server.")
 	flags.DurationVar(&o.probeInterval, "probe-interval", o.probeInterval, "The interval by which the agent periodically checks if its connections to the proxy server are ready.")
 	flags.StringVar(&o.serviceAccountTokenPath, "service-account-token-path", o.serviceAccountTokenPath, "If non-empty proxy agent uses this token to prove its identity to the proxy server.")
+	flags.StringVar(&o.agentHost, "agent-host", o.agentHost, "The host address of this agent, which can be either an IPv4/6 address or a domain name.")
 	return flags
 }
 
@@ -116,6 +119,7 @@ func (o *GrpcProxyAgentOptions) Print() {
 	klog.Warningf("SyncInterval set to %v.\n", o.syncInterval)
 	klog.Warningf("ProbeInterval set to %v.\n", o.probeInterval)
 	klog.Warningf("ServiceAccountTokenPath set to \"%s\".\n", o.serviceAccountTokenPath)
+	klog.Warningf("AgentHost set to %s.\n", o.agentHost)
 }
 
 func (o *GrpcProxyAgentOptions) Validate() error {
@@ -168,6 +172,7 @@ func newGrpcProxyAgentOptions() *GrpcProxyAgentOptions {
 		healthServerPort:        8093,
 		adminServerPort:         8094,
 		agentID:                 uuid.New().String(),
+		agentHost:               "",
 		syncInterval:            1 * time.Second,
 		probeInterval:           1 * time.Second,
 		serviceAccountTokenPath: "",

cmd/server/main.go

@@ -107,6 +107,10 @@ type ProxyRunOptions struct {
 	authenticationAudience string
 	// Path to kubeconfig (used by kubernetes client)
 	kubeconfigPath string
+
+	// Proxy strategy used by the server, can be either "destHost" or "default".
+	// If not set, the "default" strategy will be used.
+	proxyStrategy string
 }
 
 func (o *ProxyRunOptions) Flags() *pflag.FlagSet {
@@ -132,6 +136,7 @@ func (o *ProxyRunOptions) Flags() *pflag.FlagSet {
 	flags.StringVar(&o.agentServiceAccount, "agent-service-account", o.agentServiceAccount, "Expected agent's service account during agent authentication (used with agent-namespace, authentication-audience, kubeconfig).")
 	flags.StringVar(&o.kubeconfigPath, "kubeconfig", o.kubeconfigPath, "absolute path to the kubeconfig file (used with agent-namespace, agent-service-account, authentication-audience).")
 	flags.StringVar(&o.authenticationAudience, "authentication-audience", o.authenticationAudience, "Expected agent's token authentication audience (used with agent-namespace, agent-service-account, kubeconfig).")
+	flags.StringVar(&o.proxyStrategy, "proxy-strategy", o.proxyStrategy, "TODO ...")
 	return flags
 }
 
@@ -157,6 +162,7 @@ func (o *ProxyRunOptions) Print() {
 	klog.Warningf("AgentServiceAccount set to %q.\n", o.agentServiceAccount)
 	klog.Warningf("AuthenticationAudience set to %q.\n", o.authenticationAudience)
 	klog.Warningf("KubeconfigPath set to %q.\n", o.kubeconfigPath)
+	klog.Warningf("ProxyStrategy set to %s.\n", o.proxyStrategy)
 }
 
 func (o *ProxyRunOptions) Validate() error {
@@ -298,6 +304,7 @@ func newProxyRunOptions() *ProxyRunOptions {
 		agentServiceAccount:       "",
 		kubeconfigPath:            "",
 		authenticationAudience:    "",
+		proxyStrategy:             "default",
 	}
 	return &o
 }
@@ -347,7 +354,7 @@ func (p *Proxy) run(o *ProxyRunOptions) error {
 		KubernetesClient:       k8sClient,
 		AuthenticationAudience: o.authenticationAudience,
 	}
-	server := server.NewProxyServer(o.serverID, int(o.serverCount), authOpt)
+	server := server.NewProxyServer(o.serverID, o.proxyStrategy, int(o.serverCount), authOpt)
 
 	klog.Info("Starting master server for client connections.")
 	masterStop, err := p.runMasterServer(ctx, o, server)

pkg/agent/client.go

@@ -88,9 +88,10 @@ type AgentClient struct {
 
 	cs *ClientSet // the clientset that includes this AgentClient.
 
-	stream   agent.AgentService_ConnectClient
-	agentID  string
-	serverID string // the id of the proxy server this client connects to.
+	stream    agent.AgentService_ConnectClient
+	agentID   string
+	agentHost string
+	serverID  string // the id of the proxy server this client connects to.
 
 	// connect opts
 	address string
@@ -107,11 +108,12 @@ type AgentClient struct {
 	serviceAccountTokenPath string
 }
 
-func newAgentClient(address, agentID string, cs *ClientSet, opts ...grpc.DialOption) (*AgentClient, int, error) {
+func newAgentClient(address, agentID, agentHost string, cs *ClientSet, opts ...grpc.DialOption) (*AgentClient, int, error) {
 	a := &AgentClient{
 		cs:                      cs,
 		address:                 address,
 		agentID:                 agentID,
+		agentHost:               agentHost,
 		opts:                    opts,
 		probeInterval:           cs.probeInterval,
 		stopCh:                  make(chan struct{}),
@@ -132,7 +134,9 @@ func (a *AgentClient) Connect() (int, error) {
 	if err != nil {
 		return 0, err
 	}
-	ctx := metadata.AppendToOutgoingContext(context.Background(), header.AgentID, a.agentID)
+	ctx := metadata.AppendToOutgoingContext(context.Background(),
+		header.AgentID, a.agentID,
+		header.AgentHost, a.agentHost)
 	if a.serviceAccountTokenPath != "" {
 		if ctx, err = a.initializeAuthContext(ctx); err != nil {
 			conn.Close()

pkg/agent/clientset.go

@@ -48,6 +48,8 @@ type ClientSet struct {
 	serviceAccountTokenPath string
 	// channel to signal shutting down the client set. Primarily for test.
 	stopCh <-chan struct{}
+
+	agentHost string // The IP address of the agent
 }
 
 func (cs *ClientSet) ClientsCount() int {
@@ -108,6 +110,7 @@ func (cs *ClientSet) RemoveClient(serverID string) {
 type ClientSetConfig struct {
 	Address                 string
 	AgentID                 string
+	AgentHost                 string
 	SyncInterval            time.Duration
 	ProbeInterval           time.Duration
 	DialOption              grpc.DialOption
@@ -118,6 +121,7 @@ func (cc *ClientSetConfig) NewAgentClientSet(stopCh <-chan struct{}) *ClientSet
 	return &ClientSet{
 		clients:                 make(map[string]*AgentClient),
 		agentID:                 cc.AgentID,
+		agentHost:                 cc.AgentHost,
 		address:                 cc.Address,
 		syncInterval:            cc.SyncInterval,
 		probeInterval:           cc.ProbeInterval,
@@ -128,7 +132,7 @@ func (cc *ClientSetConfig) NewAgentClientSet(stopCh <-chan struct{}) *ClientSet
 }
 
 func (cs *ClientSet) newAgentClient() (*AgentClient, int, error) {
-	return newAgentClient(cs.address, cs.agentID, cs, cs.dialOption)
+	return newAgentClient(cs.address, cs.agentID, cs.agentHost, cs, cs.dialOption)
 }
 
 func (cs *ClientSet) resetBackoff() *wait.Backoff {

pkg/server/backend_manager.go

@@ -78,6 +78,7 @@ type BackendManager interface {
 	// contains multiple requests.
 	Backend(ctx context.Context) (Backend, error)
 	BackendStorage
+	ReadinessManager
 }
 
 var _ BackendManager = &DefaultBackendManager{}

pkg/server/desthost_backend_manager.go

@@ -0,0 +1,39 @@
+package server
+
+import (
+	"context"
+	"k8s.io/klog/v2"
+)
+
+type DestHostBackendManager struct {
+	*DefaultBackendStorage
+}
+
+var _ BackendManager = &DestHostBackendManager{}
+
+func NewDestHostBackendManager() *DestHostBackendManager {
+	return &DestHostBackendManager{DefaultBackendStorage: NewDefaultBackendStorage()}
+}
+
+// Backend tries to get a backend by IP. If the desired backend is not found,
+// just return a default one
+func (dibm *DestHostBackendManager) Backend(ctx context.Context) (Backend, error) {
+	dibm.mu.RLock()
+	defer dibm.mu.RUnlock()
+	if len(dibm.backends) == 0 {
+		return nil, &ErrNotFound{}
+	}
+	destHost := ctx.Value("destHost").(string)
+	if destHost != "" {
+		bes, exist := dibm.backends[destHost]
+		if exist && len(bes) > 0 {
+			return dibm.backends[destHost][0], nil
+		}
+		return nil, &ErrNotFound{}
+	}
+
+	// fallback to the default behavior, i.e., get a random backend
+	agentID := dibm.agentIDs[dibm.random.Intn(len(dibm.agentIDs))]
+	klog.Infof("pick agentID=%s as backend", agentID)
+	return dibm.backends[agentID][0], nil
+}

pkg/server/readiness_manager.go

@@ -23,9 +23,9 @@ type ReadinessManager interface {
 	Ready() (bool, string)
 }
 
-var _ ReadinessManager = &DefaultBackendManager{}
+var _ ReadinessManager = &DefaultBackendStorage{}
 
-func (s *DefaultBackendManager) Ready() (bool, string) {
+func (s *DefaultBackendStorage) Ready() (bool, string) {
 	if s.NumBackends() == 0 {
 		return false, "no connection to any proxy agent"
 	}

pkg/server/server.go

@@ -32,6 +32,7 @@ import (
 	"k8s.io/klog/v2"
 	"sigs.k8s.io/apiserver-network-proxy/konnectivity-client/proto/client"
 	"sigs.k8s.io/apiserver-network-proxy/pkg/server/metrics"
+	"sigs.k8s.io/apiserver-network-proxy/pkg/util"
 	"sigs.k8s.io/apiserver-network-proxy/proto/agent"
 	"sigs.k8s.io/apiserver-network-proxy/proto/header"
 )
@@ -125,6 +126,8 @@ type ProxyServer struct {
 
 	// agent authentication
 	AgentAuthenticationOptions *AgentTokenAuthenticationOptions
+
+	proxyStrategy string
 }
 
 // AgentTokenAuthenticationOptions contains list of parameters required for agent token based authentication
@@ -201,8 +204,15 @@ func (s *ProxyServer) getFrontendsForBackendConn(agentID string, backend Backend
 }
 
 // NewProxyServer creates a new ProxyServer instance
-func NewProxyServer(serverID string, serverCount int, agentAuthenticationOptions *AgentTokenAuthenticationOptions) *ProxyServer {
-	bm := NewDefaultBackendManager()
+func NewProxyServer(serverID, proxyStrategy string, serverCount int, agentAuthenticationOptions *AgentTokenAuthenticationOptions) *ProxyServer {
+	var bm BackendManager
+	if proxyStrategy == "destHost" {
+		bm = NewDestHostBackendManager()
+	} else {
+		// to ensure backward compatibility, use the "default" backend manager
+		// for all other proxy strategies
+		bm = NewDefaultBackendManager()
+	}
 	return &ProxyServer{
 		frontends:                  make(map[string](map[int64]*ProxyClientConnection)),
 		PendingDial:                NewPendingDialManager(),
@@ -211,6 +221,7 @@ func NewProxyServer(serverID string, serverCount int, agentAuthenticationOptions
 		BackendManager:             bm,
 		AgentAuthenticationOptions: agentAuthenticationOptions,
 		Readiness:                  bm,
+		proxyStrategy:              proxyStrategy,
 	}
 }
 
@@ -270,7 +281,12 @@ func (s *ProxyServer) serveRecvFrontend(stream client.ProxyService_ProxyServer,
 			// the address, then we can send the Dial_REQ to the
 			// same agent. That way we save the agent from creating
 			// a new connection to the address.
-			backend, err = s.BackendManager.Backend(context.TODO())
+			ctx := context.Background()
+			if s.proxyStrategy == "destHost" {
+				addr := util.RemovePortFromHost(pkt.GetDialRequest().Address)
+				ctx = context.WithValue(ctx, "destHost", addr)
+			}
+			backend, err = s.BackendManager.Backend(ctx)
 			if err != nil {
 				klog.Errorf(">>> failed to get a backend: %v", err)
 				continue
@@ -370,6 +386,21 @@ func agentID(stream agent.AgentService_ConnectServer) (string, error) {
 	return agentIDs[0], nil
 }
 
+func agentHost(stream agent.AgentService_ConnectServer) (string, error) {
+	md, ok := metadata.FromIncomingContext(stream.Context())
+	if !ok {
+		return "", fmt.Errorf("failed to get context")
+	}
+	agentHosts := md.Get(header.AgentHost)
+	if len(agentHosts) > 1 {
+		return "", fmt.Errorf("expected at most one agent IP in the context, got %v", agentHosts)
+	}
+	if len(agentHosts) == 0 {
+		return "", nil
+	}
+	return agentHosts[0], nil
+}
+
 func (s *ProxyServer) validateAuthToken(token string) error {
 	trReq := &authv1.TokenReview{
 		Spec: authv1.TokenReviewSpec{
@@ -446,6 +477,18 @@ func (s *ProxyServer) Connect(stream agent.AgentService_ConnectServer) error {
 	if err != nil {
 		return err
 	}
+	if s.proxyStrategy == "destHost" {
+		// if the proxyStrategy is "destHost", use the agentHost as the ID
+		agentHost, err := agentHost(stream)
+		if err != nil {
+			return err
+		}
+		if agentHost != "" {
+			// if the agentHost is set, assign it to the agentID
+			agentID = agentHost
+		}
+	}
+
 	klog.Infof("Connect request from agent %s", agentID)
 	backend := s.BackendManager.AddBackend(agentID, stream)
 	defer s.BackendManager.RemoveBackend(agentID, stream)

pkg/server/tunnel.go

@@ -26,6 +26,7 @@ import (
 
 	"k8s.io/klog/v2"
 	"sigs.k8s.io/apiserver-network-proxy/konnectivity-client/proto/client"
+	"sigs.k8s.io/apiserver-network-proxy/pkg/util"
 )
 
 // Tunnel implements Proxy based on HTTP Connect, which tunnels the traffic to
@@ -69,7 +70,12 @@ func (t *Tunnel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		},
 	}
 	klog.Infof("Set pending(rand=%d) to %v", random, w)
-	backend, err := t.Server.BackendManager.Backend(context.TODO())
+	ctx := context.Background()
+	if t.Server.proxyStrategy == "destHost" {
+		addr := util.RemovePortFromHost(r.Host)
+		ctx = context.WithValue(ctx, "destHost", addr)
+	}
+	backend, err := t.Server.BackendManager.Backend(ctx)
 	if err != nil {
 		http.Error(w, fmt.Sprintf("currently no tunnels available: %v", err), http.StatusInternalServerError)
 		return

pkg/util/net.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2019 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package util
+
+import (
+	"strings"
+)
+
+// containIPv6Addr checks if the given host identity contains an
+// IPv6 address
+func containIPv6Addr(host string) bool {
+	// the shortest IPv6 address is ::
+	return len(strings.Split(host, ":")) > 2
+}
+
+// containPortIPv6 checks if the host that contains an IPv6 address
+// also contains a port number
+func containPortIPv6(host string) bool {
+	// based on to RFC 3986, section 3.2.2, host identified by an
+	// IPv6 is distinguished by enclosing the IP literal within square
+	// brackets ("[" and "]")
+	return strings.ContainsRune(host, '[')
+}
+
+// RemovePortFromHost removes port number from the host address that
+// may be of the form "<host>:<port>" where the <host> can be an either
+// an IPv4/6 address or a domain name
+func RemovePortFromHost(host string) string {
+	if !containIPv6Addr(host) {
+		return strings.Split(host, ":")[0]
+	}
+	if containPortIPv6(host) {
+		host = host[:strings.LastIndexByte(host, ':')]
+	}
+	return strings.Trim(host, "[]")
+}