thrift/lib/cpp/test/ThrifttReadCheckTests.cpp

/*
 * 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.
 */

#define  MAX_MESSAGE_SIZE  2

#include <boost/test/unit_test.hpp>
#include <iostream>
#include <climits>
#include <random>
#include <vector>
#include <thrift/TConfiguration.h>
#include <thrift/protocol/TBinaryProtocol.h>
#include <thrift/protocol/TCompactProtocol.h>
#include <thrift/protocol/TJSONProtocol.h>
#include <thrift/Thrift.h>
#include <memory>
#include <thrift/transport/TTransportUtils.h>
#include <thrift/transport/TBufferTransports.h>
#ifdef THRIFT_TEST_WITH_ZLIB
#include <zlib.h>
#include <thrift/transport/THeaderTransport.h>
#include <thrift/protocol/THeaderProtocol.h>
#endif
#include <thrift/TApplicationException.h>
#include <thrift/transport/TSimpleFileTransport.h>
#include <thrift/transport/TFileTransport.h>
#include <thrift/protocol/TEnum.h>
#include <thrift/protocol/TList.h>
#include <thrift/protocol/TSet.h>
#include <thrift/protocol/TMap.h>

BOOST_AUTO_TEST_SUITE(ThriftReadCheckExceptionTest)

using apache::thrift::TConfiguration;
using apache::thrift::protocol::TBinaryProtocol;
using apache::thrift::protocol::TCompactProtocol;
using apache::thrift::protocol::TJSONProtocol;
using apache::thrift::protocol::TType;
using apache::thrift::transport::TPipedTransport;
using apache::thrift::transport::TMemoryBuffer;
using apache::thrift::transport::TSimpleFileTransport;
using apache::thrift::transport::TFileTransport;
using apache::thrift::transport::TFDTransport;
using apache::thrift::transport::TTransportException;
using apache::thrift::transport::TBufferedTransport;
using apache::thrift::transport::TFramedTransport;
using std::shared_ptr;
using std::string;
using std::memset;
using namespace apache::thrift;
using namespace apache::thrift::protocol;


BOOST_AUTO_TEST_CASE(test_tmemorybuffer_read_check_exception) {
  std::shared_ptr<TConfiguration> config(new TConfiguration(MAX_MESSAGE_SIZE));
  TMemoryBuffer trans_out(config);
  uint8_t buffer[6] = {1, 2, 3, 4, 5, 6};
  trans_out.write((const uint8_t*)buffer, sizeof(buffer));
  trans_out.close();

  TMemoryBuffer trans_in(config);
  memset(buffer, 0, sizeof(buffer));
  BOOST_CHECK_THROW(trans_in.read(buffer, sizeof(buffer)), TTransportException);
  trans_in.close();
}

BOOST_AUTO_TEST_CASE(test_tpipedtransport_read_check_exception) {
  std::shared_ptr<TConfiguration> config(new TConfiguration(MAX_MESSAGE_SIZE));
  std::shared_ptr<TMemoryBuffer> pipe(new TMemoryBuffer);
  std::shared_ptr<TMemoryBuffer> underlying(new TMemoryBuffer);
  std::shared_ptr<TPipedTransport> trans(new TPipedTransport(underlying, pipe, config));

  uint8_t buffer[4];

  underlying->write((uint8_t*)"abcd", 4);
  BOOST_CHECK_THROW(trans->read(buffer, sizeof(buffer)), TTransportException);
  BOOST_CHECK_THROW(trans->readAll(buffer, sizeof(buffer)), TTransportException);
  trans->readEnd();
  pipe->resetBuffer();
  underlying->write((uint8_t*)"ef", 2);
  BOOST_CHECK_THROW(trans->read(buffer, sizeof(buffer)), TTransportException);
  BOOST_CHECK_THROW(trans->readAll(buffer, sizeof(buffer)), TTransportException);
  trans->readEnd();
}

BOOST_AUTO_TEST_CASE(test_tsimplefiletransport_read_check_exception) {
  std::shared_ptr<TConfiguration> config(new TConfiguration(MAX_MESSAGE_SIZE));
  TSimpleFileTransport trans_out("data", false, true, config);
  uint8_t buffer[6] = {1, 2, 3, 4, 5, 6};
  trans_out.write((const uint8_t*)buffer, sizeof(buffer));
  trans_out.close();

  TSimpleFileTransport trans_in("data",true, false, config);
  memset(buffer, 0, sizeof(buffer));
  BOOST_CHECK_THROW(trans_in.read(buffer, sizeof(buffer)), TTransportException);
  trans_in.close();

  remove("./data");
}

BOOST_AUTO_TEST_CASE(test_tfiletransport_read_check_exception) {
  std::shared_ptr<TConfiguration> config(new TConfiguration(MAX_MESSAGE_SIZE));
  TFileTransport trans_out("data", false, config);
  uint8_t buffer[6] = {1, 2, 3, 4, 5, 6};
  trans_out.write((const uint8_t*)buffer, sizeof(buffer));

  TFileTransport trans_in("data", false, config);
  memset(buffer, 0, sizeof(buffer));
  BOOST_CHECK_THROW(trans_in.read(buffer, sizeof(buffer)), TTransportException);

  remove("./data");
}

BOOST_AUTO_TEST_CASE(test_tbufferedtransport_read_check_exception) {
  uint8_t arr[4] = {1, 2, 3, 4};
  std::shared_ptr<TMemoryBuffer> buffer (new TMemoryBuffer(arr, sizeof(arr)));
  std::shared_ptr<TConfiguration> config (new TConfiguration(MAX_MESSAGE_SIZE));
  std::shared_ptr<TBufferedTransport> trans (new TBufferedTransport(buffer, config));

  trans->write((const uint8_t*)arr, sizeof(arr));
  BOOST_CHECK_THROW(trans->read(arr, sizeof(arr)), TTransportException);
}

BOOST_AUTO_TEST_CASE(test_tframedtransport_read_check_exception) {
  uint8_t arr[4] = {1, 2, 3, 4};
  std::shared_ptr<TMemoryBuffer> buffer (new TMemoryBuffer(arr, sizeof(arr)));
  std::shared_ptr<TConfiguration> config (new TConfiguration(MAX_MESSAGE_SIZE));
  std::shared_ptr<TFramedTransport> trans (new TFramedTransport(buffer, config));

  trans->write((const uint8_t*)arr, sizeof(arr));
  BOOST_CHECK_THROW(trans->read(arr, sizeof(arr)), TTransportException);
}

BOOST_AUTO_TEST_CASE(test_tthriftbinaryprotocol_read_check_exception) {
  std::shared_ptr<TConfiguration> config (new TConfiguration(MAX_MESSAGE_SIZE));
  std::shared_ptr<TMemoryBuffer> transport(new TMemoryBuffer(config));
  std::shared_ptr<TBinaryProtocol> protocol(new TBinaryProtocol(transport));

  uint32_t val = 0;
  TType elemType = apache::thrift::protocol::T_STOP;
  TType elemType1 = apache::thrift::protocol::T_STOP;
  TList list(T_I32, 8);
  protocol->writeListBegin(list.elemType_, list.size_);
  protocol->writeListEnd();
  BOOST_CHECK_THROW(protocol->readListBegin(elemType, val), TTransportException);
  protocol->readListEnd();

  TSet set(T_I32, 8);
  protocol->writeSetBegin(set.elemType_, set.size_);
  protocol->writeSetEnd();
  BOOST_CHECK_THROW(protocol->readSetBegin(elemType, val), TTransportException);
  protocol->readSetEnd();

  TMap map(T_I32, T_I32, 8);
  protocol->writeMapBegin(map.keyType_, map.valueType_, map.size_);
  protocol->writeMapEnd();
  BOOST_CHECK_THROW(protocol->readMapBegin(elemType, elemType1, val), TTransportException);
  protocol->readMapEnd();
}

BOOST_AUTO_TEST_CASE(test_tthriftcompactprotocol_read_check_exception) {
  // Set Max Message Size to 11 since all structs are 12B long
  std::shared_ptr<TConfiguration> config (new TConfiguration(11));
  std::shared_ptr<TMemoryBuffer> transport(new TMemoryBuffer(config));
  std::shared_ptr<TCompactProtocol> protocol(new TCompactProtocol(transport));

  uint32_t val = 0;
  TType elemType = apache::thrift::protocol::T_STOP;
  TType elemType1 = apache::thrift::protocol::T_STOP;

  // This list needs 12B
  TList list(T_I32, 12);
  protocol->writeListBegin(list.elemType_, list.size_);
  protocol->writeListEnd();
  BOOST_CHECK_THROW(protocol->readListBegin(elemType, val), TTransportException);
  protocol->readListEnd();

  // This set needs 12B
  TSet set(T_I32, 12);
  protocol->writeSetBegin(set.elemType_, set.size_);
  protocol->writeSetEnd();
  BOOST_CHECK_THROW(protocol->readSetBegin(elemType, val), TTransportException);
  protocol->readSetEnd();


  // This map needs 12B (2x elem)
  TMap map(T_I32, T_I32, 6);
  protocol->writeMapBegin(map.keyType_, map.valueType_, map.size_);
  protocol->writeMapEnd();
  BOOST_CHECK_THROW(protocol->readMapBegin(elemType, elemType1, val), TTransportException);
  protocol->readMapEnd();

  // This string needs 12B (1 for size + str)
  string eleven = "1234567890A";
  protocol->writeString(eleven);
  BOOST_CHECK_THROW(protocol->readString(eleven), TTransportException);
}

BOOST_AUTO_TEST_CASE(test_tthriftcompactprotocol_read_check_pass) {
  // Set Max Message Size to 12 to check the edge case
  std::shared_ptr<TConfiguration> config (new TConfiguration(12));
  std::shared_ptr<TMemoryBuffer> transport(new TMemoryBuffer(config));
  std::shared_ptr<TCompactProtocol> protocol(new TCompactProtocol(transport));

  uint32_t val = 0;
  TType elemType = apache::thrift::protocol::T_STOP;
  TType elemType1 = apache::thrift::protocol::T_STOP;

  // This list needs 12B
  TList list(T_I32, 12);
  protocol->writeListBegin(list.elemType_, list.size_);
  protocol->writeListEnd();
  BOOST_CHECK_NO_THROW(protocol->readListBegin(elemType, val));
  protocol->readListEnd();

  // This set needs 12B
  TSet set(T_I32, 12);
  protocol->writeSetBegin(set.elemType_, set.size_);
  protocol->writeSetEnd();
  BOOST_CHECK_NO_THROW(protocol->readSetBegin(elemType, val));
  protocol->readSetEnd();

  // This map needs 12B (2x elem)
  TMap map(T_I32, T_I32, 6);
  protocol->writeMapBegin(map.keyType_, map.valueType_, map.size_);
  protocol->writeMapEnd();
  BOOST_CHECK_NO_THROW(protocol->readMapBegin(elemType, elemType1, val));
  protocol->readMapEnd();

  // This string needs 12B (1 for size + str)
  string eleven = "1234567890A";
  protocol->writeString(eleven);
  BOOST_CHECK_NO_THROW(protocol->readString(eleven));
}

BOOST_AUTO_TEST_CASE(test_tthriftbinaryprotocol_container_size_overflow) {
  std::shared_ptr<TConfiguration> config (new TConfiguration(1024));
  std::shared_ptr<TMemoryBuffer> transport(new TMemoryBuffer(config));
  std::shared_ptr<TBinaryProtocol> protocol(new TBinaryProtocol(transport));

  uint32_t val = 0;
  TType elemType = apache::thrift::protocol::T_STOP;
  // 0x40000000 elements of min size 4 require 4 GiB; the product wraps to 0 in
  // 32-bit math and used to slip past the MaxMessageSize check.
  TList list(T_I32, 0x40000000);
  protocol->writeListBegin(list.elemType_, list.size_);
  protocol->writeListEnd();
  BOOST_CHECK_THROW(protocol->readListBegin(elemType, val), TTransportException);
  protocol->readListEnd();
}

BOOST_AUTO_TEST_CASE(test_tthriftcompactprotocol_container_size_overflow) {
  std::shared_ptr<TConfiguration> config (new TConfiguration(1024));
  std::shared_ptr<TMemoryBuffer> transport(new TMemoryBuffer(config));
  std::shared_ptr<TCompactProtocol> protocol(new TCompactProtocol(transport));

  uint32_t val = 0;
  TType elemType = apache::thrift::protocol::T_STOP;
  // 0x10000000 elements of min size 16 (UUID) require 4 GiB; the product wraps
  // to 0 in 32-bit math and used to slip past the MaxMessageSize check.
  TList list(T_UUID, 0x10000000);
  protocol->writeListBegin(list.elemType_, list.size_);
  protocol->writeListEnd();
  BOOST_CHECK_THROW(protocol->readListBegin(elemType, val), TTransportException);
  protocol->readListEnd();
}

BOOST_AUTO_TEST_CASE(test_tthriftjsonprotocol_container_size_overflow) {
  std::shared_ptr<TConfiguration> config (new TConfiguration(1024));
  std::shared_ptr<TMemoryBuffer> transport(new TMemoryBuffer(config));
  std::shared_ptr<TJSONProtocol> protocol(new TJSONProtocol(transport));

  uint32_t val = 0;
  TType elemType = apache::thrift::protocol::T_STOP;
  // 0x10000000 elements of min size 16 (UUID) require 4 GiB; the product wraps
  // to 0 in 32-bit math and used to slip past the MaxMessageSize check.
  TList list(T_UUID, 0x10000000);
  protocol->writeListBegin(list.elemType_, list.size_);
  protocol->writeListEnd();
  BOOST_CHECK_THROW(protocol->readListBegin(elemType, val), TTransportException);
  protocol->readListEnd();
}

BOOST_AUTO_TEST_CASE(test_tthriftjsonprotocol_read_check_exception) {
  // Unlike the length-prefixed protocols, the JSON container header
  // ("[\"i32\",<n>") is itself read through the message-size-bounded string and
  // number readers, so a limit as small as MAX_MESSAGE_SIZE would already be hit
  // while reading the element type name. Use a limit that leaves room for the
  // header but stays far below the declared container payload, which is what
  // this test exercises.
  const int JSON_MAX_MESSAGE_SIZE = 64;
  const int CONTAINER_SIZE = 1024;
  std::shared_ptr<TConfiguration> config (new TConfiguration(JSON_MAX_MESSAGE_SIZE));
  std::shared_ptr<TMemoryBuffer> transport(new TMemoryBuffer(config));
  std::shared_ptr<TJSONProtocol> protocol(new TJSONProtocol(transport));

  uint32_t val = 0;
  TType elemType = apache::thrift::protocol::T_STOP;
  TType elemType1 = apache::thrift::protocol::T_STOP;
  TList list(T_I32, CONTAINER_SIZE);
  protocol->writeListBegin(list.elemType_, list.size_);
  protocol->writeListEnd();
  BOOST_CHECK_THROW(protocol->readListBegin(elemType, val), TTransportException);
  protocol->readListEnd();

  TSet set(T_I32, CONTAINER_SIZE);
  protocol->writeSetBegin(set.elemType_, set.size_);
  protocol->writeSetEnd();
  BOOST_CHECK_THROW(protocol->readSetBegin(elemType, val), TTransportException);
  protocol->readSetEnd();

  TMap map(T_I32, T_I32, CONTAINER_SIZE);
  protocol->writeMapBegin(map.keyType_, map.valueType_, map.size_);
  protocol->writeMapEnd();
  BOOST_CHECK_THROW(protocol->readMapBegin(elemType, elemType1, val), TTransportException);
  protocol->readMapEnd();
}

// THeaderTransport lives in libthriftz, which is only built with zlib.
#ifdef THRIFT_TEST_WITH_ZLIB
BOOST_AUTO_TEST_CASE(test_theadertransport_header_size_exceeds_frame) {
  using apache::thrift::transport::THeaderTransport;
  // Header-format frame whose declared header size (3 * 4 = 12) leaves fewer
  // than the 10 common-header bytes inside the 14-byte frame. The trailing
  // varint bytes are all continuation bytes, so the reader used to run off the
  // end of the receive buffer.
  uint8_t frame[] = {
      0x00, 0x00, 0x00, 0x0E, // frame length = 14
      0x0F, 0xFF, 0x00, 0x00, // header magic
      0x00, 0x00, 0x00, 0x00, // seqId
      0x00, 0x03,             // header size field (3 -> 12 bytes)
      0x02,                   // protocol id varint
      0x00,                   // num transforms = 0
      0x80, 0x80              // info-header varint, all continuation
  };
  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer(frame, sizeof(frame)));
  std::shared_ptr<THeaderTransport> trans(new THeaderTransport(buffer));

  uint8_t out[1];
  BOOST_CHECK_THROW(trans->read(out, sizeof(out)), TTransportException);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_info_header_string_overrun) {
  using apache::thrift::transport::THeaderTransport;
  // Header-format frame whose info-header key length (4) does not fit within the
  // header bytes that remain once the length varint itself is accounted for. The
  // header section (8 bytes) exactly fills the frame, so the boundary sits at the
  // buffer end; the key length has to be bounded against the remaining bytes and
  // rejected.
  uint8_t frame[] = {
      0x00, 0x00, 0x00, 0x12, // frame length = 18
      0x0F, 0xFF, 0x00, 0x00, // header magic
      0x00, 0x00, 0x00, 0x00, // seqId
      0x00, 0x02,             // header size field (2 -> 8 bytes)
      0x02,                   // protocol id varint
      0x00,                   // num transforms = 0
      0x01,                   // info id = key/value
      0x01,                   // one key/value pair
      0x04,                   // key length = 4 (only 3 bytes remain)
      0xAA, 0xBB, 0xCC        // key bytes
  };
  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer(frame, sizeof(frame)));
  std::shared_ptr<THeaderTransport> trans(new THeaderTransport(buffer));

  uint8_t out[1];
  BOOST_CHECK_THROW(trans->read(out, sizeof(out)), TTransportException);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_info_header_string_negative_length) {
  using apache::thrift::transport::THeaderTransport;
  // Header-format frame whose info-header key length varint decodes to a
  // negative int32 (top bit set). The length has to be treated as out of range
  // rather than converted to a size_t, so the read is rejected instead of
  // reaching the string assignment. The three trailing bytes only pad the
  // header section out to its declared size and are never reached.
  uint8_t frame[] = {
      0x00, 0x00, 0x00, 0x16,       // frame length = 22
      0x0F, 0xFF, 0x00, 0x00,       // header magic
      0x00, 0x00, 0x00, 0x00,       // seqId
      0x00, 0x03,                   // header size field (3 -> 12 bytes)
      0x02,                         // protocol id varint
      0x00,                         // num transforms = 0
      0x01,                         // info id = key/value
      0x01,                         // one key/value pair
      0x80, 0x80, 0x80, 0x80, 0x08, // key length varint = INT32_MIN
      0x00, 0x00, 0x00              // padding to fill the header section
  };
  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer(frame, sizeof(frame)));
  std::shared_ptr<THeaderTransport> trans(new THeaderTransport(buffer));

  uint8_t out[1];
  BOOST_CHECK_THROW(trans->read(out, sizeof(out)), TTransportException);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_zlib_roundtrip) {
  using apache::thrift::transport::THeaderTransport;
  // A run of identical bytes compresses to far fewer bytes than it occupies
  // once expanded again, so the result of the zlib transform is much larger
  // than the frame section it is read from.  This drives the full write/read
  // round trip through the zlib transform path.  Keep the payload small enough
  // to stay within the transform buffer the reader sizes from its write buffer.
  const std::size_t N = 700;
  std::vector<uint8_t> payload(N, 0x42);

  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer());
  std::shared_ptr<THeaderTransport> writer(new THeaderTransport(buffer));
  writer->setTransform(THeaderTransport::ZLIB_TRANSFORM);
  writer->write(payload.data(), static_cast<uint32_t>(payload.size()));
  writer->flush();

  std::shared_ptr<THeaderTransport> reader(new THeaderTransport(buffer));
  std::vector<uint8_t> out(N, 0x00);
  reader->readAll(out.data(), static_cast<uint32_t>(out.size()));

  BOOST_CHECK(out == payload);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_zlib_write_incompressible_roundtrip) {
  using apache::thrift::transport::THeaderTransport;
  // Round-trip a PRNG-filled (incompressible) payload through the zlib
  // transform.  The payload stays within the transform buffer the reader sizes
  // from its own write buffer, so the read side can inflate it.
  const std::size_t N = 512;
  std::vector<uint8_t> payload(N);
  std::mt19937 rng(0xC0FFEEu);
  for (auto& b : payload) {
    b = static_cast<uint8_t>(rng());
  }

  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer());
  std::shared_ptr<THeaderTransport> writer(new THeaderTransport(buffer));
  writer->setTransform(THeaderTransport::ZLIB_TRANSFORM);
  writer->write(payload.data(), static_cast<uint32_t>(payload.size()));
  writer->flush();

  std::shared_ptr<THeaderTransport> reader(new THeaderTransport(buffer));
  std::vector<uint8_t> out(N, 0x00);
  reader->readAll(out.data(), static_cast<uint32_t>(out.size()));

  BOOST_CHECK(out == payload);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_zlib_write_large_incompressible) {
  using apache::thrift::transport::THeaderTransport;
  // Write a large incompressible frame through the zlib transform and check
  // that the emitted payload inflates to the whole frame.  The reader cannot
  // inflate a frame this large, so decode the emitted header frame and inflate
  // the payload directly.
  const std::size_t N = 2u * 1024u * 1024u;
  std::vector<uint8_t> payload(N);
  std::mt19937 rng(0xC0FFEEu);
  for (auto& b : payload) {
    b = static_cast<uint8_t>(rng());
  }

  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer());
  std::shared_ptr<THeaderTransport> writer(new THeaderTransport(buffer));
  writer->setTransform(THeaderTransport::ZLIB_TRANSFORM);
  writer->write(payload.data(), static_cast<uint32_t>(payload.size()));
  writer->flush();

  // Emitted header frame: 4B size, 2B magic, 2B flags, 4B seqId, then a 2B
  // header-word count (headerSize / 4), then that many header words, then the
  // zlib payload.
  uint8_t* frame = nullptr;
  uint32_t frameLen = 0;
  buffer->getBuffer(&frame, &frameLen);
  BOOST_REQUIRE(frameLen > 14);
  const uint32_t headerBytes =
      ((static_cast<uint32_t>(frame[12]) << 8) | frame[13]) * 4u;
  const uint32_t payloadOffset = 14u + headerBytes;
  BOOST_REQUIRE(frameLen > payloadOffset);

  std::vector<uint8_t> out(N, 0x00);
  uLongf outLen = static_cast<uLongf>(out.size());
  int rc = uncompress(out.data(), &outLen, frame + payloadOffset,
                      static_cast<uLong>(frameLen - payloadOffset));
  BOOST_CHECK_EQUAL(rc, Z_OK);
  BOOST_CHECK_EQUAL(outLen, static_cast<uLongf>(N));
  BOOST_CHECK(out == payload);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_framed_size_below_magic_word) {
  using apache::thrift::transport::THeaderTransport;
  // The framed-binary and framed-compact branches of readFrame() read a 4-byte
  // magic word out of the frame and then read the remainder as `sz - 4`. `sz`
  // is unsigned, so a declared frame size below 4 wraps that subtraction to
  // roughly 4 GB while the receive buffer holds only the four bytes
  // ensureReadBuffer(4) reserved, and the read runs past the end of it. Only
  // the header-format branch had a lower bound of its own, so every declared
  // size below the magic word has to be rejected before the subtraction.
  const uint8_t magics[][4] = {
      {0x80, 0x01, 0x00, 0x01}, // TBinaryProtocol::VERSION_1
      {0x82, 0x21, 0x00, 0x00}, // TCompactProtocol PROTOCOL_ID / VERSION_N
  };

  for (const auto& magic : magics) {
    for (uint32_t declared = 0; declared < 4; ++declared) {
      std::vector<uint8_t> frame = {
          static_cast<uint8_t>((declared >> 24) & 0xFF),
          static_cast<uint8_t>((declared >> 16) & 0xFF),
          static_cast<uint8_t>((declared >> 8) & 0xFF),
          static_cast<uint8_t>(declared & 0xFF),
          magic[0],
          magic[1],
          magic[2],
          magic[3],
      };
      // Payload the wrapped-around read would copy past the receive buffer.
      frame.resize(frame.size() + 64, 0x41);

      std::shared_ptr<TMemoryBuffer> buffer(
          new TMemoryBuffer(frame.data(), static_cast<uint32_t>(frame.size())));
      std::shared_ptr<THeaderTransport> trans(new THeaderTransport(buffer));

      uint8_t out[1];
      bool caught = false;
      try {
        trans->read(out, sizeof(out));
      } catch (const TTransportException& e) {
        caught = true;
        BOOST_CHECK_EQUAL(e.getType(), TTransportException::CORRUPTED_DATA);
      }
      BOOST_CHECK(caught);
    }
  }
}
#endif // THRIFT_TEST_WITH_ZLIB

BOOST_AUTO_TEST_CASE(test_prealloc_size_caps_reserved_capacity) {
  using apache::thrift::protocol::preallocSize;
  using apache::thrift::protocol::MAX_PREALLOC_SIZE;

  // A count that fits under the cap reserves exactly that count.
  BOOST_CHECK_EQUAL(preallocSize(0), 0u);
  BOOST_CHECK_EQUAL(preallocSize(1), 1u);
  BOOST_CHECK_EQUAL(preallocSize(500), 500u);
  BOOST_CHECK_EQUAL(preallocSize(MAX_PREALLOC_SIZE), MAX_PREALLOC_SIZE);

  // A count above the cap reserves only the cap; the container grows as elements are added.
  BOOST_CHECK_EQUAL(preallocSize(MAX_PREALLOC_SIZE + 1), MAX_PREALLOC_SIZE);
  BOOST_CHECK_EQUAL(preallocSize(1000000u), MAX_PREALLOC_SIZE);
  BOOST_CHECK_EQUAL(preallocSize(0xffffffffu), MAX_PREALLOC_SIZE);
}

#ifdef THRIFT_TEST_WITH_ZLIB
BOOST_AUTO_TEST_CASE(test_theadertransport_framed_size_equal_to_magic_word) {
  using apache::thrift::transport::THeaderTransport;
  // Boundary control for the case above: a declared size of exactly 4 covers
  // the magic word, `sz - 4` is zero and nothing is read past it. The frame
  // carries no payload beyond the magic, so the transport delivers exactly
  // those four bytes. This pins the lower bound to the subtraction it protects
  // rather than rejecting one size too many.
  uint8_t frame[] = {
      0x00, 0x00, 0x00, 0x04, // frame length = 4, the magic word alone
      0x80, 0x01, 0x00, 0x01  // TBinaryProtocol::VERSION_1
  };
  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer(frame, sizeof(frame)));
  std::shared_ptr<THeaderTransport> trans(new THeaderTransport(buffer));

  uint8_t out[4];
  BOOST_CHECK_NO_THROW(trans->readAll(out, sizeof(out)));
  BOOST_CHECK_EQUAL(out[0], 0x80);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_zlib_read_failure_releases_stream) {
  using apache::thrift::transport::THeaderTransport;
  using apache::thrift::TApplicationException;
  // A ZLIB-transformed frame whose inflated size exceeds the transform buffer
  // the reader sizes from its own (default) write buffer cannot be inflated, so
  // the read fails. The zlib stream inflateInit acquired for it has to be
  // released on that failure path too; otherwise it is leaked once per frame.
  // A plain build only checks that the failure still surfaces cleanly here --
  // the leak itself is observable under LeakSanitizer (build the library and
  // this test with -fsanitize=address and run: unfixed leaks ~40 KB per frame,
  // fixed exits clean).
  const std::size_t N = 4096; // > the ~1 KB reader transform buffer
  std::vector<uint8_t> payload(N, 0x41);

  std::shared_ptr<TMemoryBuffer> buffer(new TMemoryBuffer());
  std::shared_ptr<THeaderTransport> writer(new THeaderTransport(buffer));
  writer->setTransform(THeaderTransport::ZLIB_TRANSFORM);
  writer->write(payload.data(), static_cast<uint32_t>(payload.size()));
  writer->flush();

  std::shared_ptr<THeaderTransport> reader(new THeaderTransport(buffer));
  uint8_t out[16];
  BOOST_CHECK_THROW(reader->read(out, sizeof(out)), TApplicationException);
}

BOOST_AUTO_TEST_CASE(test_theadertransport_unknown_protocol_id_error_reply) {
  using apache::thrift::transport::THeaderTransport;
  using apache::thrift::protocol::THeaderProtocol;
  using apache::thrift::protocol::T_BINARY_PROTOCOL;
  using apache::thrift::TApplicationException;
  using apache::thrift::protocol::TMessageType;
  // Header-format frame whose protocol id varint decodes to 0xFFFF. The
  // transport stores that in its signed 16-bit protocol id as -1. readMessageBegin
  // rejects the unknown protocol and writes an error reply back on the same
  // transport, which re-encodes the protocol id varint. That encoder has to
  // terminate on a negative value rather than run off its fixed output buffer,
  // so the call surfaces a clean TApplicationException instead of overrunning.
  uint8_t frame[] = {
      0x00, 0x00, 0x00, 0x0E, // frame length = 14
      0x0F, 0xFF, 0x00, 0x00, // header magic + flags
      0x00, 0x00, 0x00, 0x00, // seqId
      0x00, 0x01,             // header size field (1 -> 4 bytes)
      0xFF, 0xFF, 0x03,       // protocol id varint = 0xFFFF
      0x00                    // num transforms = 0
  };
  std::shared_ptr<TMemoryBuffer> inBuf(new TMemoryBuffer(frame, sizeof(frame)));
  std::shared_ptr<TMemoryBuffer> outBuf(new TMemoryBuffer());
  THeaderProtocol proto(inBuf, outBuf, T_BINARY_PROTOCOL);

  std::string name;
  TMessageType messageType;
  int32_t seqId = 0;
  BOOST_CHECK_THROW(proto.readMessageBegin(name, messageType, seqId), TApplicationException);
}
#endif // THRIFT_TEST_WITH_ZLIB

BOOST_AUTO_TEST_SUITE_END()