Use struct embedding to avoid casts and reuse methods

- Removed (*testbench.Connection)(&conn) like casts
- Removed redundant definition of Drain, Close and ExpectFrame

PiperOrigin-RevId: 360727788

5 files changed, 24 insertions, 30 deletions

diff --git a/test/packetimpact/tests/icmpv6_param_problem_test.go b/test/packetimpact/tests/icmpv6_param_problem_test.go
index 40d7a491d..1beccb6cf 100644
--- a/test/packetimpact/tests/icmpv6_param_problem_test.go
+++ b/test/packetimpact/tests/icmpv6_param_problem_test.go
@@ -44,8 +44,8 @@ func TestICMPv6ParamProblemTest(t *testing.T) {
 		Payload: []byte("hello world"),
 	}
 
-	toSend := (*testbench.Connection)(&conn).CreateFrame(t, testbench.Layers{&ipv6}, &icmpv6)
-	(*testbench.Connection)(&conn).SendFrame(t, toSend)
+	toSend := conn.CreateFrame(t, testbench.Layers{&ipv6}, &icmpv6)
+	conn.SendFrame(t, toSend)
 
 	// Build the expected ICMPv6 payload, which includes an index to the
 	// problematic byte and also the problematic packet as described in
diff --git a/test/packetimpact/tests/ipv6_fragment_icmp_error_test.go b/test/packetimpact/tests/ipv6_fragment_icmp_error_test.go
index a37867e85..e0b2a2178 100644
--- a/test/packetimpact/tests/ipv6_fragment_icmp_error_test.go
+++ b/test/packetimpact/tests/ipv6_fragment_icmp_error_test.go
@@ -37,7 +37,7 @@ func init() {
 	testbench.Initialize(flag.CommandLine)
 }
 
-func fragmentedICMPEchoRequest(t *testing.T, n *testbench.DUTTestNet, conn *testbench.Connection, firstPayloadLength uint16, payload []byte, secondFragmentOffset uint16) ([]testbench.Layers, [][]byte) {
+func fragmentedICMPEchoRequest(t *testing.T, n *testbench.DUTTestNet, conn *testbench.IPv6Conn, firstPayloadLength uint16, payload []byte, secondFragmentOffset uint16) ([]testbench.Layers, [][]byte) {
 	t.Helper()
 
 	icmpv6Header := header.ICMPv6(make([]byte, header.ICMPv6EchoMinimumSize))
@@ -121,17 +121,16 @@ func TestIPv6ICMPEchoRequestFragmentReassembly(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 			dut := testbench.NewDUT(t)
-			ipv6Conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
-			conn := (*testbench.Connection)(&ipv6Conn)
-			defer ipv6Conn.Close(t)
+			conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
+			defer conn.Close(t)
 
-			fragments, _ := fragmentedICMPEchoRequest(t, dut.Net, conn, test.firstPayloadLength, test.payload, test.secondFragmentOffset)
+			fragments, _ := fragmentedICMPEchoRequest(t, dut.Net, &conn, test.firstPayloadLength, test.payload, test.secondFragmentOffset)
 
 			for _, i := range test.sendFrameOrder {
 				conn.SendFrame(t, fragments[i-1])
 			}
 
-			gotEchoReply, err := ipv6Conn.ExpectFrame(t, testbench.Layers{
+			gotEchoReply, err := conn.ExpectFrame(t, testbench.Layers{
 				&testbench.Ether{},
 				&testbench.IPv6{},
 				&testbench.ICMPv6{
@@ -223,17 +222,16 @@ func TestIPv6FragmentReassemblyTimeout(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 			dut := testbench.NewDUT(t)
-			ipv6Conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
-			conn := (*testbench.Connection)(&ipv6Conn)
-			defer ipv6Conn.Close(t)
+			conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
+			defer conn.Close(t)
 
-			fragments, ipv6Bytes := fragmentedICMPEchoRequest(t, dut.Net, conn, test.firstPayloadLength, test.payload, test.secondFragmentOffset)
+			fragments, ipv6Bytes := fragmentedICMPEchoRequest(t, dut.Net, &conn, test.firstPayloadLength, test.payload, test.secondFragmentOffset)
 
 			for _, i := range test.sendFrameOrder {
 				conn.SendFrame(t, fragments[i-1])
 			}
 
-			gotErrorMessage, err := ipv6Conn.ExpectFrame(t, testbench.Layers{
+			gotErrorMessage, err := conn.ExpectFrame(t, testbench.Layers{
 				&testbench.Ether{},
 				&testbench.IPv6{},
 				&testbench.ICMPv6{
@@ -319,17 +317,16 @@ func TestIPv6FragmentParamProblem(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 			dut := testbench.NewDUT(t)
-			ipv6Conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
-			conn := (*testbench.Connection)(&ipv6Conn)
-			defer ipv6Conn.Close(t)
+			conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
+			defer conn.Close(t)
 
-			fragments, ipv6Bytes := fragmentedICMPEchoRequest(t, dut.Net, conn, test.firstPayloadLength, test.payload, test.secondFragmentOffset)
+			fragments, ipv6Bytes := fragmentedICMPEchoRequest(t, dut.Net, &conn, test.firstPayloadLength, test.payload, test.secondFragmentOffset)
 
 			for _, i := range test.sendFrameOrder {
 				conn.SendFrame(t, fragments[i-1])
 			}
 
-			gotErrorMessage, err := ipv6Conn.ExpectFrame(t, testbench.Layers{
+			gotErrorMessage, err := conn.ExpectFrame(t, testbench.Layers{
 				&testbench.Ether{},
 				&testbench.IPv6{},
 				&testbench.ICMPv6{
diff --git a/test/packetimpact/tests/ipv6_unknown_options_action_test.go b/test/packetimpact/tests/ipv6_unknown_options_action_test.go
index cb5396417..f999d13d2 100644
--- a/test/packetimpact/tests/ipv6_unknown_options_action_test.go
+++ b/test/packetimpact/tests/ipv6_unknown_options_action_test.go
@@ -141,9 +141,8 @@ func TestIPv6UnknownOptionAction(t *testing.T) {
 	} {
 		t.Run(tt.description, func(t *testing.T) {
 			dut := testbench.NewDUT(t)
-			ipv6Conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
-			conn := (*testbench.Connection)(&ipv6Conn)
-			defer ipv6Conn.Close(t)
+			conn := dut.Net.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
+			defer conn.Close(t)
 
 			outgoingOverride := testbench.Layers{}
 			if tt.multicastDst {
@@ -166,7 +165,7 @@ func TestIPv6UnknownOptionAction(t *testing.T) {
 			// after the IPv6 header (after NextHeader and ExtHdrLen).
 			binary.BigEndian.PutUint32(icmpv6Payload, header.IPv6MinimumSize+2)
 			icmpv6Payload = append(icmpv6Payload, invokingPacket...)
-			gotICMPv6, err := ipv6Conn.ExpectFrame(t, testbench.Layers{
+			gotICMPv6, err := conn.ExpectFrame(t, testbench.Layers{
 				&testbench.Ether{},
 				&testbench.IPv6{},
 				&testbench.ICMPv6{
diff --git a/test/packetimpact/tests/tcp_network_unreachable_test.go b/test/packetimpact/tests/tcp_network_unreachable_test.go
index b1690aaaf..53dc903e4 100644
--- a/test/packetimpact/tests/tcp_network_unreachable_test.go
+++ b/test/packetimpact/tests/tcp_network_unreachable_test.go
@@ -57,8 +57,7 @@ func TestTCPSynSentUnreachable(t *testing.T) {
 	}
 
 	// Send a host unreachable message.
-	rawConn := (*testbench.Connection)(&conn)
-	layers := rawConn.CreateFrame(t, nil)
+	layers := conn.CreateFrame(t, nil)
 	layers = layers[:len(layers)-1]
 	const ipLayer = 1
 	const tcpLayer = ipLayer + 1
@@ -76,7 +75,7 @@ func TestTCPSynSentUnreachable(t *testing.T) {
 	}
 
 	layers = append(layers, &icmpv4, ip, tcp)
-	rawConn.SendFrameStateless(t, layers)
+	conn.SendFrameStateless(t, layers)
 
 	if err := getConnectError(t, &dut, clientFD); err != unix.EHOSTUNREACH {
 		t.Errorf("got connect() = %v, want EHOSTUNREACH", err)
@@ -112,8 +111,7 @@ func TestTCPSynSentUnreachable6(t *testing.T) {
 	}
 
 	// Send a host unreachable message.
-	rawConn := (*testbench.Connection)(&conn)
-	layers := rawConn.CreateFrame(t, nil)
+	layers := conn.CreateFrame(t, nil)
 	layers = layers[:len(layers)-1]
 	const ipLayer = 1
 	const tcpLayer = ipLayer + 1
@@ -132,7 +130,7 @@ func TestTCPSynSentUnreachable6(t *testing.T) {
 		Payload: []byte{0, 0, 0, 0},
 	}
 	layers = append(layers, &icmpv6, ip, tcp)
-	rawConn.SendFrameStateless(t, layers)
+	conn.SendFrameStateless(t, layers)
 
 	if err := getConnectError(t, &dut, clientFD); err != unix.ENETUNREACH {
 		t.Errorf("got connect() = %v, want EHOSTUNREACH", err)
diff --git a/test/packetimpact/tests/udp_icmp_error_propagation_test.go b/test/packetimpact/tests/udp_icmp_error_propagation_test.go
index 58d49d31a..3fca8c7a3 100644
--- a/test/packetimpact/tests/udp_icmp_error_propagation_test.go
+++ b/test/packetimpact/tests/udp_icmp_error_propagation_test.go
@@ -102,7 +102,7 @@ func wantErrno(c connectionMode, icmpErr icmpError) syscall.Errno {
 func sendICMPError(t *testing.T, conn *testbench.UDPIPv4, icmpErr icmpError, udp *testbench.UDP) {
 	t.Helper()
 
-	layers := (*testbench.Connection)(conn).CreateFrame(t, nil)
+	layers := conn.CreateFrame(t, nil)
 	layers = layers[:len(layers)-1]
 	ip, ok := udp.Prev().(*testbench.IPv4)
 	if !ok {
@@ -120,7 +120,7 @@ func sendICMPError(t *testing.T, conn *testbench.UDPIPv4, icmpErr icmpError, udp
 	// resulting in a mal-formed packet.
 	layers = append(layers, icmpErr.ToICMPv4(), ip, udp)
 
-	(*testbench.Connection)(conn).SendFrameStateless(t, layers)
+	conn.SendFrameStateless(t, layers)
 }
 
 // testRecv tests observing the ICMP error through the recv syscall. A packet