Add getUnderlyingTransport to AsyncTransportWrapper.

Summary: Inspired by getSocketFromTransport from proxygen.

Reviewed By: siyengar

Differential Revision: D2663937

fb-gh-sync-id: f076215907cd06d6da3de033c57eec8a6a6ce320

diff --git a/folly/io/async/AsyncTransport.h b/folly/io/async/AsyncTransport.h
index 031b88e41e733b06d4f72e1b9198c66328683da2..f9925d841c9485c071914996a4103ef538441cbb 100644 (file)
--- 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 (file)
index 0000000..0145f87
--- /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 9202cd9ceea46836e9110cafada4634d8500b514..1e987d3db92a934a4cdf894ee45d599ed4dbe5a6 100644 (file)
--- 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*());
 
 };