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 {