thrift/transport/THeaderTransport.h

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one
0003  * or more contributor license agreements. See the NOTICE file
0004  * distributed with this work for additional information
0005  * regarding copyright ownership. The ASF licenses this file
0006  * to you under the Apache License, Version 2.0 (the
0007  * "License"); you may not use this file except in compliance
0008  * with the License. You may obtain a copy of the License at
0009  *
0010  *   http://www.apache.org/licenses/LICENSE-2.0
0011  *
0012  * Unless required by applicable law or agreed to in writing,
0013  * software distributed under the License is distributed on an
0014  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0015  * KIND, either express or implied. See the License for the
0016  * specific language governing permissions and limitations
0017  * under the License.
0018  */
0019
0020 #ifndef THRIFT_TRANSPORT_THEADERTRANSPORT_H_
0021 #define THRIFT_TRANSPORT_THEADERTRANSPORT_H_ 1
0022
0023 #include <bitset>
0024 #include <limits>
0025 #include <vector>
0026 #include <stdexcept>
0027 #include <string>
0028 #include <map>
0029
0030 #ifdef HAVE_STDINT_H
0031 #include <stdint.h>
0032 #elif HAVE_INTTYPES_H
0033 #include <inttypes.h>
0034 #endif
0035
0036 #include <thrift/protocol/TProtocolTypes.h>
0037 #include <thrift/transport/TBufferTransports.h>
0038 #include <thrift/transport/TTransport.h>
0039 #include <thrift/transport/TVirtualTransport.h>
0040
0041 enum CLIENT_TYPE {
0042   THRIFT_HEADER_CLIENT_TYPE = 0,
0043   THRIFT_FRAMED_BINARY = 1,
0044   THRIFT_UNFRAMED_BINARY = 2,
0045   THRIFT_FRAMED_COMPACT = 3,
0046   THRIFT_UNFRAMED_COMPACT = 4,
0047   THRIFT_UNKNOWN_CLIENT_TYPE = 5,
0048 };
0049
0050 namespace apache {
0051 namespace thrift {
0052 namespace transport {
0053
0054 using apache::thrift::protocol::T_COMPACT_PROTOCOL;
0055
0056 /**
0057  * Header transport. All writes go into an in-memory buffer until flush is
0058  * called, at which point the transport writes the length of the entire
0059  * binary chunk followed by the data payload. This allows the receiver on the
0060  * other end to always do fixed-length reads.
0061  *
0062  * Subclass TFramedTransport because most of the read/write methods are similar
0063  * and need similar buffers.  Major changes are readFrame & flush.
0064  *
0065  * Header Transport *must* be the same transport for both input and
0066  * output when used on the server side - client responses should be
0067  * the same protocol as those in the request.
0068  */
0069 class THeaderTransport : public TVirtualTransport<THeaderTransport, TFramedTransport> {
0070 public:
0071   static const int DEFAULT_BUFFER_SIZE = 512u;
0072   static const int THRIFT_MAX_VARINT32_BYTES = 5;
0073
0074   /// Use default buffer sizes.
0075   explicit THeaderTransport(const std::shared_ptr<TTransport>& transport,
0076                             std::shared_ptr<TConfiguration> config = nullptr)
0077     : TVirtualTransport(transport, config),
0078       outTransport_(transport),
0079       protoId(T_COMPACT_PROTOCOL),
0080       clientType(THRIFT_HEADER_CLIENT_TYPE),
0081       seqId(0),
0082       flags(0),
0083       tBufSize_(0),
0084       tBuf_(nullptr) {
0085     if (!transport_) throw std::invalid_argument("transport is empty");
0086     initBuffers();
0087   }
0088
0089   THeaderTransport(const std::shared_ptr<TTransport> inTransport,
0090                    const std::shared_ptr<TTransport> outTransport,
0091                    std::shared_ptr<TConfiguration> config = nullptr)
0092     : TVirtualTransport(inTransport, config),
0093       outTransport_(outTransport),
0094       protoId(T_COMPACT_PROTOCOL),
0095       clientType(THRIFT_HEADER_CLIENT_TYPE),
0096       seqId(0),
0097       flags(0),
0098       tBufSize_(0),
0099       tBuf_(nullptr) {
0100     if (!transport_) throw std::invalid_argument("inTransport is empty");
0101     if (!outTransport_) throw std::invalid_argument("outTransport is empty");
0102     initBuffers();
0103   }
0104
0105   uint32_t readSlow(uint8_t* buf, uint32_t len) override;
0106   void flush() override;
0107
0108   void resizeTransformBuffer(uint32_t additionalSize = 0);
0109
0110   uint16_t getProtocolId() const;
0111   void setProtocolId(uint16_t protoId) { this->protoId = protoId; }
0112
0113   void resetProtocol();
0114
0115   /**
0116    * We know we got a packet in header format here, try to parse the header
0117    *
0118    * @param headerSize size of the header portion
0119    * @param sz Size of the whole message, including header
0120    */
0121   void readHeaderFormat(uint16_t headerSize, uint32_t sz);
0122
0123   /**
0124    * Untransform the data based on the received header flags
0125    * On conclusion of function, setReadBuffer is called with the
0126    * untransformed data.
0127    *
0128    * @param ptr ptr to data
0129    * @param size of data
0130    */
0131   void untransform(uint8_t* ptr, uint32_t sz);
0132
0133   /**
0134    * Transform the data based on our write transform flags
0135    * At conclusion of function the write buffer is set to the
0136    * transformed data.
0137    *
0138    * @param ptr Ptr to data to transform
0139    * @param sz Size of data buffer
0140    */
0141   void transform(uint8_t* ptr, uint32_t sz);
0142
0143   uint16_t getNumTransforms() const {
0144     return safe_numeric_cast<uint16_t>(writeTrans_.size());
0145   }
0146
0147   void setTransform(uint16_t transId) { writeTrans_.push_back(transId); }
0148
0149   // Info headers
0150
0151   typedef std::map<std::string, std::string> StringToStringMap;
0152
0153   // these work with write headers
0154   void setHeader(const std::string& key, const std::string& value);
0155
0156   void clearHeaders();
0157
0158   StringToStringMap& getWriteHeaders() { return writeHeaders_; }
0159
0160   // these work with read headers
0161   const StringToStringMap& getHeaders() const { return readHeaders_; }
0162
0163   // accessors for seqId
0164   int32_t getSequenceNumber() const { return seqId; }
0165   void setSequenceNumber(int32_t seqId) { this->seqId = seqId; }
0166
0167   enum TRANSFORMS {
0168     ZLIB_TRANSFORM = 0x01,
0169   };
0170
0171 protected:
0172   /**
0173    * Reads a frame of input from the underlying stream.
0174    *
0175    * Returns true if a frame was read successfully, or false on EOF.
0176    * (Raises a TTransportException if EOF occurs after a partial frame.)
0177    */
0178   bool readFrame() override;
0179
0180   void ensureReadBuffer(uint32_t sz);
0181   uint32_t getWriteBytes();
0182
0183   void initBuffers() {
0184     setReadBuffer(nullptr, 0);
0185     setWriteBuffer(wBuf_.get(), wBufSize_);
0186   }
0187
0188   std::shared_ptr<TTransport> outTransport_;
0189
0190   // 0 and 16th bits must be 0 to differentiate from framed & unframed
0191   static const uint32_t HEADER_MAGIC = 0x0FFF0000;
0192   static const uint32_t HEADER_MASK = 0xFFFF0000;
0193   static const uint32_t FLAGS_MASK = 0x0000FFFF;
0194
0195   static const uint32_t MAX_FRAME_SIZE = 0x3FFFFFFF;
0196
0197   int16_t protoId;
0198   uint16_t clientType;
0199   uint32_t seqId;
0200   uint16_t flags;
0201
0202   std::vector<uint16_t> readTrans_;
0203   std::vector<uint16_t> writeTrans_;
0204
0205   // Map to use for headers
0206   StringToStringMap readHeaders_;
0207   StringToStringMap writeHeaders_;
0208
0209   /**
0210    * Returns the maximum number of bytes that write k/v headers can take
0211    */
0212   uint32_t getMaxWriteHeadersSize() const;
0213
0214   struct infoIdType {
0215     enum idType {
0216       // start at 1 to avoid confusing header padding for an infoId
0217       KEYVALUE = 1,
0218       END // signal the end of infoIds we can handle
0219     };
0220   };
0221
0222   // Buffers to use for transform processing
0223   uint32_t tBufSize_;
0224   std::unique_ptr<uint8_t[]> tBuf_;
0225
0226   void readString(uint8_t*& ptr, /* out */ std::string& str, uint8_t const* headerBoundary);
0227
0228   void writeString(uint8_t*& ptr, const std::string& str);
0229
0230   // Varint utils
0231   /**
0232    * Read an i16 from the wire as a varint. The MSB of each byte is set
0233    * if there is another byte to follow. This can read up to 3 bytes.
0234    */
0235   uint32_t readVarint16(uint8_t const* ptr, int16_t* i16, uint8_t const* boundary);
0236
0237   /**
0238    * Read an i32 from the wire as a varint. The MSB of each byte is set
0239    * if there is another byte to follow. This can read up to 5 bytes.
0240    */
0241   uint32_t readVarint32(uint8_t const* ptr, int32_t* i32, uint8_t const* boundary);
0242
0243   /**
0244    * Write an i32 as a varint. Results in 1-5 bytes on the wire.
0245    */
0246   uint32_t writeVarint32(int32_t n, uint8_t* pkt);
0247
0248   /**
0249    * Write an i16 as a varint. Results in 1-3 bytes on the wire.
0250    */
0251   uint32_t writeVarint16(int16_t n, uint8_t* pkt);
0252 };
0253
0254 /**
0255  * Wraps a transport into a header one.
0256  *
0257  */
0258 class THeaderTransportFactory : public TTransportFactory {
0259 public:
0260   THeaderTransportFactory() = default;
0261
0262   ~THeaderTransportFactory() override = default;
0263
0264   /**
0265    * Wraps the transport into a header one.
0266    */
0267   std::shared_ptr<TTransport> getTransport(std::shared_ptr<TTransport> trans) override {
0268     return std::shared_ptr<TTransport>(new THeaderTransport(trans));
0269   }
0270 };
0271 }
0272 }
0273 } // apache::thrift::transport
0274
0275 #endif // #ifndef THRIFT_TRANSPORT_THEADERTRANSPORT_H_