/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you 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. */ #ifndef THRIFT_TRANSPORT_THEADER_H_ #define THRIFT_TRANSPORT_THEADER_H_ 1 #include #include "thrift/lib/cpp/protocol/TBinaryProtocol.h" #include "thrift/lib/cpp/protocol/TCompactProtocol.h" #include "thrift/lib/cpp/protocol/TProtocolTypes.h" #include "folly/experimental/io/IOBuf.h" #include "folly/experimental/io/IOBufQueue.h" #include #include "boost/scoped_array.hpp" #include #include // Don't include the unknown client. #define CLIENT_TYPES_LEN 5 enum CLIENT_TYPE { THRIFT_HEADER_CLIENT_TYPE = 0, THRIFT_FRAMED_DEPRECATED = 1, THRIFT_UNFRAMED_DEPRECATED = 2, THRIFT_HTTP_CLIENT_TYPE = 3, THRIFT_FRAMED_COMPACT = 4, THRIFT_UNKNOWN_CLIENT_TYPE = 5, }; namespace apache { namespace thrift { namespace transport { using apache::thrift::protocol::T_COMPACT_PROTOCOL; /** * Class that will take an IOBuf and wrap it in some thrift headers. * see thrift/doc/HeaderFormat.txt for details. * * Supports transforms: zlib snappy hmac * Supports headers: http-style key/value per request and per connection * other: Protocol Id and seq ID in header. * * Backwards compatible with TFramed format, and unframed format, assuming * your server transport is compatible (some server types require 4-byte size * at the start, such as TNonblockingServer). */ class THeader { public: virtual ~THeader() {} explicit THeader() : queue_(new folly::IOBufQueue) , protoId(T_COMPACT_PROTOCOL) , clientType(THRIFT_HEADER_CLIENT_TYPE) , seqId(0) , flags(0) , identity(s_identity) { initSupportedClients(NULL); } explicit THeader(std::bitset const* clientTypes) : queue_(new folly::IOBufQueue) , protoId(T_COMPACT_PROTOCOL) , clientType(THRIFT_HEADER_CLIENT_TYPE) , seqId(0) , flags(0) , identity(s_identity) { initSupportedClients(clientTypes); } uint16_t getProtocolId() const; void setProtocolId(uint16_t protoId) { this->protoId = protoId; } virtual void resetProtocol(); /** * We know we got a packet in header format here, try to parse the header * * @param IObuf of the header + data. Untransforms the data as appropriate. * @return Just the data section in an IOBuf */ std::unique_ptr readHeaderFormat(std::unique_ptr); /** * Untransform the data based on the received header flags * On conclusion of function, setReadBuffer is called with the * untransformed data. * * @param IOBuf input data section * @return IOBuf output data section */ std::unique_ptr untransform(std::unique_ptr); /** * Transform the data based on our write transform flags * At conclusion of function the write buffer is set to the * transformed data. * * @param IOBuf to transform. Returns transformed IOBuf (or chain) * @return transformed data IOBuf */ std::unique_ptr transform(std::unique_ptr); uint16_t getNumTransforms() const { int trans = writeTrans_.size(); if (macCallback_) { trans += 1; } return trans; } void setTransform(uint16_t transId) { writeTrans_.push_back(transId); } // Info headers typedef std::map StringToStringMap; // these work with write headers void setHeader(const std::string& key, const std::string& value); void setPersistentHeader(const std::string& key, const std::string& value); void clearHeaders(); /** * this function only clears the local persistent * header. does not affect the persistent header * that already set */ void clearPersistentHeaders(); StringToStringMap& getWriteHeaders() { return writeHeaders_; } StringToStringMap& getPersistentWriteHeaders() { return persisWriteHeaders_; } // these work with read headers const StringToStringMap& getHeaders() const { return readHeaders_; } std::string getPeerIdentity(); void setIdentity(const std::string& identity); // accessors for seqId int32_t getSequenceNumber() const { return seqId; } void setSequenceNumber(int32_t seqId) { this->seqId = seqId; } enum TRANSFORMS { ZLIB_TRANSFORM = 0x01, HMAC_TRANSFORM = 0x02, SNAPPY_TRANSFORM = 0x03, }; /** * Callbacks to get and verify a mac transform. * * If a mac callback is provided, it will be called with the outgoing packet, * with the returned string appended at the end of the data. * * If a verify callback is provided, all incoming packets will be called with * their mac data and packet data to verify. If false is returned, an * exception is thrown. Packets without any mac also throw an exception if a * verify function is provided. * * If no verify callback is provided, and an incoming packet contains a mac, * the mac is ignored. * **/ typedef std::tr1::function MacCallback; typedef std::tr1::function< bool(const std::string&, const std::string)> VerifyMacCallback; void setHmac(MacCallback macCb, VerifyMacCallback verifyCb) { macCallback_ = macCb; verifyCallback_ = verifyCb; } /* IOBuf interface */ /** * Adds the header based on the type of transport: * unframed - does nothing. * framed - prepends frame size * header - prepends header, optionally appends mac * http - only supported for sync case, prepends http header. * * @return IOBuf chain with header _and_ data. Data is not copied */ std::unique_ptr addHeader(std::unique_ptr); /** * Given an IOBuf Chain, remove the header. Supports unframed (sync * only), framed, header, and http (sync case only). * * @param IOBufQueue - queue to try to read message from. * * @param needed - if the return is NULL (i.e. we didn't read a full * message), needed is set to the number of bytes needed * before you should call removeHeader again. * * @return IOBuf - the message chain. May be shared, may be chained. * If NULL, we didn't get enough data for a whole message, * call removeHeader again after reading needed more bytes. */ std::unique_ptr removeHeader(folly::IOBufQueue*, size_t& needed); protected: std::bitset supported_clients; void initSupportedClients(std::bitset const*); std::unique_ptr queue_; // 0 and 16th bits must be 0 to differentiate from framed & unframed static const uint32_t HEADER_MAGIC = 0x0FFF0000; static const uint32_t HEADER_MASK = 0xFFFF0000; static const uint32_t FLAGS_MASK = 0x0000FFFF; static const uint32_t HTTP_MAGIC = 0x504F5354; // POST static const uint32_t MAX_FRAME_SIZE = 0x3FFFFFFF; int16_t protoId; uint16_t clientType; uint32_t seqId; uint16_t flags; std::string identity; std::vector readTrans_; std::vector writeTrans_; // Map to use for headers StringToStringMap readHeaders_; StringToStringMap writeHeaders_; // Map to use for persistent headers StringToStringMap persisReadHeaders_; StringToStringMap persisWriteHeaders_; static const std::string IDENTITY_HEADER; static const std::string ID_VERSION_HEADER; static const std::string ID_VERSION; static std::string s_identity; MacCallback macCallback_; VerifyMacCallback verifyCallback_; /** * Returns the maximum number of bytes that write k/v headers can take */ size_t getMaxWriteHeadersSize() const; /** * Returns whether the 1st byte of the protocol payload should be hadled * as compact framed. */ bool compactFramed(uint32_t magic); struct infoIdType { enum idType { // start at 1 to avoid confusing header padding for an infoId KEYVALUE = 1, // for persistent header PKEYVALUE = 2, END // signal the end of infoIds we can handle }; }; }; }}} // apache::thrift::transport #endif // #ifndef THRIFT_TRANSPORT_THEADER_H_