From a74ba58b99f13b0390ea901d0e7aa60a0e241cee Mon Sep 17 00:00:00 2001
From: Kyle Nekritz <knekritz@fb.com>
Date: Wed, 18 Nov 2015 10:20:14 -0800
Subject: [PATCH] Add getUnderlyingTransport to AsyncTransportWrapper.

Summary: Inspired by getSocketFromTransport from proxygen.

Reviewed By: siyengar

Differential Revision: D2663937

fb-gh-sync-id: f076215907cd06d6da3de033c57eec8a6a6ce320
---
 folly/io/async/AsyncTransport.h            | 18 ++++++++
 folly/io/async/test/AsyncTransportTest.cpp | 49 ++++++++++++++++++++++
 folly/io/async/test/MockAsyncTransport.h   |  1 +
 3 files changed, 68 insertions(+)
 create mode 100644 folly/io/async/test/AsyncTransportTest.cpp

diff --git a/folly/io/async/AsyncTransport.h b/folly/io/async/AsyncTransport.h
index 031b88e4..f9925d84 100644
--- a/folly/io/async/AsyncTransport.h
+++ b/folly/io/async/AsyncTransport.h
@@ -547,6 +547,24 @@ class AsyncTransportWrapper : virtual public AsyncTransport,
     return nullptr;
   }
 
+  /**
+   * In many cases when we need to set socket properties or otherwise access the
+   * underlying transport from a wrapped transport. This method allows access to
+   * the derived classes of the underlying transport.
+   */
+  template <class T>
+  T* getUnderlyingTransport() {
+    AsyncTransportWrapper* current = this;
+    while (current) {
+      auto sock = dynamic_cast<T*>(current);
+      if (sock) {
+        return sock;
+      }
+      current = current->getWrappedTransport();
+    }
+    return nullptr;
+  }
+
   /**
    * Return the application protocol being used by the underlying transport
    * protocol. This is useful for transports which are used to tunnel other
diff --git a/folly/io/async/test/AsyncTransportTest.cpp b/folly/io/async/test/AsyncTransportTest.cpp
new file mode 100644
index 00000000..0145f877
--- /dev/null
+++ b/folly/io/async/test/AsyncTransportTest.cpp
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2015 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <folly/io/async/AsyncTransport.h>
+#include <folly/io/async/AsyncSocket.h>
+
+#include <folly/io/async/test/MockAsyncTransport.h>
+
+#include <gtest/gtest.h>
+
+using namespace testing;
+
+namespace folly {
+
+TEST(AsyncTransportTest, getSocketFromSocket) {
+  AsyncSocket::UniquePtr transport(new AsyncSocket());
+  auto sock = transport->getUnderlyingTransport<AsyncSocket>();
+  ASSERT_EQ(transport.get(), sock);
+}
+
+TEST(AsyncTransportTest, getSocketFromWrappedTransport) {
+  AsyncSocket::UniquePtr transport(new AsyncSocket());
+  auto transportAddr = transport.get();
+
+  test::MockAsyncTransport wrapped1;
+  test::MockAsyncTransport wrapped2;
+
+  EXPECT_CALL(wrapped2, getWrappedTransport())
+    .WillOnce(Return(&wrapped1));
+  EXPECT_CALL(wrapped1, getWrappedTransport())
+    .WillOnce(Return(transportAddr));
+
+  auto sock = wrapped2.getUnderlyingTransport<AsyncSocket>();
+  ASSERT_EQ(transportAddr, sock);
+}
+
+} // namespace
diff --git a/folly/io/async/test/MockAsyncTransport.h b/folly/io/async/test/MockAsyncTransport.h
index 9202cd9c..1e987d3d 100644
--- a/folly/io/async/test/MockAsyncTransport.h
+++ b/folly/io/async/test/MockAsyncTransport.h
@@ -77,6 +77,7 @@ class MockAsyncTransport: public AsyncTransportWrapper {
   MOCK_CONST_METHOD0(getRawBytesReceived, size_t());
   MOCK_CONST_METHOD0(isEorTrackingEnabled, bool());
   MOCK_METHOD1(setEorTracking, void(bool));
+  MOCK_METHOD0(getWrappedTransport, AsyncTransportWrapper*());
 
 };
 
-- 
2.34.1