Add a few utility methods and helpers to Thrift C++ Sockets

Summary: Submitted by Dave Simpson @ Powerset.

Reviewed By: mcslee

Test Plan: test/cpp

git-svn-id: https://svn.apache.org/repos/asf/incubator/thrift/trunk@665370 13f79535-47bb-0310-9956-ffa450edef68

diff --git a/lib/cpp/src/transport/TSocket.cpp b/lib/cpp/src/transport/TSocket.cpp
index acf5288..0114fdc 100644 (file)
--- a/lib/cpp/src/transport/TSocket.cpp
+++ b/lib/cpp/src/transport/TSocket.cpp
@@ -486,4 +486,65 @@ string TSocket::getSocketInfo() {
   return oss.str();
 }
 
+std::string TSocket::getPeerHost() {
+  if (peerHost_.empty()) {
+    struct sockaddr_storage addr;
+    socklen_t addrLen = sizeof(addr);
+
+    if (socket_ < 0) {
+      return host_;
+    }
+
+    int rv = getpeername(socket_, (sockaddr*) &addr, &addrLen);
+
+    if (rv != 0) {
+      return peerHost_;
+    }
+
+    char clienthost[NI_MAXHOST];
+    char clientservice[NI_MAXSERV];
+
+    getnameinfo((sockaddr*) &addr, addrLen,
+                clienthost, sizeof(clienthost),
+                clientservice, sizeof(clientservice), 0);
+
+    peerHost_ = clienthost;
+  }
+  return peerHost_;
+}
+
+std::string TSocket::getPeerAddress() {
+  if (peerAddress_.empty()) {
+    struct sockaddr_storage addr;
+    socklen_t addrLen = sizeof(addr);
+
+    if (socket_ < 0) {
+      return peerAddress_;
+    }
+
+    int rv = getpeername(socket_, (sockaddr*) &addr, &addrLen);
+
+    if (rv != 0) {
+      return peerAddress_;
+    }
+
+    char clienthost[NI_MAXHOST];
+    char clientservice[NI_MAXSERV];
+
+    getnameinfo((sockaddr*) &addr, addrLen,
+                clienthost, sizeof(clienthost),
+                clientservice, sizeof(clientservice),
+                NI_NUMERICHOST|NI_NUMERICSERV);
+
+    peerAddress_ = clienthost;
+    peerPort_ = std::atoi(clientservice);
+  }
+  return peerAddress_;
+}
+
+int TSocket::getPeerPort() {
+  getPeerAddress();
+  return peerPort_;
+}
+
 }}} // facebook::thrift::transport

diff --git a/lib/cpp/src/transport/TSocket.h b/lib/cpp/src/transport/TSocket.h
index 57ff127..f0495ec 100644 (file)
--- a/lib/cpp/src/transport/TSocket.h
+++ b/lib/cpp/src/transport/TSocket.h
@@ -13,7 +13,7 @@
 #include "TTransport.h"
 #include "TServerSocket.h"
 
-namespace facebook { namespace thrift { namespace transport { 
+namespace facebook { namespace thrift { namespace transport {
 
 /**
  * TCP Socket implementation of the TTransport interface.
@@ -37,7 +37,7 @@ class TSocket : public TTransport {
    *
    */
   TSocket();
-  
+
   /**
    * Constructs a new socket. Note that this does NOT actually connect the
    * socket.
@@ -126,7 +126,7 @@ class TSocket : public TTransport {
    * Whether to enable/disable Nagle's algorithm.
    *
    * @param noDelay Whether or not to disable the algorithm.
-   * @return 
+   * @return
    */
   void setNoDelay(bool noDelay);
 
@@ -151,9 +151,27 @@ class TSocket : public TTransport {
    */
   void setMaxRecvRetries(int maxRecvRetries);
 
-  /** get socket information */
+  /**
+   * Get socket information formated as a string <Host: x Port: x>
+   */
   std::string getSocketInfo();
 
+  /**
+   * Returns the DNS name of the host to which the socket is connected
+   */
+  std::string getPeerHost();
+
+  /**
+   * Returns the address of the host to which the socket is connected
+   */
+  std::string getPeerAddress();
+
+  /**
+   * Returns the port of the host to which the socket is connected
+   **/
+  int getPeerPort();
+
+
  protected:
   /**
    * Constructor to create socket from raw UNIX handle. Never called directly
@@ -167,6 +185,15 @@ class TSocket : public TTransport {
   /** Host to connect to */
   std::string host_;
 
+  /** Peer hostname */
+  std::string peerHost_;
+
+  /** Peer address */
+  std::string peerAddress_;
+
+  /** Peer port */
+  int peerPort_;
+
   /** Port number to connect on */
   int port_;
 
@@ -184,7 +211,7 @@ class TSocket : public TTransport {
 
   /** Linger on */
   bool lingerOn_;
-  
+
   /** Linger val */
   int lingerVal_;
 

diff --git a/lib/cpp/src/transport/TTransportUtils.h b/lib/cpp/src/transport/TTransportUtils.h
index 1ebbfbd..38d2787 100644 (file)
--- a/lib/cpp/src/transport/TTransportUtils.h
+++ b/lib/cpp/src/transport/TTransportUtils.h
@@ -110,6 +110,10 @@ class TBufferedTransport : public TTransport {
 
   void consume(uint32_t len);
 
+  boost::shared_ptr<TTransport> getUnderlyingTransport() {
+    return transport_;
+  }
+
  protected:
   boost::shared_ptr<TTransport> transport_;
   uint8_t* rBuf_;
@@ -225,6 +229,10 @@ class TFramedTransport : public TTransport {
 
   void consume(uint32_t len);
 
+  boost::shared_ptr<TTransport> getUnderlyingTransport() {
+    return transport_;
+  }
+
  protected:
   boost::shared_ptr<TTransport> transport_;
   uint8_t* rBuf_;
@@ -243,6 +251,27 @@ class TFramedTransport : public TTransport {
   void readFrame();
 };
 
+/**
+ * Wraps a transport into a framed one.
+ *
+ * @author Dave Simpson <dave@powerset.com>
+ */
+class TFramedTransportFactory : public TTransportFactory {
+ public:
+  TFramedTransportFactory() {}
+
+  virtual ~TFramedTransportFactory() {}
+
+  /**
+   * Wraps the transport into a framed one.
+   */
+  virtual boost::shared_ptr<TTransport> getTransport(boost::shared_ptr<TTransport> trans) {
+    return boost::shared_ptr<TTransport>(new TFramedTransport(trans));
+  }
+
+};
+
+
 /**
  * A memory buffer is a tranpsort that simply reads from and writes to an
  * in memory buffer. Anytime you call write on it, the data is simply placed