Enable leak checking for network tests.

Updates #6910

PiperOrigin-RevId: 429127828

pkg/tcpip/network/BUILD

@@ -7,9 +7,12 @@ go_test(
     size = "small",
     srcs = [
         "ip_test.go",
+        "main_test.go",
         "multicast_group_test.go",
     ],
     deps = [
+        "//pkg/refs",
+        "//pkg/refsvfs2",
         "//pkg/sync",
         "//pkg/tcpip",
         "//pkg/tcpip/buffer",

pkg/tcpip/network/ip_test.go

@@ -222,11 +222,26 @@ func (*testObject) AddHeader(*stack.PacketBuffer) {
 	panic("not implemented")
 }
 
-func buildIPv4Route(local, remote tcpip.Address) (*stack.Route, tcpip.Error) {
+type testContext struct {
+	s *stack.Stack
+}
+
+func newTestContext() *testContext {
 	s := stack.New(stack.Options{
-		NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4.NewProtocol},
+		NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4.NewProtocol, ipv6.NewProtocol},
 		TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol, tcp.NewProtocol},
+		RawFactory:         raw.EndpointFactory{},
 	})
+	return &testContext{s: s}
+}
+
+func (ctx *testContext) cleanup() {
+	ctx.s.Close()
+	ctx.s.Wait()
+}
+
+func buildIPv4Route(ctx *testContext, local, remote tcpip.Address) (*stack.Route, tcpip.Error) {
+	s := ctx.s
 	s.CreateNIC(nicID, loopback.New())
 	protocolAddr := tcpip.ProtocolAddress{
 		Protocol:          ipv4.ProtocolNumber,
@@ -244,11 +259,8 @@ func buildIPv4Route(local, remote tcpip.Address) (*stack.Route, tcpip.Error) {
 	return s.FindRoute(nicID, local, remote, ipv4.ProtocolNumber, false /* multicastLoop */)
 }
 
-func buildIPv6Route(local, remote tcpip.Address) (*stack.Route, tcpip.Error) {
-	s := stack.New(stack.Options{
-		NetworkProtocols:   []stack.NetworkProtocolFactory{ipv6.NewProtocol},
-		TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol, tcp.NewProtocol},
-	})
+func buildIPv6Route(ctx *testContext, local, remote tcpip.Address) (*stack.Route, tcpip.Error) {
+	s := ctx.s
 	s.CreateNIC(nicID, loopback.New())
 	protocolAddr := tcpip.ProtocolAddress{
 		Protocol:          ipv6.ProtocolNumber,
@@ -266,13 +278,8 @@ func buildIPv6Route(local, remote tcpip.Address) (*stack.Route, tcpip.Error) {
 	return s.FindRoute(nicID, local, remote, ipv6.ProtocolNumber, false /* multicastLoop */)
 }
 
-func buildDummyStackWithLinkEndpoint(t *testing.T, mtu uint32) (*stack.Stack, *channel.Endpoint) {
+func addLinkEndpointToStackWithMTU(t *testing.T, s *stack.Stack, mtu uint32) *channel.Endpoint {
 	t.Helper()
-
-	s := stack.New(stack.Options{
-		NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4.NewProtocol, ipv6.NewProtocol},
-		TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol, tcp.NewProtocol},
-	})
 	e := channel.New(1, mtu, "")
 	if err := s.CreateNIC(nicID, e); err != nil {
 		t.Fatalf("CreateNIC(%d, _) = %s", nicID, err)
@@ -288,14 +295,12 @@ func buildDummyStackWithLinkEndpoint(t *testing.T, mtu uint32) (*stack.Stack, *c
 		t.Fatalf("AddProtocolAddress(%d, %+v, {}) = %s", nicID, v6Addr, err)
 	}
 
-	return s, e
+	return e
 }
 
-func buildDummyStack(t *testing.T) *stack.Stack {
+func addLinkEndpointToStack(t *testing.T, s *stack.Stack) *channel.Endpoint {
 	t.Helper()
-
-	s, _ := buildDummyStackWithLinkEndpoint(t, header.IPv6MinimumMTU)
-	return s
+	return addLinkEndpointToStackWithMTU(t, s, header.IPv6MinimumMTU)
 }
 
 var _ stack.NetworkInterface = (*testInterface)(nil)
@@ -380,9 +385,11 @@ func TestSourceAddressValidation(t *testing.T) {
 		})
 		ip.SetChecksum(^ip.CalculateChecksum())
 
-		e.InjectInbound(header.IPv4ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{
+		pktBuf := stack.NewPacketBuffer(stack.PacketBufferOptions{
 			Data: hdr.View().ToVectorisedView(),
-		}))
+		})
+		e.InjectInbound(header.IPv4ProtocolNumber, pktBuf)
+		pktBuf.DecRef()
 	}
 
 	rxIPv6ICMP := func(e *channel.Endpoint, src tcpip.Address) {
@@ -405,9 +412,11 @@ func TestSourceAddressValidation(t *testing.T) {
 			SrcAddr:           src,
 			DstAddr:           localIPv6Addr,
 		})
-		e.InjectInbound(header.IPv6ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{
+		pktBuf := stack.NewPacketBuffer(stack.PacketBufferOptions{
 			Data: hdr.View().ToVectorisedView(),
-		}))
+		})
+		e.InjectInbound(header.IPv6ProtocolNumber, pktBuf)
+		pktBuf.DecRef()
 	}
 
 	tests := []struct {
@@ -471,7 +480,11 @@ func TestSourceAddressValidation(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			s, e := buildDummyStackWithLinkEndpoint(t, header.IPv6MinimumMTU)
+			ctx := newTestContext()
+			defer ctx.cleanup()
+			s := ctx.s
+
+			e := addLinkEndpointToStack(t, s)
 			test.rxICMP(e, test.srcAddress)
 
 			var wantValid uint64
@@ -515,6 +528,11 @@ func TestEnableWhenNICDisabled(t *testing.T) {
 			s := stack.New(stack.Options{
 				NetworkProtocols: []stack.NetworkProtocolFactory{test.protocolFactory},
 			})
+			defer func() {
+				s.Close()
+				s.Wait()
+			}()
+
 			p := s.NetworkProtocolInstance(test.protoNum)
 
 			// We pass nil for all parameters except the NetworkInterface and Stack
@@ -573,7 +591,10 @@ func TestEnableWhenNICDisabled(t *testing.T) {
 }
 
 func TestIPv4Send(t *testing.T) {
-	s := buildDummyStack(t)
+	ctx := newTestContext()
+	defer ctx.cleanup()
+	s := ctx.s
+
 	proto := s.NetworkProtocolInstance(ipv4.ProtocolNumber)
 	nic := testInterface{
 		testObject: testObject{
@@ -595,6 +616,7 @@ func TestIPv4Send(t *testing.T) {
 		ReserveHeaderBytes: int(ep.MaxHeaderLength()),
 		Data:               payload.ToVectorisedView(),
 	})
+	defer pkt.DecRef()
 
 	// Issue the write.
 	nic.testObject.protocol = 123
@@ -602,7 +624,7 @@ func TestIPv4Send(t *testing.T) {
 	nic.testObject.dstAddr = remoteIPv4Addr
 	nic.testObject.contents = payload
 
-	r, err := buildIPv4Route(localIPv4Addr, remoteIPv4Addr)
+	r, err := buildIPv4Route(ctx, localIPv4Addr, remoteIPv4Addr)
 	if err != nil {
 		t.Fatalf("could not find route: %v", err)
 	}
@@ -659,6 +681,7 @@ func TestReceive(t *testing.T) {
 					Data: view.ToVectorisedView(),
 				})
 				ep.HandlePacket(pkt)
+				pkt.DecRef()
 			},
 		},
 		{
@@ -694,6 +717,7 @@ func TestReceive(t *testing.T) {
 					Data: view.ToVectorisedView(),
 				})
 				ep.HandlePacket(pkt)
+				pkt.DecRef()
 			},
 		},
 	}
@@ -703,6 +727,11 @@ func TestReceive(t *testing.T) {
 			s := stack.New(stack.Options{
 				NetworkProtocols: []stack.NetworkProtocolFactory{test.protoFactory},
 			})
+			defer func() {
+				s.Close()
+				s.Wait()
+			}()
+
 			nic := testInterface{
 				testObject: testObject{
 					t:  t,
@@ -823,7 +852,10 @@ func TestIPv4ReceiveControl(t *testing.T) {
 	}
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
-			s := buildDummyStack(t)
+			ctx := newTestContext()
+			defer ctx.cleanup()
+			s := ctx.s
+
 			proto := s.NetworkProtocolInstance(ipv4.ProtocolNumber)
 			nic := testInterface{
 				testObject: testObject{
@@ -900,6 +932,7 @@ func TestIPv4ReceiveControl(t *testing.T) {
 
 			pkt := truncatedPacket(view, c.trunc, header.IPv4MinimumSize)
 			ep.HandlePacket(pkt)
+			pkt.DecRef()
 			if want := c.expectedCount; nic.testObject.controlCalls != want {
 				t.Fatalf("Bad number of control calls for %q case: got %v, want %v", c.name, nic.testObject.controlCalls, want)
 			}
@@ -908,9 +941,10 @@ func TestIPv4ReceiveControl(t *testing.T) {
 }
 
 func TestIPv4FragmentationReceive(t *testing.T) {
-	s := stack.New(stack.Options{
-		NetworkProtocols: []stack.NetworkProtocolFactory{ipv4.NewProtocol},
-	})
+	ctx := newTestContext()
+	defer ctx.cleanup()
+	s := ctx.s
+
 	proto := s.NetworkProtocolInstance(ipv4.ProtocolNumber)
 	nic := testInterface{
 		testObject: testObject{
@@ -968,11 +1002,6 @@ func TestIPv4FragmentationReceive(t *testing.T) {
 	nic.testObject.dstAddr = localIPv4Addr
 	nic.testObject.contents = append(frag1[header.IPv4MinimumSize:totalLen], frag2[header.IPv4MinimumSize:totalLen]...)
 
-	// Send first segment.
-	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		Data: frag1.ToVectorisedView(),
-	})
-
 	addressableEndpoint, ok := ep.(stack.AddressableEndpoint)
 	if !ok {
 		t.Fatal("expected IPv4 network endpoint to implement stack.AddressableEndpoint")
@@ -984,7 +1013,13 @@ func TestIPv4FragmentationReceive(t *testing.T) {
 		ep.DecRef()
 	}
 
+	// Send first segment.
+	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
+		Data: frag1.ToVectorisedView(),
+	})
 	ep.HandlePacket(pkt)
+	pkt.DecRef()
+
 	if nic.testObject.dataCalls != 0 {
 		t.Fatalf("Bad number of data calls: got %d, want 0", nic.testObject.dataCalls)
 	}
@@ -997,6 +1032,8 @@ func TestIPv4FragmentationReceive(t *testing.T) {
 		Data: frag2.ToVectorisedView(),
 	})
 	ep.HandlePacket(pkt)
+	pkt.DecRef()
+
 	if nic.testObject.dataCalls != 1 {
 		t.Fatalf("Bad number of data calls: got %d, want 1", nic.testObject.dataCalls)
 	}
@@ -1006,7 +1043,10 @@ func TestIPv4FragmentationReceive(t *testing.T) {
 }
 
 func TestIPv6Send(t *testing.T) {
-	s := buildDummyStack(t)
+	ctx := newTestContext()
+	defer ctx.cleanup()
+	s := ctx.s
+
 	proto := s.NetworkProtocolInstance(ipv6.ProtocolNumber)
 	nic := testInterface{
 		testObject: testObject{
@@ -1031,14 +1071,14 @@ func TestIPv6Send(t *testing.T) {
 		ReserveHeaderBytes: int(ep.MaxHeaderLength()),
 		Data:               payload.ToVectorisedView(),
 	})
-
+	defer pkt.DecRef()
 	// Issue the write.
 	nic.testObject.protocol = 123
 	nic.testObject.srcAddr = localIPv6Addr
 	nic.testObject.dstAddr = remoteIPv6Addr
 	nic.testObject.contents = payload
 
-	r, err := buildIPv6Route(localIPv6Addr, remoteIPv6Addr)
+	r, err := buildIPv6Route(ctx, localIPv6Addr, remoteIPv6Addr)
 	if err != nil {
 		t.Fatalf("could not find route: %v", err)
 	}
@@ -1161,7 +1201,10 @@ func TestIPv6ReceiveControl(t *testing.T) {
 	}
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
-			s := buildDummyStack(t)
+			ctx := newTestContext()
+			defer ctx.cleanup()
+			s := ctx.s
+
 			proto := s.NetworkProtocolInstance(ipv6.ProtocolNumber)
 			nic := testInterface{
 				testObject: testObject{
@@ -1251,6 +1294,7 @@ func TestIPv6ReceiveControl(t *testing.T) {
 			}
 			pkt := truncatedPacket(view, c.trunc, header.IPv6MinimumSize)
 			ep.HandlePacket(pkt)
+			pkt.DecRef()
 			if want := c.expectedCount; nic.testObject.controlCalls != want {
 				t.Fatalf("Bad number of control calls for %q case: got %v, want %v", c.name, nic.testObject.controlCalls, want)
 			}
@@ -1683,6 +1727,11 @@ func TestWriteHeaderIncludedPacket(t *testing.T) {
 					s := stack.New(stack.Options{
 						NetworkProtocols: []stack.NetworkProtocolFactory{test.protoFactory},
 					})
+					defer func() {
+						s.Close()
+						s.Wait()
+					}()
+
 					e := channel.New(1, header.IPv6MinimumMTU, "")
 					if err := s.CreateNIC(nicID, e); err != nil {
 						t.Fatalf("s.CreateNIC(%d, _): %s", nicID, err)
@@ -1704,9 +1753,11 @@ func TestWriteHeaderIncludedPacket(t *testing.T) {
 					defer r.Release()
 
 					{
-						err := r.WriteHeaderIncludedPacket(stack.NewPacketBuffer(stack.PacketBufferOptions{
+						pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 							Data: test.pktGen(t, subTest.srcAddr),
-						}))
+						})
+						err := r.WriteHeaderIncludedPacket(pkt)
+						pkt.DecRef()
 						if diff := cmp.Diff(test.expectedErr, err); diff != "" {
 							t.Fatalf("unexpected error from r.WriteHeaderIncludedPacket(_), (-want, +got):\n%s", diff)
 						}
@@ -1759,9 +1810,11 @@ func TestICMPInclusionSize(t *testing.T) {
 		// Take a copy before InjectInbound takes ownership of vv
 		// as vv may be changed during the call.
 		v := vv.ToView()
-		e.InjectInbound(header.IPv4ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{
+		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 			Data: vv,
-		}))
+		})
+		e.InjectInbound(header.IPv4ProtocolNumber, pkt)
+		pkt.DecRef()
 		return v
 	}
 
@@ -1786,9 +1839,11 @@ func TestICMPInclusionSize(t *testing.T) {
 		// as vv may be changed during the call.
 		v := vv.ToView()
 
-		e.InjectInbound(header.IPv6ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{
+		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 			Data: vv,
-		}))
+		})
+		e.InjectInbound(header.IPv6ProtocolNumber, pkt)
+		pkt.DecRef()
 		return v
 	}
 
@@ -1939,7 +1994,11 @@ func TestICMPInclusionSize(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			s, e := buildDummyStackWithLinkEndpoint(t, test.linkMTU)
+			ctx := newTestContext()
+			defer ctx.cleanup()
+			s := ctx.s
+
+			e := addLinkEndpointToStackWithMTU(t, s, test.linkMTU)
 			// Allocate and initialize the payload view.
 			payload := buffer.NewView(test.payloadLength)
 			for i := 0; i < len(payload); i++ {
@@ -2009,10 +2068,10 @@ func TestJoinLeaveAllRoutersGroup(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			for _, nicDisabled := range [...]bool{true, false} {
 				t.Run(fmt.Sprintf("NIC Disabled = %t", nicDisabled), func(t *testing.T) {
-					s := stack.New(stack.Options{
-						NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4.NewProtocol, ipv6.NewProtocol},
-						TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol, tcp.NewProtocol},
-					})
+					ctx := newTestContext()
+					defer ctx.cleanup()
+					s := ctx.s
+
 					opts := stack.NICOptions{Disabled: nicDisabled}
 					if err := s.CreateNICWithOptions(nicID, channel.New(0, 0, ""), opts); err != nil {
 						t.Fatalf("CreateNICWithOptions(%d, _, %#v) = %s", nicID, opts, err)
@@ -2072,14 +2131,10 @@ func TestSetNICIDBeforeDeliveringToRawEndpoint(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			s := stack.New(stack.Options{
-				NetworkProtocols: []stack.NetworkProtocolFactory{
-					ipv4.NewProtocol,
-					ipv6.NewProtocol,
-				},
-				TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol},
-				RawFactory:         raw.EndpointFactory{},
-			})
+			ctx := newTestContext()
+			defer ctx.cleanup()
+			s := ctx.s
+
 			if err := s.CreateNIC(nicID, loopback.New()); err != nil {
 				t.Fatalf("CreateNIC(%d, _): %s", nicID, err)
 			}

pkg/tcpip/network/main_test.go

@@ -0,0 +1,30 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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 ip_test
+
+import (
+	"os"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/refsvfs2"
+)
+
+func TestMain(m *testing.M) {
+	refs.SetLeakMode(refs.LeaksPanic)
+	code := m.Run()
+	refsvfs2.DoLeakCheck()
+	os.Exit(code)
+}

pkg/tcpip/network/multicast_group_test.go

@@ -118,12 +118,27 @@ func validateIGMPPacket(t *testing.T, p *stack.PacketBuffer, remoteAddress tcpip
 	)
 }
 
-func createStack(t *testing.T, v4, mgpEnabled bool) (*channel.Endpoint, *stack.Stack, *faketime.ManualClock) {
+type multicastTestContext struct {
+	s     *stack.Stack
+	e     *channel.Endpoint
+	clock *faketime.ManualClock
+}
+
+func newMulticastTestContext(t *testing.T, v4, mgpEnabled bool) *multicastTestContext {
 	t.Helper()
 
 	e := channel.New(maxUnsolicitedReports, header.IPv6MinimumMTU, linkAddr)
 	s, clock := createStackWithLinkEndpoint(t, v4, mgpEnabled, e)
-	return e, s, clock
+	return &multicastTestContext{
+		s:     s,
+		e:     e,
+		clock: clock,
+	}
+}
+
+func (ctx *multicastTestContext) cleanup() {
+	ctx.s.Close()
+	ctx.s.Wait()
 }
 
 func createStackWithLinkEndpoint(t *testing.T, v4, mgpEnabled bool, e stack.LinkEndpoint) (*stack.Stack, *faketime.ManualClock) {
@@ -241,9 +256,11 @@ func createAndInjectIGMPPacket(e *channel.Endpoint, igmpType byte, maxRespTime b
 	igmp.SetGroupAddress(groupAddress)
 	igmp.SetChecksum(header.IGMPCalculateChecksum(igmp))
 
-	e.InjectInbound(ipv4.ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{
+	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 		Data: buf.ToVectorisedView(),
-	}))
+	})
+	e.InjectInbound(ipv4.ProtocolNumber, pkt)
+	pkt.DecRef()
 }
 
 // createAndInjectMLDPacket creates and injects an MLD packet with the
@@ -280,9 +297,11 @@ func createAndInjectMLDPacket(e *channel.Endpoint, mldType uint8, maxRespDelay b
 		Dst:    header.IPv6AllNodesMulticastAddress,
 	}))
 
-	e.InjectInbound(ipv6.ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{
+	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 		Data: buf.ToVectorisedView(),
-	}))
+	})
+	e.InjectInbound(ipv6.ProtocolNumber, pkt)
+	pkt.DecRef()
 }
 
 // TestMGPDisabled tests that the multicast group protocol is not enabled by
@@ -328,7 +347,11 @@ func TestMGPDisabled(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			e, s, clock := createStack(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, false /* mgpEnabled */)
+			ctx := newMulticastTestContext(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, false /* mgpEnabled */)
+			defer ctx.cleanup()
+			s := ctx.s
+			e := ctx.e
+			clock := ctx.clock
 
 			// This NIC may join multicast groups when it is enabled but since MGP is
 			// disabled, no reports should be sent.
@@ -451,10 +474,11 @@ func TestMGPReceiveCounters(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			e, s, _ := createStack(t, len(test.groupAddress) == header.IPv4AddressSize /* v4 */, true /* mgpEnabled */)
+			ctx := newMulticastTestContext(t, len(test.groupAddress) == header.IPv4AddressSize /* v4 */, true /* mgpEnabled */)
+			defer ctx.cleanup()
 
-			test.rxMGPkt(e, test.headerType, test.maxRespTime, test.groupAddress)
-			if got := test.statCounter(s).Value(); got != 1 {
+			test.rxMGPkt(ctx.e, test.headerType, test.maxRespTime, test.groupAddress)
+			if got := test.statCounter(ctx.s).Value(); got != 1 {
 				t.Fatalf("got %s received = %d, want = 1", test.name, got)
 			}
 		})
@@ -513,7 +537,9 @@ func TestMGPJoinGroup(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			e, s, clock := createStack(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			ctx := newMulticastTestContext(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			defer ctx.cleanup()
+			s, e, clock := ctx.s, ctx.e, ctx.clock
 
 			var reportCounter uint64
 			if test.checkInitialGroups != nil {
@@ -625,7 +651,9 @@ func TestMGPLeaveGroup(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			e, s, clock := createStack(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			ctx := newMulticastTestContext(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			defer ctx.cleanup()
+			s, e, clock := ctx.s, ctx.e, ctx.clock
 
 			var reportCounter uint64
 			var leaveCounter uint64
@@ -764,7 +792,9 @@ func TestMGPQueryMessages(t *testing.T) {
 
 			for _, subTest := range subTests {
 				t.Run(subTest.name, func(t *testing.T) {
-					e, s, clock := createStack(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+					ctx := newMulticastTestContext(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+					defer ctx.cleanup()
+					s, e, clock := ctx.s, ctx.e, ctx.clock
 
 					var reportCounter uint64
 					if test.checkInitialGroups != nil {
@@ -892,7 +922,9 @@ func TestMGPReportMessages(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			e, s, clock := createStack(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			ctx := newMulticastTestContext(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			defer ctx.cleanup()
+			s, e, clock := ctx.s, ctx.e, ctx.clock
 
 			var reportCounter uint64
 			var leaveCounter uint64
@@ -1076,7 +1108,9 @@ func TestMGPWithNICLifecycle(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			e, s, clock := createStack(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			ctx := newMulticastTestContext(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */)
+			defer ctx.cleanup()
+			s, e, clock := ctx.s, ctx.e, ctx.clock
 
 			var reportCounter uint64
 			var leaveCounter uint64
@@ -1258,7 +1292,10 @@ func TestMGPDisabledOnLoopback(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			s, clock := createStackWithLinkEndpoint(t, test.protoNum == ipv4.ProtocolNumber /* v4 */, true /* mgpEnabled */, loopback.New())
-
+			defer func() {
+				s.Close()
+				s.Wait()
+			}()
 			sentReportStat := test.sentReportStat(s)
 			if got := sentReportStat.Value(); got != 0 {
 				t.Fatalf("got sentReportStat.Value() = %d, want = 0", got)