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

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
 	}

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 {