From d9d51621337300ccdc2a42c049a9f9ec56916060 Mon Sep 17 00:00:00 2001 From: Nino Kodabande Date: Wed, 11 Dec 2024 10:06:44 -0800 Subject: [PATCH] Address code review feedbacks Signed-off-by: Nino Kodabande --- src/go/guestagent/main.go | 2 +- src/go/guestagent/pkg/iptables/iptables.go | 50 +++++-------- .../guestagent/pkg/iptables/iptables_test.go | 72 ++++++++++--------- 3 files changed, 58 insertions(+), 66 deletions(-) diff --git a/src/go/guestagent/main.go b/src/go/guestagent/main.go index 9c01f318863..278ba15cf83 100644 --- a/src/go/guestagent/main.go +++ b/src/go/guestagent/main.go @@ -177,7 +177,7 @@ func main() { if k8sServiceListenerIP == nil || !(k8sServiceListenerIP.Equal(net.IPv4zero) || k8sServiceListenerIP.Equal(net.IPv4(127, 0, 0, 1))) { - log.Fatalf("empty or none valid input for Kubernetes service listener IP address %s. "+ + log.Fatalf("empty or invalid input for Kubernetes service listener IP address %s. "+ "Valid options are 0.0.0.0 and 127.0.0.1.", *k8sServiceListenerAddr) } diff --git a/src/go/guestagent/pkg/iptables/iptables.go b/src/go/guestagent/pkg/iptables/iptables.go index 4c75f7bf283..598cc2806a8 100644 --- a/src/go/guestagent/pkg/iptables/iptables.go +++ b/src/go/guestagent/pkg/iptables/iptables.go @@ -16,6 +16,7 @@ package iptables import ( "context" + "fmt" "net" "strconv" "strings" @@ -23,7 +24,7 @@ import ( "github.com/Masterminds/log-go" "github.com/docker/go-connections/nat" - "github.com/lima-vm/lima/pkg/guestagent/iptables" + limaiptables "github.com/lima-vm/lima/pkg/guestagent/iptables" "github.com/rancher-sandbox/rancher-desktop/src/go/guestagent/pkg/tracker" "github.com/rancher-sandbox/rancher-desktop/src/go/guestagent/pkg/utils" ) @@ -34,21 +35,21 @@ import ( // the k8sServiceListenerAddr setting for the hostIP property to create a port mapping and // forwards them to both the API tracker and the WSL Proxy for proper routing and handling. type Iptables struct { - context context.Context - apiTracker tracker.Tracker - IptablesScanner Scanner - listenerIP net.IP + context context.Context + apiTracker tracker.Tracker + scanner Scanner + listenerIP net.IP // time, in seconds, to wait between updating. updateInterval time.Duration } func New(ctx context.Context, tracker tracker.Tracker, iptablesScanner Scanner, listenerIP net.IP, updateInterval time.Duration) *Iptables { return &Iptables{ - context: ctx, - apiTracker: tracker, - IptablesScanner: iptablesScanner, - listenerIP: listenerIP, - updateInterval: updateInterval, + context: ctx, + apiTracker: tracker, + scanner: iptablesScanner, + listenerIP: listenerIP, + updateInterval: updateInterval, } } @@ -58,7 +59,7 @@ func New(ctx context.Context, tracker tracker.Tracker, iptablesScanner Scanner, // as part of the normal forwarding system. This function detects those ports // and binds them to k8sServiceListenerAddr so that they are picked up. func (i *Iptables) ForwardPorts() error { - var ports []iptables.Entry + var ports []limaiptables.Entry ticker := time.NewTicker(i.updateInterval) defer ticker.Stop() @@ -70,7 +71,7 @@ func (i *Iptables) ForwardPorts() error { case <-ticker.C: } // Detect ports for forward - newPorts, err := i.IptablesScanner.GetPorts() + newPorts, err := i.scanner.GetPorts() if err != nil { // iptables exiting with an exit status of 4 means there // is a resource problem. For example, something else is @@ -113,11 +114,7 @@ func (i *Iptables) ForwardPorts() error { HostIP: i.listenerIP.String(), HostPort: port, } - if pb, ok := portMap[portMapKey]; ok { - if !portExist(port, pb) { - portMap[portMapKey] = append(pb, portBinding) - } - } else { + if _, ok := portMap[portMapKey]; !ok { portMap[portMapKey] = []nat.PortBinding{portBinding} } name := entryToString(p) @@ -131,26 +128,13 @@ func (i *Iptables) ForwardPorts() error { } } -// portExist checks if the given port is already present in the list of port bindings. -// Since we always use the k8sServiceListenerAddr for the HostIP, the actual IP -// returned by GetPorts is irrelevant, and we only care about whether the HostPort is -// already mapped. This avoids adding duplicate entries to the nat.PortMap. -func portExist(port string, portBindings []nat.PortBinding) bool { - for _, p := range portBindings { - if port == p.HostPort { - return true - } - } - return false -} - // comparePorts compares the old and new ports to find those added or removed. // This function is mostly lifted from lima (github.com/lima-vm/lima) which is // licensed under the Apache 2. // //nolint:nonamedreturns -func comparePorts(oldPorts, newPorts []iptables.Entry) (added, removed []iptables.Entry) { - oldPortMap := make(map[string]iptables.Entry, len(oldPorts)) +func comparePorts(oldPorts, newPorts []limaiptables.Entry) (added, removed []limaiptables.Entry) { + oldPortMap := make(map[string]limaiptables.Entry, len(oldPorts)) portExistMap := make(map[string]bool, len(oldPorts)) for _, oldPort := range oldPorts { key := entryToString(oldPort) @@ -174,6 +158,6 @@ func comparePorts(oldPorts, newPorts []iptables.Entry) (added, removed []iptable return } -func entryToString(ip iptables.Entry) string { +func entryToString(ip limaiptables.Entry) string { return net.JoinHostPort(ip.IP.String(), strconv.Itoa(ip.Port)) } diff --git a/src/go/guestagent/pkg/iptables/iptables_test.go b/src/go/guestagent/pkg/iptables/iptables_test.go index 31a2aeb3a4a..d5c0489a799 100644 --- a/src/go/guestagent/pkg/iptables/iptables_test.go +++ b/src/go/guestagent/pkg/iptables/iptables_test.go @@ -21,7 +21,7 @@ import ( "time" "github.com/docker/go-connections/nat" - limaIptables "github.com/lima-vm/lima/pkg/guestagent/iptables" + limaiptables "github.com/lima-vm/lima/pkg/guestagent/iptables" "github.com/rancher-sandbox/rancher-desktop/src/go/guestagent/pkg/iptables" "github.com/rancher-sandbox/rancher-desktop/src/go/guestagent/pkg/utils" "github.com/stretchr/testify/require" @@ -29,17 +29,18 @@ import ( func TestForwardPorts(t *testing.T) { tests := []struct { - name string - remove bool - listenerIP net.IP - expectedEntries []limaIptables.Entry - expectedEntriesAfterRemove []limaIptables.Entry - expectedAddFuncErr error + name string + remove bool + listenerIP net.IP + expectedEntries []limaiptables.Entry + removedEntries []limaiptables.Entry + updateEntries []limaiptables.Entry + expectedAddFuncErr error }{ { name: "With localhost listener and valid port mappings", listenerIP: net.IPv4(127, 0, 0, 1), - expectedEntries: []limaIptables.Entry{ + expectedEntries: []limaiptables.Entry{ {TCP: true, IP: net.IPv4(192, 168, 20, 10), Port: 1080}, {TCP: true, IP: net.IPv4(192, 168, 20, 11), Port: 1081}, {TCP: true, IP: net.IPv4(192, 168, 20, 12), Port: 1082}, @@ -48,7 +49,7 @@ func TestForwardPorts(t *testing.T) { { name: "With wildcard listener and valid port mappings", listenerIP: net.IPv4(0, 0, 0, 0), - expectedEntries: []limaIptables.Entry{ + expectedEntries: []limaiptables.Entry{ {TCP: true, IP: net.IPv4(192, 168, 21, 10), Port: 1080}, {TCP: true, IP: net.IPv4(192, 168, 21, 11), Port: 1081}, {TCP: true, IP: net.IPv4(192, 168, 21, 12), Port: 1082}, @@ -58,13 +59,23 @@ func TestForwardPorts(t *testing.T) { name: "With entries removed", remove: true, listenerIP: net.IPv4(0, 0, 0, 0), - expectedEntries: []limaIptables.Entry{ + expectedEntries: []limaiptables.Entry{ {TCP: true, IP: net.IPv4(192, 168, 22, 10), Port: 1080}, {TCP: true, IP: net.IPv4(192, 168, 22, 11), Port: 1081}, {TCP: true, IP: net.IPv4(192, 168, 22, 12), Port: 1082}, {TCP: true, IP: net.IPv4(192, 168, 22, 13), Port: 1083}, {TCP: true, IP: net.IPv4(192, 168, 22, 14), Port: 1084}, }, + removedEntries: []limaiptables.Entry{ + {TCP: true, IP: net.IPv4(192, 168, 22, 11), Port: 1081}, + {TCP: true, IP: net.IPv4(192, 168, 22, 12), Port: 1082}, + }, + updateEntries: []limaiptables.Entry{ + {TCP: true, IP: net.IPv4(192, 168, 22, 10), Port: 1080}, + {TCP: true, IP: net.IPv4(192, 168, 22, 13), Port: 1083}, + {TCP: true, IP: net.IPv4(192, 168, 22, 14), Port: 1084}, + {TCP: true, IP: net.IPv4(192, 168, 22, 15), Port: 1085}, + }, }, } @@ -109,26 +120,23 @@ func TestForwardPorts(t *testing.T) { } if tt.remove { - removedEntries := []limaIptables.Entry{ - {TCP: true, IP: net.IPv4(192, 168, 22, 11), Port: 1081}, - {TCP: true, IP: net.IPv4(192, 168, 22, 12), Port: 1082}, - } - // update the entries - updatedEntries := []limaIptables.Entry{ - {TCP: true, IP: net.IPv4(192, 168, 22, 10), Port: 1080}, - {TCP: true, IP: net.IPv4(192, 168, 22, 13), Port: 1083}, - {TCP: true, IP: net.IPv4(192, 168, 22, 14), Port: 1084}, - {TCP: true, IP: net.IPv4(192, 168, 22, 15), Port: 1085}, + iptablesScanner.expectedEntries = tt.updateEntries + + // Collect all removed IDs. + var actualRemovedIDs []string + for i := 0; i < len(tt.removedEntries); i++ { + select { + case id := <-testTracker.receivedRemoveID: + actualRemovedIDs = append(actualRemovedIDs, id) + case <-time.After(5 * time.Second): + t.Fatalf("Timeout waiting for remove ID for entry %v", tt.removedEntries[i]) + } } - iptablesScanner.expectedEntries = updatedEntries - for i := 0; i < len(removedEntries); i++ { - id := <-testTracker.receivedRemoveID - expectedID := utils.GenerateID(entryToString(removedEntries[i])) - require.Equal(t, expectedID, id) + for _, removedEntry := range tt.removedEntries { + require.Contains(t, actualRemovedIDs, utils.GenerateID(entryToString(removedEntry))) } - - addedElement := updatedEntries[len(updatedEntries)-1] + addedElement := tt.updateEntries[len(tt.updateEntries)-1] id := <-testTracker.receivedID expectedID := utils.GenerateID(entryToString(addedElement)) require.Equal(t, expectedID, id) @@ -156,13 +164,13 @@ func TestForwardPortsSamePortDifferentIP(t *testing.T) { tests := []struct { name string listenerIP net.IP - expectedEntries []limaIptables.Entry + expectedEntries []limaiptables.Entry expectedAddFuncErr error }{ { name: "Same Port with different IP", listenerIP: net.IPv4(0, 0, 0, 0), - expectedEntries: []limaIptables.Entry{ + expectedEntries: []limaiptables.Entry{ {TCP: true, IP: net.IPv4(192, 168, 22, 10), Port: 1080}, {TCP: true, IP: net.IPv4(192, 168, 22, 11), Port: 1081}, {TCP: true, IP: net.IPv4(192, 168, 22, 12), Port: 1082}, @@ -253,15 +261,15 @@ func (f *fakeTracker) RemoveAll() error { // Fake Scanner to simulate iptables entries type fakeScanner struct { - expectedEntries []limaIptables.Entry + expectedEntries []limaiptables.Entry expectedErr error } -func (f *fakeScanner) GetPorts() ([]limaIptables.Entry, error) { +func (f *fakeScanner) GetPorts() ([]limaiptables.Entry, error) { return f.expectedEntries, f.expectedErr } // Utility function to convert iptables entry to string -func entryToString(ip limaIptables.Entry) string { +func entryToString(ip limaiptables.Entry) string { return net.JoinHostPort(ip.IP.String(), strconv.Itoa(ip.Port)) }