Cache neighbor entries in route struct.

This will avoid allocating closures and having to retrieve the entry from
the table every time.

PiperOrigin-RevId: 478551631

pkg/tcpip/stack/neighbor_cache.go

@@ -122,9 +122,7 @@ func (n *neighborCache) getOrCreateEntry(remoteAddr tcpip.Address) *neighborEntr
 // If specified, the local address must be an address local to the interface the
 // neighbor cache belongs to. The local address is the source address of a
 // packet prompting NUD/link address resolution.
-//
-// TODO(gvisor.dev/issue/5151): Don't return the neighbor entry.
-func (n *neighborCache) entry(remoteAddr, localAddr tcpip.Address, onResolve func(LinkResolutionResult)) (NeighborEntry, <-chan struct{}, tcpip.Error) {
+func (n *neighborCache) entry(remoteAddr, localAddr tcpip.Address, onResolve func(LinkResolutionResult)) (*neighborEntry, <-chan struct{}, tcpip.Error) {
 	entry := n.getOrCreateEntry(remoteAddr)
 	entry.mu.Lock()
 	defer entry.mu.Unlock()
@@ -142,7 +140,7 @@ func (n *neighborCache) entry(remoteAddr, localAddr tcpip.Address, onResolve fun
 		if onResolve != nil {
 			onResolve(LinkResolutionResult{LinkAddress: entry.mu.neigh.LinkAddr, Err: nil})
 		}
-		return entry.mu.neigh, nil, nil
+		return entry, nil, nil
 	case Unknown, Incomplete, Unreachable:
 		if onResolve != nil {
 			entry.mu.onResolve = append(entry.mu.onResolve, onResolve)
@@ -152,7 +150,7 @@ func (n *neighborCache) entry(remoteAddr, localAddr tcpip.Address, onResolve fun
 			entry.mu.done = make(chan struct{})
 		}
 		entry.handlePacketQueuedLocked(localAddr)
-		return entry.mu.neigh, entry.mu.done, &tcpip.ErrWouldBlock{}
+		return entry, entry.mu.done, &tcpip.ErrWouldBlock{}
 	default:
 		panic(fmt.Sprintf("Invalid cache entry state: %s", s))
 	}

pkg/tcpip/stack/neighbor_cache_test.go

@@ -874,7 +874,7 @@ func TestNeighborCacheAddStaticEntryThenOverflow(t *testing.T) {
 		State:     Static,
 		UpdatedAt: c.clock.NowMonotonic(),
 	}
-	if diff := cmp.Diff(want, e, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
+	if diff := cmp.Diff(want, e.mu.neigh, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
 		t.Errorf("c.linkRes.neigh.entry(%s, \"\", nil) mismatch (-want, +got):\n%s", entry.Addr, diff)
 	}
 
@@ -1243,7 +1243,7 @@ func TestNeighborCacheReplace(t *testing.T) {
 			State:     Delay,
 			UpdatedAt: clock.NowMonotonic(),
 		}
-		if diff := cmp.Diff(want, e, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
+		if diff := cmp.Diff(want, e.mu.neigh, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
 			t.Errorf("linkRes.neigh.entry(%s, '', nil) mismatch (-want, +got):\n%s", entry.Addr, diff)
 		}
 	}
@@ -1262,7 +1262,7 @@ func TestNeighborCacheReplace(t *testing.T) {
 			State:     Reachable,
 			UpdatedAt: clock.NowMonotonic(),
 		}
-		if diff := cmp.Diff(want, e, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
+		if diff := cmp.Diff(want, e.mu.neigh, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
 			t.Errorf("linkRes.neigh.entry(%s, '', nil) mismatch (-want, +got):\n%s", entry.Addr, diff)
 		}
 	}
@@ -1300,7 +1300,7 @@ func TestNeighborCacheResolutionFailed(t *testing.T) {
 		State:     Reachable,
 		UpdatedAt: clock.NowMonotonic(),
 	}
-	if diff := cmp.Diff(want, got, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
+	if diff := cmp.Diff(want, got.mu.neigh, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
 		t.Errorf("linkRes.neigh.entry(%s, '', nil) mismatch (-want, +got):\n%s", entry.Addr, diff)
 	}
 
@@ -1472,8 +1472,8 @@ func TestNeighborCacheRetryResolution(t *testing.T) {
 		if _, ok := err.(*tcpip.ErrWouldBlock); !ok {
 			t.Fatalf("got linkRes.neigh.entry(%s, '', _) = %v, want = %s", entry.Addr, err, &tcpip.ErrWouldBlock{})
 		}
-		if incompleteEntry.State != Incomplete {
-			t.Fatalf("got entry.State = %s, want = %s", incompleteEntry.State, Incomplete)
+		if incompleteEntry.mu.neigh.State != Incomplete {
+			t.Fatalf("got entry.State = %s, want = %s", incompleteEntry.mu.neigh.State, Incomplete)
 		}
 
 		{
@@ -1540,7 +1540,7 @@ func TestNeighborCacheRetryResolution(t *testing.T) {
 				State:     Reachable,
 				UpdatedAt: clock.NowMonotonic(),
 			}
-			if diff := cmp.Diff(gotEntry, wantEntry, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
+			if diff := cmp.Diff(gotEntry.mu.neigh, wantEntry, cmp.AllowUnexported(tcpip.MonotonicTime{})); diff != "" {
 				t.Fatalf("neighbor entry mismatch (-got, +want):\n%s", diff)
 			}
 		}

pkg/tcpip/stack/neighbor_entry.go

@@ -232,6 +232,8 @@ func (e *neighborEntry) cancelTimerLocked() {
 func (e *neighborEntry) removeLocked() {
 	e.mu.neigh.UpdatedAt = e.cache.nic.stack.clock.NowMonotonic()
 	e.dispatchRemoveEventLocked()
+	// Set state to unknown to invalidate this entry if it's cached in a Route.
+	e.setStateLocked(Unknown)
 	e.cancelTimerLocked()
 	// TODO(https://gvisor.dev/issues/5583): test the case where this function is
 	// called during resolution; that can happen in at least these scenarios:
@@ -607,3 +609,18 @@ func (e *neighborEntry) handleUpperLevelConfirmationLocked() {
 		panic(fmt.Sprintf("Invalid cache entry state: %s", e.mu.neigh.State))
 	}
 }
+
+// getRemoteLinkAddress returns the entry's link address and whether that link
+// address is valid.
+func (e *neighborEntry) getRemoteLinkAddress() (tcpip.LinkAddress, bool) {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+	switch e.mu.neigh.State {
+	case Reachable, Static, Delay, Probe:
+		return e.mu.neigh.LinkAddr, true
+	case Unknown, Incomplete, Unreachable, Stale:
+		return "", false
+	default:
+		panic(fmt.Sprintf("invalid state for neighbor entry %v: %v", e.mu.neigh, e.mu.neigh.State))
+	}
+}

pkg/tcpip/stack/neighbor_entry_test.go

@@ -1957,6 +1957,47 @@ func TestEntryProbeToReachableWhenSolicitedOverrideConfirmation(t *testing.T) {
 	}
 }
 
+func TestGetRemoteLinkAddressFailsWhenResolutionRequired(t *testing.T) {
+	c := DefaultNUDConfigurations()
+	c.MinRandomFactor = 1
+	c.MaxRandomFactor = 1
+	e, nudDisp, linkRes, clock := entryTestSetup(c)
+
+	if _, ok := e.getRemoteLinkAddress(); ok {
+		t.Errorf("getRemoteLinkAddress() = _, true, want false")
+	}
+	if err := unknownToStale(e, nudDisp, linkRes, clock); err != nil {
+		t.Fatalf("unknownToStale(...) = %s", err)
+	}
+	if _, ok := e.getRemoteLinkAddress(); ok {
+		t.Errorf("getRemoteLinkAddress() = _, true, want false")
+	}
+	if err := staleToDelay(e, nudDisp, linkRes, clock); err != nil {
+		t.Fatalf("staleToDelay(...) = %s", err)
+	}
+	if _, ok := e.getRemoteLinkAddress(); !ok {
+		t.Errorf("getRemoteLinkAddress() = _, false, want true")
+	}
+	if err := delayToProbe(c, e, nudDisp, linkRes, clock); err != nil {
+		t.Fatalf("delayToProbe(...) = %s", err)
+	}
+	if _, ok := e.getRemoteLinkAddress(); !ok {
+		t.Errorf("getRemoteLinkAddress() = _, false, want true")
+	}
+	if err := probeToReachable(e, nudDisp, linkRes, clock); err != nil {
+		t.Fatalf("probeToReachable(...) = %s", err)
+	}
+	if _, ok := e.getRemoteLinkAddress(); !ok {
+		t.Errorf("getRemoteLinkAddress() = _, false, want true")
+	}
+	if err := reachableToStale(c, e, nudDisp, linkRes, clock); err != nil {
+		t.Fatalf("reachableToStale(...) = %s", err)
+	}
+	if _, ok := e.getRemoteLinkAddress(); ok {
+		t.Errorf("getRemoteLinkAddress() = _, true, want false")
+	}
+}
+
 func probeToReachableWithFlags(e *neighborEntry, nudDisp *testNUDDispatcher, linkRes *entryTestLinkResolver, clock *faketime.ManualClock, linkAddr tcpip.LinkAddress, flags ReachabilityConfirmationFlags) error {
 	if err := func() error {
 		e.mu.Lock()

pkg/tcpip/stack/nic.go

@@ -30,11 +30,6 @@ type linkResolver struct {
 	neigh neighborCache
 }
 
-func (l *linkResolver) getNeighborLinkAddress(addr, localAddr tcpip.Address, onResolve func(LinkResolutionResult)) (tcpip.LinkAddress, <-chan struct{}, tcpip.Error) {
-	entry, ch, err := l.neigh.entry(addr, localAddr, onResolve)
-	return entry.LinkAddr, ch, err
-}
-
 func (l *linkResolver) confirmReachable(addr tcpip.Address) {
 	l.neigh.handleUpperLevelConfirmation(addr)
 }
@@ -625,7 +620,7 @@ func (n *nic) getLinkAddress(addr, localAddr tcpip.Address, protocol tcpip.Netwo
 		return nil
 	}
 
-	_, _, err := linkRes.getNeighborLinkAddress(addr, localAddr, onResolve)
+	_, _, err := linkRes.neigh.entry(addr, localAddr, onResolve)
 	return err
 }
 

pkg/tcpip/stack/route.go

@@ -50,6 +50,10 @@ type Route struct {
 	// linkRes is set if link address resolution is enabled for this protocol on
 	// the route's NIC.
 	linkRes *linkResolver
+
+	// neighborEntry is the cached result of fetching a neighbor entry from the
+	// neighbor cache.
+	neighborEntry *neighborEntry
 }
 
 // +stateify savable
@@ -390,22 +394,48 @@ func (r *Route) resolvedFields(afterResolve func(ResolvedFieldsResult)) (RouteIn
 		linkAddressResolutionRequestLocalAddr = r.LocalAddress()
 	}
 
+	nEntry := r.getCachedNeighborEntry()
+	if nEntry != nil {
+		if addr, ok := nEntry.getRemoteLinkAddress(); ok {
+			fields.RemoteLinkAddress = addr
+			if afterResolve != nil {
+				afterResolve(ResolvedFieldsResult{RouteInfo: fields, Err: nil})
+			}
+			return fields, nil, nil
+		}
+	}
 	afterResolveFields := fields
-	linkAddr, ch, err := r.linkRes.getNeighborLinkAddress(r.nextHop(), linkAddressResolutionRequestLocalAddr, func(r LinkResolutionResult) {
+	entry, ch, err := r.linkRes.neigh.entry(r.nextHop(), linkAddressResolutionRequestLocalAddr, func(lrr LinkResolutionResult) {
+		if lrr.Err != nil {
+			r.setCachedNeighborEntry(nil)
+		}
 		if afterResolve != nil {
-			if r.Err == nil {
-				afterResolveFields.RemoteLinkAddress = r.LinkAddress
+			if lrr.Err == nil {
+				afterResolveFields.RemoteLinkAddress = lrr.LinkAddress
 			}
 
-			afterResolve(ResolvedFieldsResult{RouteInfo: afterResolveFields, Err: r.Err})
+			afterResolve(ResolvedFieldsResult{RouteInfo: afterResolveFields, Err: lrr.Err})
 		}
 	})
 	if err == nil {
-		fields.RemoteLinkAddress = linkAddr
+		fields.RemoteLinkAddress, _ = entry.getRemoteLinkAddress()
 	}
+	r.setCachedNeighborEntry(entry)
 	return fields, ch, err
 }
 
+func (r *Route) getCachedNeighborEntry() *neighborEntry {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return r.neighborEntry
+}
+
+func (r *Route) setCachedNeighborEntry(entry *neighborEntry) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.neighborEntry = entry
+}
+
 func (r *Route) nextHop() tcpip.Address {
 	if len(r.NextHop()) == 0 {
 		return r.RemoteAddress()

pkg/tcpip/tests/integration/BUILD

@@ -58,12 +58,14 @@ go_test(
     srcs = ["link_resolution_test.go"],
     deps = [
         "//pkg/bufferv2",
+        "//pkg/sync",
         "//pkg/tcpip",
         "//pkg/tcpip/checker",
         "//pkg/tcpip/checksum",
         "//pkg/tcpip/faketime",
         "//pkg/tcpip/header",
         "//pkg/tcpip/link/channel",
+        "//pkg/tcpip/link/ethernet",
         "//pkg/tcpip/link/pipe",
         "//pkg/tcpip/network/arp",
         "//pkg/tcpip/network/ipv4",

pkg/tcpip/tests/integration/link_resolution_test.go

@@ -25,12 +25,14 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"gvisor.dev/gvisor/pkg/bufferv2"
+	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/checker"
 	"gvisor.dev/gvisor/pkg/tcpip/checksum"
 	"gvisor.dev/gvisor/pkg/tcpip/faketime"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/link/channel"
+	"gvisor.dev/gvisor/pkg/tcpip/link/ethernet"
 	"gvisor.dev/gvisor/pkg/tcpip/link/pipe"
 	"gvisor.dev/gvisor/pkg/tcpip/network/arp"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
@@ -45,10 +47,14 @@ import (
 )
 
 func setupStack(t *testing.T, stackOpts stack.Options, host1NICID, host2NICID tcpip.NICID) (*stack.Stack, *stack.Stack) {
+	return setupStackWithSeparateOpts(t, stackOpts, stackOpts, host1NICID, host2NICID)
+}
+
+func setupStackWithSeparateOpts(t *testing.T, stack1Opts stack.Options, stack2Opts stack.Options, host1NICID, host2NICID tcpip.NICID) (*stack.Stack, *stack.Stack) {
 	const maxFrameSize = header.IPv6MinimumMTU + header.EthernetMinimumSize
 
-	host1Stack := stack.New(stackOpts)
-	host2Stack := stack.New(stackOpts)
+	host1Stack := stack.New(stack1Opts)
+	host2Stack := stack.New(stack2Opts)
 
 	host1NIC, host2NIC := pipe.New(utils.LinkAddr1, utils.LinkAddr2, maxFrameSize)
 
@@ -1644,3 +1650,165 @@ func TestDAD(t *testing.T) {
 		})
 	}
 }
+
+type settableLinkEndpoint struct {
+	stack.LinkEndpoint
+	mu   sync.Mutex
+	addr tcpip.LinkAddress
+}
+
+func newSettableLinkEndpoint(e stack.LinkEndpoint) *settableLinkEndpoint {
+	return &settableLinkEndpoint{
+		LinkEndpoint: e,
+		addr:         e.LinkAddress(),
+	}
+}
+
+func (e *settableLinkEndpoint) setLinkAddress(addr tcpip.LinkAddress) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	e.addr = addr
+}
+
+func (e *settableLinkEndpoint) LinkAddress() tcpip.LinkAddress {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	return e.addr
+}
+
+type monitorableLinkEndpoint struct {
+	stack.LinkEndpoint
+	ch chan tcpip.LinkAddress
+}
+
+func newMonitorableLinkEndpoint(e stack.LinkEndpoint) *monitorableLinkEndpoint {
+	return &monitorableLinkEndpoint{e, make(chan tcpip.LinkAddress, 1)}
+}
+
+func (e *monitorableLinkEndpoint) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
+	for _, pkt := range pkts.AsSlice() {
+		dstAddr := header.Ethernet(pkt.LinkHeader().Slice()).DestinationAddress()
+		e.ch <- dstAddr
+	}
+	e.LinkEndpoint.WritePackets(pkts)
+
+	return 0, nil
+}
+
+func (e *monitorableLinkEndpoint) waitForLinkAddress(addr tcpip.LinkAddress, wait time.Duration) error {
+	c := time.After(wait)
+	for {
+		select {
+		case sentAddr := <-e.ch:
+			if addr == sentAddr {
+				return nil
+			}
+		case <-c:
+			return fmt.Errorf("timed out waiting for endpoint to send packet with destination address: %v", addr)
+		}
+	}
+}
+
+func TestUpdateCachedNeighborEntry(t *testing.T) {
+	d := []byte{1, 2}
+	params := stack.NetworkHeaderParams{
+		Protocol: udp.ProtocolNumber,
+		TTL:      64,
+		TOS:      stack.DefaultTOS,
+	}
+
+	writePacket := func(t *testing.T, r *stack.Route) {
+		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
+			ReserveHeaderBytes: header.UDPMinimumSize + int(r.MaxHeaderLength()),
+			Payload:            bufferv2.MakeWithData(d),
+		})
+		if err := r.WritePacket(params, pkt); err != nil {
+			t.Fatalf("WritePacket(...): %s", err)
+		}
+		pkt.DecRef()
+	}
+
+	const (
+		host1NICID = 1
+		host2NICID = 4
+	)
+	stackOpts := stack.Options{
+		NetworkProtocols:   []stack.NetworkProtocolFactory{arp.NewProtocol, ipv4.NewProtocol, ipv6.NewProtocol},
+		TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol},
+	}
+
+	const maxFrameSize = header.IPv6MinimumMTU + header.EthernetMinimumSize
+
+	host1Stack := stack.New(stackOpts)
+	host2Stack := stack.New(stackOpts)
+
+	host1Pipe, host2Pipe := pipe.New(utils.LinkAddr1, utils.LinkAddr2, maxFrameSize)
+
+	host1NICMonitorable := newMonitorableLinkEndpoint(ethernet.New(host1Pipe))
+	host2NICSettable := newSettableLinkEndpoint(host2Pipe)
+
+	if err := host1Stack.CreateNIC(host1NICID, host1NICMonitorable); err != nil {
+		t.Fatalf("host1Stack.CreateNIC(%d, _): %s", host1NICID, err)
+	}
+	if err := host2Stack.CreateNIC(host2NICID, ethernet.New(host2NICSettable)); err != nil {
+		t.Fatalf("host2Stack.CreateNIC(%d, _): %s", host2NICID, err)
+	}
+
+	if err := host1Stack.AddProtocolAddress(host1NICID, utils.Ipv4Addr1, stack.AddressProperties{}); err != nil {
+		t.Fatalf("host1Stack.AddProtocolAddress(%d, %+v, {}): %s", host1NICID, utils.Ipv4Addr1, err)
+	}
+	if err := host2Stack.AddProtocolAddress(host2NICID, utils.Ipv4Addr2, stack.AddressProperties{}); err != nil {
+		t.Fatalf("host2Stack.AddProtocolAddress(%d, %+v, {}): %s", host2NICID, utils.Ipv4Addr2, err)
+	}
+
+	host1Stack.SetRouteTable([]tcpip.Route{
+		{
+			Destination: utils.Ipv4Addr1.AddressWithPrefix.Subnet(),
+			NIC:         host1NICID,
+		},
+	})
+	host2Stack.SetRouteTable([]tcpip.Route{
+		{
+			Destination: utils.Ipv4Addr2.AddressWithPrefix.Subnet(),
+			NIC:         host2NICID,
+		},
+	})
+
+	localAddr := utils.Ipv4Addr1.AddressWithPrefix.Address
+	neighborAddr := utils.Ipv4Addr2.AddressWithPrefix.Address
+
+	// Obtain a route to a neighbor.
+	r, err := host1Stack.FindRoute(host1NICID, localAddr, neighborAddr, header.IPv4ProtocolNumber, false)
+	if err != nil {
+		t.Fatalf("host1Stack.FindRoute(...): %s", err)
+	}
+
+	// Send packet to neighbor (start link resolution & resolve, then send
+	// packet). Send twice to use cached address the second time.
+	for i := 0; i < 2; i++ {
+		writePacket(t, r)
+		if err := host1NICMonitorable.waitForLinkAddress(utils.LinkAddr2, time.Second); err != nil {
+			t.Fatalf("host1NIC.waitForLinkAddress(%s): %s", utils.LinkAddr2, err)
+		}
+	}
+
+	// Neighbor no longer reachable, deleted from the neighbor cache.
+	host1Stack.RemoveNeighbor(host1NICID, header.IPv4ProtocolNumber, neighborAddr)
+	host2Stack.DisableNIC(host2NICID)
+
+	// Send packet to neighbor that's no longer reachable (should fail).
+	writePacket(t, r)
+	if err := host1NICMonitorable.waitForLinkAddress(utils.LinkAddr2, time.Second); err == nil {
+		t.Fatalf("got host1NIC.waitForLinkAddress(%s) = nil, want err", utils.LinkAddr2)
+	}
+
+	// Neighbor reachable again with new MAC address.
+	host2Stack.EnableNIC(host2NICID)
+	host2NICSettable.setLinkAddress(utils.LinkAddr3)
+
+	// Send packet to neighbor (start link resolution and then send packet).
+	writePacket(t, r)
+	if err := host1NICMonitorable.waitForLinkAddress(utils.LinkAddr3, 5*time.Second); err != nil {
+		t.Fatalf("host1NIC.waitForLinkAddress(%s): %s", utils.LinkAddr3, err)
+	}
+}