From 646a0332135480e21980e5348dbe00e0a0de9c8a Mon Sep 17 00:00:00 2001
From: Andrei Vagin <avagin@google.com>
Date: Fri, 21 Jun 2024 18:47:21 -0700
Subject: [PATCH] tcpip: destroy both ends of one veth pair together

PiperOrigin-RevId: 645562424
---
 pkg/sentry/socket/netstack/stack.go           |   2 +-
 pkg/tcpip/link/channel/channel.go             |   4 +-
 pkg/tcpip/link/fdbased/endpoint.go            |   3 +
 pkg/tcpip/link/loopback/loopback.go           |   3 +
 pkg/tcpip/link/muxed/injectable.go            |   3 +
 pkg/tcpip/link/nested/nested.go               |   5 +
 .../link/packetsocket/packetsocket_test.go    |   1 +
 pkg/tcpip/link/pipe/pipe.go                   |   3 +
 pkg/tcpip/link/veth/BUILD                     |  16 ++-
 pkg/tcpip/link/veth/veth.go                   |   7 +-
 pkg/tcpip/link/veth/veth_test.go              | 110 +++++++++++++++++-
 pkg/tcpip/link/waitable/waitable.go           |   5 +
 pkg/tcpip/link/waitable/waitable_test.go      |   3 +
 pkg/tcpip/link/xdp/endpoint.go                |   3 +
 pkg/tcpip/network/ip_test.go                  |   3 +
 pkg/tcpip/network/ipv6/icmp_test.go           |   2 +
 pkg/tcpip/stack/bridge.go                     |   3 +
 pkg/tcpip/stack/forwarding_test.go            |   3 +
 pkg/tcpip/stack/nic.go                        |   1 +
 pkg/tcpip/stack/registration.go               |   3 +
 pkg/tcpip/transport/datagram_test.go          |   3 +
 21 files changed, 176 insertions(+), 10 deletions(-)

diff --git a/pkg/sentry/socket/netstack/stack.go b/pkg/sentry/socket/netstack/stack.go
index bd81375c3..ddf238e5f 100644
--- a/pkg/sentry/socket/netstack/stack.go
+++ b/pkg/sentry/socket/netstack/stack.go
@@ -262,7 +262,7 @@ func (s *Stack) newVeth(ctx context.Context, linkAttrs map[uint16]nlmsg.BytesVie
 		peerEP.Close()
 		return syserr.TranslateNetstackError(err)
 	}
-	peerEP.SetStack(peerStack.Stack, id)
+	peerEP.SetStack(peerStack.Stack, peerID)
 	if peerLinkAttrs != nil {
 		if err := s.setLink(peerID, peerLinkAttrs); err != nil {
 			peerStack.Stack.RemoveNIC(peerID)
diff --git a/pkg/tcpip/link/channel/channel.go b/pkg/tcpip/link/channel/channel.go
index 0f3337ab7..334dc542c 100644
--- a/pkg/tcpip/link/channel/channel.go
+++ b/pkg/tcpip/link/channel/channel.go
@@ -54,7 +54,9 @@ type queue struct {
 func (q *queue) Close() {
 	q.mu.Lock()
 	defer q.mu.Unlock()
-	close(q.c)
+	if !q.closed {
+		close(q.c)
+	}
 	q.closed = true
 }
 
diff --git a/pkg/tcpip/link/fdbased/endpoint.go b/pkg/tcpip/link/fdbased/endpoint.go
index 237fc0716..bfe2e7a70 100644
--- a/pkg/tcpip/link/fdbased/endpoint.go
+++ b/pkg/tcpip/link/fdbased/endpoint.go
@@ -813,6 +813,9 @@ func (e *endpoint) ARPHardwareType() header.ARPHardwareType {
 	return header.ARPHardwareNone
 }
 
+// Close implements stack.LinkEndpoint.
+func (e *endpoint) Close() {}
+
 // InjectableEndpoint is an injectable fd-based endpoint. The endpoint writes
 // to the FD, but does not read from it. All reads come from injected packets.
 //
diff --git a/pkg/tcpip/link/loopback/loopback.go b/pkg/tcpip/link/loopback/loopback.go
index de4a29799..3ebf4acbc 100644
--- a/pkg/tcpip/link/loopback/loopback.go
+++ b/pkg/tcpip/link/loopback/loopback.go
@@ -116,3 +116,6 @@ func (*endpoint) AddHeader(*stack.PacketBuffer) {}
 
 // ParseHeader implements stack.LinkEndpoint.
 func (*endpoint) ParseHeader(*stack.PacketBuffer) bool { return true }
+
+// Close implements stack.LinkEndpoint.
+func (*endpoint) Close() {}
diff --git a/pkg/tcpip/link/muxed/injectable.go b/pkg/tcpip/link/muxed/injectable.go
index b297c0fcb..1a5f3aa70 100644
--- a/pkg/tcpip/link/muxed/injectable.go
+++ b/pkg/tcpip/link/muxed/injectable.go
@@ -155,6 +155,9 @@ func (*InjectableEndpoint) AddHeader(*stack.PacketBuffer) {}
 // ParseHeader implements stack.LinkEndpoint.ParseHeader.
 func (*InjectableEndpoint) ParseHeader(*stack.PacketBuffer) bool { return true }
 
+// Close implements stack.LinkEndpoint.
+func (*InjectableEndpoint) Close() {}
+
 // NewInjectableEndpoint creates a new multi-endpoint injectable endpoint.
 func NewInjectableEndpoint(routes map[tcpip.Address]stack.InjectableLinkEndpoint) *InjectableEndpoint {
 	return &InjectableEndpoint{
diff --git a/pkg/tcpip/link/nested/nested.go b/pkg/tcpip/link/nested/nested.go
index 72af02180..9b35513d5 100644
--- a/pkg/tcpip/link/nested/nested.go
+++ b/pkg/tcpip/link/nested/nested.go
@@ -161,3 +161,8 @@ func (e *Endpoint) AddHeader(pkt *stack.PacketBuffer) {
 func (e *Endpoint) ParseHeader(pkt *stack.PacketBuffer) bool {
 	return e.child.ParseHeader(pkt)
 }
+
+// Close implements stack.LinkEndpoint.
+func (e *Endpoint) Close() {
+	e.child.Close()
+}
diff --git a/pkg/tcpip/link/packetsocket/packetsocket_test.go b/pkg/tcpip/link/packetsocket/packetsocket_test.go
index 64c2a3bf5..5c7f69a49 100644
--- a/pkg/tcpip/link/packetsocket/packetsocket_test.go
+++ b/pkg/tcpip/link/packetsocket/packetsocket_test.go
@@ -55,6 +55,7 @@ func (*nullEndpoint) Wait()                                   {}
 func (*nullEndpoint) ARPHardwareType() header.ARPHardwareType { return header.ARPHardwareNone }
 func (*nullEndpoint) AddHeader(*stack.PacketBuffer)           {}
 func (*nullEndpoint) ParseHeader(*stack.PacketBuffer) bool    { return true }
+func (*nullEndpoint) Close()                                  {}
 
 var _ stack.NetworkDispatcher = (*testNetworkDispatcher)(nil)
 
diff --git a/pkg/tcpip/link/pipe/pipe.go b/pkg/tcpip/link/pipe/pipe.go
index 7cdf422b3..06df08633 100644
--- a/pkg/tcpip/link/pipe/pipe.go
+++ b/pkg/tcpip/link/pipe/pipe.go
@@ -138,3 +138,6 @@ func (*Endpoint) AddHeader(*stack.PacketBuffer) {}
 
 // ParseHeader implements stack.LinkEndpoint.
 func (*Endpoint) ParseHeader(*stack.PacketBuffer) bool { return true }
+
+// Close implements stack.LinkEndpoint.
+func (e *Endpoint) Close() {}
diff --git a/pkg/tcpip/link/veth/BUILD b/pkg/tcpip/link/veth/BUILD
index 8911686a6..5108862b4 100644
--- a/pkg/tcpip/link/veth/BUILD
+++ b/pkg/tcpip/link/veth/BUILD
@@ -19,7 +19,17 @@ go_library(
 
 go_test(
     name = "veth_test",
-    srcs = ["veth_test.go"],
-    library = ":veth",
-    deps = ["//pkg/tcpip"],
+    size = "small",
+    srcs = [
+        "veth_test.go",
+    ],
+    deps = [
+        "//pkg/buffer",
+        "//pkg/refs",
+        "//pkg/tcpip",
+        "//pkg/tcpip/header",
+        "//pkg/tcpip/link/ethernet",
+        "//pkg/tcpip/link/veth",
+        "//pkg/tcpip/stack",
+    ],
 )
diff --git a/pkg/tcpip/link/veth/veth.go b/pkg/tcpip/link/veth/veth.go
index 3c9226229..2e8daa09b 100644
--- a/pkg/tcpip/link/veth/veth.go
+++ b/pkg/tcpip/link/veth/veth.go
@@ -107,7 +107,12 @@ func (e *Endpoint) Close() {
 	e.stack = nil
 	e.mu.Unlock()
 	if stack != nil {
-		stack.RemoveNIC(idx)
+		// The pair endpoint can live in the current stack or another one.
+		// RemoveNIC will take the stack lock, so let's run it in another
+		// goroutine to avoid lock conflicts.
+		go func() {
+			stack.RemoveNIC(idx)
+		}()
 	}
 	close(*e.backlogQueue)
 }
diff --git a/pkg/tcpip/link/veth/veth_test.go b/pkg/tcpip/link/veth/veth_test.go
index 76433e977..1da4ee0bc 100644
--- a/pkg/tcpip/link/veth/veth_test.go
+++ b/pkg/tcpip/link/veth/veth_test.go
@@ -11,20 +11,28 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-package veth
+
+package veth_test
 
 import (
+	"os"
 	"testing"
+	"time"
 
+	"gvisor.dev/gvisor/pkg/buffer"
+	"gvisor.dev/gvisor/pkg/refs"
 	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
+	"gvisor.dev/gvisor/pkg/tcpip/link/ethernet"
+	"gvisor.dev/gvisor/pkg/tcpip/link/veth"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 func TestSetLinkAddress(t *testing.T) {
 	addrs := []tcpip.LinkAddress{"abc", "def"}
-	e := &Endpoint{
-		linkAddr: tcpip.LinkAddress("xyz"),
-	}
+	e, e2 := veth.NewPair(1500)
 	defer e.Close()
+	defer e2.Close()
 	for _, addr := range addrs {
 		e.SetLinkAddress(addr)
 
@@ -33,3 +41,97 @@ func TestSetLinkAddress(t *testing.T) {
 		}
 	}
 }
+
+type testNetworkDispatcher struct {
+	ch chan *stack.PacketBuffer
+}
+
+func (d *testNetworkDispatcher) DeliverNetworkPacket(_ tcpip.NetworkProtocolNumber, pkt *stack.PacketBuffer) {
+	pkt.IncRef()
+	d.ch <- pkt
+}
+
+func (*testNetworkDispatcher) DeliverLinkPacket(tcpip.NetworkProtocolNumber, *stack.PacketBuffer) {
+	panic("not implemented")
+}
+
+func TestWritePacket(t *testing.T) {
+	const (
+		localLinkAddr  = tcpip.LinkAddress("\x02\x02\x03\x04\x05\x06")
+		remoteLinkAddr = tcpip.LinkAddress("\x02\x02\x03\x04\x05\x07")
+
+		netProto = 55
+		nicID    = 5
+	)
+
+	veth1, veth2 := veth.NewPair(1500)
+	veth1.SetLinkAddress(localLinkAddr)
+	veth2.SetLinkAddress(remoteLinkAddr)
+
+	s := stack.New(stack.Options{})
+	if err := s.CreateNIC(nicID, ethernet.New(veth1)); err != nil {
+		t.Fatalf("s.CreateNIC(%d, _): %s", nicID, err)
+	}
+
+	sink := &testNetworkDispatcher{ch: make(chan *stack.PacketBuffer, 1)}
+	veth2Ethernet := ethernet.New(veth2)
+	veth2Ethernet.Attach(sink)
+
+	if err := s.WritePacketToRemote(nicID, remoteLinkAddr, netProto, buffer.Buffer{}); err != nil {
+		t.Fatalf("s.WritePacketToRemote(%d, %s, _): %s", nicID, remoteLinkAddr, err)
+	}
+	pkt := <-sink.ch
+	if pkt == nil {
+		t.Fatal("expected to read a packet")
+	}
+
+	eth := header.Ethernet(pkt.LinkHeader().Slice())
+	pkt.DecRef()
+	if got := eth.SourceAddress(); got != localLinkAddr {
+		t.Errorf("got eth.SourceAddress() = %s, want = %s", got, localLinkAddr)
+	}
+	if got := eth.DestinationAddress(); got != remoteLinkAddr {
+		t.Errorf("got eth.DestinationAddress() = %s, want = %s", got, remoteLinkAddr)
+	}
+	if got := eth.Type(); got != netProto {
+		t.Errorf("got eth.Type() = %d, want = %d", got, netProto)
+	}
+}
+
+func TestDestroyDevices(t *testing.T) {
+	const (
+		vethFirstID  = 5
+		vethSecondID = 6
+	)
+
+	veth1, veth2 := veth.NewPair(1500)
+
+	s1 := stack.New(stack.Options{})
+	if err := s1.CreateNIC(vethFirstID, ethernet.New(veth1)); err != nil {
+		t.Fatalf("s.CreateNIC(%d, _): %s", vethFirstID, err)
+	}
+	veth1.SetStack(s1, vethFirstID)
+
+	s2 := stack.New(stack.Options{})
+	if err := s2.CreateNIC(vethSecondID, ethernet.New(veth2)); err != nil {
+		t.Fatalf("s.CreateNIC(%d, _): %s", vethSecondID, err)
+	}
+	veth2.SetStack(s2, vethSecondID)
+
+	s1.RemoveNIC(vethFirstID)
+	timeout := time.Millisecond
+	for s2.HasNIC(vethSecondID) && timeout < 5*time.Second {
+		time.Sleep(timeout)
+		timeout += timeout
+	}
+	if s2.HasNIC(vethSecondID) {
+		t.Fatalf("veth2 hasn't been destroyed")
+	}
+}
+
+func TestMain(m *testing.M) {
+	refs.SetLeakMode(refs.LeaksPanic)
+	code := m.Run()
+	refs.DoLeakCheck()
+	os.Exit(code)
+}
diff --git a/pkg/tcpip/link/waitable/waitable.go b/pkg/tcpip/link/waitable/waitable.go
index 2a8a9c659..46bd0b617 100644
--- a/pkg/tcpip/link/waitable/waitable.go
+++ b/pkg/tcpip/link/waitable/waitable.go
@@ -176,3 +176,8 @@ func (e *Endpoint) AddHeader(pkt *stack.PacketBuffer) {
 func (e *Endpoint) ParseHeader(pkt *stack.PacketBuffer) bool {
 	return e.lower.ParseHeader(pkt)
 }
+
+// Close implements stack.LinkEndpoint.
+func (e *Endpoint) Close() {
+	e.lower.Close()
+}
diff --git a/pkg/tcpip/link/waitable/waitable_test.go b/pkg/tcpip/link/waitable/waitable_test.go
index 9f8c8c169..5f6ebbeb4 100644
--- a/pkg/tcpip/link/waitable/waitable_test.go
+++ b/pkg/tcpip/link/waitable/waitable_test.go
@@ -101,6 +101,9 @@ func (*countedEndpoint) ParseHeader(*stack.PacketBuffer) bool {
 	panic("unimplemented")
 }
 
+// Close implements stack.LinkEndpoint.
+func (*countedEndpoint) Close() {}
+
 func TestWaitWrite(t *testing.T) {
 	ep := &countedEndpoint{}
 	wep := New(ep)
diff --git a/pkg/tcpip/link/xdp/endpoint.go b/pkg/tcpip/link/xdp/endpoint.go
index a053fbcf7..0ccbcf502 100644
--- a/pkg/tcpip/link/xdp/endpoint.go
+++ b/pkg/tcpip/link/xdp/endpoint.go
@@ -410,3 +410,6 @@ func (ep *endpoint) dispatch() (bool, tcpip.Error) {
 		}
 	}
 }
+
+// Close implements stack.LinkEndpoint.
+func (*endpoint) Close() {}
diff --git a/pkg/tcpip/network/ip_test.go b/pkg/tcpip/network/ip_test.go
index 307ac23a2..2ba467bfa 100644
--- a/pkg/tcpip/network/ip_test.go
+++ b/pkg/tcpip/network/ip_test.go
@@ -382,6 +382,9 @@ func (*testInterface) CheckLocalAddress(tcpip.NetworkProtocolNumber, tcpip.Addre
 	return false
 }
 
+// Close implements stack.LinkEndpoint.
+func (*testInterface) Close() {}
+
 func TestSourceAddressValidation(t *testing.T) {
 	rxIPv4ICMP := func(e *channel.Endpoint, src tcpip.Address) {
 		totalLen := header.IPv4MinimumSize + header.ICMPv4MinimumSize
diff --git a/pkg/tcpip/network/ipv6/icmp_test.go b/pkg/tcpip/network/ipv6/icmp_test.go
index 0f405c417..a2a2e2caf 100644
--- a/pkg/tcpip/network/ipv6/icmp_test.go
+++ b/pkg/tcpip/network/ipv6/icmp_test.go
@@ -61,6 +61,8 @@ type stubLinkEndpoint struct {
 	stack.LinkEndpoint
 }
 
+func (*stubLinkEndpoint) Close() {}
+
 func (*stubLinkEndpoint) MTU() uint32 {
 	return defaultMTU
 }
diff --git a/pkg/tcpip/stack/bridge.go b/pkg/tcpip/stack/bridge.go
index 9a70a9042..f7bfcacf9 100644
--- a/pkg/tcpip/stack/bridge.go
+++ b/pkg/tcpip/stack/bridge.go
@@ -211,3 +211,6 @@ func (b *BridgeEndpoint) AddHeader(pkt *PacketBuffer) {
 func (b *BridgeEndpoint) ParseHeader(*PacketBuffer) bool {
 	return true
 }
+
+// Close implements stack.LinkEndpoint.Close.
+func (b *BridgeEndpoint) Close() {}
diff --git a/pkg/tcpip/stack/forwarding_test.go b/pkg/tcpip/stack/forwarding_test.go
index 8e695fbff..07bec92dc 100644
--- a/pkg/tcpip/stack/forwarding_test.go
+++ b/pkg/tcpip/stack/forwarding_test.go
@@ -141,6 +141,7 @@ func (f *fwdTestNetworkEndpoint) WriteHeaderIncludedPacket(r *Route, pkt *Packet
 	return f.nic.WritePacket(r, pkt)
 }
 
+// Close implements stack.LinkEndpoint.
 func (f *fwdTestNetworkEndpoint) Close() {
 	f.AddressableEndpointState.Cleanup()
 }
@@ -338,6 +339,8 @@ func (*fwdTestLinkEndpoint) AddHeader(*PacketBuffer) {}
 // ParseHeader implements stack.LinkEndpoint.ParseHeader.
 func (*fwdTestLinkEndpoint) ParseHeader(*PacketBuffer) bool { return true }
 
+func (*fwdTestLinkEndpoint) Close() {}
+
 func fwdTestNetFactory(t *testing.T, proto *fwdTestNetworkProtocol) (*faketime.ManualClock, *fwdTestLinkEndpoint, *fwdTestLinkEndpoint) {
 	clock := faketime.NewManualClock()
 	// Create a stack with the network protocol and two NICs.
diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index e7979122e..e3ffae5dd 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -329,6 +329,7 @@ func (n *nic) remove() tcpip.Error {
 	// Prevent packets from going down to the link before shutting the link down.
 	n.qDisc.Close()
 	n.NetworkLinkEndpoint.Attach(nil)
+	n.NetworkLinkEndpoint.Close()
 
 	return nil
 }
diff --git a/pkg/tcpip/stack/registration.go b/pkg/tcpip/stack/registration.go
index 53e70a2ae..27e7cc270 100644
--- a/pkg/tcpip/stack/registration.go
+++ b/pkg/tcpip/stack/registration.go
@@ -1138,6 +1138,9 @@ type NetworkLinkEndpoint interface {
 
 	// ParseHeader parses the link layer header to the packet.
 	ParseHeader(*PacketBuffer) bool
+
+	// Close is called when the endpoint is removed from a stack.
+	Close()
 }
 
 // QueueingDiscipline provides a queueing strategy for outgoing packets (e.g
diff --git a/pkg/tcpip/transport/datagram_test.go b/pkg/tcpip/transport/datagram_test.go
index dd878bfba..7e73a40d3 100644
--- a/pkg/tcpip/transport/datagram_test.go
+++ b/pkg/tcpip/transport/datagram_test.go
@@ -183,6 +183,9 @@ func (e *mockEndpoint) pktsSize() int {
 	return s
 }
 
+// Close implements stack.LinkEndpoint.
+func (*mockEndpoint) Close() {}
+
 func TestSndBuf(t *testing.T) {
 	const nicID = 1