From af1b43876af66e0c0d50a22cc1604c7a67213815 Mon Sep 17 00:00:00 2001 From: Joseph Anttila Hall Date: Fri, 22 Mar 2024 10:44:40 -0700 Subject: [PATCH 1/2] Add test coverage around BackendManager and BackendStorage. Includes a trivial API refactor. --- pkg/server/backend_manager.go | 35 ++- pkg/server/backend_manager_test.go | 273 +++++++++++++++++--- pkg/server/default_route_backend_manager.go | 28 +- pkg/server/desthost_backend_manager.go | 34 ++- pkg/server/server.go | 63 +---- pkg/server/server_test.go | 254 ++++++++++++++++++ 6 files changed, 578 insertions(+), 109 deletions(-) diff --git a/pkg/server/backend_manager.go b/pkg/server/backend_manager.go index d6c62e612..52c3af118 100644 --- a/pkg/server/backend_manager.go +++ b/pkg/server/backend_manager.go @@ -182,10 +182,10 @@ func NewBackend(conn agent.AgentService_ConnectServer) (Backend, error) { // BackendStorage is an interface to manage the storage of the backend // connections, i.e., get, add and remove type BackendStorage interface { - // AddBackend adds a backend. - AddBackend(identifier string, idType header.IdentifierType, backend Backend) - // RemoveBackend removes a backend. - RemoveBackend(identifier string, idType header.IdentifierType, backend Backend) + // addBackend adds a backend. + addBackend(identifier string, idType header.IdentifierType, backend Backend) + // removeBackend removes a backend. + removeBackend(identifier string, idType header.IdentifierType, backend Backend) // NumBackends returns the number of backends. NumBackends() int } @@ -199,6 +199,10 @@ type BackendManager interface { // pick a backend for every tunnel session and each tunnel session may // contains multiple requests. Backend(ctx context.Context) (Backend, error) + // AddBackend adds a backend. + AddBackend(backend Backend) + // RemoveBackend adds a backend. + RemoveBackend(backend Backend) BackendStorage ReadinessManager } @@ -215,6 +219,18 @@ func (dbm *DefaultBackendManager) Backend(_ context.Context) (Backend, error) { return dbm.DefaultBackendStorage.GetRandomBackend() } +func (dbm *DefaultBackendManager) AddBackend(backend Backend) { + agentID := backend.GetAgentID() + klog.V(5).InfoS("Add the agent to DefaultBackendManager", "agentID", agentID) + dbm.addBackend(agentID, header.UID, backend) +} + +func (dbm *DefaultBackendManager) RemoveBackend(backend Backend) { + agentID := backend.GetAgentID() + klog.V(5).InfoS("Remove the agent from the DefaultBackendManager", "agentID", agentID) + dbm.removeBackend(agentID, header.UID, backend) +} + // DefaultBackendStorage is the default backend storage. type DefaultBackendStorage struct { mu sync.RWMutex //protects the following @@ -222,6 +238,9 @@ type DefaultBackendStorage struct { // For a given agent, ProxyServer prefers backends[agentID][0] to send // traffic, because backends[agentID][1:] are more likely to be closed // by the agent to deduplicate connections to the same server. + // + // TODO: fix documentation. This is not always agentID, e.g. in + // the case of DestHostBackendManager. backends map[string][]Backend // agentID is tracked in this slice to enable randomly picking an // agentID in the Backend() method. There is no reliable way to @@ -267,8 +286,8 @@ func containIDType(idTypes []header.IdentifierType, idType header.IdentifierType return false } -// AddBackend adds a backend. -func (s *DefaultBackendStorage) AddBackend(identifier string, idType header.IdentifierType, backend Backend) { +// addBackend adds a backend. +func (s *DefaultBackendStorage) addBackend(identifier string, idType header.IdentifierType, backend Backend) { if !containIDType(s.idTypes, idType) { klog.V(4).InfoS("fail to add backend", "backend", identifier, "error", &ErrWrongIDType{idType, s.idTypes}) return @@ -295,8 +314,8 @@ func (s *DefaultBackendStorage) AddBackend(identifier string, idType header.Iden } } -// RemoveBackend removes a backend. -func (s *DefaultBackendStorage) RemoveBackend(identifier string, idType header.IdentifierType, backend Backend) { +// removeBackend removes a backend. +func (s *DefaultBackendStorage) removeBackend(identifier string, idType header.IdentifierType, backend Backend) { if !containIDType(s.idTypes, idType) { klog.ErrorS(&ErrWrongIDType{idType, s.idTypes}, "fail to remove backend") return diff --git a/pkg/server/backend_manager_test.go b/pkg/server/backend_manager_test.go index e716ab3fc..1c90cae8f 100644 --- a/pkg/server/backend_manager_test.go +++ b/pkg/server/backend_manager_test.go @@ -25,7 +25,6 @@ import ( "google.golang.org/grpc/metadata" agentmock "sigs.k8s.io/apiserver-network-proxy/proto/agent/mocks" - "sigs.k8s.io/apiserver-network-proxy/proto/header" ) func mockAgentConn(ctrl *gomock.Controller, agentID string, agentIdentifiers []string) *agentmock.MockAgentService_ConnectServer { @@ -106,7 +105,7 @@ func TestNewBackend(t *testing.T) { } } -func TestAddRemoveBackendsWithDefaultStrategy(t *testing.T) { +func TestDefaultBackendManager_AddRemoveBackends(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -118,8 +117,8 @@ func TestAddRemoveBackendsWithDefaultStrategy(t *testing.T) { p := NewDefaultBackendManager() - p.AddBackend("agent1", header.UID, backend1) - p.RemoveBackend("agent1", header.UID, backend1) + p.AddBackend(backend1) + p.RemoveBackend(backend1) expectedBackends := make(map[string][]Backend) expectedAgentIDs := []string{} expectedDefaultRouteAgentIDs := []string(nil) @@ -134,18 +133,16 @@ func TestAddRemoveBackendsWithDefaultStrategy(t *testing.T) { } p = NewDefaultBackendManager() - p.AddBackend("agent1", header.UID, backend1) - p.AddBackend("agent1", header.UID, backend12) + p.AddBackend(backend1) + p.AddBackend(backend12) // Adding the same connection again should be a no-op. - p.AddBackend("agent1", header.UID, backend12) - p.AddBackend("agent2", header.UID, backend2) - p.AddBackend("agent2", header.UID, backend22) - p.AddBackend("agent3", header.UID, backend3) - p.RemoveBackend("agent2", header.UID, backend22) - p.RemoveBackend("agent2", header.UID, backend2) - p.RemoveBackend("agent1", header.UID, backend1) - // This is invalid. agent1 doesn't have backend3. This should be a no-op. - p.RemoveBackend("agent1", header.UID, backend3) + p.AddBackend(backend12) + p.AddBackend(backend2) + p.AddBackend(backend22) + p.AddBackend(backend3) + p.RemoveBackend(backend22) + p.RemoveBackend(backend2) + p.RemoveBackend(backend1) expectedBackends = map[string][]Backend{ "agent1": {backend12}, "agent3": {backend3}, @@ -163,20 +160,20 @@ func TestAddRemoveBackendsWithDefaultStrategy(t *testing.T) { } } -func TestAddRemoveBackendsWithDefaultRouteStrategy(t *testing.T) { +func TestDefaultRouteBackendManager_AddRemoveBackends(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() - backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"default-route"})) - backend12, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"default-route"})) - backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route"})) - backend22, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route"})) - backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"default-route"})) + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"default-route=true"})) + backend12, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"default-route=true"})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route=true"})) + backend22, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route=true"})) + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"default-route=true"})) p := NewDefaultRouteBackendManager() - p.AddBackend("agent1", header.DefaultRoute, backend1) - p.RemoveBackend("agent1", header.DefaultRoute, backend1) + p.AddBackend(backend1) + p.RemoveBackend(backend1) expectedBackends := make(map[string][]Backend) expectedAgentIDs := []string{} expectedDefaultRouteAgentIDs := []string{} @@ -191,18 +188,16 @@ func TestAddRemoveBackendsWithDefaultRouteStrategy(t *testing.T) { } p = NewDefaultRouteBackendManager() - p.AddBackend("agent1", header.DefaultRoute, backend1) - p.AddBackend("agent1", header.DefaultRoute, backend12) + p.AddBackend(backend1) + p.AddBackend(backend12) // Adding the same connection again should be a no-op. - p.AddBackend("agent1", header.DefaultRoute, backend12) - p.AddBackend("agent2", header.DefaultRoute, backend2) - p.AddBackend("agent2", header.DefaultRoute, backend22) - p.AddBackend("agent3", header.DefaultRoute, backend3) - p.RemoveBackend("agent2", header.DefaultRoute, backend22) - p.RemoveBackend("agent2", header.DefaultRoute, backend2) - p.RemoveBackend("agent1", header.DefaultRoute, backend1) - // This is invalid. agent1 doesn't have backend3. This should be a no-op. - p.RemoveBackend("agent1", header.DefaultRoute, backend3) + p.AddBackend(backend12) + p.AddBackend(backend2) + p.AddBackend(backend22) + p.AddBackend(backend3) + p.RemoveBackend(backend22) + p.RemoveBackend(backend2) + p.RemoveBackend(backend1) expectedBackends = map[string][]Backend{ "agent1": {backend12}, @@ -221,3 +216,213 @@ func TestAddRemoveBackendsWithDefaultRouteStrategy(t *testing.T) { t.Errorf("expected %v, got %v", e, a) } } + +func TestDestHostBackendManager_AddRemoveBackends(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + // backend2 has no desthost relevant identifiers + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route=true"})) + // TODO: if backend3 is given conflicting identifiers with backend1, the wrong thing happens in RemoveBackend. + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"host=node2.mydomain.com&ipv4=5.6.7.8&ipv6=::"})) + + p := NewDestHostBackendManager() + + p.AddBackend(backend1) + p.RemoveBackend(backend1) + expectedBackends := make(map[string][]Backend) + expectedAgentIDs := []string{} + expectedDefaultRouteAgentIDs := []string(nil) + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + + p = NewDestHostBackendManager() + p.AddBackend(backend1) + + expectedBackends = map[string][]Backend{ + "localhost": {backend1}, + "1.2.3.4": {backend1}, + "9878::7675:1292:9183:7562": {backend1}, + "node1.mydomain.com": {backend1}, + } + expectedAgentIDs = []string{ + "1.2.3.4", + "9878::7675:1292:9183:7562", + "localhost", + "node1.mydomain.com", + } + + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + + p.AddBackend(backend2) + p.AddBackend(backend3) + + expectedBackends = map[string][]Backend{ + "localhost": {backend1}, + "node1.mydomain.com": {backend1}, + "node2.mydomain.com": {backend3}, + "1.2.3.4": {backend1}, + "5.6.7.8": {backend3}, + "9878::7675:1292:9183:7562": {backend1}, + "::": {backend3}, + } + expectedAgentIDs = []string{ + "1.2.3.4", + "9878::7675:1292:9183:7562", + "localhost", + "node1.mydomain.com", + "5.6.7.8", + "::", + "node2.mydomain.com", + } + expectedDefaultRouteAgentIDs = []string(nil) + + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + + p.RemoveBackend(backend2) + p.RemoveBackend(backend1) + + expectedBackends = map[string][]Backend{ + "node2.mydomain.com": {backend3}, + "5.6.7.8": {backend3}, + "::": {backend3}, + } + expectedAgentIDs = []string{ + "node2.mydomain.com", + "::", + "5.6.7.8", + } + + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + + p.RemoveBackend(backend3) + expectedBackends = map[string][]Backend{} + expectedAgentIDs = []string{} + + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } +} + +func TestDestHostBackendManager_WithDuplicateIdents(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"host=localhost&host=node2.mydomain.com&ipv4=5.6.7.8&ipv6=::"})) + + p := NewDestHostBackendManager() + + p.AddBackend(backend1) + p.AddBackend(backend2) + p.AddBackend(backend3) + + expectedBackends := map[string][]Backend{ + "localhost": {backend1, backend2, backend3}, + "1.2.3.4": {backend1, backend2}, + "5.6.7.8": {backend3}, + "9878::7675:1292:9183:7562": {backend1, backend2}, + "::": {backend3}, + "node1.mydomain.com": {backend1, backend2}, + "node2.mydomain.com": {backend3}, + } + expectedAgentIDs := []string{ + "1.2.3.4", + "9878::7675:1292:9183:7562", + "localhost", + "node1.mydomain.com", + "5.6.7.8", + "::", + "node2.mydomain.com", + } + expectedDefaultRouteAgentIDs := []string(nil) + + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + + p.RemoveBackend(backend1) + p.RemoveBackend(backend3) + + expectedBackends = map[string][]Backend{ + "localhost": {backend2}, + "1.2.3.4": {backend2}, + "9878::7675:1292:9183:7562": {backend2}, + "node1.mydomain.com": {backend2}, + } + expectedAgentIDs = []string{ + "1.2.3.4", + "9878::7675:1292:9183:7562", + "localhost", + "node1.mydomain.com", + } + + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + + p.RemoveBackend(backend2) + expectedBackends = map[string][]Backend{} + expectedAgentIDs = []string{} + + if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedAgentIDs, p.agentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } + if e, a := expectedDefaultRouteAgentIDs, p.defaultRouteAgentIDs; !reflect.DeepEqual(e, a) { + t.Errorf("expected %v, got %v", e, a) + } +} diff --git a/pkg/server/default_route_backend_manager.go b/pkg/server/default_route_backend_manager.go index cfbd3b3bc..fe67bf5dc 100644 --- a/pkg/server/default_route_backend_manager.go +++ b/pkg/server/default_route_backend_manager.go @@ -35,17 +35,25 @@ func NewDefaultRouteBackendManager() *DefaultRouteBackendManager { []header.IdentifierType{header.DefaultRoute})} } -// Backend tries to get a backend associating to the request destination host. +// Backend tries to get a backend that advertises default route, with random selection. func (dibm *DefaultRouteBackendManager) Backend(ctx context.Context) (Backend, error) { - dibm.mu.RLock() - defer dibm.mu.RUnlock() - if len(dibm.backends) == 0 { - return nil, &ErrNotFound{} + return dibm.GetRandomBackend() +} + +func (dibm *DefaultRouteBackendManager) AddBackend(backend Backend) { + agentID := backend.GetAgentID() + agentIdentifiers := backend.GetAgentIdentifiers() + if agentIdentifiers.DefaultRoute { + klog.V(5).InfoS("Add the agent to DefaultRouteBackendManager", "agentID", agentID) + dibm.addBackend(agentID, header.DefaultRoute, backend) } - if len(dibm.defaultRouteAgentIDs) == 0 { - return nil, &ErrNotFound{} +} + +func (dibm *DefaultRouteBackendManager) RemoveBackend(backend Backend) { + agentID := backend.GetAgentID() + agentIdentifiers := backend.GetAgentIdentifiers() + if agentIdentifiers.DefaultRoute { + klog.V(5).InfoS("Remove the agent from the DefaultRouteBackendManager", "agentID", agentID) + dibm.removeBackend(agentID, header.DefaultRoute, backend) } - agentID := dibm.defaultRouteAgentIDs[dibm.random.Intn(len(dibm.defaultRouteAgentIDs))] - klog.V(4).InfoS("Picked agent as backend", "agentID", agentID) - return dibm.backends[agentID][0], nil } diff --git a/pkg/server/desthost_backend_manager.go b/pkg/server/desthost_backend_manager.go index 304c8ccd0..2857659c5 100644 --- a/pkg/server/desthost_backend_manager.go +++ b/pkg/server/desthost_backend_manager.go @@ -35,6 +35,38 @@ func NewDestHostBackendManager() *DestHostBackendManager { []header.IdentifierType{header.IPv4, header.IPv6, header.Host})} } +func (dibm *DestHostBackendManager) AddBackend(backend Backend) { + agentIdentifiers := backend.GetAgentIdentifiers() + for _, ipv4 := range agentIdentifiers.IPv4 { + klog.V(5).InfoS("Add the agent to DestHostBackendManager", "agent address", ipv4) + dibm.addBackend(ipv4, header.IPv4, backend) + } + for _, ipv6 := range agentIdentifiers.IPv6 { + klog.V(5).InfoS("Add the agent to DestHostBackendManager", "agent address", ipv6) + dibm.addBackend(ipv6, header.IPv6, backend) + } + for _, host := range agentIdentifiers.Host { + klog.V(5).InfoS("Add the agent to DestHostBackendManager", "agent address", host) + dibm.addBackend(host, header.Host, backend) + } +} + +func (dibm *DestHostBackendManager) RemoveBackend(backend Backend) { + agentIdentifiers := backend.GetAgentIdentifiers() + for _, ipv4 := range agentIdentifiers.IPv4 { + klog.V(5).InfoS("Remove the agent from the DestHostBackendManager", "agentHost", ipv4) + dibm.removeBackend(ipv4, header.IPv4, backend) + } + for _, ipv6 := range agentIdentifiers.IPv6 { + klog.V(5).InfoS("Remove the agent from the DestHostBackendManager", "agentHost", ipv6) + dibm.removeBackend(ipv6, header.IPv6, backend) + } + for _, host := range agentIdentifiers.Host { + klog.V(5).InfoS("Remove the agent from the DestHostBackendManager", "agentHost", host) + dibm.removeBackend(host, header.Host, backend) + } +} + // Backend tries to get a backend associating to the request destination host. func (dibm *DestHostBackendManager) Backend(ctx context.Context) (Backend, error) { dibm.mu.RLock() @@ -42,7 +74,7 @@ func (dibm *DestHostBackendManager) Backend(ctx context.Context) (Backend, error if len(dibm.backends) == 0 { return nil, &ErrNotFound{} } - destHost := ctx.Value(destHost).(string) + destHost := ctx.Value(destHostKey).(string) if destHost != "" { bes, exist := dibm.backends[destHost] if exist && len(bes) > 0 { diff --git a/pkg/server/server.go b/pkg/server/server.go index d65d5c4be..94a76feec 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -107,7 +107,7 @@ type ProxyClientConnection struct { } const ( - destHost key = iota + destHostKey key = iota ) func (c *ProxyClientConnection) send(pkt *client.Packet) error { @@ -216,6 +216,7 @@ type ProxyServer struct { // agent authentication AgentAuthenticationOptions *AgentTokenAuthenticationOptions + // TODO: move strategies into BackendStorage proxyStrategies []ProxyStrategy } @@ -238,7 +239,7 @@ func genContext(proxyStrategies []ProxyStrategy, reqHost string) context.Context switch ps { case ProxyStrategyDestHost: addr := util.RemovePortFromHost(reqHost) - ctx = context.WithValue(ctx, destHost, addr) + ctx = context.WithValue(ctx, destHostKey, addr) } } return ctx @@ -261,65 +262,15 @@ func (s *ProxyServer) getBackend(reqHost string) (Backend, error) { } func (s *ProxyServer) addBackend(backend Backend) { - agentID := backend.GetAgentID() - for i := 0; i < len(s.BackendManagers); i++ { - switch s.BackendManagers[i].(type) { - case *DestHostBackendManager: - agentIdentifiers := backend.GetAgentIdentifiers() - for _, ipv4 := range agentIdentifiers.IPv4 { - klog.V(5).InfoS("Add the agent to DestHostBackendManager", "agent address", ipv4) - s.BackendManagers[i].AddBackend(ipv4, header.IPv4, backend) - } - for _, ipv6 := range agentIdentifiers.IPv6 { - klog.V(5).InfoS("Add the agent to DestHostBackendManager", "agent address", ipv6) - s.BackendManagers[i].AddBackend(ipv6, header.IPv6, backend) - } - for _, host := range agentIdentifiers.Host { - klog.V(5).InfoS("Add the agent to DestHostBackendManager", "agent address", host) - s.BackendManagers[i].AddBackend(host, header.Host, backend) - } - case *DefaultRouteBackendManager: - agentIdentifiers := backend.GetAgentIdentifiers() - if agentIdentifiers.DefaultRoute { - klog.V(5).InfoS("Add the agent to DefaultRouteBackendManager", "agentID", agentID) - s.BackendManagers[i].AddBackend(agentID, header.DefaultRoute, backend) - } - default: - klog.V(5).InfoS("Add the agent to DefaultBackendManager", "agentID", agentID) - s.BackendManagers[i].AddBackend(agentID, header.UID, backend) - } + // TODO: refactor BackendStorage to acquire lock once, not up to 3 times. + for _, bm := range s.BackendManagers { + bm.AddBackend(backend) } - return } func (s *ProxyServer) removeBackend(backend Backend) { - agentID := backend.GetAgentID() for _, bm := range s.BackendManagers { - switch bm.(type) { - case *DestHostBackendManager: - agentIdentifiers := backend.GetAgentIdentifiers() - for _, ipv4 := range agentIdentifiers.IPv4 { - klog.V(5).InfoS("Remove the agent from the DestHostBackendManager", "agentHost", ipv4) - bm.RemoveBackend(ipv4, header.IPv4, backend) - } - for _, ipv6 := range agentIdentifiers.IPv6 { - klog.V(5).InfoS("Remove the agent from the DestHostBackendManager", "agentHost", ipv6) - bm.RemoveBackend(ipv6, header.IPv6, backend) - } - for _, host := range agentIdentifiers.Host { - klog.V(5).InfoS("Remove the agent from the DestHostBackendManager", "agentHost", host) - bm.RemoveBackend(host, header.Host, backend) - } - case *DefaultRouteBackendManager: - agentIdentifiers := backend.GetAgentIdentifiers() - if agentIdentifiers.DefaultRoute { - klog.V(5).InfoS("Remove the agent from the DefaultRouteBackendManager", "agentID", agentID) - bm.RemoveBackend(agentID, header.DefaultRoute, backend) - } - default: - klog.V(5).InfoS("Remove the agent from the DefaultBackendManager", "agentID", agentID) - bm.RemoveBackend(agentID, header.UID, backend) - } + bm.RemoveBackend(backend) } } diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go index d12e0e01f..b96fdfe13 100644 --- a/pkg/server/server_test.go +++ b/pkg/server/server_test.go @@ -253,6 +253,260 @@ func TestAddRemoveFrontends(t *testing.T) { } } +func TestAddRemoveBackends_DefaultStrategy(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{})) + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{})) + + p := NewProxyServer("", []ProxyStrategy{ProxyStrategyDefault}, 1, nil) + + p.addBackend(backend1) + + if got, _ := p.getBackend("127.0.0.1"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + + p.addBackend(backend2) + p.addBackend(backend3) + p.removeBackend(backend1) + p.removeBackend(backend2) + + if got, _ := p.getBackend("127.0.0.1"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + + p.removeBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } +} + +func TestAddRemoveBackends_DefaultRouteStrategy(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route=false"})) + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"default-route=true"})) + + p := NewProxyServer("", []ProxyStrategy{ProxyStrategyDefaultRoute}, 1, nil) + + p.addBackend(backend1) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } + + p.addBackend(backend2) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } + + p.addBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + + p.removeBackend(backend1) + p.removeBackend(backend2) + + if got, _ := p.getBackend("127.0.0.1"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + + p.removeBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } +} + +func TestAddRemoveBackends_DestHostStrategy(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route=true"})) + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"host=node2.mydomain.com&ipv4=5.6.7.8&ipv6=::"})) + + p := NewProxyServer("", []ProxyStrategy{ProxyStrategyDestHost}, 1, nil) + + p.addBackend(backend1) + p.addBackend(backend2) + p.addBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } + if got, _ := p.getBackend("localhost"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("node1.mydomain.com"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("1.2.3.4"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("9878::7675:1292:9183:7562"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("node2.mydomain.com"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + if got, _ := p.getBackend("5.6.7.8"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + if got, _ := p.getBackend("::"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + + p.removeBackend(backend1) + p.removeBackend(backend2) + p.removeBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } +} + +func TestAddRemoveBackends_DestHostSanitizeRequest(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"host=node2.mydomain.com&ipv4=5.6.7.8&ipv6=::"})) + + p := NewProxyServer("", []ProxyStrategy{ProxyStrategyDestHost}, 1, nil) + + p.addBackend(backend1) + p.addBackend(backend2) + + if got, _ := p.getBackend("127.0.0.1:443"); got != nil { + t.Errorf("expected nil, got %v", got) + } + if got, _ := p.getBackend("node1.mydomain.com:443"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("node2.mydomain.com:443"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } +} + +func TestAddRemoveBackends_DestHostWithDefault(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"default-route=false"})) + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"host=node2.mydomain.com&ipv4=5.6.7.8&ipv6=::"})) + + p := NewProxyServer("", []ProxyStrategy{ProxyStrategyDestHost, ProxyStrategyDefault}, 1, nil) + + p.addBackend(backend1) + p.addBackend(backend2) + p.addBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got == nil { + t.Errorf("expected random fallback, got nil") + } + if got, _ := p.getBackend("localhost"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("node1.mydomain.com"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("1.2.3.4"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("9878::7675:1292:9183:7562"); got != backend1 { + t.Errorf("expected %v, got %v", backend1, got) + } + if got, _ := p.getBackend("node2.mydomain.com"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + if got, _ := p.getBackend("5.6.7.8"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + if got, _ := p.getBackend("::"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + + p.removeBackend(backend1) + p.removeBackend(backend2) + + if got, _ := p.getBackend("127.0.0.1"); got != backend3 { + t.Errorf("expected %v, got %v", backend3, got) + } + + p.removeBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } +} + +func TestAddRemoveBackends_DestHostWithDuplicateIdents(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backend1, _ := NewBackend(mockAgentConn(ctrl, "agent1", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + backend2, _ := NewBackend(mockAgentConn(ctrl, "agent2", []string{"host=localhost&host=node1.mydomain.com&ipv4=1.2.3.4&ipv6=9878::7675:1292:9183:7562"})) + backend3, _ := NewBackend(mockAgentConn(ctrl, "agent3", []string{"host=localhost&host=node2.mydomain.com&ipv4=5.6.7.8&ipv6=::"})) + + p := NewProxyServer("", []ProxyStrategy{ProxyStrategyDestHost, ProxyStrategyDefault}, 1, nil) + + p.addBackend(backend1) + p.addBackend(backend2) + p.addBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got == nil { + t.Errorf("expected random fallback, got nil") + } + if got, _ := p.getBackend("localhost"); got == nil { + t.Errorf("expected any backend, got nil") + } + + p.removeBackend(backend1) + p.removeBackend(backend3) + + if got, _ := p.getBackend("127.0.0.1"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + if got, _ := p.getBackend("localhost"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + if got, _ := p.getBackend("node1.mydomain.com"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + if got, _ := p.getBackend("1.2.3.4"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + if got, _ := p.getBackend("9878::7675:1292:9183:7562"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + if got, _ := p.getBackend("node2.mydomain.com"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + if got, _ := p.getBackend("5.6.7.8"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + if got, _ := p.getBackend("::"); got != backend2 { + t.Errorf("expected %v, got %v", backend2, got) + } + + p.removeBackend(backend2) + + if got, _ := p.getBackend("127.0.0.1"); got != nil { + t.Errorf("expected nil, got %v", got) + } +} + func TestEstablishedConnsMetric(t *testing.T) { metrics.Metrics.Reset() From d1fd7dc793495ae2ca6a61d0aa4eba7620943aef Mon Sep 17 00:00:00 2001 From: Joseph Anttila Hall Date: Fri, 29 Mar 2024 12:38:10 -0700 Subject: [PATCH 2/2] Makefile add 'fast-test' target for development. --- Makefile | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/Makefile b/Makefile index 69f4036ce..1e564ef9d 100644 --- a/Makefile +++ b/Makefile @@ -64,6 +64,13 @@ mock_gen: cat hack/go-license-header.txt proto/agent/mocks/agent_mock.go > proto/agent/mocks/agent_mock.licensed.go mv proto/agent/mocks/agent_mock.licensed.go proto/agent/mocks/agent_mock.go +# Unit tests with faster execution (nicer for development). +.PHONY: fast-test +fast-test: + go test -mod=vendor -race ./... + cd konnectivity-client && go test -race ./... + +# Unit tests with fuller coverage, invoked by CI system. .PHONY: test test: go test -mod=vendor -race -covermode=atomic -coverprofile=konnectivity.out ./... && go tool cover -html=konnectivity.out -o=konnectivity.html