3 files changed, 44 insertions, 15 deletions
diff --git a/test/packetimpact/testbench/dut.go b/test/packetimpact/testbench/dut.go
index aedcf6013..81634b5f4 100644
--- a/test/packetimpact/testbench/dut.go
+++ b/test/packetimpact/testbench/dut.go
@@ -486,6 +486,23 @@ func (dut *DUT) ListenWithErrno(ctx context.Context, t *testing.T, sockfd, backl
 	return resp.GetRet(), syscall.Errno(resp.GetErrno_())
 }
 
+// PollOne calls poll on the DUT and asserts that the expected event must be
+// signaled on the given fd within the given timeout.
+func (dut *DUT) PollOne(t *testing.T, fd int32, events int16, timeout time.Duration) {
+	t.Helper()
+
+	pfds := dut.Poll(t, []unix.PollFd{{Fd: fd, Events: events}}, timeout)
+	if n := len(pfds); n != 1 {
+		t.Fatalf("Poll returned %d ready file descriptors, expected 1", n)
+	}
+	if readyFd := pfds[0].Fd; readyFd != fd {
+		t.Fatalf("Poll returned an fd %d that was not requested (%d)", readyFd, fd)
+	}
+	if got, want := pfds[0].Revents, int16(events); got&want == 0 {
+		t.Fatalf("Poll returned no events in our interest, got: %#b, want: %#b", got, want)
+	}
+}
+
 // Poll calls poll on the DUT and causes a fatal test failure if it doesn't
 // succeed. If more control over error handling is needed, use PollWithErrno.
 // Only pollfds with non-empty revents are returned, the only way to tie the
diff --git a/test/packetimpact/tests/tcp_network_unreachable_test.go b/test/packetimpact/tests/tcp_network_unreachable_test.go
index 6cd6d2edf..4eb11f7a1 100644
--- a/test/packetimpact/tests/tcp_network_unreachable_test.go
+++ b/test/packetimpact/tests/tcp_network_unreachable_test.go
@@ -78,8 +78,8 @@ func TestTCPSynSentUnreachable(t *testing.T) {
 	layers = append(layers, &icmpv4, ip, tcp)
 	rawConn.SendFrameStateless(t, layers)
 
-	if _, err = dut.ConnectWithErrno(ctx, t, clientFD, &sa); err != syscall.Errno(unix.EHOSTUNREACH) {
-		t.Errorf("expected connect to fail with EHOSTUNREACH, but got %v", err)
+	if err := getConnectError(ctx, t, &dut, clientFD); err != unix.EHOSTUNREACH {
+		t.Errorf("Expected connect to fail with EHOSTUNREACH, but got %s", err)
 	}
 }
 
@@ -134,7 +134,29 @@ func TestTCPSynSentUnreachable6(t *testing.T) {
 	layers = append(layers, &icmpv6, ip, tcp)
 	rawConn.SendFrameStateless(t, layers)
 
-	if _, err = dut.ConnectWithErrno(ctx, t, clientFD, &sa); err != syscall.Errno(unix.ENETUNREACH) {
-		t.Errorf("expected connect to fail with ENETUNREACH, but got %v", err)
+	if err := getConnectError(ctx, t, &dut, clientFD); err != unix.ENETUNREACH {
+		t.Errorf("Expected connect to fail with ENETUNREACH, but got %s", err)
 	}
 }
+
+// getConnectError gets the errno generated by the on-going connect attempt on
+// fd. fd must be non-blocking and there must be a connect call to fd which
+// returned EINPROGRESS before. These conditions are guaranteed in this test.
+func getConnectError(ctx context.Context, t *testing.T, dut *testbench.DUT, fd int32) error {
+	t.Helper()
+	// We previously got EINPROGRESS form the connect call. We can
+	// handle it as explained by connect(2):
+	// EINPROGRESS:
+	//     The socket is nonblocking and the connection cannot be
+	// completed immediately. It is possible to select(2) or poll(2)
+	// for completion by selecting the socket for writing.  After
+	// select(2) indicates writability, use getsockopt(2) to read
+	// the SO_ERROR option at level SOL_SOCKET to determine
+	// whether connect() completed successfully (SO_ERROR is
+	// zero) or unsuccessfully (SO_ERROR is one of the usual
+	// error codes listed here, explaining the reason for the
+	// failure).
+	dut.PollOne(t, fd, unix.POLLOUT, 10*time.Second)
+	errno := dut.GetSockOptInt(t, fd, unix.SOL_SOCKET, unix.SO_ERROR)
+	return syscall.Errno(errno)
+}
diff --git a/test/packetimpact/tests/tcp_noaccept_close_rst_test.go b/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
index c874a8912..d2871df08 100644
--- a/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
+++ b/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
@@ -38,17 +38,7 @@ func TestTcpNoAcceptCloseReset(t *testing.T) {
 	// established. Otherwise there could be a race when we issue the Close
 	// command prior to the DUT receiving the last ack of the handshake and
 	// it will only respond RST instead of RST+ACK.
-	timeout := time.Second
-	pfds := dut.Poll(t, []unix.PollFd{{Fd: listenFd, Events: unix.POLLIN}}, timeout)
-	if n := len(pfds); n != 1 {
-		t.Fatalf("poll returned %d ready file descriptors, expected 1", n)
-	}
-	if readyFd := pfds[0].Fd; readyFd != listenFd {
-		t.Fatalf("poll returned an fd %d that was not requested (%d)", readyFd, listenFd)
-	}
-	if got, want := pfds[0].Revents, int16(unix.POLLIN); got&want == 0 {
-		t.Fatalf("poll returned no events in our interest, got: %#b, want: %#b", got, want)
-	}
+	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 {
 		t.Fatalf("expected a RST-ACK packet but got none: %s", err)