Optional peerAddress argument for sockets created with fd

Summary: For TCP sockets, it may not be possible to query the peer address from the fd if the endpoint is no longer connected (closed with reset).  In most server cases, the peer address is known when the connection is accepted, so it can be prepopulated.

Reviewed By: mjoras

Differential Revision: D27985942

fbshipit-source-id: 630098ac00930b978fe8a9cad1894834c44d4d93

folly/io/async/AsyncSSLSocket.cpp

@@ -271,8 +271,9 @@ AsyncSSLSocket::AsyncSSLSocket(
     EventBase* evb,
     NetworkSocket fd,
     bool server,
-    bool deferSecurityNegotiation)
-    : AsyncSocket(evb, fd),
+    bool deferSecurityNegotiation,
+    const SocketAddress* peerAddress)
+    : AsyncSocket(evb, fd, 0, peerAddress),
       server_(server),
       ctx_(std::move(ctx)),
       handshakeTimeout_(this, evb),
@@ -340,8 +341,10 @@ AsyncSSLSocket::AsyncSSLSocket(
     EventBase* evb,
     NetworkSocket fd,
     const std::string& serverName,
-    bool deferSecurityNegotiation)
-    : AsyncSSLSocket(ctx, evb, fd, false, deferSecurityNegotiation) {
+    bool deferSecurityNegotiation,
+    const SocketAddress* peerAddress)
+    : AsyncSSLSocket(
+          ctx, evb, fd, false, deferSecurityNegotiation, peerAddress) {
   tlsextHostname_ = serverName;
 }
 #endif // FOLLY_OPENSSL_HAS_SNI

folly/io/async/AsyncSSLSocket.h

@@ -252,13 +252,17 @@ class AsyncSSLSocket : public AsyncSocket {
    * @param server Is socket in server mode?
    * @param deferSecurityNegotiation
    *          unencrypted data can be sent before sslConn/Accept
+   * @param peerAddress optional peer address (eg: returned from accept).  If
+   *          nullptr, AsyncSocket will lazily attempt to determine it from fd
+   *          via a system call
    */
   AsyncSSLSocket(
       std::shared_ptr<folly::SSLContext> ctx,
       EventBase* evb,
       NetworkSocket fd,
       bool server = true,
-      bool deferSecurityNegotiation = false);
+      bool deferSecurityNegotiation = false,
+      const SocketAddress* peerAddress = nullptr);
 
   /**
    * Create a server/client AsyncSSLSocket from an already connected
@@ -288,9 +292,10 @@ class AsyncSSLSocket : public AsyncSocket {
       EventBase* evb,
       NetworkSocket fd,
       bool server = true,
-      bool deferSecurityNegotiation = false) {
-    return AsyncSSLSocket::UniquePtr(
-        new AsyncSSLSocket(ctx, evb, fd, server, deferSecurityNegotiation));
+      bool deferSecurityNegotiation = false,
+      const folly::SocketAddress* peerAddress = nullptr) {
+    return AsyncSSLSocket::UniquePtr(new AsyncSSLSocket(
+        ctx, evb, fd, server, deferSecurityNegotiation, peerAddress));
   }
 
   /**
@@ -329,13 +334,19 @@ class AsyncSSLSocket : public AsyncSocket {
    * @param evb  EventBase that will manage this socket.
    * @param fd   File descriptor to take over (should be a connected socket).
    * @param serverName tlsext_hostname that will be sent in ClientHello.
+   * @param deferSecurityNegotiation
+   *          unencrypted data can be sent before sslConn/Accept
+   * @param peerAddress optional peer address (eg: returned from accept).  If
+   *          nullptr, AsyncSocket will lazily attempt to determine it from fd
+   *          via a system call
    */
   AsyncSSLSocket(
       const std::shared_ptr<folly::SSLContext>& ctx,
       EventBase* evb,
       NetworkSocket fd,
       const std::string& serverName,
-      bool deferSecurityNegotiation = false);
+      bool deferSecurityNegotiation = false,
+      const SocketAddress* peerAddr = nullptr);
 
   static UniquePtr newSocket(
       const std::shared_ptr<folly::SSLContext>& ctx,

folly/io/async/AsyncSocket.cpp

@@ -563,7 +563,10 @@ AsyncSocket::AsyncSocket(
 }
 
 AsyncSocket::AsyncSocket(
-    EventBase* evb, NetworkSocket fd, uint32_t zeroCopyBufId)
+    EventBase* evb,
+    NetworkSocket fd,
+    uint32_t zeroCopyBufId,
+    const SocketAddress* peerAddress)
     : zeroCopyBufId_(zeroCopyBufId),
       eventBase_(evb),
       writeTimeout_(this, evb),
@@ -576,13 +579,17 @@ AsyncSocket::AsyncSocket(
   disableTransparentFunctions(fd_, noTransparentTls_, noTSocks_);
   setCloseOnExec();
   state_ = StateEnum::ESTABLISHED;
+  if (peerAddress) {
+    addr_ = *peerAddress;
+  }
 }
 
 AsyncSocket::AsyncSocket(AsyncSocket* oldAsyncSocket)
     : AsyncSocket(
           oldAsyncSocket->getEventBase(),
           oldAsyncSocket->detachNetworkSocket(),
-          oldAsyncSocket->getZeroCopyBufId()) {
+          oldAsyncSocket->getZeroCopyBufId(),
+          &oldAsyncSocket->addr_) {
   appBytesWritten_ = oldAsyncSocket->appBytesWritten_;
   rawBytesWritten_ = oldAsyncSocket->rawBytesWritten_;
   byteEventHelper_ = std::move(oldAsyncSocket->byteEventHelper_);

folly/io/async/AsyncSocket.h

@@ -376,8 +376,15 @@ class AsyncSocket : public AsyncTransport {
    * @param evb EventBase that will manage this socket.
    * @param fd  File descriptor to take over (should be a connected socket).
    * @param zeroCopyBufId Zerocopy buf id to start with.
+   * @param peerAddress optional peer address (eg: returned from accept).  If
+   *        nullptr, AsyncSocket will lazily attempt to determine it from fd
+   *        via a system call
    */
-  AsyncSocket(EventBase* evb, NetworkSocket fd, uint32_t zeroCopyBufId = 0);
+  AsyncSocket(
+      EventBase* evb,
+      NetworkSocket fd,
+      uint32_t zeroCopyBufId = 0,
+      const SocketAddress* peerAddress = nullptr);
 
   /**
    * Create an AsyncSocket from a different, already connected AsyncSocket.
@@ -434,8 +441,11 @@ class AsyncSocket : public AsyncTransport {
   /**
    * Helper function to create an AsyncSocket.
    */
-  static UniquePtr newSocket(EventBase* evb, NetworkSocket fd) {
-    return UniquePtr{new AsyncSocket(evb, fd)};
+  static UniquePtr newSocket(
+      EventBase* evb,
+      NetworkSocket fd,
+      const SocketAddress* peerAddress = nullptr) {
+    return UniquePtr{new AsyncSocket(evb, fd, 0, peerAddress)};
   }
 
   /**

folly/io/async/test/AsyncSocketTest2.cpp

@@ -5731,8 +5731,11 @@ TEST(AsyncSocket, PreReceivedDataTakeover) {
 
   socket->writeChain(nullptr, IOBuf::copyBuffer("hello"));
 
+  auto fd = server.acceptFD();
+  SocketAddress peerAddress;
+  peerAddress.setFromPeerAddress(fd);
   auto acceptedSocket =
-      AsyncSocket::UniquePtr(new AsyncSocket(&evb, server.acceptFD()));
+      AsyncSocket::UniquePtr(new AsyncSocket(&evb, fd, 0, &peerAddress));
   AsyncSocket::UniquePtr takeoverSocket;
 
   ReadCallback peekCallback(3);
@@ -5753,6 +5756,12 @@ TEST(AsyncSocket, PreReceivedDataTakeover) {
   acceptedSocket->setReadCB(&peekCallback);
 
   evb.loop();
+  // Verify we can still get the peer address after the peer socket is reset.
+  socket->closeWithReset();
+  evb.loopOnce();
+  SocketAddress socketPeerAddress;
+  takeoverSocket->getPeerAddress(&socketPeerAddress);
+  EXPECT_EQ(socketPeerAddress, peerAddress);
 }
 
 #ifdef MSG_NOSIGNAL

folly/io/async/test/TestSSLServer.h

@@ -63,7 +63,7 @@ class SSLServerAcceptCallbackBase : public AsyncServerSocket::AcceptCallback {
 
   void connectionAccepted(
       folly::NetworkSocket fd,
-      const SocketAddress& /* clientAddr */) noexcept override {
+      const SocketAddress& clientAddr) noexcept override {
     if (socket_) {
       socket_->detachEventBase();
     }
@@ -71,7 +71,13 @@ class SSLServerAcceptCallbackBase : public AsyncServerSocket::AcceptCallback {
     try {
       // Create a AsyncSSLSocket object with the fd. The socket should be
       // added to the event base and in the state of accepting SSL connection.
-      socket_ = AsyncSSLSocket::newSocket(ctx_, base_, fd);
+      socket_ = AsyncSSLSocket::newSocket(
+          ctx_,
+          base_,
+          fd,
+          /*server=*/true,
+          /*deferSecurityNegotiation=*/false,
+          &clientAddr);
     } catch (const std::exception& e) {
       LOG(ERROR) << "Exception %s caught while creating a AsyncSSLSocket "
                     "object with socket "