Better network proxy tests

Fix sending error on `connect` failure.

PiperOrigin-RevId: 562693682
Change-Id: I70c710a9001f22e172cbe4df328983bfa7188d3d

sandboxed_api/sandbox2/BUILD.bazel

@@ -1108,6 +1108,7 @@ cc_test(
         "//sandboxed_api/sandbox2/network_proxy:testing",
         "//sandboxed_api/util:status_matchers",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
     ],
 )

sandboxed_api/sandbox2/CMakeLists.txt

@@ -1206,6 +1206,7 @@ if(BUILD_TESTING AND SAPI_BUILD_TESTING)
   )
   target_link_libraries(sandbox2_network_proxy_test
     PRIVATE absl::strings
+            absl::time
             sandbox2::sandbox2
             sandbox2::network_proxy_testing
             sapi::status_matchers

sandboxed_api/sandbox2/examples/network/network_sandbox.cc

@@ -126,9 +126,10 @@ int main(int argc, char* argv[]) {
   absl::ParseCommandLine(argc, argv);
   absl::InitializeLog();
 
-  absl::StatusOr<int> port = sandbox2::StartNetworkProxyTestServer();
-  if (!port.ok()) {
-    LOG(ERROR) << port.status();
+  absl::StatusOr<std::unique_ptr<sandbox2::NetworkProxyTestServer>> server =
+      sandbox2::NetworkProxyTestServer::Start();
+  if (!server.ok()) {
+    LOG(ERROR) << server.status();
     return EXIT_FAILURE;
   }
 
@@ -165,7 +166,7 @@ int main(int argc, char* argv[]) {
     return 2;
   }
 
-  if (!HandleSandboxee(comms, *port)) {
+  if (!HandleSandboxee(comms, (*server)->port())) {
     if (!s2.IsTerminated()) {
       // Kill the sandboxee, because failure to receive the data over the Comms
       // channel doesn't automatically mean that the sandboxee itself had

sandboxed_api/sandbox2/examples/network_proxy/networkproxy_sandbox.cc

@@ -74,9 +74,10 @@ int main(int argc, char* argv[]) {
   absl::ParseCommandLine(argc, argv);
   absl::InitializeLog();
 
-  absl::StatusOr<int> port = sandbox2::StartNetworkProxyTestServer();
-  if (!port.ok()) {
-    LOG(ERROR) << port.status();
+  absl::StatusOr<std::unique_ptr<sandbox2::NetworkProxyTestServer>> server =
+      sandbox2::NetworkProxyTestServer::Start();
+  if (!server.ok()) {
+    LOG(ERROR) << server.status();
     return EXIT_FAILURE;
   }
 
@@ -118,7 +119,7 @@ int main(int argc, char* argv[]) {
   }
 
   // Send the port number via comms
-  if (!comms->SendInt32(*port)) {
+  if (!comms->SendInt32((*server)->port())) {
     LOG(ERROR) << "sandboxee_comms->SendInt32() failed";
     return 3;
   }

sandboxed_api/sandbox2/network_proxy/BUILD.bazel

@@ -84,7 +84,10 @@ cc_library(
     hdrs = ["testing.h"],
     deps = [
         "//sandboxed_api/util:fileops",
-        "@com_google_absl//absl/log",
+        "//sandboxed_api/util:status",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

sandboxed_api/sandbox2/network_proxy/CMakeLists.txt

@@ -67,11 +67,14 @@ add_library(sandbox2_network_proxy_testing ${SAPI_LIB_TYPE}
 )
 add_library(sandbox2::network_proxy_testing ALIAS sandbox2_network_proxy_testing)
 target_link_libraries(sandbox2_network_proxy_testing PRIVATE
-  absl::log
+  absl::check
+  absl::core_headers
+  absl::memory
   absl::status
   absl::statusor
   absl::strings
   sapi::fileops
+  sapi::status
   sapi::base
 )
 

sandboxed_api/sandbox2/network_proxy/server.cc

@@ -22,6 +22,7 @@
 
 #include <atomic>
 #include <cerrno>
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
@@ -30,6 +31,8 @@
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "sandboxed_api/sandbox2/comms.h"
+#include "sandboxed_api/sandbox2/network_proxy/filtering.h"
 #include "sandboxed_api/util/fileops.h"
 
 namespace sandbox2 {
@@ -77,14 +80,15 @@ void NetworkProxyServer::ProcessConnectRequest() {
   int result = connect(
       new_socket, reinterpret_cast<const sockaddr*>(addr.data()), addr.size());
 
-  if (result == 0) {
-    NotifySuccess();
-    if (!fatal_error_) {
-      if (!comms_->SendFD(new_socket)) {
-        fatal_error_ = true;
-        return;
-      }
-    }
+  if (result < 0) {
+    SendError(errno);
+    return;
+  }
+
+  NotifySuccess();
+  if (!fatal_error_ && !comms_->SendFD(new_socket)) {
+    fatal_error_ = true;
+    return;
   }
 }
 

sandboxed_api/sandbox2/network_proxy/testing.cc

@@ -16,76 +16,106 @@
 
 #include <arpa/inet.h>
 #include <netinet/in.h>
+#include <poll.h>
+#include <sys/eventfd.h>
 #include <sys/socket.h>
 
+#include <cstddef>
+#include <cstdint>
+#include <memory>
 #include <thread>
+#include <utility>
 
-#include "absl/log/log.h"
+#include "absl/base/macros.h"
+#include "absl/log/check.h"
+#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "sandboxed_api/util/fileops.h"
+#include "sandboxed_api/util/status_macros.h"
 
 namespace sandbox2 {
 namespace {
-void ServerThread(int port) {
-  sapi::file_util::fileops::FDCloser s{
-      socket(AF_INET6, SOCK_STREAM | SOCK_CLOEXEC, 0)};
+
+using sapi::file_util::fileops::FDCloser;
+
+absl::StatusOr<FDCloser> CreateServerSocket(int port, bool ipv6 = true) {
+  // Listen to localhost only
+  sockaddr_in6 addr6;
+  sockaddr_in addr4;
+  addr6.sin6_family = AF_INET6;
+  addr4.sin_family = AF_INET;
+  addr4.sin_port = addr6.sin6_port = htons(port);
+  PCHECK(inet_pton(AF_INET, "127.0.0.1", &addr4.sin_addr.s_addr) == 1);
+  PCHECK(inet_pton(AF_INET6, "::1", &addr6.sin6_addr.s6_addr) == 1);
+  sockaddr* addr = ipv6 ? reinterpret_cast<sockaddr*>(&addr6)
+                        : reinterpret_cast<sockaddr*>(&addr4);
+  size_t addr_size = ipv6 ? sizeof(addr6) : sizeof(addr4);
+  sapi::file_util::fileops::FDCloser s(
+      socket(addr->sa_family, SOCK_STREAM | SOCK_CLOEXEC, 0));
   if (s.get() < 0) {
-    PLOG(ERROR) << "socket() failed";
-    return;
+    return absl::InternalError("socket() failed");
   }
-
-  {
-    int enable = 1;
-    if (setsockopt(s.get(), SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(enable)) <
-        0) {
-      PLOG(ERROR) << "setsockopt(SO_REUSEADDR) failed";
-      return;
-    }
+  int enable = 1;
+  if (setsockopt(s.get(), SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(enable)) <
+      0) {
+    return absl::InternalError("setsockopt() failed");
   }
-
-  // Listen to localhost only.
-  struct sockaddr_in6 addr = {};
-  addr.sin6_family = AF_INET6;
-  addr.sin6_port = htons(port);
-
-  int err = inet_pton(AF_INET6, "::1", &addr.sin6_addr.s6_addr);
-  if (err == 0) {
-    LOG(ERROR) << "inet_pton() failed";
-    return;
-  } else if (err == -1) {
-    PLOG(ERROR) << "inet_pton() failed";
-    return;
+  if (bind(s.get(), addr, addr_size) < 0) {
+    return absl::InternalError("bind() failed");
   }
-
-  if (bind(s.get(), (struct sockaddr*)&addr, sizeof(addr)) < 0) {
-    PLOG(ERROR) << "bind() failed";
-    return;
-  }
-
   if (listen(s.get(), 1) < 0) {
-    PLOG(ERROR) << "listen() failed";
-    return;
-  }
-
-  sapi::file_util::fileops::FDCloser client{accept(s.get(), 0, 0)};
-  if (client.get() < 0) {
-    PLOG(ERROR) << "accept() failed";
-    return;
-  }
-
-  constexpr absl::string_view kMsg = "Hello World\n";
-  if (write(client.get(), kMsg.data(), kMsg.size()) < 0) {
-    PLOG(ERROR) << "write() failed";
+    return absl::InternalError("listen() failed");
   }
+  return s;
 }
+
 }  // namespace
 
-absl::StatusOr<int> StartNetworkProxyTestServer() {
-  static int port = 8085;
-  std::thread server_thread([] { ServerThread(++port); });
-  server_thread.detach();
-  return port;
+absl::StatusOr<std::unique_ptr<NetworkProxyTestServer>>
+NetworkProxyTestServer::Start(bool ipv6) {
+   static int port = 8085;
+  FDCloser event_fd(eventfd(0, 0));
+  if (event_fd.get() < 0) {
+    return absl::InternalError("eventfd() failed");
+  }
+  SAPI_ASSIGN_OR_RETURN(FDCloser server_socket, CreateServerSocket(port, ipv6));
+  auto server = absl::WrapUnique(new NetworkProxyTestServer(
+      port, std::move(server_socket), std::move(event_fd)));
+  server->Spawn();
+  return server;
+}
+
+void NetworkProxyTestServer::Stop() {
+  if (event_fd_.get() < 0) {
+    return;
+  }
+  uint64_t value = 1;
+  PCHECK(write(event_fd_.get(), &value, sizeof(value)) == sizeof(value));
+  thread_.join();
+  event_fd_.Close();
+  server_socket_.Close();
+}
+
+void NetworkProxyTestServer::Run() {
+  struct pollfd pfds[] = {
+      {.fd = server_socket_.get(), .events = POLLIN},
+      {.fd = event_fd_.get(), .events = POLLIN},
+  };
+  do {
+    PCHECK(poll(pfds, ABSL_ARRAYSIZE(pfds), -1) > 0);
+    if (pfds[1].revents & POLLIN) {
+      return;
+    }
+  } while (!(pfds[0].revents & POLLIN));
+  FDCloser client(accept(server_socket_.get(), 0, 0));
+  PCHECK(client.get() >= 0);
+  constexpr absl::string_view kMsg = "Hello World\n";
+  PCHECK(write(client.get(), kMsg.data(), kMsg.size()) == kMsg.size());
+}
+
+void NetworkProxyTestServer::Spawn() {
+  thread_ = std::thread([this] { Run(); });
 }
 
 }  // namespace sandbox2

sandboxed_api/sandbox2/network_proxy/testing.h

@@ -15,12 +15,42 @@
 #ifndef SANDBOXED_API_SANDBOX2_EXAMPLES_NETWORK_PROXY_NETWORKPROXY_LIB_H_
 #define SANDBOXED_API_SANDBOX2_EXAMPLES_NETWORK_PROXY_NETWORKPROXY_LIB_H_
 
+#include <memory>
+#include <thread>
+#include <utility>
+
 #include "absl/status/statusor.h"
+#include "sandboxed_api/util/fileops.h"
 
 namespace sandbox2 {
 
-// Returns the port the server was started on.
-absl::StatusOr<int> StartNetworkProxyTestServer();
+class NetworkProxyTestServer {
+ public:
+  static absl::StatusOr<std::unique_ptr<NetworkProxyTestServer>> Start(
+      bool ipv6 = true);
+
+  NetworkProxyTestServer(NetworkProxyTestServer&&) = delete;
+  NetworkProxyTestServer& operator=(NetworkProxyTestServer&&) = delete;
+
+  ~NetworkProxyTestServer() { Stop(); }
+
+  int port() { return port_; }
+  void Stop();
+
+ private:
+  NetworkProxyTestServer(int port,
+                         sapi::file_util::fileops::FDCloser server_socket,
+                         sapi::file_util::fileops::FDCloser event_fd)
+      : port_(port),
+        server_socket_(std::move(server_socket)),
+        event_fd_(std::move(event_fd)) {}
+  void Spawn();
+  void Run();
+  std::thread thread_;
+  int port_;
+  sapi::file_util::fileops::FDCloser server_socket_;
+  sapi::file_util::fileops::FDCloser event_fd_;
+};
 
 }  // namespace sandbox2
 

sandboxed_api/sandbox2/network_proxy_test.cc

@@ -22,6 +22,7 @@
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/time/time.h"
 #include "sandboxed_api/sandbox2/executor.h"
 #include "sandboxed_api/sandbox2/network_proxy/testing.h"
 #include "sandboxed_api/sandbox2/policybuilder.h"
@@ -36,12 +37,19 @@ namespace {
 using ::sapi::GetTestSourcePath;
 using ::testing::Eq;
 
-TEST(NetworkProxy, ProxyWithHandlerAllowed) {
+using NetworkProxyTest = ::testing::TestWithParam<bool>;
+
+TEST_P(NetworkProxyTest, ProxyWithHandlerAllowed) {
   SKIP_SANITIZERS;
+  const bool ipv6 = GetParam();
   const std::string path =
       GetTestSourcePath("sandbox2/testcases/network_proxy");
   std::vector<std::string> args = {"network_proxy"};
+  if (ipv6) {
+    args.push_back("--ipv6");
+  }
   auto executor = std::make_unique<Executor>(path, args);
+  executor->limits()->set_walltime_limit(absl::Seconds(3));
 
   PolicyBuilder builder;
   builder.AllowDynamicStartup()
@@ -52,28 +60,38 @@ TEST(NetworkProxy, ProxyWithHandlerAllowed) {
       .AllowTcMalloc()
       .AddNetworkProxyHandlerPolicy()
       .AllowLlvmCoverage()
-      .AllowIPv6("::1")
       .AddLibrariesForBinary(path);
 
+  if (ipv6) {
+    builder.AllowIPv6("::1");
+  } else {
+    builder.AllowIPv4("127.0.0.1");
+  }
+
   SAPI_ASSERT_OK_AND_ASSIGN(auto policy, builder.TryBuild());
 
   Sandbox2 s2(std::move(executor), std::move(policy));
   ASSERT_TRUE(s2.RunAsync());
 
-  SAPI_ASSERT_OK_AND_ASSIGN(int port, StartNetworkProxyTestServer());
-  ASSERT_TRUE(s2.comms()->SendInt32(port));
+  SAPI_ASSERT_OK_AND_ASSIGN(auto server, NetworkProxyTestServer::Start(ipv6));
+  ASSERT_TRUE(s2.comms()->SendInt32(server->port()));
 
   sandbox2::Result result = s2.AwaitResult();
   ASSERT_THAT(result.final_status(), Eq(Result::OK));
   EXPECT_THAT(result.reason_code(), Eq(EXIT_SUCCESS));
 }
 
-TEST(NetworkProxy, ProxyWithHandlerNotAllowed) {
+TEST_P(NetworkProxyTest, ProxyWithHandlerNotAllowed) {
   SKIP_SANITIZERS;
+  const bool ipv6 = GetParam();
   const std::string path =
       GetTestSourcePath("sandbox2/testcases/network_proxy");
   std::vector<std::string> args = {"network_proxy"};
+  if (ipv6) {
+    args.push_back("--ipv6");
+  }
   auto executor = std::make_unique<Executor>(path, args);
+  executor->limits()->set_walltime_limit(absl::Seconds(3));
 
   PolicyBuilder builder;
   builder.AllowDynamicStartup()
@@ -91,20 +109,25 @@ TEST(NetworkProxy, ProxyWithHandlerNotAllowed) {
   Sandbox2 s2(std::move(executor), std::move(policy));
   ASSERT_TRUE(s2.RunAsync());
 
-  SAPI_ASSERT_OK_AND_ASSIGN(int port, StartNetworkProxyTestServer());
-  ASSERT_TRUE(s2.comms()->SendInt32(port));
+  SAPI_ASSERT_OK_AND_ASSIGN(auto server, NetworkProxyTestServer::Start(ipv6));
+  ASSERT_TRUE(s2.comms()->SendInt32(server->port()));
 
   sandbox2::Result result = s2.AwaitResult();
   ASSERT_THAT(result.final_status(), Eq(Result::VIOLATION));
   EXPECT_THAT(result.reason_code(), Eq(Result::VIOLATION_NETWORK));
 }
 
-TEST(NetworkProxy, ProxyWithoutHandlerAllowed) {
+TEST_P(NetworkProxyTest, ProxyWithoutHandlerAllowed) {
   SKIP_SANITIZERS;
+  const bool ipv6 = GetParam();
   const std::string path =
       GetTestSourcePath("sandbox2/testcases/network_proxy");
   std::vector<std::string> args = {"network_proxy", "--noconnect_with_handler"};
+  if (ipv6) {
+    args.push_back("--ipv6");
+  }
   auto executor = std::make_unique<Executor>(path, args);
+  executor->limits()->set_walltime_limit(absl::Seconds(3));
 
   PolicyBuilder builder;
   builder.AllowDynamicStartup()
@@ -115,21 +138,64 @@ TEST(NetworkProxy, ProxyWithoutHandlerAllowed) {
       .AllowTcMalloc()
       .AddNetworkProxyHandlerPolicy()
       .AllowLlvmCoverage()
-      .AllowIPv6("::1")
       .AddLibrariesForBinary(path);
 
+  if (ipv6) {
+    builder.AllowIPv6("::1");
+  } else {
+    builder.AllowIPv4("127.0.0.1");
+  }
+
   SAPI_ASSERT_OK_AND_ASSIGN(auto policy, builder.TryBuild());
 
   Sandbox2 s2(std::move(executor), std::move(policy));
   ASSERT_TRUE(s2.RunAsync());
 
-  SAPI_ASSERT_OK_AND_ASSIGN(int port, StartNetworkProxyTestServer());
-  ASSERT_TRUE(s2.comms()->SendInt32(port));
+  SAPI_ASSERT_OK_AND_ASSIGN(auto server, NetworkProxyTestServer::Start(ipv6));
+  ASSERT_TRUE(s2.comms()->SendInt32(server->port()));
 
   sandbox2::Result result = s2.AwaitResult();
   ASSERT_THAT(result.final_status(), Eq(Result::OK));
   EXPECT_THAT(result.reason_code(), Eq(EXIT_SUCCESS));
 }
 
+TEST(NetworkProxyTest, ProxyNonExistantAddress) {
+  // Creates a IPv6 server tries to connect with IPv4
+  SKIP_SANITIZERS;
+  const std::string path =
+      GetTestSourcePath("sandbox2/testcases/network_proxy");
+  std::vector<std::string> args = {"network_proxy", "--noconnect_with_handler"};
+  auto executor = std::make_unique<Executor>(path, args);
+  executor->limits()->set_walltime_limit(absl::Seconds(3));
+
+  PolicyBuilder builder;
+  builder.AllowDynamicStartup()
+      .AllowExit()
+      .AllowWrite()
+      .AllowRead()
+      .AllowSyscall(__NR_sendto)
+      .AllowTcMalloc()
+      .AddNetworkProxyHandlerPolicy()
+      .AllowLlvmCoverage()
+      .AddLibrariesForBinary(path)
+      .AllowIPv4("127.0.0.1");
+
+  SAPI_ASSERT_OK_AND_ASSIGN(auto policy, builder.TryBuild());
+
+  Sandbox2 s2(std::move(executor), std::move(policy));
+  ASSERT_TRUE(s2.RunAsync());
+
+  SAPI_ASSERT_OK_AND_ASSIGN(auto server,
+                            NetworkProxyTestServer::Start(/*ipv6=*/true));
+  ASSERT_TRUE(s2.comms()->SendInt32(server->port()));
+
+  sandbox2::Result result = s2.AwaitResult();
+  ASSERT_THAT(result.final_status(), Eq(Result::OK));
+  EXPECT_THAT(result.reason_code(), Eq(3));
+}
+
+INSTANTIATE_TEST_SUITE_P(NetworkProxyTest, NetworkProxyTest,
+                         ::testing::Values(true, false));
+
 }  // namespace
 }  // namespace sandbox2

sandboxed_api/sandbox2/testcases/BUILD.bazel

@@ -252,6 +252,7 @@ cc_binary(
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/flags:parse",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/log:globals",
         "@com_google_absl//absl/log:initialize",
         "@com_google_absl//absl/status",

sandboxed_api/sandbox2/testcases/CMakeLists.txt

@@ -302,6 +302,7 @@ set_target_properties(sandbox2_testcase_network_proxy PROPERTIES
   OUTPUT_NAME network_proxy
 )
 target_link_libraries(sandbox2_testcase_network_proxy PRIVATE
+  absl::check
   absl::flags
   absl::flags_parse
   absl::log

sandboxed_api/sandbox2/testcases/network_proxy.cc

@@ -27,10 +27,12 @@
 #include <cerrno>
 #include <cstdint>
 #include <cstring>
+#include <variant>
 
 #include "absl/base/log_severity.h"
 #include "absl/flags/flag.h"
 #include "absl/flags/parse.h"
+#include "absl/log/check.h"
 #include "absl/log/globals.h"
 #include "absl/log/initialize.h"
 #include "absl/log/log.h"
@@ -45,11 +47,24 @@
 #include "sandboxed_api/util/status_macros.h"
 
 ABSL_FLAG(bool, connect_with_handler, true, "Connect using automatic mode.");
+ABSL_FLAG(bool, ipv6, false, "Use IPv6.");
 
 namespace {
 
 using ::sapi::file_util::fileops::FDCloser;
 
+struct IPAddr {
+  size_t GetSize() const {
+    return addr.index() == 0 ? sizeof(sockaddr_in) : sizeof(sockaddr_in6);
+  }
+  const sockaddr* GetPtr() const {
+    return addr.index() == 0
+               ? reinterpret_cast<const sockaddr*>(&std::get<0>(addr))
+               : reinterpret_cast<const sockaddr*>(&std::get<1>(addr));
+  }
+  std::variant<sockaddr_in, sockaddr_in6> addr;
+};
+
 static sandbox2::NetworkProxyClient* g_proxy_client;
 
 ssize_t ReadFromFd(int fd, uint8_t* buf, size_t size) {
@@ -82,25 +97,24 @@ absl::Status CommunicationTest(int sock) {
   return absl::OkStatus();
 }
 
-absl::StatusOr<struct sockaddr_in6> CreateAddres(int port) {
-  static struct sockaddr_in6 saddr {};
-  saddr.sin6_family = AF_INET6;
-  saddr.sin6_port = htons(port);
-
-  if (int err = inet_pton(AF_INET6, "::1", &saddr.sin6_addr); err <= 0) {
-    return absl::ErrnoToStatus(errno, "socket()");
-  }
-  return saddr;
+IPAddr CreateAddress(int port) {
+  static struct sockaddr_in saddr4 {};
+  static struct sockaddr_in6 saddr6 {};
+  saddr4.sin_family = AF_INET;
+  saddr6.sin6_family = AF_INET6;
+  saddr4.sin_port = saddr6.sin6_port = htons(port);
+  PCHECK(inet_pton(AF_INET6, "::1", &saddr6.sin6_addr) == 1);
+  PCHECK(inet_pton(AF_INET, "127.0.0.1", &saddr4.sin_addr) == 1);
+  return absl::GetFlag(FLAGS_ipv6) ? IPAddr{.addr = saddr6}
+                                   : IPAddr{.addr = saddr4};
 }
 
-absl::Status ConnectWithoutHandler(int s, const struct sockaddr_in6& saddr) {
-  return g_proxy_client->Connect(
-      s, reinterpret_cast<const struct sockaddr*>(&saddr), sizeof(saddr));
+absl::Status ConnectWithoutHandler(int s, IPAddr saddr) {
+  return g_proxy_client->Connect(s, saddr.GetPtr(), saddr.GetSize());
 }
 
-absl::Status ConnectWithHandler(int s, const struct sockaddr_in6& saddr) {
-  int err = connect(s, reinterpret_cast<const struct sockaddr*>(&saddr),
-                    sizeof(saddr));
+absl::Status ConnectWithHandler(int s, IPAddr saddr) {
+  int err = connect(s, saddr.GetPtr(), saddr.GetSize());
   if (err != 0) {
     return absl::InternalError("connect() failed");
   }
@@ -109,17 +123,18 @@ absl::Status ConnectWithHandler(int s, const struct sockaddr_in6& saddr) {
 }
 
 absl::StatusOr<FDCloser> ConnectToServer(int port) {
-  SAPI_ASSIGN_OR_RETURN(struct sockaddr_in6 saddr, CreateAddres(port));
+  IPAddr addr = CreateAddress(port);
 
-  FDCloser s(socket(AF_INET6, SOCK_STREAM, 0));
+  FDCloser s(
+      socket(absl::GetFlag(FLAGS_ipv6) ? AF_INET6 : AF_INET, SOCK_STREAM, 0));
   if (s.get() < 0) {
     return absl::ErrnoToStatus(errno, "socket()");
   }
 
   if (absl::GetFlag(FLAGS_connect_with_handler)) {
-    SAPI_RETURN_IF_ERROR(ConnectWithHandler(s.get(), saddr));
+    SAPI_RETURN_IF_ERROR(ConnectWithHandler(s.get(), addr));
   } else {
-    SAPI_RETURN_IF_ERROR(ConnectWithoutHandler(s.get(), saddr));
+    SAPI_RETURN_IF_ERROR(ConnectWithoutHandler(s.get(), addr));
   }
 
   LOG(INFO) << "Connected to the server";