1 files changed, 9 insertions, 13 deletions
diff --git a/pkg/tcpip/network/ipv4/ipv4_test.go b/pkg/tcpip/network/ipv4/ipv4_test.go
index 1b5a55bea..ae827ca27 100644
--- a/pkg/tcpip/network/ipv4/ipv4_test.go
+++ b/pkg/tcpip/network/ipv4/ipv4_test.go
@@ -36,11 +36,11 @@ func TestExcludeBroadcast(t *testing.T) {
 	s := stack.New([]string{ipv4.ProtocolName}, []string{udp.ProtocolName}, stack.Options{})
 
 	const defaultMTU = 65536
-	id, _ := channel.New(256, defaultMTU, "")
+	ep := stack.LinkEndpoint(channel.New(256, defaultMTU, ""))
 	if testing.Verbose() {
-		id = sniffer.New(id)
+		ep = sniffer.New(ep)
 	}
-	if err := s.CreateNIC(1, id); err != nil {
+	if err := s.CreateNIC(1, ep); err != nil {
 		t.Fatalf("CreateNIC failed: %v", err)
 	}
 
@@ -184,15 +184,12 @@ type errorChannel struct {
 // newErrorChannel creates a new errorChannel endpoint. Each call to WritePacket
 // will return successive errors from packetCollectorErrors until the list is
 // empty and then return nil each time.
-func newErrorChannel(size int, mtu uint32, linkAddr tcpip.LinkAddress, packetCollectorErrors []*tcpip.Error) (tcpip.LinkEndpointID, *errorChannel) {
-	_, e := channel.New(size, mtu, linkAddr)
-	ec := errorChannel{
-		Endpoint:              e,
+func newErrorChannel(size int, mtu uint32, linkAddr tcpip.LinkAddress, packetCollectorErrors []*tcpip.Error) *errorChannel {
+	return &errorChannel{
+		Endpoint:              channel.New(size, mtu, linkAddr),
 		Ch:                    make(chan packetInfo, size),
 		packetCollectorErrors: packetCollectorErrors,
 	}
-
-	return stack.RegisterLinkEndpoint(e), &ec
 }
 
 // packetInfo holds all the information about an outbound packet.
@@ -242,9 +239,8 @@ type context struct {
 func buildContext(t *testing.T, packetCollectorErrors []*tcpip.Error, mtu uint32) context {
 	// Make the packet and write it.
 	s := stack.New([]string{ipv4.ProtocolName}, []string{}, stack.Options{})
-	_, linkEP := newErrorChannel(100 /* Enough for all tests. */, mtu, "", packetCollectorErrors)
-	linkEPId := stack.RegisterLinkEndpoint(linkEP)
-	s.CreateNIC(1, linkEPId)
+	ep := newErrorChannel(100 /* Enough for all tests. */, mtu, "", packetCollectorErrors)
+	s.CreateNIC(1, ep)
 	const (
 		src = "\x10\x00\x00\x01"
 		dst = "\x10\x00\x00\x02"
@@ -266,7 +262,7 @@ func buildContext(t *testing.T, packetCollectorErrors []*tcpip.Error, mtu uint32
 	}
 	return context{
 		Route:  r,
-		linkEP: linkEP,
+		linkEP: ep,
 	}
 }