Make tcp_noaccept_close_rst more robust

There used to be a race condition where we may call Close before the connection
is established. Adding poll support so that we can eliminate this kind of race.

Startblock:
  has LGTM from iyerm
  and then
  add reviewer tamird
PiperOrigin-RevId: 354369130

5 files changed, 130 insertions, 0 deletions

diff --git a/test/packetimpact/dut/BUILD b/test/packetimpact/dut/BUILD
index ccf1c735f..0be14ca3e 100644
--- a/test/packetimpact/dut/BUILD
+++ b/test/packetimpact/dut/BUILD
@@ -14,6 +14,7 @@ cc_binary(
         grpcpp,
         "//test/packetimpact/proto:posix_server_cc_grpc_proto",
         "//test/packetimpact/proto:posix_server_cc_proto",
+        "@com_google_absl//absl/strings:str_format",
     ],
 )
 
@@ -24,5 +25,6 @@ cc_binary(
         grpcpp,
         "//test/packetimpact/proto:posix_server_cc_grpc_proto",
         "//test/packetimpact/proto:posix_server_cc_proto",
+        "@com_google_absl//absl/strings:str_format",
     ],
 )
diff --git a/test/packetimpact/dut/posix_server.cc b/test/packetimpact/dut/posix_server.cc
index 4de8540f6..eba21df12 100644
--- a/test/packetimpact/dut/posix_server.cc
+++ b/test/packetimpact/dut/posix_server.cc
@@ -16,6 +16,7 @@
 #include <getopt.h>
 #include <netdb.h>
 #include <netinet/in.h>
+#include <poll.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -30,6 +31,7 @@
 #include "include/grpcpp/security/server_credentials.h"
 #include "include/grpcpp/server_builder.h"
 #include "include/grpcpp/server_context.h"
+#include "absl/strings/str_format.h"
 #include "test/packetimpact/proto/posix_server.grpc.pb.h"
 #include "test/packetimpact/proto/posix_server.pb.h"
 
@@ -256,6 +258,44 @@ class PosixImpl final : public posix_server::Posix::Service {
     return ::grpc::Status::OK;
   }
 
+  ::grpc::Status Poll(::grpc::ServerContext *context,
+                      const ::posix_server::PollRequest *request,
+                      ::posix_server::PollResponse *response) override {
+    std::vector<struct pollfd> pfds;
+    pfds.reserve(request->pfds_size());
+    for (const auto &pfd : request->pfds()) {
+      pfds.push_back({
+          .fd = pfd.fd(),
+          .events = static_cast<short>(pfd.events()),
+      });
+    }
+    int ret = ::poll(pfds.data(), pfds.size(), request->timeout_millis());
+
+    response->set_ret(ret);
+    if (ret < 0) {
+      response->set_errno_(errno);
+    } else {
+      // Only pollfds that have non-empty revents are returned, the client can't
+      // rely on indexes of the request array.
+      for (const auto &pfd : pfds) {
+        if (pfd.revents) {
+          auto *proto_pfd = response->add_pfds();
+          proto_pfd->set_fd(pfd.fd);
+          proto_pfd->set_events(pfd.revents);
+        }
+      }
+      if (int ready = response->pfds_size(); ret != ready) {
+        return ::grpc::Status(
+            ::grpc::StatusCode::INTERNAL,
+            absl::StrFormat(
+                "poll's return value(%d) doesn't match the number of "
+                "file descriptors that are actually ready(%d)",
+                ret, ready));
+      }
+    }
+    return ::grpc::Status::OK;
+  }
+
   ::grpc::Status Send(::grpc::ServerContext *context,
                       const ::posix_server::SendRequest *request,
                       ::posix_server::SendResponse *response) override {
diff --git a/test/packetimpact/proto/posix_server.proto b/test/packetimpact/proto/posix_server.proto
index f32ed54ef..b4c68764a 100644
--- a/test/packetimpact/proto/posix_server.proto
+++ b/test/packetimpact/proto/posix_server.proto
@@ -142,6 +142,25 @@ message ListenResponse {
   int32 errno_ = 2;  // "errno" may fail to compile in c++.
 }
 
+// The events field is overloaded: when used for request, it is copied into the
+// events field of posix struct pollfd; when used for response, it is filled by
+// the revents field from the posix struct pollfd.
+message PollFd {
+  int32 fd = 1;
+  uint32 events = 2;
+}
+
+message PollRequest {
+  repeated PollFd pfds = 1;
+  int32 timeout_millis = 2;
+}
+
+message PollResponse {
+  int32 ret = 1;
+  int32 errno_ = 2;  // "errno" may fail to compile in c++.
+  repeated PollFd pfds = 3;
+}
+
 message SendRequest {
   int32 sockfd = 1;
   bytes buf = 2;
@@ -226,6 +245,10 @@ service Posix {
   rpc GetSockOpt(GetSockOptRequest) returns (GetSockOptResponse);
   // Call listen() on the DUT.
   rpc Listen(ListenRequest) returns (ListenResponse);
+  // Call poll() on the DUT. Only pollfds that have non-empty revents are
+  // returned, the only way to tie the response back to the original request
+  // is using the fd number.
+  rpc Poll(PollRequest) returns (PollResponse);
   // Call send() on the DUT.
   rpc Send(SendRequest) returns (SendResponse);
   // Call sendto() on the DUT.
diff --git a/test/packetimpact/testbench/dut.go b/test/packetimpact/testbench/dut.go
index d40890e12..aedcf6013 100644
--- a/test/packetimpact/testbench/dut.go
+++ b/test/packetimpact/testbench/dut.go
@@ -486,6 +486,56 @@ func (dut *DUT) ListenWithErrno(ctx context.Context, t *testing.T, sockfd, backl
 	return resp.GetRet(), syscall.Errno(resp.GetErrno_())
 }
 
+// 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
+// response back to the original request is using the fd number.
+func (dut *DUT) Poll(t *testing.T, pfds []unix.PollFd, timeout time.Duration) []unix.PollFd {
+	t.Helper()
+
+	ctx := context.Background()
+	var cancel context.CancelFunc
+	if timeout >= 0 {
+		ctx, cancel = context.WithTimeout(ctx, timeout+RPCTimeout)
+		defer cancel()
+	}
+	ret, result, err := dut.PollWithErrno(ctx, t, pfds, timeout)
+	if ret < 0 {
+		t.Fatalf("failed to poll: %s", err)
+	}
+	return result
+}
+
+// PollWithErrno calls poll on the DUT.
+func (dut *DUT) PollWithErrno(ctx context.Context, t *testing.T, pfds []unix.PollFd, timeout time.Duration) (int32, []unix.PollFd, error) {
+	t.Helper()
+
+	req := pb.PollRequest{
+		TimeoutMillis: int32(timeout.Milliseconds()),
+	}
+	for _, pfd := range pfds {
+		req.Pfds = append(req.Pfds, &pb.PollFd{
+			Fd:     pfd.Fd,
+			Events: uint32(pfd.Events),
+		})
+	}
+	resp, err := dut.posixServer.Poll(ctx, &req)
+	if err != nil {
+		t.Fatalf("failed to call Poll: %s", err)
+	}
+	if ret, npfds := resp.GetRet(), len(resp.GetPfds()); ret >= 0 && int(ret) != npfds {
+		t.Fatalf("nonsensical poll response: ret(%d) != len(pfds)(%d)", ret, npfds)
+	}
+	var result []unix.PollFd
+	for _, protoPfd := range resp.GetPfds() {
+		result = append(result, unix.PollFd{
+			Fd:      protoPfd.GetFd(),
+			Revents: int16(protoPfd.GetEvents()),
+		})
+	}
+	return resp.GetRet(), result, syscall.Errno(resp.GetErrno_())
+}
+
 // Send calls send on the DUT and causes a fatal test failure if it doesn't
 // succeed. If more control over the timeout or error handling is needed, use
 // SendWithErrno.
diff --git a/test/packetimpact/tests/tcp_noaccept_close_rst_test.go b/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
index f0af5352d..c874a8912 100644
--- a/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
+++ b/test/packetimpact/tests/tcp_noaccept_close_rst_test.go
@@ -34,6 +34,21 @@ func TestTcpNoAcceptCloseReset(t *testing.T) {
 	conn := dut.Net.NewTCPIPv4(t, testbench.TCP{DstPort: &remotePort}, testbench.TCP{SrcPort: &remotePort})
 	conn.Connect(t)
 	defer conn.Close(t)
+	// We need to wait for POLLIN event on listenFd to know the connection is
+	// 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.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)