From 4c7a736d6529f22451a0ec965e093e7e318695e3 Mon Sep 17 00:00:00 2001
From: Kyle Nekritz <knekritz@fb.com>
Date: Fri, 23 Jun 2017 15:51:53 -0700
Subject: [PATCH] Move address caching logic from AsyncSSLSocket to
 AsyncSocket.

Summary: So that it is available on other transports.

Reviewed By: Orvid

Differential Revision: D5302039

fbshipit-source-id: cbfdadd158061ed9a2b3ed3e0960ce66f0d545fd
---
 folly/io/async/AsyncSSLSocket.cpp | 21 ++++-----------------
 folly/io/async/AsyncSSLSocket.h   |  2 --
 folly/io/async/AsyncSocket.cpp    | 28 +++++++++++++++++++++++++---
 folly/io/async/AsyncSocket.h      | 10 ++++++++++
 4 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/folly/io/async/AsyncSSLSocket.cpp b/folly/io/async/AsyncSSLSocket.cpp
index bb2452e6..e6513f82 100644
--- a/folly/io/async/AsyncSSLSocket.cpp
+++ b/folly/io/async/AsyncSSLSocket.cpp
@@ -462,8 +462,8 @@ void AsyncSSLSocket::sslAccept(
 
   // Cache local and remote socket addresses to keep them available
   // after socket file descriptor is closed.
-  if (cacheAddrOnFailure_ && -1 != getFd()) {
-    cacheLocalPeerAddr();
+  if (cacheAddrOnFailure_) {
+    cacheAddresses();
   }
 
   handshakeStartTime_ = std::chrono::steady_clock::now();
@@ -665,19 +665,6 @@ void AsyncSSLSocket::invokeHandshakeCB() {
   }
 }
 
-void AsyncSSLSocket::cacheLocalPeerAddr() {
-  SocketAddress address;
-  try {
-    getLocalAddress(&address);
-    getPeerAddress(&address);
-  } catch (const std::system_error& e) {
-    // The handle can be still valid while the connection is already closed.
-    if (e.code() != std::error_code(ENOTCONN, std::system_category())) {
-      throw;
-    }
-  }
-}
-
 void AsyncSSLSocket::connect(
     ConnectCallback* callback,
     const folly::SocketAddress& address,
@@ -755,8 +742,8 @@ void AsyncSSLSocket::sslConn(
 
   // Cache local and remote socket addresses to keep them available
   // after socket file descriptor is closed.
-  if (cacheAddrOnFailure_ && -1 != getFd()) {
-    cacheLocalPeerAddr();
+  if (cacheAddrOnFailure_) {
+    cacheAddresses();
   }
 
   verifyPeer_ = verifyPeer;
diff --git a/folly/io/async/AsyncSSLSocket.h b/folly/io/async/AsyncSSLSocket.h
index 516e07c2..474225c0 100644
--- a/folly/io/async/AsyncSSLSocket.h
+++ b/folly/io/async/AsyncSSLSocket.h
@@ -811,8 +811,6 @@ class AsyncSSLSocket : public virtual AsyncSocket {
   void invokeConnectSuccess() override;
   void scheduleConnectTimeout() override;
 
-  void cacheLocalPeerAddr();
-
   void startSSLConnect();
 
   static void sslInfoCallback(const SSL *ssl, int type, int val);
diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index 6cd879d9..f32da3a7 100644
--- a/folly/io/async/AsyncSocket.cpp
+++ b/folly/io/async/AsyncSocket.cpp
@@ -1285,17 +1285,39 @@ bool AsyncSocket::isDetachable() const {
   return !ioHandler_.isHandlerRegistered() && !writeTimeout_.isScheduled();
 }
 
-void AsyncSocket::getLocalAddress(folly::SocketAddress* address) const {
+void AsyncSocket::cacheAddresses() {
+  if (fd_ >= 0) {
+    try {
+      cacheLocalAddress();
+      cachePeerAddress();
+    } catch (const std::system_error& e) {
+      if (e.code() != std::error_code(ENOTCONN, std::system_category())) {
+        VLOG(1) << "Error caching addresses: " << e.code().value() << ", "
+                << e.code().message();
+      }
+    }
+  }
+}
+
+void AsyncSocket::cacheLocalAddress() const {
   if (!localAddr_.isInitialized()) {
     localAddr_.setFromLocalAddress(fd_);
   }
-  *address = localAddr_;
 }
 
-void AsyncSocket::getPeerAddress(folly::SocketAddress* address) const {
+void AsyncSocket::cachePeerAddress() const {
   if (!addr_.isInitialized()) {
     addr_.setFromPeerAddress(fd_);
   }
+}
+
+void AsyncSocket::getLocalAddress(folly::SocketAddress* address) const {
+  cacheLocalAddress();
+  *address = localAddr_;
+}
+
+void AsyncSocket::getPeerAddress(folly::SocketAddress* address) const {
+  cachePeerAddress();
   *address = addr_;
 }
 
diff --git a/folly/io/async/AsyncSocket.h b/folly/io/async/AsyncSocket.h
index 4b0fd069..6462f115 100644
--- a/folly/io/async/AsyncSocket.h
+++ b/folly/io/async/AsyncSocket.h
@@ -786,6 +786,13 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     evbChangeCb_ = std::move(cb);
   }
 
+  /**
+   * Attempt to cache the current local and peer addresses (if not already
+   * cached) so that they are available from getPeerAddress() and
+   * getLocalAddress() even after the socket is closed.
+   */
+  void cacheAddresses();
+
   /**
    * writeReturn is the total number of bytes written, or WRITE_ERROR on error.
    * If no data has been written, 0 is returned.
@@ -1123,6 +1130,9 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
 
   std::string withAddr(const std::string& s);
 
+  void cacheLocalAddress() const;
+  void cachePeerAddress() const;
+
   StateEnum state_;                      ///< StateEnum describing current state
   uint8_t shutdownFlags_;                ///< Shutdown state (ShutdownFlags)
   uint16_t eventFlags_;                  ///< EventBase::HandlerFlags settings
-- 
2.34.1