ceph/src/jaegertracing/thrift/lib/cpp/src/thrift/transport/THeaderTransport.h

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one
   3  * or more contributor license agreements. See the NOTICE file
   4  * distributed with this work for additional information
   5  * regarding copyright ownership. The ASF licenses this file
   6  * to you under the Apache License, Version 2.0 (the
   7  * "License"); you may not use this file except in compliance
   8  * with the License. You may obtain a copy of the License at
   9  *
  10  *   http://www.apache.org/licenses/LICENSE-2.0
  11  *
  12  * Unless required by applicable law or agreed to in writing,
  13  * software distributed under the License is distributed on an
  14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15  * KIND, either express or implied. See the License for the
  16  * specific language governing permissions and limitations
  17  * under the License.
  18  */
  19
  20 #ifndef THRIFT_TRANSPORT_THEADERTRANSPORT_H_
  21 #define THRIFT_TRANSPORT_THEADERTRANSPORT_H_ 1
  22
  23 #include <bitset>
  24 #include <limits>
  25 #include <vector>
  26 #include <stdexcept>
  27 #include <string>
  28 #include <map>
  29
  30 #ifdef HAVE_STDINT_H
  31 #include <stdint.h>
  32 #elif HAVE_INTTYPES_H
  33 #include <inttypes.h>
  34 #endif
  35
  36 #include <boost/scoped_array.hpp>
  37
  38 #include <thrift/protocol/TProtocolTypes.h>
  39 #include <thrift/transport/TBufferTransports.h>
  40 #include <thrift/transport/TTransport.h>
  41 #include <thrift/transport/TVirtualTransport.h>
  42
  43 enum CLIENT_TYPE {
  44   THRIFT_HEADER_CLIENT_TYPE = 0,
  45   THRIFT_FRAMED_BINARY = 1,
  46   THRIFT_UNFRAMED_BINARY = 2,
  47   THRIFT_FRAMED_COMPACT = 3,
  48   THRIFT_UNFRAMED_COMPACT = 4,
  49   THRIFT_UNKNOWN_CLIENT_TYPE = 5,
  50 };
  51
  52 namespace apache {
  53 namespace thrift {
  54 namespace transport {
  55
  56 using apache::thrift::protocol::T_COMPACT_PROTOCOL;
  57
  58 /**
  59  * Header transport. All writes go into an in-memory buffer until flush is
  60  * called, at which point the transport writes the length of the entire
  61  * binary chunk followed by the data payload. This allows the receiver on the
  62  * other end to always do fixed-length reads.
  63  *
  64  * Subclass TFramedTransport because most of the read/write methods are similar
  65  * and need similar buffers.  Major changes are readFrame & flush.
  66  *
  67  * Header Transport *must* be the same transport for both input and
  68  * output when used on the server side - client responses should be
  69  * the same protocol as those in the request.
  70  */
  71 class THeaderTransport : public TVirtualTransport<THeaderTransport, TFramedTransport> {
  72 public:
  73   static const int DEFAULT_BUFFER_SIZE = 512u;
  74   static const int THRIFT_MAX_VARINT32_BYTES = 5;
  75
  76   /// Use default buffer sizes.
  77   explicit THeaderTransport(const std::shared_ptr<TTransport>& transport)
  78     : TVirtualTransport(transport),
  79       outTransport_(transport),
  80       protoId(T_COMPACT_PROTOCOL),
  81       clientType(THRIFT_HEADER_CLIENT_TYPE),
  82       seqId(0),
  83       flags(0),
  84       tBufSize_(0),
  85       tBuf_(nullptr) {
  86     if (!transport_) throw std::invalid_argument("transport is empty");
  87     initBuffers();
  88   }
  89
  90   THeaderTransport(const std::shared_ptr<TTransport> inTransport,
  91                    const std::shared_ptr<TTransport> outTransport)
  92     : TVirtualTransport(inTransport),
  93       outTransport_(outTransport),
  94       protoId(T_COMPACT_PROTOCOL),
  95       clientType(THRIFT_HEADER_CLIENT_TYPE),
  96       seqId(0),
  97       flags(0),
  98       tBufSize_(0),
  99       tBuf_(nullptr) {
 100     if (!transport_) throw std::invalid_argument("inTransport is empty");
 101     if (!outTransport_) throw std::invalid_argument("outTransport is empty");
 102     initBuffers();
 103   }
 104
 105   uint32_t readSlow(uint8_t* buf, uint32_t len) override;
 106   void flush() override;
 107
 108   void resizeTransformBuffer(uint32_t additionalSize = 0);
 109
 110   uint16_t getProtocolId() const;
 111   void setProtocolId(uint16_t protoId) { this->protoId = protoId; }
 112
 113   void resetProtocol();
 114
 115   /**
 116    * We know we got a packet in header format here, try to parse the header
 117    *
 118    * @param headerSize size of the header portion
 119    * @param sz Size of the whole message, including header
 120    */
 121   void readHeaderFormat(uint16_t headerSize, uint32_t sz);
 122
 123   /**
 124    * Untransform the data based on the received header flags
 125    * On conclusion of function, setReadBuffer is called with the
 126    * untransformed data.
 127    *
 128    * @param ptr ptr to data
 129    * @param size of data
 130    */
 131   void untransform(uint8_t* ptr, uint32_t sz);
 132
 133   /**
 134    * Transform the data based on our write transform flags
 135    * At conclusion of function the write buffer is set to the
 136    * transformed data.
 137    *
 138    * @param ptr Ptr to data to transform
 139    * @param sz Size of data buffer
 140    */
 141   void transform(uint8_t* ptr, uint32_t sz);
 142
 143   uint16_t getNumTransforms() const {
 144     return safe_numeric_cast<uint16_t>(writeTrans_.size());
 145   }
 146
 147   void setTransform(uint16_t transId) { writeTrans_.push_back(transId); }
 148
 149   // Info headers
 150
 151   typedef std::map<std::string, std::string> StringToStringMap;
 152
 153   // these work with write headers
 154   void setHeader(const std::string& key, const std::string& value);
 155
 156   void clearHeaders();
 157
 158   StringToStringMap& getWriteHeaders() { return writeHeaders_; }
 159
 160   // these work with read headers
 161   const StringToStringMap& getHeaders() const { return readHeaders_; }
 162
 163   // accessors for seqId
 164   int32_t getSequenceNumber() const { return seqId; }
 165   void setSequenceNumber(int32_t seqId) { this->seqId = seqId; }
 166
 167   enum TRANSFORMS {
 168     ZLIB_TRANSFORM = 0x01,
 169   };
 170
 171 protected:
 172   /**
 173    * Reads a frame of input from the underlying stream.
 174    *
 175    * Returns true if a frame was read successfully, or false on EOF.
 176    * (Raises a TTransportException if EOF occurs after a partial frame.)
 177    */
 178   bool readFrame() override;
 179
 180   void ensureReadBuffer(uint32_t sz);
 181   uint32_t getWriteBytes();
 182
 183   void initBuffers() {
 184     setReadBuffer(nullptr, 0);
 185     setWriteBuffer(wBuf_.get(), wBufSize_);
 186   }
 187
 188   std::shared_ptr<TTransport> outTransport_;
 189
 190   // 0 and 16th bits must be 0 to differentiate from framed & unframed
 191   static const uint32_t HEADER_MAGIC = 0x0FFF0000;
 192   static const uint32_t HEADER_MASK = 0xFFFF0000;
 193   static const uint32_t FLAGS_MASK = 0x0000FFFF;
 194
 195   static const uint32_t MAX_FRAME_SIZE = 0x3FFFFFFF;
 196
 197   int16_t protoId;
 198   uint16_t clientType;
 199   uint32_t seqId;
 200   uint16_t flags;
 201
 202   std::vector<uint16_t> readTrans_;
 203   std::vector<uint16_t> writeTrans_;
 204
 205   // Map to use for headers
 206   StringToStringMap readHeaders_;
 207   StringToStringMap writeHeaders_;
 208
 209   /**
 210    * Returns the maximum number of bytes that write k/v headers can take
 211    */
 212   uint32_t getMaxWriteHeadersSize() const;
 213
 214   struct infoIdType {
 215     enum idType {
 216       // start at 1 to avoid confusing header padding for an infoId
 217       KEYVALUE = 1,
 218       END // signal the end of infoIds we can handle
 219     };
 220   };
 221
 222   // Buffers to use for transform processing
 223   uint32_t tBufSize_;
 224   boost::scoped_array<uint8_t> tBuf_;
 225
 226   void readString(uint8_t*& ptr, /* out */ std::string& str, uint8_t const* headerBoundary);
 227
 228   void writeString(uint8_t*& ptr, const std::string& str);
 229
 230   // Varint utils
 231   /**
 232    * Read an i16 from the wire as a varint. The MSB of each byte is set
 233    * if there is another byte to follow. This can read up to 3 bytes.
 234    */
 235   uint32_t readVarint16(uint8_t const* ptr, int16_t* i16, uint8_t const* boundary);
 236
 237   /**
 238    * Read an i32 from the wire as a varint. The MSB of each byte is set
 239    * if there is another byte to follow. This can read up to 5 bytes.
 240    */
 241   uint32_t readVarint32(uint8_t const* ptr, int32_t* i32, uint8_t const* boundary);
 242
 243   /**
 244    * Write an i32 as a varint. Results in 1-5 bytes on the wire.
 245    */
 246   uint32_t writeVarint32(int32_t n, uint8_t* pkt);
 247
 248   /**
 249    * Write an i16 as a varint. Results in 1-3 bytes on the wire.
 250    */
 251   uint32_t writeVarint16(int16_t n, uint8_t* pkt);
 252 };
 253
 254 /**
 255  * Wraps a transport into a header one.
 256  *
 257  */
 258 class THeaderTransportFactory : public TTransportFactory {
 259 public:
 260   THeaderTransportFactory() = default;
 261
 262   ~THeaderTransportFactory() override = default;
 263
 264   /**
 265    * Wraps the transport into a header one.
 266    */
 267   std::shared_ptr<TTransport> getTransport(std::shared_ptr<TTransport> trans) override {
 268     return std::shared_ptr<TTransport>(new THeaderTransport(trans));
 269   }
 270 };
 271 }
 272 }
 273 } // apache::thrift::transport
 274
 275 #endif // #ifndef THRIFT_TRANSPORT_THEADERTRANSPORT_H_