Prevent TLVWrite from possibly performing buffer overrun if trying to reserve

src/lib/core/TLVBackingStore.h

@@ -155,6 +155,8 @@ class DLL_EXPORT TLVBackingStore
      *
      */
     virtual CHIP_ERROR FinalizeBuffer(TLVWriter & writer, uint8_t * bufStart, uint32_t bufLen) = 0;
+
+    virtual bool IsSafeToReserve() { return false; }
 };
 
 } // namespace TLV

src/lib/core/TLVWriter.cpp

@@ -97,6 +97,22 @@ CHIP_ERROR TLVWriter::Finalize()
     return err;
 }
 
+CHIP_ERROR TLVWriter::ReserveBuffer(uint32_t aBufferSize)
+{
+    VerifyOrReturnError(mRemainingLen >= aBufferSize, CHIP_ERROR_NO_MEMORY);
+
+    uint32_t spaceLeftInCurrentBuffer        = mReservedSize + mRemainingLen;
+    bool canPossiblyAllocateAdditionalMemory = spaceLeftInCurrentBuffer < mMaxLen;
+
+    if (mBackingStore && canPossiblyAllocateAdditionalMemory)
+    {
+        VerifyOrReturnError(mBackingStore->IsSafeToReserve(), CHIP_ERROR_INCORRECT_STATE);
+    }
+    mReservedSize += aBufferSize;
+    mRemainingLen -= aBufferSize;
+    return CHIP_NO_ERROR;
+}
+
 CHIP_ERROR TLVWriter::PutBoolean(Tag tag, bool v)
 {
     return WriteElementHead((v) ? TLVElementType::BooleanTrue : TLVElementType::BooleanFalse, tag, 0);

src/lib/core/TLVWriter.h

@@ -128,14 +128,12 @@ class DLL_EXPORT TLVWriter
      *
      * @retval #CHIP_NO_ERROR        Successfully reserved required buffer size.
      * @retval #CHIP_ERROR_NO_MEMORY The reserved buffer size cannot fits into the remaining buffer size.
+     * @retval #CHIP_ERROR_INCORRECT_STATE
+     *                               Uses TLVBackingStore and is in a state where it might allocate
+     *                               additional non-contigious memory, thus making it difficult/impossible
+     *                               to properly reserve space.
      */
-    CHIP_ERROR ReserveBuffer(uint32_t aBufferSize)
-    {
-        VerifyOrReturnError(mRemainingLen >= aBufferSize, CHIP_ERROR_NO_MEMORY);
-        mReservedSize += aBufferSize;
-        mRemainingLen -= aBufferSize;
-        return CHIP_NO_ERROR;
-    }
+    CHIP_ERROR ReserveBuffer(uint32_t aBufferSize);
 
     /**
      * Release previously reserved buffer.

src/system/TLVPacketBufferBackingStore.h

@@ -80,6 +80,11 @@ class TLVPacketBufferBackingStore : public chip::TLV::TLVBackingStore
     CHIP_ERROR OnInit(chip::TLV::TLVWriter & writer, uint8_t *& bufStart, uint32_t & bufLen) override;
     CHIP_ERROR GetNewBuffer(chip::TLV::TLVWriter & writer, uint8_t *& bufStart, uint32_t & bufLen) override;
     CHIP_ERROR FinalizeBuffer(chip::TLV::TLVWriter & writer, uint8_t * bufStart, uint32_t bufLen) override;
+    virtual bool IsSafeToReserve() override
+    {
+        // It is safe to reserve if there is no chained buffers.
+        return !mUseChainedBuffers;
+    }
 
 protected:
     chip::System::PacketBufferHandle mHeadBuffer;

src/system/tests/TestTLVPacketBufferBackingStore.cpp

@@ -35,6 +35,18 @@ using namespace ::chip;
 
 namespace {
 
+void WriteUntilRemainingLessThan(nlTestSuite * inSuite, PacketBufferTLVWriter & writer, const uint32_t & remainingSize)
+{
+    CHIP_ERROR error;
+    uint32_t lengthRemaining = writer.GetRemainingFreeLength();
+    while (lengthRemaining >= remainingSize)
+    {
+        error = writer.Put(TLV::AnonymousTag(), static_cast<uint8_t>(7));
+        NL_TEST_ASSERT(inSuite, error == CHIP_NO_ERROR);
+        lengthRemaining = writer.GetRemainingFreeLength();
+    }
+}
+
 class TLVPacketBufferBackingStoreTest
 {
 public:
@@ -43,6 +55,9 @@ class TLVPacketBufferBackingStoreTest
 
     static void BasicEncodeDecode(nlTestSuite * inSuite, void * inContext);
     static void MultiBufferEncode(nlTestSuite * inSuite, void * inContext);
+    static void NonChainedBufferCanReserve(nlTestSuite * inSuite, void * inContext);
+    static void TestWriterReserveUnreserveDoesNotOverflow(nlTestSuite * inSuite, void * inContext);
+    static void TestWriterReserve(nlTestSuite * inSuite, void * inContext);
 };
 
 int TLVPacketBufferBackingStoreTest::TestSetup(void * inContext)
@@ -238,14 +253,111 @@ void TLVPacketBufferBackingStoreTest::MultiBufferEncode(nlTestSuite * inSuite, v
     NL_TEST_ASSERT(inSuite, error == CHIP_END_OF_TLV);
 }
 
+void TLVPacketBufferBackingStoreTest::NonChainedBufferCanReserve(nlTestSuite * inSuite, void * inContext)
+{
+    // Start with a too-small buffer.
+    uint32_t smallSize             = 5;
+    uint32_t smallerSizeToReserver = smallSize - 1;
+
+    auto buffer = PacketBufferHandle::New(smallSize, 0);
+
+    PacketBufferTLVWriter writer;
+    writer.Init(std::move(buffer), /* useChainedBuffers = */ false);
+
+    CHIP_ERROR error = writer.ReserveBuffer(smallerSizeToReserver);
+    NL_TEST_ASSERT(inSuite, error == CHIP_NO_ERROR);
+}
+
+// This test previously was created to show that there was an overflow bug, now this test mainly
+// just checks that you cannot reserve this type of TLVBackingStorage buffer.
+void TLVPacketBufferBackingStoreTest::TestWriterReserveUnreserveDoesNotOverflow(nlTestSuite * inSuite, void * inContext)
+{
+    // Start with a too-small buffer.
+    uint32_t smallSize             = 100;
+    uint32_t smallerSizeToReserver = smallSize - 1;
+
+    auto buffer = PacketBufferHandle::New(smallSize, 0);
+
+    PacketBufferTLVWriter writer;
+    writer.Init(std::move(buffer), /* useChainedBuffers = */ true);
+
+    CHIP_ERROR error = writer.ReserveBuffer(smallerSizeToReserver);
+    if (error == CHIP_NO_ERROR)
+    {
+        uint32_t lengthRemaining = writer.GetRemainingFreeLength();
+        NL_TEST_ASSERT(inSuite, lengthRemaining == 1);
+        // Lets try to overflow by getting next buffer in the chain,
+        // unreserving then writting until the end of the current buffer.
+        error = writer.Put(TLV::AnonymousTag(), static_cast<uint8_t>(7));
+        NL_TEST_ASSERT(inSuite, error == CHIP_NO_ERROR);
+
+        lengthRemaining = writer.GetRemainingFreeLength();
+        NL_TEST_ASSERT(inSuite, lengthRemaining > smallerSizeToReserver);
+
+        WriteUntilRemainingLessThan(inSuite, writer, 2);
+
+        lengthRemaining = writer.GetRemainingFreeLength();
+        NL_TEST_ASSERT(inSuite, lengthRemaining != 0);
+        NL_TEST_ASSERT(inSuite, lengthRemaining < smallerSizeToReserver);
+
+        error = writer.UnreserveBuffer(smallerSizeToReserver);
+        NL_TEST_ASSERT(inSuite, error == CHIP_NO_ERROR);
+
+        lengthRemaining = writer.GetRemainingFreeLength();
+        NL_TEST_ASSERT(inSuite, lengthRemaining > smallerSizeToReserver);
+
+        // This is where we get overflow.
+        WriteUntilRemainingLessThan(inSuite, writer, 2);
+
+        // If we get here then the overflow condition we were expecting did not happen. If that is the case,
+        // either we have fixed reservation for chained buffers, or expected failure didn't hit on this
+        // platform.
+        //
+        // If there is a fix please add reservation for chained buffers, please make sure you account for
+        // what happens if TLVWriter::WriteData fails to get a new buffer but we are not at max size, do
+        // you actually have space for what was supposed to be reserved.
+        NL_TEST_ASSERT(inSuite, false);
+    }
+
+    // We no longer allow non-contigous buffers to be reserved.
+    NL_TEST_ASSERT(inSuite, error == CHIP_ERROR_INCORRECT_STATE);
+}
+
+void TLVPacketBufferBackingStoreTest::TestWriterReserve(nlTestSuite * inSuite, void * inContext)
+{
+    // Start with a too-small buffer.
+    uint32_t smallSize             = 5;
+    uint32_t smallerSizeToReserver = smallSize - 1;
+
+    auto buffer = PacketBufferHandle::New(smallSize, 0);
+
+    PacketBufferTLVWriter writer;
+    writer.Init(std::move(buffer), /* useChainedBuffers = */ false);
+
+    CHIP_ERROR error = writer.ReserveBuffer(smallerSizeToReserver);
+    NL_TEST_ASSERT(inSuite, error == CHIP_NO_ERROR);
+
+    error = writer.Put(TLV::AnonymousTag(), static_cast<uint8_t>(7));
+    NL_TEST_ASSERT(inSuite, error == CHIP_ERROR_NO_MEMORY);
+
+    error = writer.UnreserveBuffer(smallerSizeToReserver);
+    NL_TEST_ASSERT(inSuite, error == CHIP_NO_ERROR);
+
+    error = writer.Put(TLV::AnonymousTag(), static_cast<uint8_t>(7));
+    NL_TEST_ASSERT(inSuite, error == CHIP_NO_ERROR);
+}
+
 /**
  *   Test Suite. It lists all the test functions.
  */
 // clang-format off
 const nlTest sTests[] =
 {
-    NL_TEST_DEF("BasicEncodeDecode",                    TLVPacketBufferBackingStoreTest::BasicEncodeDecode),
-    NL_TEST_DEF("MultiBufferEncode",                    TLVPacketBufferBackingStoreTest::MultiBufferEncode),
+    NL_TEST_DEF("BasicEncodeDecode",                         TLVPacketBufferBackingStoreTest::BasicEncodeDecode),
+    NL_TEST_DEF("MultiBufferEncode",                         TLVPacketBufferBackingStoreTest::MultiBufferEncode),
+    NL_TEST_DEF("NonChainedBufferCanReserve",                TLVPacketBufferBackingStoreTest::NonChainedBufferCanReserve),
+    NL_TEST_DEF("TestWriterReserveUnreserveDoesNotOverflow", TLVPacketBufferBackingStoreTest::TestWriterReserveUnreserveDoesNotOverflow),
+    NL_TEST_DEF("TestWriterReserve",                         TLVPacketBufferBackingStoreTest::TestWriterReserve),
 
     NL_TEST_SENTINEL()
 };