Refactor connections.go to make it easier to add new connection types.

Rather than have a struct for the state of each type of connection, such as
TCP/IPv4, UDP/IPv4, TCP/IPv6, etc, have a state for each layer, such as UDP,
TCP, IPv4, IPv6.  Those states can be composed into connections.

Tested:
  Existing unit tests still pass/fail as expected.
PiperOrigin-RevId: 306703180

3 files changed, 24 insertions, 8 deletions

diff --git a/test/packetimpact/tests/tcp_should_piggyback_test.go b/test/packetimpact/tests/tcp_should_piggyback_test.go
index f2ab49e51..b0be6ba23 100644
--- a/test/packetimpact/tests/tcp_should_piggyback_test.go
+++ b/test/packetimpact/tests/tcp_should_piggyback_test.go
@@ -40,7 +40,11 @@ func TestPiggyback(t *testing.T) {
 	sampleData := []byte("Sample Data")
 
 	dut.Send(acceptFd, sampleData, 0)
-	conn.ExpectData(tb.TCP{Flags: tb.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, sampleData, time.Second)
+	expectedTCP := tb.TCP{Flags: tb.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}
+	expectedPayload := tb.Payload{Bytes: sampleData}
+	if _, err := conn.ExpectData(&expectedTCP, &expectedPayload, time.Second); err != nil {
+		t.Fatalf("Expected %v but didn't get one: %s", tb.Layers{&expectedTCP, &expectedPayload}, err)
+	}
 
 	// Cause DUT to send us more data as soon as we ACK their first data segment because we have
 	// a small window.
@@ -48,6 +52,8 @@ func TestPiggyback(t *testing.T) {
 
 	// DUT should ACK our segment by piggybacking ACK to their outstanding data segment instead of
 	// sending a separate ACK packet.
-	conn.Send(tb.TCP{Flags: tb.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, &tb.Payload{Bytes: sampleData})
-	conn.ExpectData(tb.TCP{Flags: tb.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, sampleData, time.Second)
+	conn.Send(expectedTCP, &expectedPayload)
+	if _, err := conn.ExpectData(&expectedTCP, &expectedPayload, time.Second); err != nil {
+		t.Fatalf("Expected %v but didn't get one: %s", tb.Layers{&expectedTCP, &expectedPayload}, err)
+	}
 }
diff --git a/test/packetimpact/tests/tcp_window_shrink_test.go b/test/packetimpact/tests/tcp_window_shrink_test.go
index b48cc6491..c9354074e 100644
--- a/test/packetimpact/tests/tcp_window_shrink_test.go
+++ b/test/packetimpact/tests/tcp_window_shrink_test.go
@@ -38,15 +38,22 @@ func TestWindowShrink(t *testing.T) {
 	dut.SetSockOptInt(acceptFd, unix.IPPROTO_TCP, unix.TCP_NODELAY, 1)
 
 	sampleData := []byte("Sample Data")
+	samplePayload := &tb.Payload{Bytes: sampleData}
 
 	dut.Send(acceptFd, sampleData, 0)
-	conn.ExpectData(tb.TCP{}, sampleData, time.Second)
+	if _, err := conn.ExpectData(&tb.TCP{}, samplePayload, time.Second); err != nil {
+		t.Fatalf("expected a packet with payload %v: %s", samplePayload, err)
+	}
 	conn.Send(tb.TCP{Flags: tb.Uint8(header.TCPFlagAck)})
 
 	dut.Send(acceptFd, sampleData, 0)
 	dut.Send(acceptFd, sampleData, 0)
-	conn.ExpectData(tb.TCP{}, sampleData, time.Second)
-	conn.ExpectData(tb.TCP{}, sampleData, time.Second)
+	if _, err := conn.ExpectData(&tb.TCP{}, samplePayload, time.Second); err != nil {
+		t.Fatalf("expected a packet with payload %v: %s", samplePayload, err)
+	}
+	if _, err := conn.ExpectData(&tb.TCP{}, samplePayload, time.Second); err != nil {
+		t.Fatalf("expected a packet with payload %v: %s", samplePayload, err)
+	}
 	// We close our receiving window here
 	conn.Send(tb.TCP{Flags: tb.Uint8(header.TCPFlagAck), WindowSize: tb.Uint16(0)})
 
@@ -54,5 +61,8 @@ func TestWindowShrink(t *testing.T) {
 	// Note: There is another kind of zero-window probing which Windows uses (by sending one
 	// new byte at `RemoteSeqNum`), if netstack wants to go that way, we may want to change
 	// the following lines.
-	conn.ExpectData(tb.TCP{SeqNum: tb.Uint32(uint32(conn.RemoteSeqNum - 1))}, nil, time.Second)
+	expectedRemoteSeqNum := *conn.RemoteSeqNum() - 1
+	if _, err := conn.ExpectData(&tb.TCP{SeqNum: tb.Uint32(uint32(expectedRemoteSeqNum))}, nil, time.Second); err != nil {
+		t.Fatalf("expected a packet with sequence number %v: %s", expectedRemoteSeqNum, err)
+	}
 }
diff --git a/test/packetimpact/tests/udp_recv_multicast_test.go b/test/packetimpact/tests/udp_recv_multicast_test.go
index bc1b0be49..61fd17050 100644
--- a/test/packetimpact/tests/udp_recv_multicast_test.go
+++ b/test/packetimpact/tests/udp_recv_multicast_test.go
@@ -30,7 +30,7 @@ func TestUDPRecvMulticast(t *testing.T) {
 	defer dut.Close(boundFD)
 	conn := tb.NewUDPIPv4(t, tb.UDP{DstPort: &remotePort}, tb.UDP{SrcPort: &remotePort})
 	defer conn.Close()
-	frame := conn.CreateFrame(tb.UDP{}, &tb.Payload{Bytes: []byte("hello world")})
+	frame := conn.CreateFrame(&tb.UDP{}, &tb.Payload{Bytes: []byte("hello world")})
 	frame[1].(*tb.IPv4).DstAddr = tb.Address(tcpip.Address(net.ParseIP("224.0.0.1").To4()))
 	conn.SendFrame(frame)
 	dut.Recv(boundFD, 100, 0)