Implement LinkAddressResolver on NetworkEndpoints

This removes the need to provide the link address request with the NIC
the request is being performed on since the NetworkEndpoints already
have a reference to the NIC.

PiperOrigin-RevId: 354721940

7 files changed, 60 insertions, 191 deletions

diff --git a/pkg/tcpip/network/arp/BUILD b/pkg/tcpip/network/arp/BUILD
index c7ab876bf..933845269 100644
--- a/pkg/tcpip/network/arp/BUILD
+++ b/pkg/tcpip/network/arp/BUILD
@@ -10,7 +10,6 @@ go_library(
     ],
     visibility = ["//visibility:public"],
     deps = [
-        "//pkg/sync",
         "//pkg/tcpip",
         "//pkg/tcpip/buffer",
         "//pkg/tcpip/header",
diff --git a/pkg/tcpip/network/arp/arp.go b/pkg/tcpip/network/arp/arp.go
index 7838cc753..5c79d6485 100644
--- a/pkg/tcpip/network/arp/arp.go
+++ b/pkg/tcpip/network/arp/arp.go
@@ -22,7 +22,6 @@ import (
 	"reflect"
 	"sync/atomic"
 
-	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
@@ -35,6 +34,8 @@ const (
 	ProtocolNumber = header.ARPProtocolNumber
 )
 
+var _ stack.LinkAddressResolver = (*endpoint)(nil)
+
 // ARP endpoints need to implement stack.NetworkEndpoint because the stack
 // considers the layer above the link-layer a network layer; the only
 // facility provided by the stack to deliver packets to a layer above
@@ -101,9 +102,7 @@ func (e *endpoint) MaxHeaderLength() uint16 {
 	return e.nic.MaxHeaderLength() + header.ARPSize
 }
 
-func (e *endpoint) Close() {
-	e.protocol.forgetEndpoint(e.nic.ID())
-}
+func (*endpoint) Close() {}
 
 func (*endpoint) WritePacket(*stack.Route, *stack.GSO, stack.NetworkHeaderParams, *stack.PacketBuffer) tcpip.Error {
 	return &tcpip.ErrNotSupported{}
@@ -154,7 +153,7 @@ func (e *endpoint) HandlePacket(pkt *stack.PacketBuffer) {
 		if e.nud == nil {
 			e.linkAddrCache.AddLinkAddress(remoteAddr, remoteLinkAddr)
 		} else {
-			e.nud.HandleProbe(remoteAddr, ProtocolNumber, remoteLinkAddr, e.protocol)
+			e.nud.HandleProbe(remoteAddr, ProtocolNumber, remoteLinkAddr, e)
 		}
 
 		respPkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
@@ -221,19 +220,10 @@ func (e *endpoint) Stats() stack.NetworkEndpointStats {
 }
 
 var _ stack.NetworkProtocol = (*protocol)(nil)
-var _ stack.LinkAddressResolver = (*protocol)(nil)
 
 // protocol implements stack.NetworkProtocol and stack.LinkAddressResolver.
 type protocol struct {
 	stack *stack.Stack
-
-	mu struct {
-		sync.RWMutex
-
-		// eps is keyed by NICID to allow protocol methods to retrieve the correct
-		// endpoint depending on the NIC.
-		eps map[tcpip.NICID]*endpoint
-	}
 }
 
 func (p *protocol) Number() tcpip.NetworkProtocolNumber { return ProtocolNumber }
@@ -257,43 +247,26 @@ func (p *protocol) NewEndpoint(nic stack.NetworkInterface, linkAddrCache stack.L
 	stackStats := p.stack.Stats()
 	e.stats.arp.init(&e.stats.localStats.ARP, &stackStats.ARP)
 
-	p.mu.Lock()
-	p.mu.eps[nic.ID()] = e
-	p.mu.Unlock()
-
 	return e
 }
 
-func (p *protocol) forgetEndpoint(nicID tcpip.NICID) {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	delete(p.mu.eps, nicID)
-}
-
 // LinkAddressProtocol implements stack.LinkAddressResolver.LinkAddressProtocol.
-func (*protocol) LinkAddressProtocol() tcpip.NetworkProtocolNumber {
+func (*endpoint) LinkAddressProtocol() tcpip.NetworkProtocolNumber {
 	return header.IPv4ProtocolNumber
 }
 
 // LinkAddressRequest implements stack.LinkAddressResolver.LinkAddressRequest.
-func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remoteLinkAddr tcpip.LinkAddress, nic stack.NetworkInterface) tcpip.Error {
-	nicID := nic.ID()
-
-	p.mu.Lock()
-	netEP, ok := p.mu.eps[nicID]
-	p.mu.Unlock()
-	if !ok {
-		return &tcpip.ErrNotConnected{}
-	}
+func (e *endpoint) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remoteLinkAddr tcpip.LinkAddress) tcpip.Error {
+	nicID := e.nic.ID()
 
-	stats := netEP.stats.arp
+	stats := e.stats.arp
 
 	if len(remoteLinkAddr) == 0 {
 		remoteLinkAddr = header.EthernetBroadcastAddress
 	}
 
 	if len(localAddr) == 0 {
-		addr, ok := p.stack.GetMainNICAddress(nicID, header.IPv4ProtocolNumber)
+		addr, ok := e.protocol.stack.GetMainNICAddress(nicID, header.IPv4ProtocolNumber)
 		if !ok {
 			return &tcpip.ErrUnknownNICID{}
 		}
@@ -304,13 +277,13 @@ func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remot
 		}
 
 		localAddr = addr.Address
-	} else if p.stack.CheckLocalAddress(nicID, header.IPv4ProtocolNumber, localAddr) == 0 {
+	} else if e.protocol.stack.CheckLocalAddress(nicID, header.IPv4ProtocolNumber, localAddr) == 0 {
 		stats.outgoingRequestBadLocalAddressErrors.Increment()
 		return &tcpip.ErrBadLocalAddress{}
 	}
 
 	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		ReserveHeaderBytes: int(nic.MaxHeaderLength()) + header.ARPSize,
+		ReserveHeaderBytes: int(e.nic.MaxHeaderLength()) + header.ARPSize,
 	})
 	h := header.ARP(pkt.NetworkHeader().Push(header.ARPSize))
 	pkt.NetworkProtocolNumber = ProtocolNumber
@@ -318,14 +291,14 @@ func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remot
 	h.SetOp(header.ARPRequest)
 	// TODO(gvisor.dev/issue/4582): check copied length once TAP devices have a
 	// link address.
-	_ = copy(h.HardwareAddressSender(), nic.LinkAddress())
+	_ = copy(h.HardwareAddressSender(), e.nic.LinkAddress())
 	if n := copy(h.ProtocolAddressSender(), localAddr); n != header.IPv4AddressSize {
 		panic(fmt.Sprintf("copied %d bytes, expected %d bytes", n, header.IPv4AddressSize))
 	}
 	if n := copy(h.ProtocolAddressTarget(), targetAddr); n != header.IPv4AddressSize {
 		panic(fmt.Sprintf("copied %d bytes, expected %d bytes", n, header.IPv4AddressSize))
 	}
-	if err := nic.WritePacketToRemote(remoteLinkAddr, nil /* gso */, ProtocolNumber, pkt); err != nil {
+	if err := e.nic.WritePacketToRemote(remoteLinkAddr, nil /* gso */, ProtocolNumber, pkt); err != nil {
 		stats.outgoingRequestsDropped.Increment()
 		return err
 	}
@@ -334,7 +307,7 @@ func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remot
 }
 
 // ResolveStaticAddress implements stack.LinkAddressResolver.ResolveStaticAddress.
-func (*protocol) ResolveStaticAddress(addr tcpip.Address) (tcpip.LinkAddress, bool) {
+func (*endpoint) ResolveStaticAddress(addr tcpip.Address) (tcpip.LinkAddress, bool) {
 	if addr == header.IPv4Broadcast {
 		return header.EthernetBroadcastAddress, true
 	}
@@ -369,9 +342,5 @@ func (*protocol) Parse(pkt *stack.PacketBuffer) (proto tcpip.TransportProtocolNu
 func NewProtocol(s *stack.Stack) stack.NetworkProtocol {
 	return &protocol{
 		stack: s,
-		mu: struct {
-			sync.RWMutex
-			eps map[tcpip.NICID]*endpoint
-		}{eps: make(map[tcpip.NICID]*endpoint)},
 	}
 }
diff --git a/pkg/tcpip/network/arp/arp_test.go b/pkg/tcpip/network/arp/arp_test.go
index b0f07aa44..d753a97af 100644
--- a/pkg/tcpip/network/arp/arp_test.go
+++ b/pkg/tcpip/network/arp/arp_test.go
@@ -530,52 +530,19 @@ func TestDirectRequestWithNeighborCache(t *testing.T) {
 	}
 }
 
-var _ stack.NetworkInterface = (*testInterface)(nil)
+var _ stack.LinkEndpoint = (*testLinkEndpoint)(nil)
 
-type testInterface struct {
+type testLinkEndpoint struct {
 	stack.LinkEndpoint
 
-	nicID tcpip.NICID
-
 	writeErr tcpip.Error
 }
 
-func (t *testInterface) ID() tcpip.NICID {
-	return t.nicID
-}
-
-func (*testInterface) IsLoopback() bool {
-	return false
-}
-
-func (*testInterface) Name() string {
-	return ""
-}
-
-func (*testInterface) Enabled() bool {
-	return true
-}
-
-func (*testInterface) Promiscuous() bool {
-	return false
-}
-
-func (t *testInterface) WritePacket(r *stack.Route, gso *stack.GSO, protocol tcpip.NetworkProtocolNumber, pkt *stack.PacketBuffer) tcpip.Error {
-	return t.LinkEndpoint.WritePacket(r.Fields(), gso, protocol, pkt)
-}
-
-func (t *testInterface) WritePackets(r *stack.Route, gso *stack.GSO, pkts stack.PacketBufferList, protocol tcpip.NetworkProtocolNumber) (int, tcpip.Error) {
-	return t.LinkEndpoint.WritePackets(r.Fields(), gso, pkts, protocol)
-}
-
-func (t *testInterface) WritePacketToRemote(remoteLinkAddr tcpip.LinkAddress, gso *stack.GSO, protocol tcpip.NetworkProtocolNumber, pkt *stack.PacketBuffer) tcpip.Error {
+func (t *testLinkEndpoint) WritePacket(r stack.RouteInfo, gso *stack.GSO, protocol tcpip.NetworkProtocolNumber, pkt *stack.PacketBuffer) tcpip.Error {
 	if t.writeErr != nil {
 		return t.writeErr
 	}
 
-	var r stack.RouteInfo
-	r.NetProto = protocol
-	r.RemoteLinkAddress = remoteLinkAddr
 	return t.LinkEndpoint.WritePacket(r, gso, protocol, pkt)
 }
 
@@ -709,32 +676,31 @@ func TestLinkAddressRequest(t *testing.T) {
 			s := stack.New(stack.Options{
 				NetworkProtocols: []stack.NetworkProtocolFactory{arp.NewProtocol, ipv4.NewProtocol},
 			})
-			p := s.NetworkProtocolInstance(arp.ProtocolNumber)
-			linkRes, ok := p.(stack.LinkAddressResolver)
-			if !ok {
-				t.Fatal("expected ARP protocol to implement stack.LinkAddressResolver")
-			}
-
 			linkEP := channel.New(defaultChannelSize, defaultMTU, stackLinkAddr)
-			if err := s.CreateNIC(nicID, linkEP); err != nil {
+			if err := s.CreateNIC(nicID, &testLinkEndpoint{LinkEndpoint: linkEP, writeErr: test.linkErr}); err != nil {
 				t.Fatalf("s.CreateNIC(%d, _): %s", nicID, err)
 			}
 
+			ep, err := s.GetNetworkEndpoint(nicID, arp.ProtocolNumber)
+			if err != nil {
+				t.Fatalf("s.GetNetworkEndpoint(%d, %d): %s", nicID, arp.ProtocolNumber, err)
+			}
+			linkRes, ok := ep.(stack.LinkAddressResolver)
+			if !ok {
+				t.Fatalf("expected %T to implement stack.LinkAddressResolver", ep)
+			}
+
 			if len(test.nicAddr) != 0 {
 				if err := s.AddAddress(nicID, ipv4.ProtocolNumber, test.nicAddr); err != nil {
 					t.Fatalf("s.AddAddress(%d, %d, %s): %s", nicID, ipv4.ProtocolNumber, test.nicAddr, err)
 				}
 			}
 
-			// We pass a test network interface to LinkAddressRequest with the same
-			// NIC ID and link endpoint used by the NIC we created earlier so that we
-			// can mock a link address request and observe the packets sent to the
-			// link endpoint even though the stack uses the real NIC to validate the
-			// local address.
-			iface := testInterface{LinkEndpoint: linkEP, nicID: nicID, writeErr: test.linkErr}
-			err := linkRes.LinkAddressRequest(remoteAddr, test.localAddr, test.remoteLinkAddr, &iface)
-			if diff := cmp.Diff(test.expectedErr, err); diff != "" {
-				t.Fatalf("unexpected error from p.LinkAddressRequest(%s, %s, %s, _), (-want, +got):\n%s", remoteAddr, test.localAddr, test.remoteLinkAddr, diff)
+			{
+				err := linkRes.LinkAddressRequest(remoteAddr, test.localAddr, test.remoteLinkAddr)
+				if diff := cmp.Diff(test.expectedErr, err); diff != "" {
+					t.Fatalf("unexpected error from p.LinkAddressRequest(%s, %s, %s, _), (-want, +got):\n%s", remoteAddr, test.localAddr, test.remoteLinkAddr, diff)
+				}
 			}
 
 			if got := s.Stats().ARP.OutgoingRequestsSent.Value(); got != test.expectedRequestsSent {
@@ -782,19 +748,3 @@ func TestLinkAddressRequest(t *testing.T) {
 		})
 	}
 }
-
-func TestLinkAddressRequestWithoutNIC(t *testing.T) {
-	s := stack.New(stack.Options{
-		NetworkProtocols: []stack.NetworkProtocolFactory{arp.NewProtocol, ipv4.NewProtocol},
-	})
-	p := s.NetworkProtocolInstance(arp.ProtocolNumber)
-	linkRes, ok := p.(stack.LinkAddressResolver)
-	if !ok {
-		t.Fatal("expected ARP protocol to implement stack.LinkAddressResolver")
-	}
-
-	err := linkRes.LinkAddressRequest(remoteAddr, "", remoteLinkAddr, &testInterface{nicID: nicID})
-	if _, ok := err.(*tcpip.ErrNotConnected); !ok {
-		t.Fatalf("got p.LinkAddressRequest(%s, %s, %s, _) = %s, want = %s", remoteAddr, "", remoteLinkAddr, err, &tcpip.ErrNotConnected{})
-	}
-}
diff --git a/pkg/tcpip/network/arp/stats_test.go b/pkg/tcpip/network/arp/stats_test.go
index 036fdf739..d3b56c635 100644
--- a/pkg/tcpip/network/arp/stats_test.go
+++ b/pkg/tcpip/network/arp/stats_test.go
@@ -34,48 +34,6 @@ func (t *testInterface) ID() tcpip.NICID {
 	return t.nicID
 }
 
-func knownNICIDs(proto *protocol) []tcpip.NICID {
-	var nicIDs []tcpip.NICID
-
-	for k := range proto.mu.eps {
-		nicIDs = append(nicIDs, k)
-	}
-
-	return nicIDs
-}
-
-func TestClearEndpointFromProtocolOnClose(t *testing.T) {
-	s := stack.New(stack.Options{
-		NetworkProtocols: []stack.NetworkProtocolFactory{NewProtocol},
-	})
-	proto := s.NetworkProtocolInstance(ProtocolNumber).(*protocol)
-	nic := testInterface{nicID: 1}
-	ep := proto.NewEndpoint(&nic, nil, nil, nil).(*endpoint)
-	var nicIDs []tcpip.NICID
-
-	proto.mu.Lock()
-	foundEP, hasEndpointBeforeClose := proto.mu.eps[nic.ID()]
-	nicIDs = knownNICIDs(proto)
-	proto.mu.Unlock()
-
-	if !hasEndpointBeforeClose {
-		t.Fatalf("expected to find the nic id %d in the protocol's endpoint map (%v)", nic.ID(), nicIDs)
-	}
-	if foundEP != ep {
-		t.Fatalf("found an incorrect endpoint mapped to nic id %d", nic.ID())
-	}
-
-	ep.Close()
-
-	proto.mu.Lock()
-	_, hasEndpointAfterClose := proto.mu.eps[nic.ID()]
-	nicIDs = knownNICIDs(proto)
-	proto.mu.Unlock()
-	if hasEndpointAfterClose {
-		t.Fatalf("unexpectedly found an endpoint mapped to the nic id %d in the protocol's known nic ids (%v)", nic.ID(), nicIDs)
-	}
-}
-
 func TestMultiCounterStatsInitialization(t *testing.T) {
 	s := stack.New(stack.Options{
 		NetworkProtocols: []stack.NetworkProtocolFactory{NewProtocol},
diff --git a/pkg/tcpip/network/ipv6/icmp.go b/pkg/tcpip/network/ipv6/icmp.go
index 7298bd061..8db2454d3 100644
--- a/pkg/tcpip/network/ipv6/icmp.go
+++ b/pkg/tcpip/network/ipv6/icmp.go
@@ -290,7 +290,7 @@ func (e *endpoint) handleICMP(pkt *stack.PacketBuffer, hasFragmentHeader bool) {
 			received.invalid.Increment()
 			return
 		} else if e.nud != nil {
-			e.nud.HandleProbe(srcAddr, header.IPv6ProtocolNumber, sourceLinkAddr, e.protocol)
+			e.nud.HandleProbe(srcAddr, header.IPv6ProtocolNumber, sourceLinkAddr, e)
 		} else {
 			e.linkAddrCache.AddLinkAddress(srcAddr, sourceLinkAddr)
 		}
@@ -578,7 +578,7 @@ func (e *endpoint) handleICMP(pkt *stack.PacketBuffer, hasFragmentHeader bool) {
 			if e.nud != nil {
 				// A RS with a specified source IP address modifies the NUD state
 				// machine in the same way a reachability probe would.
-				e.nud.HandleProbe(srcAddr, ProtocolNumber, sourceLinkAddr, e.protocol)
+				e.nud.HandleProbe(srcAddr, ProtocolNumber, sourceLinkAddr, e)
 			}
 		}
 
@@ -628,7 +628,7 @@ func (e *endpoint) handleICMP(pkt *stack.PacketBuffer, hasFragmentHeader bool) {
 		// If the RA has the source link layer option, update the link address
 		// cache with the link address for the advertised router.
 		if len(sourceLinkAddr) != 0 && e.nud != nil {
-			e.nud.HandleProbe(routerAddr, ProtocolNumber, sourceLinkAddr, e.protocol)
+			e.nud.HandleProbe(routerAddr, ProtocolNumber, sourceLinkAddr, e)
 		}
 
 		e.mu.Lock()
@@ -694,24 +694,13 @@ func (e *endpoint) handleICMP(pkt *stack.PacketBuffer, hasFragmentHeader bool) {
 	}
 }
 
-var _ stack.LinkAddressResolver = (*protocol)(nil)
-
 // LinkAddressProtocol implements stack.LinkAddressResolver.
-func (*protocol) LinkAddressProtocol() tcpip.NetworkProtocolNumber {
+func (*endpoint) LinkAddressProtocol() tcpip.NetworkProtocolNumber {
 	return header.IPv6ProtocolNumber
 }
 
 // LinkAddressRequest implements stack.LinkAddressResolver.
-func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remoteLinkAddr tcpip.LinkAddress, nic stack.NetworkInterface) tcpip.Error {
-	nicID := nic.ID()
-
-	p.mu.Lock()
-	netEP, ok := p.mu.eps[nicID]
-	p.mu.Unlock()
-	if !ok {
-		return &tcpip.ErrNotConnected{}
-	}
-
+func (e *endpoint) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remoteLinkAddr tcpip.LinkAddress) tcpip.Error {
 	remoteAddr := targetAddr
 	if len(remoteLinkAddr) == 0 {
 		remoteAddr = header.SolicitedNodeAddr(targetAddr)
@@ -719,22 +708,22 @@ func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remot
 	}
 
 	if len(localAddr) == 0 {
-		addressEndpoint := netEP.AcquireOutgoingPrimaryAddress(remoteAddr, false /* allowExpired */)
+		addressEndpoint := e.AcquireOutgoingPrimaryAddress(remoteAddr, false /* allowExpired */)
 		if addressEndpoint == nil {
 			return &tcpip.ErrNetworkUnreachable{}
 		}
 
 		localAddr = addressEndpoint.AddressWithPrefix().Address
-	} else if p.stack.CheckLocalAddress(nicID, ProtocolNumber, localAddr) == 0 {
+	} else if e.protocol.stack.CheckLocalAddress(e.nic.ID(), ProtocolNumber, localAddr) == 0 {
 		return &tcpip.ErrBadLocalAddress{}
 	}
 
 	optsSerializer := header.NDPOptionsSerializer{
-		header.NDPSourceLinkLayerAddressOption(nic.LinkAddress()),
+		header.NDPSourceLinkLayerAddressOption(e.nic.LinkAddress()),
 	}
 	neighborSolicitSize := header.ICMPv6NeighborSolicitMinimumSize + optsSerializer.Length()
 	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		ReserveHeaderBytes: int(nic.MaxHeaderLength()) + header.IPv6FixedHeaderSize + neighborSolicitSize,
+		ReserveHeaderBytes: int(e.nic.MaxHeaderLength()) + header.IPv6FixedHeaderSize + neighborSolicitSize,
 	})
 	pkt.TransportProtocolNumber = header.ICMPv6ProtocolNumber
 	packet := header.ICMPv6(pkt.TransportHeader().Push(neighborSolicitSize))
@@ -751,9 +740,9 @@ func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remot
 		panic(fmt.Sprintf("failed to add IP header: %s", err))
 	}
 
-	stat := netEP.stats.icmp.packetsSent
+	stat := e.stats.icmp.packetsSent
 
-	if err := nic.WritePacketToRemote(remoteLinkAddr, nil /* gso */, ProtocolNumber, pkt); err != nil {
+	if err := e.nic.WritePacketToRemote(remoteLinkAddr, nil /* gso */, ProtocolNumber, pkt); err != nil {
 		stat.dropped.Increment()
 		return err
 	}
@@ -763,7 +752,7 @@ func (p *protocol) LinkAddressRequest(targetAddr, localAddr tcpip.Address, remot
 }
 
 // ResolveStaticAddress implements stack.LinkAddressResolver.
-func (*protocol) ResolveStaticAddress(addr tcpip.Address) (tcpip.LinkAddress, bool) {
+func (*endpoint) ResolveStaticAddress(addr tcpip.Address) (tcpip.LinkAddress, bool) {
 	if header.IsV6MulticastAddress(addr) {
 		return header.EthernetAddressFromMulticastIPv6Address(addr), true
 	}
diff --git a/pkg/tcpip/network/ipv6/icmp_test.go b/pkg/tcpip/network/ipv6/icmp_test.go
index db1c2e663..a5c88444e 100644
--- a/pkg/tcpip/network/ipv6/icmp_test.go
+++ b/pkg/tcpip/network/ipv6/icmp_test.go
@@ -1346,29 +1346,32 @@ func TestLinkAddressRequest(t *testing.T) {
 			s := stack.New(stack.Options{
 				NetworkProtocols: []stack.NetworkProtocolFactory{NewProtocol},
 			})
-			p := s.NetworkProtocolInstance(ProtocolNumber)
-			linkRes, ok := p.(stack.LinkAddressResolver)
-			if !ok {
-				t.Fatalf("expected IPv6 protocol to implement stack.LinkAddressResolver")
-			}
 
 			linkEP := channel.New(defaultChannelSize, defaultMTU, linkAddr0)
 			if err := s.CreateNIC(nicID, linkEP); err != nil {
 				t.Fatalf("s.CreateNIC(%d, _): %s", nicID, err)
 			}
+
+			ep, err := s.GetNetworkEndpoint(nicID, ProtocolNumber)
+			if err != nil {
+				t.Fatalf("s.GetNetworkEndpoint(%d, %d): %s", nicID, ProtocolNumber, err)
+			}
+			linkRes, ok := ep.(stack.LinkAddressResolver)
+			if !ok {
+				t.Fatalf("expected %T to implement stack.LinkAddressResolver", ep)
+			}
+
 			if len(test.nicAddr) != 0 {
 				if err := s.AddAddress(nicID, ProtocolNumber, test.nicAddr); err != nil {
 					t.Fatalf("s.AddAddress(%d, %d, %s): %s", nicID, ProtocolNumber, test.nicAddr, err)
 				}
 			}
 
-			// We pass a test network interface to LinkAddressRequest with the same NIC
-			// ID and link endpoint used by the NIC we created earlier so that we can
-			// mock a link address request and observe the packets sent to the link
-			// endpoint even though the stack uses the real NIC.
-			err := linkRes.LinkAddressRequest(lladdr0, test.localAddr, test.remoteLinkAddr, &testInterface{LinkEndpoint: linkEP, nicID: nicID})
-			if diff := cmp.Diff(test.expectedErr, err); diff != "" {
-				t.Fatalf("unexpected error from p.LinkAddressRequest(%s, %s, %s, _), (-want, +got):\n%s", lladdr0, test.localAddr, test.remoteLinkAddr, diff)
+			{
+				err := linkRes.LinkAddressRequest(lladdr0, test.localAddr, test.remoteLinkAddr)
+				if diff := cmp.Diff(test.expectedErr, err); diff != "" {
+					t.Fatalf("unexpected error from p.LinkAddressRequest(%s, %s, %s, _), (-want, +got):\n%s", lladdr0, test.localAddr, test.remoteLinkAddr, diff)
+				}
 			}
 
 			if test.expectedErr != nil {
diff --git a/pkg/tcpip/network/ipv6/ipv6.go b/pkg/tcpip/network/ipv6/ipv6.go
index caa62b3a2..b55a35525 100644
--- a/pkg/tcpip/network/ipv6/ipv6.go
+++ b/pkg/tcpip/network/ipv6/ipv6.go
@@ -164,6 +164,7 @@ func getLabel(addr tcpip.Address) uint8 {
 	panic(fmt.Sprintf("should have a label for address = %s", addr))
 }
 
+var _ stack.LinkAddressResolver = (*endpoint)(nil)
 var _ stack.LinkResolvableNetworkEndpoint = (*endpoint)(nil)
 var _ stack.GroupAddressableEndpoint = (*endpoint)(nil)
 var _ stack.AddressableEndpoint = (*endpoint)(nil)