Give TCP flags a dedicated type

- Implement Stringer for it so that we can improve error messages.
- Use TCPFlags through the code base. There used to be a mixed usage of byte,
  uint8 and int as TCP flags.

PiperOrigin-RevId: 361940150

28 files changed, 155 insertions, 149 deletions

diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go
index 15e1a51de..bafc45d3b 100644
--- a/test/packetimpact/testbench/connections.go
+++ b/test/packetimpact/testbench/connections.go
@@ -677,17 +677,17 @@ func (conn *TCPIPv4) Connect(t *testing.T) {
 	t.Helper()
 
 	// Send the SYN.
-	conn.Send(t, TCP{Flags: Uint8(header.TCPFlagSyn)})
+	conn.Send(t, TCP{Flags: TCPFlags(header.TCPFlagSyn)})
 
 	// Wait for the SYN-ACK.
-	synAck, err := conn.Expect(t, TCP{Flags: Uint8(header.TCPFlagSyn | header.TCPFlagAck)}, time.Second)
+	synAck, err := conn.Expect(t, TCP{Flags: TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)}, time.Second)
 	if err != nil {
 		t.Fatalf("didn't get synack during handshake: %s", err)
 	}
 	conn.layerStates[len(conn.layerStates)-1].(*tcpState).synAck = synAck
 
 	// Send an ACK.
-	conn.Send(t, TCP{Flags: Uint8(header.TCPFlagAck)})
+	conn.Send(t, TCP{Flags: TCPFlags(header.TCPFlagAck)})
 }
 
 // ConnectWithOptions performs a TCP 3-way handshake with given TCP options.
@@ -696,17 +696,17 @@ func (conn *TCPIPv4) ConnectWithOptions(t *testing.T, options []byte) {
 	t.Helper()
 
 	// Send the SYN.
-	conn.Send(t, TCP{Flags: Uint8(header.TCPFlagSyn), Options: options})
+	conn.Send(t, TCP{Flags: TCPFlags(header.TCPFlagSyn), Options: options})
 
 	// Wait for the SYN-ACK.
-	synAck, err := conn.Expect(t, TCP{Flags: Uint8(header.TCPFlagSyn | header.TCPFlagAck)}, time.Second)
+	synAck, err := conn.Expect(t, TCP{Flags: TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)}, time.Second)
 	if err != nil {
 		t.Fatalf("didn't get synack during handshake: %s", err)
 	}
 	conn.layerStates[len(conn.layerStates)-1].(*tcpState).synAck = synAck
 
 	// Send an ACK.
-	conn.Send(t, TCP{Flags: Uint8(header.TCPFlagAck)})
+	conn.Send(t, TCP{Flags: TCPFlags(header.TCPFlagAck)})
 }
 
 // ExpectData is a convenient method that expects a Layer and the Layer after
diff --git a/test/packetimpact/testbench/layers.go b/test/packetimpact/testbench/layers.go
index 64a7a171a..2311f7686 100644
--- a/test/packetimpact/testbench/layers.go
+++ b/test/packetimpact/testbench/layers.go
@@ -407,6 +407,12 @@ func Uint8(v uint8) *uint8 {
 	return &v
 }
 
+// TCPFlags is a helper routine that allocates a new
+// header.TCPFlags value to store v and returns a pointer to it.
+func TCPFlags(v header.TCPFlags) *header.TCPFlags {
+	return &v
+}
+
 // Address is a helper routine that allocates a new tcpip.Address value to
 // store v and returns a pointer to it.
 func Address(v tcpip.Address) *tcpip.Address {
@@ -1030,7 +1036,7 @@ type TCP struct {
 	SeqNum        *uint32
 	AckNum        *uint32
 	DataOffset    *uint8
-	Flags         *uint8
+	Flags         *header.TCPFlags
 	WindowSize    *uint16
 	Checksum      *uint16
 	UrgentPointer *uint16
@@ -1063,7 +1069,7 @@ func (l *TCP) ToBytes() ([]byte, error) {
 		h.SetDataOffset(uint8(l.length()))
 	}
 	if l.Flags != nil {
-		h.SetFlags(*l.Flags)
+		h.SetFlags(uint8(*l.Flags))
 	}
 	if l.WindowSize != nil {
 		h.SetWindowSize(*l.WindowSize)
@@ -1157,7 +1163,7 @@ func parseTCP(b []byte) (Layer, layerParser) {
 		SeqNum:        Uint32(h.SequenceNumber()),
 		AckNum:        Uint32(h.AckNumber()),
 		DataOffset:    Uint8(h.DataOffset()),
-		Flags:         Uint8(h.Flags()),
+		Flags:         TCPFlags(h.Flags()),
 		WindowSize:    Uint16(h.WindowSize()),
 		Checksum:      Uint16(h.Checksum()),
 		UrgentPointer: Uint16(h.UrgentPointer()),
diff --git a/test/packetimpact/testbench/layers_test.go b/test/packetimpact/testbench/layers_test.go
index eca0780b5..614a5de1e 100644
--- a/test/packetimpact/testbench/layers_test.go
+++ b/test/packetimpact/testbench/layers_test.go
@@ -178,7 +178,7 @@ func TestLayerStringFormat(t *testing.T) {
 				SeqNum:     Uint32(3452155723),
 				AckNum:     Uint32(2596996163),
 				DataOffset: Uint8(5),
-				Flags:      Uint8(20),
+				Flags:      TCPFlags(header.TCPFlagRst | header.TCPFlagAck),
 				WindowSize: Uint16(64240),
 				Checksum:   Uint16(0x2e2b),
 			},
@@ -188,7 +188,7 @@ func TestLayerStringFormat(t *testing.T) {
 				"SeqNum:3452155723 " +
 				"AckNum:2596996163 " +
 				"DataOffset:5 " +
-				"Flags:20 " +
+				"Flags:  R A  " +
 				"WindowSize:64240 " +
 				"Checksum:11819" +
 				"}",
@@ -436,7 +436,7 @@ func TestTCPOptions(t *testing.T) {
 					DstPort:       Uint16(54321),
 					SeqNum:        Uint32(0),
 					AckNum:        Uint32(0),
-					Flags:         Uint8(header.TCPFlagSyn),
+					Flags:         TCPFlags(header.TCPFlagSyn),
 					WindowSize:    Uint16(8192),
 					Checksum:      Uint16(0xf51c),
 					UrgentPointer: Uint16(0),
@@ -480,7 +480,7 @@ func TestTCPOptions(t *testing.T) {
 					DstPort:       Uint16(54321),
 					SeqNum:        Uint32(0),
 					AckNum:        Uint32(0),
-					Flags:         Uint8(header.TCPFlagSyn),
+					Flags:         TCPFlags(header.TCPFlagSyn),
 					WindowSize:    Uint16(8192),
 					Checksum:      Uint16(0xe521),
 					UrgentPointer: Uint16(0),
diff --git a/test/packetimpact/tests/fin_wait2_timeout_test.go b/test/packetimpact/tests/fin_wait2_timeout_test.go
index 11f0fcd1e..cff8ca51d 100644
--- a/test/packetimpact/tests/fin_wait2_timeout_test.go
+++ b/test/packetimpact/tests/fin_wait2_timeout_test.go
@@ -51,21 +51,21 @@ func TestFinWait2Timeout(t *testing.T) {
 			}
 			dut.Close(t, acceptFd)
 
-			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
+			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
 				t.Fatalf("expected a FIN-ACK within 1 second but got none: %s", err)
 			}
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 
 			time.Sleep(5 * time.Second)
 			conn.Drain(t)
 
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 			if tt.linger2 {
-				if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, time.Second); err != nil {
+				if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, time.Second); err != nil {
 					t.Fatalf("expected a RST packet within a second but got none: %s", err)
 				}
 			} else {
-				if got, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, 10*time.Second); got != nil || err == nil {
+				if got, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, 10*time.Second); got != nil || err == nil {
 					t.Fatalf("expected no RST packets within ten seconds but got one: %s", got)
 				}
 			}
diff --git a/test/packetimpact/tests/ipv4_id_uniqueness_test.go b/test/packetimpact/tests/ipv4_id_uniqueness_test.go
index a63b41366..2b69ceecb 100644
--- a/test/packetimpact/tests/ipv4_id_uniqueness_test.go
+++ b/test/packetimpact/tests/ipv4_id_uniqueness_test.go
@@ -100,7 +100,7 @@ func TestIPv4RetransmitIdentificationUniqueness(t *testing.T) {
 			// Let the DUT estimate RTO with RTT from the DATA-ACK.
 			// TODO(gvisor.dev/issue/2685) Estimate RTO during handshake, after which
 			// we can skip sending this ACK.
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 
 			dut.Send(t, remoteFD, tc.payload, 0)
 			expectTCP := &testbench.TCP{SeqNum: testbench.Uint32(uint32(*conn.RemoteSeqNum(t)))}
diff --git a/test/packetimpact/tests/tcp_cork_mss_test.go b/test/packetimpact/tests/tcp_cork_mss_test.go
index a7ba5035e..1db3c9883 100644
--- a/test/packetimpact/tests/tcp_cork_mss_test.go
+++ b/test/packetimpact/tests/tcp_cork_mss_test.go
@@ -60,24 +60,24 @@ func TestTCPCorkMSS(t *testing.T) {
 
 	// Expect the segments to be coalesced and sent and capped to MSS.
 	expectedPayload := testbench.Payload{Bytes: expectedData[:mss]}
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, &expectedPayload, time.Second); err != nil {
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, &expectedPayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 	// Expect the coalesced segment to be split and transmitted.
 	expectedPayload = testbench.Payload{Bytes: expectedData[mss:]}
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, &expectedPayload, time.Second); err != nil {
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, &expectedPayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
 
 	// Check for segments to *not* be held up because of TCP_CORK when
 	// the current send window is less than MSS.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), WindowSize: testbench.Uint16(uint16(2 * len(sampleData)))})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), WindowSize: testbench.Uint16(uint16(2 * len(sampleData)))})
 	dut.Send(t, acceptFD, sampleData, 0)
 	dut.Send(t, acceptFD, sampleData, 0)
 	expectedPayload = testbench.Payload{Bytes: append(sampleData, sampleData...)}
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, &expectedPayload, time.Second); err != nil {
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, &expectedPayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 }
diff --git a/test/packetimpact/tests/tcp_handshake_window_size_test.go b/test/packetimpact/tests/tcp_handshake_window_size_test.go
index 5d1266f3c..668e0275c 100644
--- a/test/packetimpact/tests/tcp_handshake_window_size_test.go
+++ b/test/packetimpact/tests/tcp_handshake_window_size_test.go
@@ -38,8 +38,8 @@ func TestTCPHandshakeWindowSize(t *testing.T) {
 	defer conn.Close(t)
 
 	// Start handshake with zero window size.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn), WindowSize: testbench.Uint16(uint16(0))})
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn | header.TCPFlagAck)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn), WindowSize: testbench.Uint16(uint16(0))})
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)}, nil, time.Second); err != nil {
 		t.Fatalf("expected SYN-ACK: %s", err)
 	}
 	// Update the advertised window size to a non-zero value with the ACK that
@@ -47,7 +47,7 @@ func TestTCPHandshakeWindowSize(t *testing.T) {
 	//
 	// Set the window size with MSB set and expect the dut to treat it as
 	// an unsigned value.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), WindowSize: testbench.Uint16(uint16(1 << 15))})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), WindowSize: testbench.Uint16(uint16(1 << 15))})
 
 	acceptFd, _ := dut.Accept(t, listenFD)
 	defer dut.Close(t, acceptFd)
@@ -59,7 +59,7 @@ func TestTCPHandshakeWindowSize(t *testing.T) {
 	// expect the dut to honor the recently advertised non-zero window
 	// and actually send out the data instead of probing for zero window.
 	dut.Send(t, acceptFd, sampleData, 0)
-	if _, err := conn.ExpectNextData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload, time.Second); err != nil {
+	if _, err := conn.ExpectNextData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
 }
diff --git a/test/packetimpact/tests/tcp_info_test.go b/test/packetimpact/tests/tcp_info_test.go
index 69275e54b..3fc2c7fe5 100644
--- a/test/packetimpact/tests/tcp_info_test.go
+++ b/test/packetimpact/tests/tcp_info_test.go
@@ -51,7 +51,7 @@ func TestTCPInfo(t *testing.T) {
 	if _, err := conn.ExpectData(t, &testbench.TCP{}, samplePayload, time.Second); err != nil {
 		t.Fatalf("expected a packet with payload %v: %s", samplePayload, err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 
 	info := linux.TCPInfo{}
 	infoBytes := dut.GetSockOpt(t, acceptFD, unix.SOL_TCP, unix.TCP_INFO, int32(linux.SizeOfTCPInfo))
diff --git a/test/packetimpact/tests/tcp_linger_test.go b/test/packetimpact/tests/tcp_linger_test.go
index 6c57b6d37..88942904d 100644
--- a/test/packetimpact/tests/tcp_linger_test.go
+++ b/test/packetimpact/tests/tcp_linger_test.go
@@ -57,10 +57,10 @@ func TestTCPLingerZeroTimeout(t *testing.T) {
 	dut.Close(t, acceptFD)
 
 	// If the linger timeout is set to zero, the DUT should send a RST.
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst | header.TCPFlagAck)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst | header.TCPFlagAck)}, time.Second); err != nil {
 		t.Errorf("expected RST-ACK packet within a second but got none: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 }
 
 // TestTCPLingerOff tests when SO_LINGER is not set. DUT should send FIN-ACK
@@ -74,10 +74,10 @@ func TestTCPLingerOff(t *testing.T) {
 	dut.Close(t, acceptFD)
 
 	// If SO_LINGER is not set, DUT should send a FIN-ACK.
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
 		t.Errorf("expected FIN-ACK packet within a second but got none: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 }
 
 // TestTCPLingerNonZeroTimeout tests when SO_LINGER is set with non-zero timeout.
@@ -114,10 +114,10 @@ func TestTCPLingerNonZeroTimeout(t *testing.T) {
 				t.Errorf("expected close to return within a second, but returned later")
 			}
 
-			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
+			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
 				t.Errorf("expected FIN-ACK packet within a second but got none: %s", err)
 			}
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 		})
 	}
 }
@@ -165,10 +165,10 @@ func TestTCPLingerSendNonZeroTimeout(t *testing.T) {
 				t.Fatalf("expected a packet with payload %v: %s", samplePayload, err)
 			}
 
-			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
+			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
 				t.Errorf("expected FIN-ACK packet within a second but got none: %s", err)
 			}
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 		})
 	}
 }
@@ -186,15 +186,15 @@ func TestTCPLingerShutdownZeroTimeout(t *testing.T) {
 	dut.Close(t, acceptFD)
 
 	// Shutdown will send FIN-ACK with read/write option.
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
 		t.Errorf("expected FIN-ACK packet within a second but got none: %s", err)
 	}
 
 	// If the linger timeout is set to zero, the DUT should send a RST.
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst | header.TCPFlagAck)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst | header.TCPFlagAck)}, time.Second); err != nil {
 		t.Errorf("expected RST-ACK packet within a second but got none: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 }
 
 // TestTCPLingerShutdownSendNonZeroTimeout tests SO_LINGER with shutdown() and
@@ -242,10 +242,10 @@ func TestTCPLingerShutdownSendNonZeroTimeout(t *testing.T) {
 				t.Fatalf("expected a packet with payload %v: %s", samplePayload, err)
 			}
 
-			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
+			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second); err != nil {
 				t.Errorf("expected FIN-ACK packet within a second but got none: %s", err)
 			}
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 		})
 	}
 }
diff --git a/test/packetimpact/tests/tcp_network_unreachable_test.go b/test/packetimpact/tests/tcp_network_unreachable_test.go
index 4e5b418e4..5168450ad 100644
--- a/test/packetimpact/tests/tcp_network_unreachable_test.go
+++ b/test/packetimpact/tests/tcp_network_unreachable_test.go
@@ -50,7 +50,7 @@ func TestTCPSynSentUnreachable(t *testing.T) {
 	}
 
 	// Get the SYN.
-	tcpLayers, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)}, nil, time.Second)
+	tcpLayers, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)}, nil, time.Second)
 	if err != nil {
 		t.Fatalf("expected SYN: %s", err)
 	}
@@ -104,7 +104,7 @@ func TestTCPSynSentUnreachable6(t *testing.T) {
 	}
 
 	// Get the SYN.
-	tcpLayers, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)}, nil, time.Second)
+	tcpLayers, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)}, nil, time.Second)
 	if err != nil {
 		t.Fatalf("expected SYN: %s", err)
 	}
diff --git a/test/packetimpact/tests/tcp_noaccept_close_rst_test.go b/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
index d2871df08..14eb7d93b 100644
--- a/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
+++ b/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
@@ -40,7 +40,7 @@ func TestTcpNoAcceptCloseReset(t *testing.T) {
 	// it will only respond RST instead of RST+ACK.
 	dut.PollOne(t, listenFd, unix.POLLIN, time.Second)
 	dut.Close(t, listenFd)
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst | header.TCPFlagAck)}, 1*time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst | header.TCPFlagAck)}, 1*time.Second); err != nil {
 		t.Fatalf("expected a RST-ACK packet but got none: %s", err)
 	}
 }
diff --git a/test/packetimpact/tests/tcp_outside_the_window_test.go b/test/packetimpact/tests/tcp_outside_the_window_test.go
index 8909a348e..e442268be 100644
--- a/test/packetimpact/tests/tcp_outside_the_window_test.go
+++ b/test/packetimpact/tests/tcp_outside_the_window_test.go
@@ -37,7 +37,7 @@ func init() {
 func TestTCPOutsideTheWindow(t *testing.T) {
 	for _, tt := range []struct {
 		description  string
-		tcpFlags     uint8
+		tcpFlags     header.TCPFlags
 		payload      []testbench.Layer
 		seqNumOffset seqnum.Size
 		expectACK    bool
@@ -76,11 +76,11 @@ func TestTCPOutsideTheWindow(t *testing.T) {
 			// to the AckNum.
 			localSeqNum := testbench.Uint32(uint32(*conn.LocalSeqNum(t)))
 			conn.Send(t, testbench.TCP{
-				Flags:  testbench.Uint8(tt.tcpFlags),
+				Flags:  testbench.TCPFlags(tt.tcpFlags),
 				SeqNum: testbench.Uint32(uint32(conn.LocalSeqNum(t).Add(windowSize))),
 			}, tt.payload...)
 			timeout := 3 * time.Second
-			gotACK, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: localSeqNum}, timeout)
+			gotACK, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: localSeqNum}, timeout)
 			if tt.expectACK && err != nil {
 				t.Fatalf("expected an ACK packet within %s but got none: %s", timeout, err)
 			}
@@ -93,11 +93,11 @@ func TestTCPOutsideTheWindow(t *testing.T) {
 				// has passed since the last ACK was sent.
 				t.Logf("sending another segment")
 				conn.Send(t, testbench.TCP{
-					Flags:  testbench.Uint8(tt.tcpFlags),
+					Flags:  testbench.TCPFlags(tt.tcpFlags),
 					SeqNum: testbench.Uint32(uint32(conn.LocalSeqNum(t).Add(windowSize))),
 				}, tt.payload...)
 				timeout := 3 * time.Second
-				gotACK, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: localSeqNum}, timeout)
+				gotACK, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: localSeqNum}, timeout)
 				if err == nil {
 					t.Fatalf("expected no ACK packet but got one: %s", gotACK)
 				}
diff --git a/test/packetimpact/tests/tcp_paws_mechanism_test.go b/test/packetimpact/tests/tcp_paws_mechanism_test.go
index 24d9ef4ec..9054955ea 100644
--- a/test/packetimpact/tests/tcp_paws_mechanism_test.go
+++ b/test/packetimpact/tests/tcp_paws_mechanism_test.go
@@ -38,8 +38,8 @@ func TestPAWSMechanism(t *testing.T) {
 
 	options := make([]byte, header.TCPOptionTSLength)
 	header.EncodeTSOption(currentTS(), 0, options)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn), Options: options})
-	synAck, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn | header.TCPFlagAck)}, time.Second)
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn), Options: options})
+	synAck, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)}, time.Second)
 	if err != nil {
 		t.Fatalf("didn't get synack during handshake: %s", err)
 	}
@@ -49,7 +49,7 @@ func TestPAWSMechanism(t *testing.T) {
 	}
 	tsecr := parsedSynOpts.TSVal
 	header.EncodeTSOption(currentTS(), tsecr, options)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), Options: options})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), Options: options})
 	acceptFD, _ := dut.Accept(t, listenFD)
 	defer dut.Close(t, acceptFD)
 
@@ -60,9 +60,9 @@ func TestPAWSMechanism(t *testing.T) {
 	// every time we send one, it should not cause any flakiness because timestamps
 	// only need to be non-decreasing.
 	time.Sleep(3 * time.Millisecond)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), Options: options}, &testbench.Payload{Bytes: sampleData})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), Options: options}, &testbench.Payload{Bytes: sampleData})
 
-	gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+	gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 	if err != nil {
 		t.Fatalf("expected an ACK but got none: %s", err)
 	}
@@ -85,9 +85,9 @@ func TestPAWSMechanism(t *testing.T) {
 	// 3ms here is chosen arbitrarily and this time.Sleep() should not cause flakiness
 	// due to the exact same reasoning discussed above.
 	time.Sleep(3 * time.Millisecond)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), Options: options}, &testbench.Payload{Bytes: sampleData})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), Options: options}, &testbench.Payload{Bytes: sampleData})
 
-	gotTCP, err = conn.Expect(t, testbench.TCP{AckNum: lastAckNum, Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+	gotTCP, err = conn.Expect(t, testbench.TCP{AckNum: lastAckNum, Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 	if err != nil {
 		t.Fatalf("expected segment with AckNum %d but got none: %s", lastAckNum, err)
 	}
diff --git a/test/packetimpact/tests/tcp_queue_send_recv_in_syn_sent_test.go b/test/packetimpact/tests/tcp_queue_send_recv_in_syn_sent_test.go
index 74616a54c..1c8b72ebe 100644
--- a/test/packetimpact/tests/tcp_queue_send_recv_in_syn_sent_test.go
+++ b/test/packetimpact/tests/tcp_queue_send_recv_in_syn_sent_test.go
@@ -47,7 +47,7 @@ func TestQueueSendInSynSentHandshake(t *testing.T) {
 	if _, err := dut.ConnectWithErrno(context.Background(), t, socket, conn.LocalAddr(t)); !errors.Is(err, unix.EINPROGRESS) {
 		t.Fatalf("failed to bring DUT to SYN-SENT, got: %s, want EINPROGRESS", err)
 	}
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)}, time.Second); err != nil {
 		t.Fatalf("expected a SYN from DUT, but got none: %s", err)
 	}
 
@@ -84,19 +84,19 @@ func TestQueueSendInSynSentHandshake(t *testing.T) {
 	time.Sleep(100 * time.Millisecond)
 
 	// Bring the connection to Established.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn | header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)})
 	// Expect the data from the DUT's enqueued send request.
 	//
 	// On Linux, this can be piggybacked with the ACK completing the
 	// handshake. On gVisor, getting such a piggyback is a bit more
 	// complicated because the actual data enqueuing occurs in the
 	// callers of endpoint Write.
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagPsh | header.TCPFlagAck)}, &testbench.Payload{Bytes: sampleData}, time.Second); err != nil {
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagPsh | header.TCPFlagAck)}, &testbench.Payload{Bytes: sampleData}, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
 
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagPsh | header.TCPFlagAck)}, &testbench.Payload{Bytes: sampleData})
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagPsh | header.TCPFlagAck)}, &testbench.Payload{Bytes: sampleData})
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second); err != nil {
 		t.Fatalf("expected an ACK from DUT, but got none: %s", err)
 	}
 }
@@ -115,7 +115,7 @@ func TestQueueRecvInSynSentHandshake(t *testing.T) {
 	if _, err := dut.ConnectWithErrno(context.Background(), t, socket, conn.LocalAddr(t)); !errors.Is(err, unix.EINPROGRESS) {
 		t.Fatalf("failed to bring DUT to SYN-SENT, got: %s, want EINPROGRESS", err)
 	}
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)}, time.Second); err != nil {
 		t.Fatalf("expected a SYN from DUT, but got none: %s", err)
 	}
 
@@ -159,14 +159,14 @@ func TestQueueRecvInSynSentHandshake(t *testing.T) {
 	time.Sleep(100 * time.Millisecond)
 
 	// Bring the connection to Established.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn | header.TCPFlagAck)})
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)})
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second); err != nil {
 		t.Fatalf("expected an ACK from DUT, but got none: %s", err)
 	}
 
 	// Send sample payload so that DUT can recv.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagPsh | header.TCPFlagAck)}, &testbench.Payload{Bytes: sampleData})
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagPsh | header.TCPFlagAck)}, &testbench.Payload{Bytes: sampleData})
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second); err != nil {
 		t.Fatalf("expected an ACK from DUT, but got none: %s", err)
 	}
 }
@@ -185,7 +185,7 @@ func TestQueueSendInSynSentRST(t *testing.T) {
 	if _, err := dut.ConnectWithErrno(context.Background(), t, socket, conn.LocalAddr(t)); !errors.Is(err, unix.EINPROGRESS) {
 		t.Fatalf("failed to bring DUT to SYN-SENT, got: %s, want EINPROGRESS", err)
 	}
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)}, time.Second); err != nil {
 		t.Fatalf("expected a SYN from DUT, but got none: %s", err)
 	}
 
@@ -223,7 +223,7 @@ func TestQueueSendInSynSentRST(t *testing.T) {
 	// request and the system actually being blocked.
 	time.Sleep(100 * time.Millisecond)
 
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst | header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst | header.TCPFlagAck)})
 }
 
 // TestQueueRecvInSynSentRST tests recv behavior when the TCP state
@@ -240,7 +240,7 @@ func TestQueueRecvInSynSentRST(t *testing.T) {
 	if _, err := dut.ConnectWithErrno(context.Background(), t, socket, conn.LocalAddr(t)); !errors.Is(err, unix.EINPROGRESS) {
 		t.Fatalf("failed to bring DUT to SYN-SENT, got: %s, want EINPROGRESS", err)
 	}
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)}, time.Second); err != nil {
 		t.Fatalf("expected a SYN from DUT, but got none: %s", err)
 	}
 
@@ -282,5 +282,5 @@ func TestQueueRecvInSynSentRST(t *testing.T) {
 	// request and the system actually being blocked.
 	time.Sleep(100 * time.Millisecond)
 
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst | header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst | header.TCPFlagAck)})
 }
diff --git a/test/packetimpact/tests/tcp_rack_test.go b/test/packetimpact/tests/tcp_rack_test.go
index ef902c54d..0a5b0f12b 100644
--- a/test/packetimpact/tests/tcp_rack_test.go
+++ b/test/packetimpact/tests/tcp_rack_test.go
@@ -97,7 +97,7 @@ func sendAndReceive(t *testing.T, dut testbench.DUT, conn testbench.TCPIPv4, num
 
 		if sendACK {
 			time.Sleep(simulatedRTT)
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(sn))})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(sn))})
 		}
 	}
 	return lastSent
@@ -149,7 +149,7 @@ func TestRACKTLPLost(t *testing.T) {
 
 	// Cumulative ACK for #[1-5] packets.
 	ackNum := seqNum1.Add(seqnum.Size(6 * payloadSize))
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(ackNum))})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(ackNum))})
 
 	// Probe Timeout (PTO) should be two times RTT. Check that the last
 	// packet is retransmitted after probe timeout.
@@ -194,7 +194,7 @@ func TestRACKWithSACK(t *testing.T) {
 	sbOff += header.EncodeSACKBlocks([]header.SACKBlock{{
 		start, end,
 	}}, sackBlock[sbOff:])
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
 
 	rtt, _ := getRTTAndRTO(t, dut, acceptFd)
 	timeout := 2 * rtt
@@ -206,7 +206,7 @@ func TestRACKWithSACK(t *testing.T) {
 
 	time.Sleep(simulatedRTT)
 	// ACK for #1 packet.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(end))})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(end))})
 
 	// RACK considers transmission times of the packets to mark them lost.
 	// As the 3rd packet was sent before the retransmitted 1st packet, RACK
@@ -243,7 +243,7 @@ func TestRACKWithoutReorder(t *testing.T) {
 		start, end,
 	}}, sackBlock[sbOff:])
 	time.Sleep(simulatedRTT)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
 
 	// RACK marks #1 and #2 packets as lost and retransmits both after
 	// RTT + reorderWindow. The reorderWindow initially will be a small
@@ -289,7 +289,7 @@ func TestRACKWithReorder(t *testing.T) {
 		sbOff += header.EncodeSACKBlocks([]header.SACKBlock{{
 			start, end,
 		}}, sackBlock[sbOff:])
-		conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
+		conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
 	}
 
 	// Send a DSACK block indicating both original and retransmitted
@@ -304,7 +304,7 @@ func TestRACKWithReorder(t *testing.T) {
 	dbOff += header.EncodeSACKBlocks([]header.SACKBlock{{
 		start, end,
 	}}, dsackBlock[dbOff:])
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1 + numPkts*payloadSize)), Options: dsackBlock[:dbOff]})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1 + numPkts*payloadSize)), Options: dsackBlock[:dbOff]})
 
 	seqNum1.UpdateForward(seqnum.Size(numPkts * payloadSize))
 	sendTime := time.Now()
@@ -321,7 +321,7 @@ func TestRACKWithReorder(t *testing.T) {
 	sbOff += header.EncodeSACKBlocks([]header.SACKBlock{{
 		start, end,
 	}}, sackBlock[sbOff:])
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
 
 	// Expect the retransmission of #1 packet after RTT+ReorderWindow.
 	if _, err := conn.Expect(t, testbench.TCP{SeqNum: testbench.Uint32(uint32(seqNum1))}, time.Second); err != nil {
@@ -361,7 +361,7 @@ func TestRACKWithLostRetransmission(t *testing.T) {
 		start, end,
 	}}, sackBlock[sbOff:])
 	time.Sleep(simulatedRTT)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock[:sbOff]})
 
 	// RACK marks #1 packet as lost and retransmits it after
 	// RTT + reorderWindow. The reorderWindow is bounded between a small
@@ -394,7 +394,7 @@ func TestRACKWithLostRetransmission(t *testing.T) {
 		start, end,
 	}}, sackBlock1[sbOff1:])
 	time.Sleep(simulatedRTT)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock1[:sbOff1]})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: testbench.Uint32(uint32(seqNum1)), Options: sackBlock1[:sbOff1]})
 
 	// Expect re-retransmission of #1 packet without entering an RTO.
 	if _, err := conn.Expect(t, testbench.TCP{SeqNum: testbench.Uint32(uint32(seqNum1))}, timeout); err != nil {
diff --git a/test/packetimpact/tests/tcp_rcv_buf_space_test.go b/test/packetimpact/tests/tcp_rcv_buf_space_test.go
index 042da8fee..f121d44eb 100644
--- a/test/packetimpact/tests/tcp_rcv_buf_space_test.go
+++ b/test/packetimpact/tests/tcp_rcv_buf_space_test.go
@@ -60,7 +60,7 @@ func TestReduceRecvBuf(t *testing.T) {
 			payloadBytes = l
 		}
 
-		conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, []testbench.Layer{&testbench.Payload{Bytes: payload[:payloadBytes]}}...)
+		conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, []testbench.Layer{&testbench.Payload{Bytes: payload[:payloadBytes]}}...)
 		payload = payload[payloadBytes:]
 	}
 
diff --git a/test/packetimpact/tests/tcp_retransmits_test.go b/test/packetimpact/tests/tcp_retransmits_test.go
index ba79fbf55..c2611c2a6 100644
--- a/test/packetimpact/tests/tcp_retransmits_test.go
+++ b/test/packetimpact/tests/tcp_retransmits_test.go
@@ -66,7 +66,7 @@ func TestRetransmits(t *testing.T) {
 	// Give a chance for the dut to estimate RTO with RTT from the DATA-ACK.
 	// TODO(gvisor.dev/issue/2685) Estimate RTO during handshake, after which
 	// we can skip sending this ACK.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 
 	const timeoutCorrection = time.Second
 	const diffCorrection = time.Millisecond
diff --git a/test/packetimpact/tests/tcp_send_window_sizes_piggyback_test.go b/test/packetimpact/tests/tcp_send_window_sizes_piggyback_test.go
index 418393796..64b7288fb 100644
--- a/test/packetimpact/tests/tcp_send_window_sizes_piggyback_test.go
+++ b/test/packetimpact/tests/tcp_send_window_sizes_piggyback_test.go
@@ -71,7 +71,7 @@ func TestSendWindowSizesPiggyback(t *testing.T) {
 
 			dut.SetSockOptInt(t, acceptFd, unix.IPPROTO_TCP, unix.TCP_NODELAY, 1)
 
-			expectedTCP := testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}
+			expectedTCP := testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}
 
 			dut.Send(t, acceptFd, sampleData, 0)
 			expectedPayload := testbench.Payload{Bytes: tt.expectedPayload1}
@@ -90,7 +90,7 @@ func TestSendWindowSizesPiggyback(t *testing.T) {
 			// Send ACK for the previous segment along with data for the dut to
 			// receive and ACK back. Sending this ACK would make room for the dut
 			// to transmit any enqueued segment.
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh), WindowSize: testbench.Uint16(tt.windowSize)}, &testbench.Payload{Bytes: sampleData})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh), WindowSize: testbench.Uint16(tt.windowSize)}, &testbench.Payload{Bytes: sampleData})
 
 			// Expect the dut to piggyback the ACK for received data along with
 			// the segment enqueued for transmit.
diff --git a/test/packetimpact/tests/tcp_synrcvd_reset_test.go b/test/packetimpact/tests/tcp_synrcvd_reset_test.go
index 32271d7b2..3346d43c4 100644
--- a/test/packetimpact/tests/tcp_synrcvd_reset_test.go
+++ b/test/packetimpact/tests/tcp_synrcvd_reset_test.go
@@ -37,11 +37,11 @@ func TestTCPSynRcvdReset(t *testing.T) {
 	defer conn.Close(t)
 
 	// Expect dut connection to have transitioned to SYN-RCVD state.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)})
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn | header.TCPFlagAck)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)})
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)}, nil, time.Second); err != nil {
 		t.Fatalf("expected SYN-ACK %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)})
 
 	// Expect the connection to have transitioned SYN-RCVD to CLOSED.
 	//
@@ -49,8 +49,8 @@ func TestTCPSynRcvdReset(t *testing.T) {
 	// CLOSED. We cannot use TCP_INFO to lookup the state as this is a passive
 	// DUT connection.
 	for i := 0; i < 5; i++ {
-		conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
-		if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, nil, time.Second); err != nil {
+		conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
+		if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, nil, time.Second); err != nil {
 			t.Logf("retransmit%d ACK as we did not get the expected RST, %s", i, err)
 			continue
 		}
diff --git a/test/packetimpact/tests/tcp_synsent_reset_test.go b/test/packetimpact/tests/tcp_synsent_reset_test.go
index 2c8bb101b..cccb0abc6 100644
--- a/test/packetimpact/tests/tcp_synsent_reset_test.go
+++ b/test/packetimpact/tests/tcp_synsent_reset_test.go
@@ -42,7 +42,7 @@ func dutSynSentState(t *testing.T) (*testbench.DUT, *testbench.TCPIPv4, uint16,
 	copy(sa.Addr[:], dut.Net.LocalIPv4)
 	// Bring the dut to SYN-SENT state with a non-blocking connect.
 	dut.Connect(t, clientFD, &sa)
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)}, nil, time.Second); err != nil {
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)}, nil, time.Second); err != nil {
 		t.Fatalf("expected SYN\n")
 	}
 
@@ -53,11 +53,11 @@ func dutSynSentState(t *testing.T) (*testbench.DUT, *testbench.TCPIPv4, uint16,
 func TestTCPSynSentReset(t *testing.T) {
 	_, conn, _, _ := dutSynSentState(t)
 	defer conn.Close(t)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst | header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst | header.TCPFlagAck)})
 	// Expect the connection to have closed.
 	// TODO(gvisor.dev/issue/478): Check for TCP_INFO on the dut side.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, nil, time.Second); err != nil {
 		t.Fatalf("expected a TCP RST")
 	}
 }
@@ -73,15 +73,15 @@ func TestTCPSynSentRcvdReset(t *testing.T) {
 	// Initiate new SYN connection with the same port pair
 	// (simultaneous open case), expect the dut connection to move to
 	// SYN-RCVD state
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn)})
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagSyn | header.TCPFlagAck)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn)})
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagSyn | header.TCPFlagAck)}, nil, time.Second); err != nil {
 		t.Fatalf("expected SYN-ACK %s\n", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)})
 	// Expect the connection to have transitioned SYN-RCVD to CLOSED.
 	// TODO(gvisor.dev/issue/478): Check for TCP_INFO on the dut side.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, nil, time.Second); err != nil {
 		t.Fatalf("expected a TCP RST")
 	}
 }
diff --git a/test/packetimpact/tests/tcp_timewait_reset_test.go b/test/packetimpact/tests/tcp_timewait_reset_test.go
index d1d2fb83d..89037f0a4 100644
--- a/test/packetimpact/tests/tcp_timewait_reset_test.go
+++ b/test/packetimpact/tests/tcp_timewait_reset_test.go
@@ -42,26 +42,26 @@ func TestTimeWaitReset(t *testing.T) {
 	// Trigger active close.
 	dut.Close(t, acceptFD)
 
-	_, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second)
+	_, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second)
 	if err != nil {
 		t.Fatalf("expected a FIN: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 	// Send a FIN, DUT should transition to TIME_WAIT from FIN_WAIT2.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)})
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)})
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second); err != nil {
 		t.Fatalf("expected an ACK for our FIN: %s", err)
 	}
 
 	// Send a RST, the DUT should transition to CLOSED from TIME_WAIT.
 	// This is the default Linux behavior, it can be changed to ignore RSTs via
 	// sysctl net.ipv4.tcp_rfc1337.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)})
 
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 	// The DUT should reply with RST to our ACK as the state should have
 	// transitioned to CLOSED.
-	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, time.Second); err != nil {
+	if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, time.Second); err != nil {
 		t.Fatalf("expected a RST: %s", err)
 	}
 }
diff --git a/test/packetimpact/tests/tcp_unacc_seq_ack_test.go b/test/packetimpact/tests/tcp_unacc_seq_ack_test.go
index b76369dbe..92b54b72e 100644
--- a/test/packetimpact/tests/tcp_unacc_seq_ack_test.go
+++ b/test/packetimpact/tests/tcp_unacc_seq_ack_test.go
@@ -58,8 +58,8 @@ func TestEstablishedUnaccSeqAck(t *testing.T) {
 			sampleData := []byte("Sample Data")
 			samplePayload := &testbench.Payload{Bytes: sampleData}
 
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
-			gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
+			gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 			if err != nil {
 				t.Fatalf("expected ack %s", err)
 			}
@@ -73,7 +73,7 @@ func TestEstablishedUnaccSeqAck(t *testing.T) {
 				// ACK matches the TCP layer state.
 				*conn.LocalSeqNum(t) = origSeq
 			}
-			gotAck, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+			gotAck, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 			if tt.expectAck && err != nil {
 				t.Fatalf("expected an ack but got none: %s", err)
 			}
@@ -109,8 +109,8 @@ func TestPassiveCloseUnaccSeqAck(t *testing.T) {
 			acceptFD, _ := dut.Accept(t, listenFD)
 
 			// Send a FIN to DUT to intiate the passive close.
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagFin)})
-			gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagFin)})
+			gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 			if err != nil {
 				t.Fatalf("expected an ACK for our fin and DUT should enter CLOSE_WAIT: %s", err)
 			}
@@ -121,7 +121,7 @@ func TestPassiveCloseUnaccSeqAck(t *testing.T) {
 
 			// Send a segment with OTW Seq / unacc ACK.
 			conn.Send(t, tt.makeTestingTCP(t, &conn, tt.seqNumOffset, windowSize), samplePayload)
-			gotAck, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+			gotAck, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 			if tt.expectAck && err != nil {
 				t.Errorf("expected an ack but got none: %s", err)
 			}
@@ -131,14 +131,14 @@ func TestPassiveCloseUnaccSeqAck(t *testing.T) {
 
 			// Now let's verify DUT is indeed in CLOSE_WAIT
 			dut.Close(t, acceptFD)
-			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagFin)}, time.Second); err != nil {
+			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagFin)}, time.Second); err != nil {
 				t.Fatalf("expected DUT to send a FIN: %s", err)
 			}
 			// Ack the FIN from DUT
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 			// Send some extra data to DUT
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, samplePayload)
-			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, time.Second); err != nil {
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, samplePayload)
+			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, time.Second); err != nil {
 				t.Fatalf("expected DUT to send an RST: %s", err)
 			}
 		})
@@ -173,11 +173,11 @@ func TestActiveCloseUnaccpSeqAck(t *testing.T) {
 			dut.Shutdown(t, acceptFD, unix.SHUT_WR)
 
 			// Get to FIN_WAIT2
-			gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)}, time.Second)
+			gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)}, time.Second)
 			if err != nil {
 				t.Fatalf("expected a FIN: %s", err)
 			}
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 
 			sendUnaccSeqAck := func(state string) {
 				t.Helper()
@@ -192,7 +192,7 @@ func TestActiveCloseUnaccpSeqAck(t *testing.T) {
 					// incoming ACK matches the TCP layer state.
 					*conn.LocalSeqNum(t) = origSeq
 				}
-				if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second); err != nil {
+				if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second); err != nil {
 					t.Errorf("expected an ack in %s state, but got none: %s", state, err)
 				}
 			}
@@ -200,8 +200,8 @@ func TestActiveCloseUnaccpSeqAck(t *testing.T) {
 			sendUnaccSeqAck("FIN_WAIT2")
 
 			// Send a FIN to DUT to get to TIME_WAIT
-			conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagFin | header.TCPFlagAck)})
-			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second); err != nil {
+			conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagFin | header.TCPFlagAck)})
+			if _, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second); err != nil {
 				t.Fatalf("expected an ACK for our fin and DUT should enter TIME_WAIT: %s", err)
 			}
 
@@ -217,7 +217,7 @@ func TestActiveCloseUnaccpSeqAck(t *testing.T) {
 func generateOTWSeqSegment(t *testing.T, conn *testbench.TCPIPv4, seqNumOffset seqnum.Size, windowSize seqnum.Size) testbench.TCP {
 	lastAcceptable := conn.LocalSeqNum(t).Add(windowSize)
 	otwSeq := uint32(lastAcceptable.Add(seqNumOffset))
-	return testbench.TCP{SeqNum: testbench.Uint32(otwSeq), Flags: testbench.Uint8(header.TCPFlagAck)}
+	return testbench.TCP{SeqNum: testbench.Uint32(otwSeq), Flags: testbench.TCPFlags(header.TCPFlagAck)}
 }
 
 // generateUnaccACKSegment generates an segment with
@@ -226,5 +226,5 @@ func generateOTWSeqSegment(t *testing.T, conn *testbench.TCPIPv4, seqNumOffset s
 func generateUnaccACKSegment(t *testing.T, conn *testbench.TCPIPv4, seqNumOffset seqnum.Size, windowSize seqnum.Size) testbench.TCP {
 	lastAcceptable := conn.RemoteSeqNum(t)
 	unaccAck := uint32(lastAcceptable.Add(seqNumOffset))
-	return testbench.TCP{AckNum: testbench.Uint32(unaccAck), Flags: testbench.Uint8(header.TCPFlagAck)}
+	return testbench.TCP{AckNum: testbench.Uint32(unaccAck), Flags: testbench.TCPFlags(header.TCPFlagAck)}
 }
diff --git a/test/packetimpact/tests/tcp_user_timeout_test.go b/test/packetimpact/tests/tcp_user_timeout_test.go
index b16e65366..ef38bd738 100644
--- a/test/packetimpact/tests/tcp_user_timeout_test.go
+++ b/test/packetimpact/tests/tcp_user_timeout_test.go
@@ -35,7 +35,7 @@ func sendPayload(t *testing.T, conn *testbench.TCPIPv4, dut *testbench.DUT, fd i
 	}
 	conn.Drain(t)
 	dut.Send(t, fd, sampleData, 0)
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, &testbench.Payload{Bytes: sampleData}, time.Second); err != nil {
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, &testbench.Payload{Bytes: sampleData}, time.Second); err != nil {
 		t.Fatalf("expected data but got none: %w", err)
 	}
 }
@@ -79,14 +79,14 @@ func TestTCPUserTimeout(t *testing.T) {
 
 				time.Sleep(tt.sendDelay)
 				conn.Drain(t)
-				conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+				conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 
 				// If TCP_USER_TIMEOUT was set and the above delay was longer than the
 				// TCP_USER_TIMEOUT then the DUT should send a RST in response to the
 				// testbench's packet.
 				expectRST := tt.userTimeout != 0 && tt.sendDelay > tt.userTimeout
 				expectTimeout := 5 * time.Second
-				got, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, expectTimeout)
+				got, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, expectTimeout)
 				if expectRST && err != nil {
 					t.Errorf("expected RST packet within %s but got none: %s", expectTimeout, err)
 				}
diff --git a/test/packetimpact/tests/tcp_window_shrink_test.go b/test/packetimpact/tests/tcp_window_shrink_test.go
index 093484721..0d65a2ea2 100644
--- a/test/packetimpact/tests/tcp_window_shrink_test.go
+++ b/test/packetimpact/tests/tcp_window_shrink_test.go
@@ -48,7 +48,7 @@ func TestWindowShrink(t *testing.T) {
 	if _, err := conn.ExpectData(t, &testbench.TCP{}, samplePayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
 
 	dut.Send(t, acceptFd, sampleData, 0)
 	dut.Send(t, acceptFd, sampleData, 0)
@@ -59,7 +59,7 @@ func TestWindowShrink(t *testing.T) {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
 	// We close our receiving window here
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
 
 	dut.Send(t, acceptFd, []byte("Sample Data"), 0)
 	// Note: There is another kind of zero-window probing which Windows uses (by sending one
diff --git a/test/packetimpact/tests/tcp_zero_receive_window_test.go b/test/packetimpact/tests/tcp_zero_receive_window_test.go
index d06690705..d73495454 100644
--- a/test/packetimpact/tests/tcp_zero_receive_window_test.go
+++ b/test/packetimpact/tests/tcp_zero_receive_window_test.go
@@ -49,8 +49,8 @@ func TestZeroReceiveWindow(t *testing.T) {
 			// Expect the DUT to eventually advertise zero receive window.
 			// The test would timeout otherwise.
 			for readOnce := false; ; {
-				conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
-				gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+				conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
+				gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 				if err != nil {
 					t.Fatalf("expected packet was not received: %s", err)
 				}
@@ -100,8 +100,8 @@ func TestNonZeroReceiveWindow(t *testing.T) {
 			// we sent. Once we have received ACKs with non-zero receive windows, we break
 			// the loop.
 			for {
-				conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
-				gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, time.Second)
+				conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
+				gotTCP, err := conn.Expect(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, time.Second)
 				if err != nil {
 					t.Fatalf("expected packet was not received: %s", err)
 				}
diff --git a/test/packetimpact/tests/tcp_zero_window_probe_retransmit_test.go b/test/packetimpact/tests/tcp_zero_window_probe_retransmit_test.go
index 413cfa236..22b17a39e 100644
--- a/test/packetimpact/tests/tcp_zero_window_probe_retransmit_test.go
+++ b/test/packetimpact/tests/tcp_zero_window_probe_retransmit_test.go
@@ -59,8 +59,8 @@ func TestZeroWindowProbeRetransmit(t *testing.T) {
 	// of the recorded first zero probe transmission duration.
 	//
 	// Advertize zero receive window along with a payload.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh), WindowSize: testbench.Uint16(0)}, samplePayload)
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh), WindowSize: testbench.Uint16(0)}, samplePayload)
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, nil, time.Second); err != nil {
 		t.Fatalf("expected packet was not received: %s", err)
 	}
 	// Wait for the payload to be received by the DUT, which is also an
@@ -104,10 +104,10 @@ func TestZeroWindowProbeRetransmit(t *testing.T) {
 		}
 		prev = got
 		// Acknowledge the zero-window probes from the dut.
-		conn.Send(t, testbench.TCP{AckNum: ackProbe, Flags: testbench.Uint8(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
+		conn.Send(t, testbench.TCP{AckNum: ackProbe, Flags: testbench.TCPFlags(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
 	}
 	// Advertize non-zero window.
-	conn.Send(t, testbench.TCP{AckNum: ackProbe, Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{AckNum: ackProbe, Flags: testbench.TCPFlags(header.TCPFlagAck)})
 	// Expect the dut to recover and transmit data.
 	if _, err := conn.ExpectData(t, &testbench.TCP{SeqNum: ackProbe}, samplePayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
diff --git a/test/packetimpact/tests/tcp_zero_window_probe_test.go b/test/packetimpact/tests/tcp_zero_window_probe_test.go
index 650a569cc..8b90fcbe9 100644
--- a/test/packetimpact/tests/tcp_zero_window_probe_test.go
+++ b/test/packetimpact/tests/tcp_zero_window_probe_test.go
@@ -53,8 +53,8 @@ func TestZeroWindowProbe(t *testing.T) {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
 	sendTime := time.Now().Sub(start)
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, nil, time.Second); err != nil {
 		t.Fatalf("expected packet was not received: %s", err)
 	}
 
@@ -62,7 +62,7 @@ func TestZeroWindowProbe(t *testing.T) {
 	//         probe to be sent.
 	//
 	// Advertize zero window to the dut.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
 
 	// Expected sequence number of the zero window probe.
 	probeSeq := testbench.Uint32(uint32(*conn.RemoteSeqNum(t) - 1))
@@ -93,7 +93,7 @@ func TestZeroWindowProbe(t *testing.T) {
 	//         and sends out the sample payload after the send window opens.
 	//
 	// Advertize non-zero window to the dut and ack the zero window probe.
-	conn.Send(t, testbench.TCP{AckNum: ackProbe, Flags: testbench.Uint8(header.TCPFlagAck)})
+	conn.Send(t, testbench.TCP{AckNum: ackProbe, Flags: testbench.TCPFlags(header.TCPFlagAck)})
 	// Expect the dut to recover and transmit data.
 	if _, err := conn.ExpectData(t, &testbench.TCP{SeqNum: ackProbe}, samplePayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
@@ -104,8 +104,8 @@ func TestZeroWindowProbe(t *testing.T) {
 	//         Basically with sequence number to one byte behind the unacknowledged
 	//         sequence number.
 	p := testbench.Uint32(uint32(*conn.LocalSeqNum(t)))
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), SeqNum: testbench.Uint32(uint32(*conn.LocalSeqNum(t) - 1))})
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), AckNum: p}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), SeqNum: testbench.Uint32(uint32(*conn.LocalSeqNum(t) - 1))})
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), AckNum: p}, nil, time.Second); err != nil {
 		t.Fatalf("expected a packet with ack number: %d: %s", p, err)
 	}
 }
diff --git a/test/packetimpact/tests/tcp_zero_window_probe_usertimeout_test.go b/test/packetimpact/tests/tcp_zero_window_probe_usertimeout_test.go
index 079fea68c..1ce4d22b7 100644
--- a/test/packetimpact/tests/tcp_zero_window_probe_usertimeout_test.go
+++ b/test/packetimpact/tests/tcp_zero_window_probe_usertimeout_test.go
@@ -51,8 +51,8 @@ func TestZeroWindowProbeUserTimeout(t *testing.T) {
 	if _, err := conn.ExpectData(t, &testbench.TCP{}, samplePayload, time.Second); err != nil {
 		t.Fatalf("expected payload was not received: %s", err)
 	}
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck | header.TCPFlagPsh)}, samplePayload)
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)}, nil, time.Second); err != nil {
 		t.Fatalf("expected packet was not received: %s", err)
 	}
 
@@ -60,7 +60,7 @@ func TestZeroWindowProbeUserTimeout(t *testing.T) {
 	//         probe to be sent.
 	//
 	// Advertize zero window to the dut.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
 
 	// Expected sequence number of the zero window probe.
 	probeSeq := testbench.Uint32(uint32(*conn.RemoteSeqNum(t) - 1))
@@ -81,7 +81,7 @@ func TestZeroWindowProbeUserTimeout(t *testing.T) {
 	// Reduce the retransmit timeout.
 	dut.SetSockOptInt(t, acceptFd, unix.IPPROTO_TCP, unix.TCP_USER_TIMEOUT, int32(startProbeDuration.Milliseconds()))
 	// Advertize zero window again.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck), WindowSize: testbench.Uint16(0)})
 	// Ask the dut to send out data that would trigger zero window probe retransmissions.
 	dut.Send(t, acceptFd, sampleData, 0)
 
@@ -90,8 +90,8 @@ func TestZeroWindowProbeUserTimeout(t *testing.T) {
 
 	// Expect the connection to have timed out and closed which would cause the dut
 	// to reply with a RST to the ACK we send.
-	conn.Send(t, testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
-	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.Uint8(header.TCPFlagRst)}, nil, time.Second); err != nil {
+	conn.Send(t, testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagAck)})
+	if _, err := conn.ExpectData(t, &testbench.TCP{Flags: testbench.TCPFlags(header.TCPFlagRst)}, nil, time.Second); err != nil {
 		t.Fatalf("expected a TCP RST")
 	}
 }