3 files changed, 37 insertions, 0 deletions
diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go
index 79c0ccf5c..2b8e2f005 100644
--- a/test/packetimpact/testbench/connections.go
+++ b/test/packetimpact/testbench/connections.go
@@ -228,6 +228,12 @@ func (conn *TCPIPv4) RecvFrame(timeout time.Duration) Layers {
 	return nil
 }
 
+// Drain drains the sniffer's receive buffer by receiving packets until there's
+// nothing else to receive.
+func (conn *TCPIPv4) Drain() {
+	conn.sniffer.Drain()
+}
+
 // Expect a packet that matches the provided tcp within the timeout specified.
 // If it doesn't arrive in time, it returns nil.
 func (conn *TCPIPv4) Expect(tcp TCP, timeout time.Duration) (*TCP, error) {
@@ -423,6 +429,12 @@ func (conn *UDPIPv4) Recv(timeout time.Duration) *UDP {
 	return nil
 }
 
+// Drain drains the sniffer's receive buffer by receiving packets until there's
+// nothing else to receive.
+func (conn *UDPIPv4) Drain() {
+	conn.sniffer.Drain()
+}
+
 // Expect a packet that matches the provided udp within the timeout specified.
 // If it doesn't arrive in time, the test fails.
 func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) {
diff --git a/test/packetimpact/testbench/rawsockets.go b/test/packetimpact/testbench/rawsockets.go
index 0074484f7..09bfa43c5 100644
--- a/test/packetimpact/testbench/rawsockets.go
+++ b/test/packetimpact/testbench/rawsockets.go
@@ -97,6 +97,29 @@ func (s *Sniffer) Recv(timeout time.Duration) []byte {
 	}
 }
 
+// Drain drains the Sniffer's socket receive buffer by receiving until there's
+// nothing else to receive.
+func (s *Sniffer) Drain() {
+	s.t.Helper()
+	flags, err := unix.FcntlInt(uintptr(s.fd), unix.F_GETFL, 0)
+	if err != nil {
+		s.t.Fatalf("failed to get sniffer socket fd flags: %s", err)
+	}
+	if _, err := unix.FcntlInt(uintptr(s.fd), unix.F_SETFL, flags|unix.O_NONBLOCK); err != nil {
+		s.t.Fatalf("failed to make sniffer socket non-blocking: %s", err)
+	}
+	for {
+		buf := make([]byte, maxReadSize)
+		_, _, err := unix.Recvfrom(s.fd, buf, unix.MSG_TRUNC)
+		if err == unix.EINTR || err == unix.EAGAIN || err == unix.EWOULDBLOCK {
+			break
+		}
+	}
+	if _, err := unix.FcntlInt(uintptr(s.fd), unix.F_SETFL, flags); err != nil {
+		s.t.Fatalf("failed to restore sniffer socket fd flags: %s", err)
+	}
+}
+
 // Close the socket that Sniffer is using.
 func (s *Sniffer) Close() {
 	if err := unix.Close(s.fd); err != nil {
diff --git a/test/packetimpact/tests/fin_wait2_timeout_test.go b/test/packetimpact/tests/fin_wait2_timeout_test.go
index 90e16ef65..b98594f94 100644
--- a/test/packetimpact/tests/fin_wait2_timeout_test.go
+++ b/test/packetimpact/tests/fin_wait2_timeout_test.go
@@ -53,6 +53,8 @@ func TestFinWait2Timeout(t *testing.T) {
 			conn.Send(tb.TCP{Flags: tb.Uint8(header.TCPFlagAck)})
 
 			time.Sleep(5 * time.Second)
+			conn.Drain()
+
 			conn.Send(tb.TCP{Flags: tb.Uint8(header.TCPFlagAck)})
 			if tt.linger2 {
 				if _, err := conn.Expect(tb.TCP{Flags: tb.Uint8(header.TCPFlagRst)}, time.Second); err != nil {