From 93d9f9c346df9efcccf9e76fdd0d6d81ae40ed5c Mon Sep 17 00:00:00 2001
From: Andrei Vagin <avagin@google.com>
Date: Tue, 9 Jul 2024 15:37:19 -0700
Subject: [PATCH] netstack: call NetworkLink.Close synchronously

PiperOrigin-RevId: 650779522
---
 pkg/tcpip/stack/nic.go   | 12 +++++++-----
 pkg/tcpip/stack/stack.go | 37 +++++++++++++++++++++++++++----------
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index a2520d663..a4fcc8fba 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -311,7 +311,10 @@ func (n *nic) enable() tcpip.Error {
 // remove detaches NIC from the link endpoint and releases network endpoint
 // resources. This guarantees no packets between this NIC and the network
 // stack.
-func (n *nic) remove(closeLinkEndpoint bool) tcpip.Error {
+//
+// It returns an action that has to be excuted after releasing the Stack lock
+// and any error encountered.
+func (n *nic) remove(closeLinkEndpoint bool) (func(), tcpip.Error) {
 	n.enableDisableMu.Lock()
 
 	n.disableLocked()
@@ -326,6 +329,7 @@ func (n *nic) remove(closeLinkEndpoint bool) tcpip.Error {
 	// We must not hold n.enableDisableMu here.
 	n.linkResQueue.cancel()
 
+	var deferAct func()
 	// Prevent packets from going down to the link before shutting the link down.
 	n.qDisc.Close()
 	n.NetworkLinkEndpoint.Attach(nil)
@@ -335,12 +339,10 @@ func (n *nic) remove(closeLinkEndpoint bool) tcpip.Error {
 		// The link endpoint has to be closed without holding a
 		// netstack lock, because it can trigger other netstack
 		// operations.
-		go func() {
-			ep.Close()
-		}()
+		deferAct = ep.Close
 	}
 
-	return nil
+	return deferAct, nil
 }
 
 // setPromiscuousMode enables or disables promiscuous mode.
diff --git a/pkg/tcpip/stack/stack.go b/pkg/tcpip/stack/stack.go
index 83079ef97..2a9c16245 100644
--- a/pkg/tcpip/stack/stack.go
+++ b/pkg/tcpip/stack/stack.go
@@ -982,25 +982,28 @@ func (s *Stack) CheckNIC(id tcpip.NICID) bool {
 // RemoveNIC removes NIC and all related routes from the network stack.
 func (s *Stack) RemoveNIC(id tcpip.NICID) tcpip.Error {
 	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	return s.removeNICLocked(id)
+	deferAct, err := s.removeNICLocked(id)
+	s.mu.Unlock()
+	if deferAct != nil {
+		deferAct()
+	}
+	return err
 }
 
 // removeNICLocked removes NIC and all related routes from the network stack.
 //
 // +checklocks:s.mu
-func (s *Stack) removeNICLocked(id tcpip.NICID) tcpip.Error {
+func (s *Stack) removeNICLocked(id tcpip.NICID) (func(), tcpip.Error) {
 	nic, ok := s.nics[id]
 	if !ok {
-		return &tcpip.ErrUnknownNICID{}
+		return nil, &tcpip.ErrUnknownNICID{}
 	}
 	delete(s.nics, id)
 
 	if nic.Primary != nil {
 		b := nic.Primary.NetworkLinkEndpoint.(CoordinatorNIC)
 		if err := b.DelNIC(nic); err != nil {
-			return err
+			return nil, err
 		}
 	}
 
@@ -1905,14 +1908,22 @@ func (s *Stack) Wait() {
 		p.Wait()
 	}
 
-	s.mu.Lock()
-	defer s.mu.Unlock()
+	deferActs := make([]func(), 0)
 
+	s.mu.Lock()
 	for id, n := range s.nics {
 		// Remove NIC to ensure that qDisc goroutines are correctly
 		// terminated on stack teardown.
-		s.removeNICLocked(id)
+		act, _ := s.removeNICLocked(id)
 		n.NetworkLinkEndpoint.Wait()
+		if act != nil {
+			deferActs = append(deferActs, act)
+		}
+	}
+	s.mu.Unlock()
+
+	for _, act := range deferActs {
+		act()
 	}
 }
 
@@ -2357,8 +2368,14 @@ func (s *Stack) SetNICStack(id tcpip.NICID, peer *Stack) (tcpip.NICID, tcpip.Err
 	// Remove routes in-place. n tracks the number of routes written.
 	s.RemoveRoutes(func(r tcpip.Route) bool { return r.NIC == id })
 	ne := nic.NetworkLinkEndpoint.(LinkEndpoint)
-	nic.remove(false /* closeLinkEndpoint */)
+	deferAct, err := nic.remove(false /* closeLinkEndpoint */)
 	s.mu.Unlock()
+	if deferAct != nil {
+		deferAct()
+	}
+	if err != nil {
+		return 0, err
+	}
 
 	id = tcpip.NICID(peer.NextNICID())
 	return id, peer.CreateNICWithOptions(id, ne, NICOptions{Name: nic.Name()})