2 files changed, 35 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 {