Add leak checking to ethernet and muxed tests.

PiperOrigin-RevId: 422698888
gristlabs · Jan 19, 2022 · 3d578af · 3d578af
1 parent 4ce5c43
commit 3d578af
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 3 deletions.
diff --git a/pkg/tcpip/link/ethernet/BUILD b/pkg/tcpip/link/ethernet/BUILD
@@ -20,6 +20,8 @@ go_test(
     srcs = ["ethernet_test.go"],
     deps = [
         ":ethernet",
+        "//pkg/refs",
+        "//pkg/refsvfs2",
         "//pkg/tcpip",
         "//pkg/tcpip/buffer",
         "//pkg/tcpip/header",

diff --git a/pkg/tcpip/link/ethernet/ethernet_test.go b/pkg/tcpip/link/ethernet/ethernet_test.go
@@ -16,8 +16,11 @@ package ethernet_test
 
 import (
 	"fmt"
+	"os"
 	"testing"
 
+	"gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
@@ -63,9 +66,9 @@ func TestDeliverNetworkPacket(t *testing.T) {
 		DstAddr: otherLinkAddr2,
 		Type:    header.IPv4ProtocolNumber,
 	})
-	e.DeliverNetworkPacket("", "", 0, stack.NewPacketBuffer(stack.PacketBufferOptions{
-		Data: eth.ToVectorisedView(),
-	}))
+	p := stack.NewPacketBuffer(stack.PacketBufferOptions{Data: eth.ToVectorisedView()})
+	defer p.DecRef()
+	e.DeliverNetworkPacket("", "", 0, p)
 	if networkDispatcher.networkPackets != 1 {
 		t.Fatalf("got networkDispatcher.networkPackets = %d, want = 1", networkDispatcher.networkPackets)
 	}
@@ -135,6 +138,7 @@ func TestWritePacketsAddHeader(t *testing.T) {
 		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 			ReserveHeaderBytes: int(e.MaxHeaderLength()),
 		})
+		defer pkt.DecRef()
 		pkt.NetworkProtocolNumber = netProto
 		pkt.EgressRoute.RemoteLinkAddress = remoteLinkAddr
 
@@ -165,3 +169,10 @@ func TestWritePacketsAddHeader(t *testing.T) {
 		}
 	}
 }
+
+func TestMain(m *testing.M) {
+	refs.SetLeakMode(refs.LeaksPanic)
+	code := m.Run()
+	refsvfs2.DoLeakCheck()
+	os.Exit(code)
+}
diff --git a/pkg/tcpip/link/muxed/BUILD b/pkg/tcpip/link/muxed/BUILD
@@ -19,6 +19,8 @@ go_test(
     srcs = ["injectable_test.go"],
     library = ":muxed",
     deps = [
+        "//pkg/refs",
+        "//pkg/refsvfs2",
         "//pkg/tcpip",
         "//pkg/tcpip/buffer",
         "//pkg/tcpip/link/fdbased",

diff --git a/pkg/tcpip/link/muxed/injectable_test.go b/pkg/tcpip/link/muxed/injectable_test.go
@@ -21,6 +21,8 @@ import (
 	"testing"
 
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/link/fdbased"
@@ -50,6 +52,7 @@ func TestInjectableEndpointDispatch(t *testing.T) {
 		ReserveHeaderBytes: 1,
 		Data:               buffer.NewViewFromBytes([]byte{0xFB}).ToVectorisedView(),
 	})
+	defer pkt.DecRef()
 	pkt.TransportHeader().Push(1)[0] = 0xFA
 	pkt.EgressRoute.RemoteAddress = dstIP
 	pkt.NetworkProtocolNumber = ipv4.ProtocolNumber
@@ -77,6 +80,7 @@ func TestInjectableEndpointDispatchHdrOnly(t *testing.T) {
 		ReserveHeaderBytes: 1,
 		Data:               buffer.NewView(0).ToVectorisedView(),
 	})
+	defer pkt.DecRef()
 	pkt.TransportHeader().Push(1)[0] = 0xFA
 	pkt.EgressRoute.RemoteAddress = dstIP
 	pkt.NetworkProtocolNumber = ipv4.ProtocolNumber
@@ -108,3 +112,10 @@ func makeTestInjectableEndpoint(t *testing.T) (*InjectableEndpoint, *os.File, tc
 	endpoint := NewInjectableEndpoint(routes)
 	return endpoint, os.NewFile(uintptr(pair[0]), "test route end"), dstIP
 }
+
+func TestMain(m *testing.M) {
+	refs.SetLeakMode(refs.LeaksPanic)
+	code := m.Run()
+	refsvfs2.DoLeakCheck()
+	os.Exit(code)
+}