Add getOrigin() function to TTransport

getOrigin returns the origin of a request, the value depends on the transport used

7 files changed, 55 insertions, 1 deletions

diff --git a/lib/cpp/src/thrift/transport/TBufferTransports.h b/lib/cpp/src/thrift/transport/TBufferTransports.h
index b669a5968..af9412c6f 100644
--- a/lib/cpp/src/thrift/transport/TBufferTransports.h
+++ b/lib/cpp/src/thrift/transport/TBufferTransports.h
@@ -126,7 +126,6 @@ class TBufferBase : public TVirtualTransport<TBufferBase> {
     }
   }
 
-
  protected:
 
   /// Slow path read.
@@ -253,6 +252,12 @@ class TBufferedTransport
 
   void flush();
 
+  /**
+   * Returns the origin of the underlying transport
+   */
+  virtual const std::string getOrigin() {
+    return transport_->getOrigin();
+  }
 
   /**
    * The following behavior is currently implemented by TBufferedTransport,
@@ -393,6 +398,13 @@ class TFramedTransport
     return TBufferBase::readAll(buf,len);
   }
 
+  /**
+   * Returns the origin of the underlying transport
+   */
+  virtual const std::string getOrigin() {
+    return transport_->getOrigin();
+  }
+
  protected:
   /**
    * Reads a frame of input from the underlying stream.
diff --git a/lib/cpp/src/thrift/transport/THttpServer.cpp b/lib/cpp/src/thrift/transport/THttpServer.cpp
index e5b33273d..620bbd2b9 100644
--- a/lib/cpp/src/thrift/transport/THttpServer.cpp
+++ b/lib/cpp/src/thrift/transport/THttpServer.cpp
@@ -49,6 +49,8 @@ void THttpServer::parseHeader(char* header) {
   } else if (strncmp(header, "Content-Length", sz) == 0) {
     chunked_ = false;
     contentLength_ = atoi(value);
+  } else if (strncmp(header, "X-Forwarded-For", sz) == 0) {
+    origin_ = value;
   }
 }
 
diff --git a/lib/cpp/src/thrift/transport/THttpTransport.cpp b/lib/cpp/src/thrift/transport/THttpTransport.cpp
index c415ddb98..79ee7d537 100644
--- a/lib/cpp/src/thrift/transport/THttpTransport.cpp
+++ b/lib/cpp/src/thrift/transport/THttpTransport.cpp
@@ -17,6 +17,8 @@
  * under the License.
  */
 
+#include <sstream>
+
 #include <thrift/transport/THttpTransport.h>
 
 namespace apache { namespace thrift { namespace transport {
@@ -29,6 +31,7 @@ const int THttpTransport::CRLF_LEN = 2;
 
 THttpTransport::THttpTransport(boost::shared_ptr<TTransport> transport) :
   transport_(transport),
+  origin_(""),
   readHeaders_(true),
   chunked_(false),
   chunkedDone_(false),
@@ -249,4 +252,13 @@ void THttpTransport::write(const uint8_t* buf, uint32_t len) {
   writeBuffer_.write(buf, len);
 }
 
+const std::string THttpTransport::getOrigin() {
+  std::ostringstream oss;
+  if ( !origin_.empty()) {
+    oss << origin_ << ", ";
+  }
+  oss << transport_->getOrigin();
+  return oss.str();
+}
+
 }}}
diff --git a/lib/cpp/src/thrift/transport/THttpTransport.h b/lib/cpp/src/thrift/transport/THttpTransport.h
index a2e883496..8967c7446 100644
--- a/lib/cpp/src/thrift/transport/THttpTransport.h
+++ b/lib/cpp/src/thrift/transport/THttpTransport.h
@@ -62,9 +62,12 @@ class THttpTransport : public TVirtualTransport<THttpTransport> {
 
   virtual void flush() = 0;
 
+  virtual const std::string getOrigin();
+
  protected:
 
   boost::shared_ptr<TTransport> transport_;
+  std::string origin_;
 
   TMemoryBuffer writeBuffer_;
   TMemoryBuffer readBuffer_;
diff --git a/lib/cpp/src/thrift/transport/TSocket.cpp b/lib/cpp/src/thrift/transport/TSocket.cpp
index e80f71287..82678bac0 100755..100644
--- a/lib/cpp/src/thrift/transport/TSocket.cpp
+++ b/lib/cpp/src/thrift/transport/TSocket.cpp
@@ -841,4 +841,10 @@ bool TSocket::getUseLowMinRto() {
   return useLowMinRto_;
 }
 
+const std::string TSocket::getOrigin() {
+  std::ostringstream oss;
+  oss << getPeerHost() << ":" << getPeerPort();
+  return oss.str();
+}
+
 }}} // apache::thrift::transport
diff --git a/lib/cpp/src/thrift/transport/TSocket.h b/lib/cpp/src/thrift/transport/TSocket.h
index 062a5b9fd..c87321884 100644
--- a/lib/cpp/src/thrift/transport/TSocket.h
+++ b/lib/cpp/src/thrift/transport/TSocket.h
@@ -235,6 +235,13 @@ class TSocket : public TVirtualTransport<TSocket> {
   static bool getUseLowMinRto();
 
   /**
+   * Get the origin the socket is connected to
+   *
+   * @return string peer host identifier and port
+   */
+  virtual const std::string getOrigin();
+
+  /**
    * Constructor to create socket from raw UNIX handle.
    */
   TSocket(THRIFT_SOCKET socket);
diff --git a/lib/cpp/src/thrift/transport/TTransport.h b/lib/cpp/src/thrift/transport/TTransport.h
index 3b552c4f9..6e9a698b5 100644
--- a/lib/cpp/src/thrift/transport/TTransport.h
+++ b/lib/cpp/src/thrift/transport/TTransport.h
@@ -237,6 +237,18 @@ class TTransport {
                               "Base TTransport cannot consume.");
   }
 
+  /**
+   * Returns the origin of the transports call. The value depends on the
+   * transport used. An IP based transport for example will return the
+   * IP address of the client making the request.
+   * If the transport doesn't know the origin Unknown is returned.
+   *
+   * The returned value can be used in a log message for example
+   */
+  virtual const std::string getOrigin() {
+    return "Unknown";
+  }
+
  protected:
   /**
    * Simple constructor.