Address code review feedbacks

Signed-off-by: Nino Kodabande <nkodabande@suse.com>

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)
 		}
 

src/go/guestagent/pkg/iptables/iptables.go

@@ -16,14 +16,15 @@ package iptables
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"strconv"
 	"strings"
 	"time"
 
 	"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))
 }

src/go/guestagent/pkg/iptables/iptables_test.go

@@ -21,25 +21,26 @@ 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"
 )
 
 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))
 }