1 files changed, 38 insertions, 6 deletions

diff --git a/test/packetimpact/testbench/rawsockets.go b/test/packetimpact/testbench/rawsockets.go
index 0c7d0f979..ff722d4a6 100644
--- a/test/packetimpact/testbench/rawsockets.go
+++ b/test/packetimpact/testbench/rawsockets.go
@@ -17,6 +17,7 @@ package testbench
 import (
 	"encoding/binary"
 	"flag"
+	"fmt"
 	"math"
 	"net"
 	"testing"
@@ -47,6 +48,12 @@ func NewSniffer(t *testing.T) (Sniffer, error) {
 	if err != nil {
 		return Sniffer{}, err
 	}
+	if err := unix.SetsockoptInt(snifferFd, unix.SOL_SOCKET, unix.SO_RCVBUFFORCE, 1); err != nil {
+		t.Fatalf("can't set sockopt SO_RCVBUFFORCE to 1: %s", err)
+	}
+	if err := unix.SetsockoptInt(snifferFd, unix.SOL_SOCKET, unix.SO_RCVBUF, 1e7); err != nil {
+		t.Fatalf("can't setsockopt SO_RCVBUF to 10M: %s", err)
+	}
 	return Sniffer{
 		t:  t,
 		fd: snifferFd,
@@ -91,12 +98,36 @@ func (s *Sniffer) Recv(timeout time.Duration) []byte {
 	}
 }
 
-// Close the socket that Sniffer is using.
-func (s *Sniffer) Close() {
+// 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() error {
 	if err := unix.Close(s.fd); err != nil {
-		s.t.Fatalf("can't close sniffer socket: %s", err)
+		return fmt.Errorf("can't close sniffer socket: %w", err)
 	}
 	s.fd = -1
+	return nil
 }
 
 // Injector can inject raw frames.
@@ -142,10 +173,11 @@ func (i *Injector) Send(b []byte) {
 	}
 }
 
-// Close the underlying socket.
-func (i *Injector) Close() {
+// close the underlying socket.
+func (i *Injector) close() error {
 	if err := unix.Close(i.fd); err != nil {
-		i.t.Fatalf("can't close sniffer socket: %s", err)
+		return fmt.Errorf("can't close sniffer socket: %w", err)
 	}
 	i.fd = -1
+	return nil
 }