Pass ProtocolAddress instead of its fields

PiperOrigin-RevId: 260803517

3 files changed, 91 insertions, 54 deletions

diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index 3e6ff4afb..dc28dc970 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -209,8 +209,13 @@ func (n *NIC) findEndpoint(protocol tcpip.NetworkProtocolNumber, address tcpip.A
 	ref = n.endpoints[id]
 	if ref == nil || !ref.tryIncRef() {
 		if netProto, ok := n.stack.networkProtocols[protocol]; ok {
-			addrWithPrefix := tcpip.AddressWithPrefix{address, netProto.DefaultPrefixLen()}
-			ref, _ = n.addAddressLocked(protocol, addrWithPrefix, peb, true)
+			ref, _ = n.addAddressLocked(tcpip.ProtocolAddress{
+				Protocol: protocol,
+				AddressWithPrefix: tcpip.AddressWithPrefix{
+					Address:   address,
+					PrefixLen: netProto.DefaultPrefixLen(),
+				},
+			}, peb, true)
 			if ref != nil {
 				ref.holdsInsertRef = false
 			}
@@ -220,14 +225,14 @@ func (n *NIC) findEndpoint(protocol tcpip.NetworkProtocolNumber, address tcpip.A
 	return ref
 }
 
-func (n *NIC) addAddressLocked(protocol tcpip.NetworkProtocolNumber, addrWithPrefix tcpip.AddressWithPrefix, peb PrimaryEndpointBehavior, replace bool) (*referencedNetworkEndpoint, *tcpip.Error) {
-	netProto, ok := n.stack.networkProtocols[protocol]
+func (n *NIC) addAddressLocked(protocolAddress tcpip.ProtocolAddress, peb PrimaryEndpointBehavior, replace bool) (*referencedNetworkEndpoint, *tcpip.Error) {
+	netProto, ok := n.stack.networkProtocols[protocolAddress.Protocol]
 	if !ok {
 		return nil, tcpip.ErrUnknownProtocol
 	}
 
 	// Create the new network endpoint.
-	ep, err := netProto.NewEndpoint(n.id, addrWithPrefix, n.stack, n, n.linkEP)
+	ep, err := netProto.NewEndpoint(n.id, protocolAddress.AddressWithPrefix, n.stack, n, n.linkEP)
 	if err != nil {
 		return nil, err
 	}
@@ -245,23 +250,23 @@ func (n *NIC) addAddressLocked(protocol tcpip.NetworkProtocolNumber, addrWithPre
 		refs:           1,
 		ep:             ep,
 		nic:            n,
-		protocol:       protocol,
+		protocol:       protocolAddress.Protocol,
 		holdsInsertRef: true,
 	}
 
 	// Set up cache if link address resolution exists for this protocol.
 	if n.linkEP.Capabilities()&CapabilityResolutionRequired != 0 {
-		if _, ok := n.stack.linkAddrResolvers[protocol]; ok {
+		if _, ok := n.stack.linkAddrResolvers[protocolAddress.Protocol]; ok {
 			ref.linkCache = n.stack
 		}
 	}
 
 	n.endpoints[id] = ref
 
-	l, ok := n.primary[protocol]
+	l, ok := n.primary[protocolAddress.Protocol]
 	if !ok {
 		l = &ilist.List{}
-		n.primary[protocol] = l
+		n.primary[protocolAddress.Protocol] = l
 	}
 
 	switch peb {
@@ -276,10 +281,10 @@ func (n *NIC) addAddressLocked(protocol tcpip.NetworkProtocolNumber, addrWithPre
 
 // AddAddress adds a new address to n, so that it starts accepting packets
 // targeted at the given address (and network protocol).
-func (n *NIC) AddAddress(protocol tcpip.NetworkProtocolNumber, addrWithPrefix tcpip.AddressWithPrefix, peb PrimaryEndpointBehavior) *tcpip.Error {
+func (n *NIC) AddAddress(protocolAddress tcpip.ProtocolAddress, peb PrimaryEndpointBehavior) *tcpip.Error {
 	// Add the endpoint.
 	n.mu.Lock()
-	_, err := n.addAddressLocked(protocol, addrWithPrefix, peb, false)
+	_, err := n.addAddressLocked(protocolAddress, peb, false)
 	n.mu.Unlock()
 
 	return err
@@ -414,8 +419,13 @@ func (n *NIC) joinGroup(protocol tcpip.NetworkProtocolNumber, addr tcpip.Address
 		if !ok {
 			return tcpip.ErrUnknownProtocol
 		}
-		addrWithPrefix := tcpip.AddressWithPrefix{addr, netProto.DefaultPrefixLen()}
-		if _, err := n.addAddressLocked(protocol, addrWithPrefix, NeverPrimaryEndpoint, false); err != nil {
+		if _, err := n.addAddressLocked(tcpip.ProtocolAddress{
+			Protocol: protocol,
+			AddressWithPrefix: tcpip.AddressWithPrefix{
+				Address:   addr,
+				PrefixLen: netProto.DefaultPrefixLen(),
+			},
+		}, NeverPrimaryEndpoint, false); err != nil {
 			return err
 		}
 	}
@@ -577,8 +587,13 @@ func (n *NIC) getRef(protocol tcpip.NetworkProtocolNumber, dst tcpip.Address) *r
 			n.mu.Unlock()
 			return nil
 		}
-		addrWithPrefix := tcpip.AddressWithPrefix{dst, netProto.DefaultPrefixLen()}
-		ref, err := n.addAddressLocked(protocol, addrWithPrefix, CanBePrimaryEndpoint, true)
+		ref, err := n.addAddressLocked(tcpip.ProtocolAddress{
+			Protocol: protocol,
+			AddressWithPrefix: tcpip.AddressWithPrefix{
+				Address:   dst,
+				PrefixLen: netProto.DefaultPrefixLen(),
+			},
+		}, CanBePrimaryEndpoint, true)
 		n.mu.Unlock()
 		if err == nil {
 			ref.holdsInsertRef = false
diff --git a/pkg/tcpip/stack/stack.go b/pkg/tcpip/stack/stack.go
index 57b8a9994..6156c3f46 100644
--- a/pkg/tcpip/stack/stack.go
+++ b/pkg/tcpip/stack/stack.go
@@ -751,10 +751,10 @@ func (s *Stack) AddAddress(id tcpip.NICID, protocol tcpip.NetworkProtocolNumber,
 	return s.AddAddressWithOptions(id, protocol, addr, CanBePrimaryEndpoint)
 }
 
-// AddAddressWithPrefix adds a new network-layer address/prefixLen to the
+// AddProtocolAddress adds a new network-layer protocol address to the
 // specified NIC.
-func (s *Stack) AddAddressWithPrefix(id tcpip.NICID, protocol tcpip.NetworkProtocolNumber, addrWithPrefix tcpip.AddressWithPrefix) *tcpip.Error {
-	return s.AddAddressWithPrefixAndOptions(id, protocol, addrWithPrefix, CanBePrimaryEndpoint)
+func (s *Stack) AddProtocolAddress(id tcpip.NICID, protocolAddress tcpip.ProtocolAddress) *tcpip.Error {
+	return s.AddProtocolAddressWithOptions(id, protocolAddress, CanBePrimaryEndpoint)
 }
 
 // AddAddressWithOptions is the same as AddAddress, but allows you to specify
@@ -764,13 +764,18 @@ func (s *Stack) AddAddressWithOptions(id tcpip.NICID, protocol tcpip.NetworkProt
 	if !ok {
 		return tcpip.ErrUnknownProtocol
 	}
-	addrWithPrefix := tcpip.AddressWithPrefix{addr, netProto.DefaultPrefixLen()}
-	return s.AddAddressWithPrefixAndOptions(id, protocol, addrWithPrefix, peb)
-}
-
-// AddAddressWithPrefixAndOptions is the same as AddAddressWithPrefixLen,
-// but allows you to specify whether the new endpoint can be primary or not.
-func (s *Stack) AddAddressWithPrefixAndOptions(id tcpip.NICID, protocol tcpip.NetworkProtocolNumber, addrWithPrefix tcpip.AddressWithPrefix, peb PrimaryEndpointBehavior) *tcpip.Error {
+	return s.AddProtocolAddressWithOptions(id, tcpip.ProtocolAddress{
+		Protocol: protocol,
+		AddressWithPrefix: tcpip.AddressWithPrefix{
+			Address:   addr,
+			PrefixLen: netProto.DefaultPrefixLen(),
+		},
+	}, peb)
+}
+
+// AddProtocolAddressWithOptions is the same as AddProtocolAddress, but allows
+// you to specify whether the new endpoint can be primary or not.
+func (s *Stack) AddProtocolAddressWithOptions(id tcpip.NICID, protocolAddress tcpip.ProtocolAddress, peb PrimaryEndpointBehavior) *tcpip.Error {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
@@ -779,7 +784,7 @@ func (s *Stack) AddAddressWithPrefixAndOptions(id tcpip.NICID, protocol tcpip.Ne
 		return tcpip.ErrUnknownNICID
 	}
 
-	return nic.AddAddress(protocol, addrWithPrefix, peb)
+	return nic.AddAddress(protocolAddress, peb)
 }
 
 // AddSubnet adds a subnet range to the specified NIC.
diff --git a/pkg/tcpip/stack/stack_test.go b/pkg/tcpip/stack/stack_test.go
index 9d082bba4..1ab9c575b 100644
--- a/pkg/tcpip/stack/stack_test.go
+++ b/pkg/tcpip/stack/stack_test.go
@@ -188,7 +188,7 @@ func (*fakeNetworkProtocol) ParseAddresses(v buffer.View) (src, dst tcpip.Addres
 func (f *fakeNetworkProtocol) NewEndpoint(nicid tcpip.NICID, addrWithPrefix tcpip.AddressWithPrefix, linkAddrCache stack.LinkAddressCache, dispatcher stack.TransportDispatcher, linkEP stack.LinkEndpoint) (stack.NetworkEndpoint, *tcpip.Error) {
 	return &fakeNetworkEndpoint{
 		nicid:      nicid,
-		id:         stack.NetworkEndpointID{addrWithPrefix.Address},
+		id:         stack.NetworkEndpointID{LocalAddress: addrWithPrefix.Address},
 		prefixLen:  addrWithPrefix.PrefixLen,
 		proto:      f,
 		dispatcher: dispatcher,
@@ -969,12 +969,18 @@ func TestGetMainNICAddressAddPrimaryNonPrimary(t *testing.T) {
 								// prefixLen.
 								address := tcpip.Address(bytes.Repeat([]byte{byte(i)}, addrLen))
 								if behavior == stack.CanBePrimaryEndpoint {
-									addressWithPrefix := tcpip.AddressWithPrefix{address, addrLen * 8}
-									if err := s.AddAddressWithPrefixAndOptions(1, fakeNetNumber, addressWithPrefix, behavior); err != nil {
-										t.Fatal("AddAddressWithPrefixAndOptions failed:", err)
+									protocolAddress := tcpip.ProtocolAddress{
+										Protocol: fakeNetNumber,
+										AddressWithPrefix: tcpip.AddressWithPrefix{
+											Address:   address,
+											PrefixLen: addrLen * 8,
+										},
+									}
+									if err := s.AddProtocolAddressWithOptions(1, protocolAddress, behavior); err != nil {
+										t.Fatal("AddProtocolAddressWithOptions failed:", err)
 									}
 									// Remember the address/prefix.
-									primaryAddrAdded[addressWithPrefix] = struct{}{}
+									primaryAddrAdded[protocolAddress.AddressWithPrefix] = struct{}{}
 								} else {
 									if err := s.AddAddressWithOptions(1, fakeNetNumber, address, behavior); err != nil {
 										t.Fatal("AddAddressWithOptions failed:", err)
@@ -1024,20 +1030,25 @@ func TestGetMainNICAddressAddRemove(t *testing.T) {
 		{"IPv6", "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01", 116},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			addressWithPrefix := tcpip.AddressWithPrefix{tc.address, tc.prefixLen}
-
-			if err := s.AddAddressWithPrefix(1, fakeNetNumber, addressWithPrefix); err != nil {
-				t.Fatal("AddAddressWithPrefix failed:", err)
+			protocolAddress := tcpip.ProtocolAddress{
+				Protocol: fakeNetNumber,
+				AddressWithPrefix: tcpip.AddressWithPrefix{
+					Address:   tc.address,
+					PrefixLen: tc.prefixLen,
+				},
+			}
+			if err := s.AddProtocolAddress(1, protocolAddress); err != nil {
+				t.Fatal("AddProtocolAddress failed:", err)
 			}
 
 			// Check that we get the right initial address and prefix length.
 			if gotAddressWithPrefix, err := s.GetMainNICAddress(1, fakeNetNumber); err != nil {
 				t.Fatal("GetMainNICAddress failed:", err)
-			} else if gotAddressWithPrefix != addressWithPrefix {
-				t.Fatalf("got GetMainNICAddress = %+v, want = %+v", gotAddressWithPrefix, addressWithPrefix)
+			} else if gotAddressWithPrefix != protocolAddress.AddressWithPrefix {
+				t.Fatalf("got GetMainNICAddress = %+v, want = %+v", gotAddressWithPrefix, protocolAddress.AddressWithPrefix)
 			}
 
-			if err := s.RemoveAddress(1, addressWithPrefix.Address); err != nil {
+			if err := s.RemoveAddress(1, protocolAddress.AddressWithPrefix.Address); err != nil {
 				t.Fatal("RemoveAddress failed:", err)
 			}
 
@@ -1102,7 +1113,7 @@ func TestAddAddress(t *testing.T) {
 	verifyAddresses(t, expectedAddresses, gotAddresses)
 }
 
-func TestAddAddressWithPrefix(t *testing.T) {
+func TestAddProtocolAddress(t *testing.T) {
 	const nicid = 1
 	s := stack.New([]string{"fakeNet"}, nil, stack.Options{})
 	id, _ := channel.New(10, defaultMTU, "")
@@ -1116,14 +1127,17 @@ func TestAddAddressWithPrefix(t *testing.T) {
 	expectedAddresses := make([]tcpip.ProtocolAddress, 0, len(addrLenRange)*len(prefixLenRange))
 	for _, addrLen := range addrLenRange {
 		for _, prefixLen := range prefixLenRange {
-			address := addrGen.next(addrLen)
-			if err := s.AddAddressWithPrefix(nicid, fakeNetNumber, tcpip.AddressWithPrefix{address, prefixLen}); err != nil {
-				t.Errorf("AddAddressWithPrefix(address=%s, prefixLen=%d) failed: %s", address, prefixLen, err)
+			protocolAddress := tcpip.ProtocolAddress{
+				Protocol: fakeNetNumber,
+				AddressWithPrefix: tcpip.AddressWithPrefix{
+					Address:   addrGen.next(addrLen),
+					PrefixLen: prefixLen,
+				},
 			}
-			expectedAddresses = append(expectedAddresses, tcpip.ProtocolAddress{
-				Protocol:          fakeNetNumber,
-				AddressWithPrefix: tcpip.AddressWithPrefix{address, prefixLen},
-			})
+			if err := s.AddProtocolAddress(nicid, protocolAddress); err != nil {
+				t.Errorf("AddProtocolAddress(%+v) failed: %s", protocolAddress, err)
+			}
+			expectedAddresses = append(expectedAddresses, protocolAddress)
 		}
 	}
 
@@ -1160,7 +1174,7 @@ func TestAddAddressWithOptions(t *testing.T) {
 	verifyAddresses(t, expectedAddresses, gotAddresses)
 }
 
-func TestAddAddressWithPrefixAndOptions(t *testing.T) {
+func TestAddProtocolAddressWithOptions(t *testing.T) {
 	const nicid = 1
 	s := stack.New([]string{"fakeNet"}, nil, stack.Options{})
 	id, _ := channel.New(10, defaultMTU, "")
@@ -1176,14 +1190,17 @@ func TestAddAddressWithPrefixAndOptions(t *testing.T) {
 	for _, addrLen := range addrLenRange {
 		for _, prefixLen := range prefixLenRange {
 			for _, behavior := range behaviorRange {
-				address := addrGen.next(addrLen)
-				if err := s.AddAddressWithPrefixAndOptions(nicid, fakeNetNumber, tcpip.AddressWithPrefix{address, prefixLen}, behavior); err != nil {
-					t.Fatalf("AddAddressWithPrefixAndOptions(address=%s, prefixLen=%d, behavior=%d) failed: %s", address, prefixLen, behavior, err)
+				protocolAddress := tcpip.ProtocolAddress{
+					Protocol: fakeNetNumber,
+					AddressWithPrefix: tcpip.AddressWithPrefix{
+						Address:   addrGen.next(addrLen),
+						PrefixLen: prefixLen,
+					},
 				}
-				expectedAddresses = append(expectedAddresses, tcpip.ProtocolAddress{
-					Protocol:          fakeNetNumber,
-					AddressWithPrefix: tcpip.AddressWithPrefix{address, prefixLen},
-				})
+				if err := s.AddProtocolAddressWithOptions(nicid, protocolAddress, behavior); err != nil {
+					t.Fatalf("AddProtocolAddressWithOptions(%+v, %d) failed: %s", protocolAddress, behavior, err)
+				}
+				expectedAddresses = append(expectedAddresses, protocolAddress)
 			}
 		}
 	}