Remove incorrect code and improve testing of Stack.GetMainNICAddress

This removes code that should have never made it in in the first place, but did so due to incomplete testing. With the new tests the original code fails, the new code passes.

PiperOrigin-RevId: 224086966
Change-Id: I646fef76977f4528f3705f497b95fad6b3ec32bc

2 files changed, 75 insertions, 11 deletions

diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index 6f3c24122..770d288cf 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -119,16 +119,6 @@ func (n *NIC) getMainNICAddress(protocol tcpip.NetworkProtocolNumber) (tcpip.Add
 
 	}
 
-	// If no primary endpoints then check for other endpoints.
-	if r == nil {
-		for _, ref := range n.endpoints {
-			if ref.holdsInsertRef && ref.tryIncRef() {
-				r = ref
-				break
-			}
-		}
-	}
-
 	if r == nil {
 		return "", tcpip.Subnet{}, tcpip.ErrNoLinkAddress
 	}
diff --git a/pkg/tcpip/stack/stack_test.go b/pkg/tcpip/stack/stack_test.go
index 74bf2c99e..8a9c9d234 100644
--- a/pkg/tcpip/stack/stack_test.go
+++ b/pkg/tcpip/stack/stack_test.go
@@ -18,6 +18,8 @@
 package stack_test
 
 import (
+	"bytes"
+	"fmt"
 	"math"
 	"strings"
 	"testing"
@@ -787,7 +789,79 @@ func TestSubnetAddRemove(t *testing.T) {
 	}
 }
 
-func TestGetMainNICAddress(t *testing.T) {
+func TestGetMainNICAddressAddPrimaryNonPrimary(t *testing.T) {
+	for _, addrLen := range []int{4, 16} {
+		t.Run(fmt.Sprintf("addrLen=%d", addrLen), func(t *testing.T) {
+			for canBe := 0; canBe < 3; canBe++ {
+				t.Run(fmt.Sprintf("canBe=%d", canBe), func(t *testing.T) {
+					for never := 0; never < 3; never++ {
+						t.Run(fmt.Sprintf("never=%d", never), func(t *testing.T) {
+							s := stack.New([]string{"fakeNet"}, nil, stack.Options{})
+							id, _ := channel.New(10, defaultMTU, "")
+							if err := s.CreateNIC(1, id); err != nil {
+								t.Fatalf("CreateNIC failed: %v", err)
+							}
+							// Insert <canBe> primary and <never> never-primary addresses.
+							// Each one will add a network endpoint to the NIC.
+							primaryAddrAdded := make(map[tcpip.Address]tcpip.Subnet)
+							for i := 0; i < canBe+never; i++ {
+								var behavior stack.PrimaryEndpointBehavior
+								if i < canBe {
+									behavior = stack.CanBePrimaryEndpoint
+								} else {
+									behavior = stack.NeverPrimaryEndpoint
+								}
+								// Add an address and in case of a primary one also add a
+								// subnet.
+								address := tcpip.Address(bytes.Repeat([]byte{byte(i)}, addrLen))
+								if err := s.AddAddressWithOptions(1, fakeNetNumber, address, behavior); err != nil {
+									t.Fatalf("AddAddressWithOptions failed: %v", err)
+								}
+								if behavior == stack.CanBePrimaryEndpoint {
+									mask := tcpip.AddressMask(strings.Repeat("\xff", len(address)))
+									subnet, err := tcpip.NewSubnet(address, mask)
+									if err != nil {
+										t.Fatalf("NewSubnet failed: %v", err)
+									}
+									if err := s.AddSubnet(1, fakeNetNumber, subnet); err != nil {
+										t.Fatalf("AddSubnet failed: %v", err)
+									}
+									// Remember the address/subnet.
+									primaryAddrAdded[address] = subnet
+								}
+							}
+							// Check that GetMainNICAddress returns an address if at least
+							// one primary address was added. In that case make sure the
+							// address/subnet matches what we added.
+							if len(primaryAddrAdded) == 0 {
+								// No primary addresses present, expect an error.
+								if _, _, err := s.GetMainNICAddress(1, fakeNetNumber); err != tcpip.ErrNoLinkAddress {
+									t.Fatalf("got s.GetMainNICAddress(...) = %v, wanted = %v", err, tcpip.ErrNoLinkAddress)
+								}
+							} else {
+								// At least one primary address was added, expect a valid
+								// address and subnet.
+								gotAddress, gotSubnet, err := s.GetMainNICAddress(1, fakeNetNumber)
+								if err != nil {
+									t.Fatalf("GetMainNICAddress failed: %v", err)
+								}
+								expectedSubnet, ok := primaryAddrAdded[gotAddress]
+								if !ok {
+									t.Fatalf("GetMainNICAddress: got address = %v, wanted any in {%v}", gotAddress, primaryAddrAdded)
+								}
+								if expectedSubnet != gotSubnet {
+									t.Fatalf("GetMainNICAddress: got subnet = %v, wanted %v", gotSubnet, expectedSubnet)
+								}
+							}
+						})
+					}
+				})
+			}
+		})
+	}
+}
+
+func TestGetMainNICAddressAddRemove(t *testing.T) {
 	s := stack.New([]string{"fakeNet"}, nil, stack.Options{})
 	id, _ := channel.New(10, defaultMTU, "")
 	if err := s.CreateNIC(1, id); err != nil {