unit test for chunking stream

Author: kai-ion

src/aws-cpp-sdk-core/include/smithy/client/features/ChunkingInterceptor.h

@@ -2,77 +2,89 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0.
  */
-
 #pragma once
-
-#include <aws/core/Core_EXPORTS.h>
-#include <aws/core/client/ClientConfiguration.h>
-#include <aws/core/utils/HashingUtils.h>
-#include <aws/core/utils/StringUtils.h>
+#include <aws/core/http/HttpRequest.h>
 #include <aws/core/utils/Array.h>
+#include <aws/core/utils/StringUtils.h>
+#include <aws/core/utils/HashingUtils.h>
+#include <aws/core/utils/logging/LogMacros.h>
+#include <aws/core/utils/memory/stl/AWSStringStream.h>
 #include <smithy/interceptor/Interceptor.h>
-#include <smithy/client/common/AwsSmithyClientUtils.h>
+#include <aws/core/client/ClientConfiguration.h>
+#include <aws/core/utils/Outcome.h>
+#include <aws/core/client/AWSError.h>
+#include <memory>
 
 namespace smithy {
 namespace client {
 namespace features {
 
-namespace {
-
-static const char* const CHECKSUM_HEADER_PREFIX = "x-amz-checksum-";
-static const char* const ALLOCATION_TAG = "ChunkingInterceptor";
-static const size_t DATA_BUFFER_SIZE = 65536;
-
-} // anonymous namespace
+static const size_t AWS_DATA_BUFFER_SIZE = 65536;
+static const char* ALLOCATION_TAG = "ChunkingInterceptor";
+static const char* CHECKSUM_HEADER_PREFIX = "x-amz-checksum-";
 
-class AwsChunkedStreamWrapper {
+template <size_t DataBufferSize = AWS_DATA_BUFFER_SIZE>
+class AwsChunkedStreamBuf : public std::streambuf {
 public:
-    AwsChunkedStreamWrapper(Aws::Http::HttpRequest* request, const std::shared_ptr<Aws::IOStream>& originalBody, size_t bufferSize = DATA_BUFFER_SIZE)
-        : m_streambuf(request, originalBody, bufferSize), m_iostream(&m_streambuf) {}
-    
-    Aws::IOStream* GetIOStream() { return &m_iostream; }
+    AwsChunkedStreamBuf(Aws::Http::HttpRequest* request,
+                        const std::shared_ptr<Aws::IOStream>& stream,
+                        size_t bufferSize = DataBufferSize)
+        : m_chunkingStream(Aws::MakeShared<Aws::StringStream>("AwsChunkedStream")),
+          m_request(request),
+          m_stream(stream),
+          m_data(bufferSize)
+    {
+        assert(m_stream != nullptr);
+        if (m_stream == nullptr) {
+            AWS_LOGSTREAM_ERROR("AwsChunkedStream", "stream is null");
+        }
+        assert(m_request != nullptr);
+        if (m_request == nullptr) {
+            AWS_LOGSTREAM_ERROR("AwsChunkedStream", "request is null");
+        }
 
-private:
-    class AwsChunkedStreamBuf : public std::streambuf {
-    public:
-        AwsChunkedStreamBuf(Aws::Http::HttpRequest* request, const std::shared_ptr<Aws::IOStream>& originalBody, size_t bufferSize)
-            : m_request(request), m_stream(originalBody), m_data(bufferSize), m_bufferSize(bufferSize),
-              m_chunkingStream(Aws::MakeShared<Aws::StringStream>(ALLOCATION_TAG)) {
-            setg(nullptr, nullptr, nullptr);
+        setg(nullptr, nullptr, nullptr);
+    }
+
+protected:
+    int_type underflow() override {
+        if (gptr() && gptr() < egptr()) {
+            return traits_type::to_int_type(*gptr());
         }
 
-    protected:
-        int_type underflow() override {
-            if (gptr() < egptr()) {
-                return traits_type::to_int_type(*gptr());
+        // only read and write to chunked stream if the underlying stream
+        // is still in a valid state
+        if (m_stream->good()) {
+            // Try to read in a 64K chunk, if we cant we know the stream is over
+            m_stream->read(m_data.GetUnderlyingData(), m_data.GetLength());
+            size_t bytesRead = static_cast<size_t>(m_stream->gcount());
+            writeChunk(bytesRead);
+
+            // if we've read everything from the stream, we want to add the trailer
+            // to the underlying stream
+            if ((m_stream->peek() == EOF || m_stream->eof()) && !m_stream->bad()) {
+                writeTrailerToUnderlyingStream();
             }
+        }
 
-            if (m_stream->good()) {
-                m_stream->read(m_data.GetUnderlyingData(), m_bufferSize);
-                size_t bytesRead = static_cast<size_t>(m_stream->gcount());
-                writeChunk(bytesRead);
-                
-                if ((m_stream->peek() == EOF || m_stream->eof()) && !m_stream->bad()) {
-                    writeTrailerToUnderlyingStream();
-                }
-            }
-            
-            if ((m_chunkingStream->peek() == EOF || m_chunkingStream->eof()) && !m_chunkingStream->bad()) {
-                return traits_type::eof();
-            }
-            
-            m_chunkingStream->read(m_buffer, sizeof(m_buffer));
-            size_t bytesRead = static_cast<size_t>(m_chunkingStream->gcount());
-            if (bytesRead == 0) {
-                return traits_type::eof();
-            }
-            
-            setg(m_buffer, m_buffer, m_buffer + bytesRead);
-            return traits_type::to_int_type(*gptr());
+        // if the underlying stream is empty there is nothing to read
+        if ((m_chunkingStream->peek() == EOF || m_chunkingStream->eof()) && !m_chunkingStream->bad()) {
+            return traits_type::eof();
         }
 
-    private:
-        void writeTrailerToUnderlyingStream() {
+        // Read from chunking stream to internal buffer
+        m_chunkingStream->read(m_buffer.GetUnderlyingData(), m_buffer.GetLength());
+        size_t bytesRead = static_cast<size_t>(m_chunkingStream->gcount());
+        if (bytesRead == 0) {
+            return traits_type::eof();
+        }
+
+        setg(m_buffer.GetUnderlyingData(), m_buffer.GetUnderlyingData(), m_buffer.GetUnderlyingData() + bytesRead);
+        return traits_type::to_int_type(*gptr());
+    }
+
+private:
+    void writeTrailerToUnderlyingStream() {
         Aws::StringStream chunkedTrailerStream;
         chunkedTrailerStream << "0\r\n";
         if (m_request->GetRequestHash().second != nullptr) {
@@ -86,13 +98,13 @@ class AwsChunkedStreamWrapper {
         }
         *m_chunkingStream << chunkedTrailer;
     }
-    
+
     void writeChunk(size_t bytesRead) {
         if (m_request->GetRequestHash().second != nullptr) {
             m_request->GetRequestHash().second->Update(reinterpret_cast<unsigned char*>(m_data.GetUnderlyingData()), bytesRead);
         }
-        
-        if (bytesRead > 0 && m_chunkingStream != nullptr && !m_chunkingStream->bad()) {
+
+        if (bytesRead > 0 && m_chunkingStream && !m_chunkingStream->bad()) {
             if (m_chunkingStream->eof()) {
                 m_chunkingStream->clear();
             }
@@ -102,16 +114,23 @@ class AwsChunkedStreamWrapper {
         }
     }
 
-        Aws::Http::HttpRequest* m_request;
-        std::shared_ptr<Aws::IOStream> m_stream;
-        Aws::Utils::Array<char> m_data;
-        size_t m_bufferSize;
-        std::shared_ptr<Aws::IOStream> m_chunkingStream;
-        char m_buffer[8192];
-    };
-    
-    AwsChunkedStreamBuf m_streambuf;
-    Aws::IOStream m_iostream;
+    std::shared_ptr<Aws::IOStream> m_chunkingStream;
+    Aws::Http::HttpRequest* m_request{nullptr};
+    std::shared_ptr<Aws::IOStream> m_stream;
+    Aws::Utils::Array<char> m_data;
+    Aws::Utils::Array<char> m_buffer{DataBufferSize};
+};
+
+class AwsChunkedIOStream : public Aws::IOStream {
+public:
+    AwsChunkedIOStream(Aws::Http::HttpRequest* request,
+                       const std::shared_ptr<Aws::IOStream>& originalBody,
+                       size_t bufferSize = AWS_DATA_BUFFER_SIZE)
+        : Aws::IOStream(&m_buf),
+          m_buf(request, originalBody, bufferSize) {}
+
+private:
+    AwsChunkedStreamBuf<> m_buf;
 };
 
 /**
@@ -159,11 +178,9 @@ class ChunkingInterceptor : public smithy::interceptor::Interceptor {
             }
         }
 
-        auto wrapper = Aws::MakeShared<AwsChunkedStreamWrapper>(
+        auto chunkedBody = Aws::MakeShared<AwsChunkedIOStream>(
             ALLOCATION_TAG, request.get(), originalBody);
-        auto chunkedBody = std::shared_ptr<Aws::IOStream>(
-            wrapper, wrapper->GetIOStream());
-        
+
         request->AddContentBody(chunkedBody);
         return request;
     }

tests/aws-cpp-sdk-core-tests/utils/stream/ChunkingInterceptorTest.cpp

@@ -12,76 +12,91 @@ using namespace Aws::Http::Standard;
 using namespace smithy::client::features;
 using namespace Aws::Utils::Crypto;
 
-class ChunkingInterceptorTest : public Aws::Testing::AwsCppSdkGTestSuite {};
-
 const char* CHUNKING_TEST_LOG_TAG = "CHUNKING_INTERCEPTOR_TEST";
 
-// Helper function to read all data from a stream
-std::string ReadAllFromStream(std::shared_ptr<Aws::IOStream> stream) {
-    std::stringstream result;
-    char buffer[1024];
-    while (stream->good() && !stream->eof()) {
-        stream->read(buffer, sizeof(buffer));
-        auto bytesRead = stream->gcount();
-        if (bytesRead > 0) {
-            result.write(buffer, bytesRead);
-        }
+class ChunkingInterceptorTest : public Aws::Testing::AwsCppSdkGTestSuite {
+protected:
+    template <typename Fn>
+    void withChunkedStream(const std::string& input, size_t bufferSize, Fn&& fn) {
+        StandardHttpRequest request{"test.com", Http::HttpMethod::HTTP_GET};
+        auto requestHash = Aws::MakeShared<CRC32>(CHUNKING_TEST_LOG_TAG);
+        request.SetRequestHash("crc32", requestHash);
+        auto inputStream = Aws::MakeShared<Aws::StringStream>(CHUNKING_TEST_LOG_TAG, input);
+        
+        AwsChunkedIOStream wrapper{&request, inputStream, bufferSize};
+        Aws::IOStream* stream = &wrapper;
+        
+        fn(*stream);
     }
-    return result.str();
-}
+};
 
 TEST_F(ChunkingInterceptorTest, ChunkedStreamShouldWork) {
-    StandardHttpRequest request{"www.elda.com/will", Http::HttpMethod::HTTP_GET};
-    auto requestHash = Aws::MakeShared<CRC32>(CHUNKING_TEST_LOG_TAG);
-    request.SetRequestHash("crc32", requestHash);
-    std::shared_ptr<Aws::IOStream> inputStream = Aws::MakeShared<Aws::StringStream>(CHUNKING_TEST_LOG_TAG, "1234567890123456789012345");
-    
-    AwsChunkedStreamWrapper wrapper{&request, inputStream, 10};
-    auto chunkedStream = wrapper.GetIOStream();
-    
-    std::string output = ReadAllFromStream(std::shared_ptr<Aws::IOStream>(chunkedStream, [](Aws::IOStream*){}));
-    auto expectedStreamWithChecksum = "A\r\n1234567890\r\nA\r\n1234567890\r\n5\r\n12345\r\n0\r\nx-amz-checksum-crc32:78DeVw==\r\n\r\n";
-    EXPECT_EQ(expectedStreamWithChecksum, output);
+    withChunkedStream("1234567890123456789012345", 10, [](Aws::IOStream& chunkedStream) {
+        char buffer[100];
+        std::stringstream output;
+        
+        // Read in 10-byte chunks like original test
+        for (int i = 0; i < 4; i++) {
+            chunkedStream.read(buffer, 10);
+            auto bytesRead = chunkedStream.gcount();
+            output.write(buffer, bytesRead);
+        }
+        
+        // Read trailing checksum (greater than 10 chars)
+        chunkedStream.read(buffer, 40);
+        auto bytesRead = static_cast<size_t>(chunkedStream.gcount());
+        EXPECT_EQ(36ul, bytesRead);
+        output.write(buffer, bytesRead);
+        
+        EXPECT_EQ("A\r\n1234567890\r\nA\r\n1234567890\r\n5\r\n12345\r\n0\r\nx-amz-checksum-crc32:78DeVw==\r\n\r\n", output.str());
+    });
 }
 
 TEST_F(ChunkingInterceptorTest, ShouldNotRequireTwoReadsOnSmallChunk) {
-    StandardHttpRequest request{"www.clemar.com/strohl", Http::HttpMethod::HTTP_GET};
-    auto requestHash = Aws::MakeShared<CRC32>(CHUNKING_TEST_LOG_TAG);
-    request.SetRequestHash("crc32", requestHash);
-    std::shared_ptr<Aws::IOStream> inputStream = Aws::MakeShared<Aws::StringStream>(CHUNKING_TEST_LOG_TAG, "12345");
-    
-    AwsChunkedStreamWrapper wrapper{&request, inputStream, 100};
-    auto chunkedStream = wrapper.GetIOStream();
-    
-    std::string output = ReadAllFromStream(std::shared_ptr<Aws::IOStream>(chunkedStream, [](Aws::IOStream*){}));
-    auto expectedStreamWithChecksum = "5\r\n12345\r\n0\r\nx-amz-checksum-crc32:y/U6HA==\r\n\r\n";
-    EXPECT_EQ(expectedStreamWithChecksum, output);
+    withChunkedStream("12345", 100, [](Aws::IOStream& chunkedStream) {
+        char buffer[100];
+        chunkedStream.read(buffer, 100);
+        auto bytesRead = static_cast<size_t>(chunkedStream.gcount());
+        EXPECT_EQ(46ul, bytesRead);
+        
+        std::string output(buffer, bytesRead);
+        EXPECT_EQ("5\r\n12345\r\n0\r\nx-amz-checksum-crc32:y/U6HA==\r\n\r\n", output);
+    });
 }
 
 TEST_F(ChunkingInterceptorTest, ShouldWorkOnSmallBuffer) {
-    StandardHttpRequest request{"www.eugief.com/hesimay", Http::HttpMethod::HTTP_GET};
-    auto requestHash = Aws::MakeShared<CRC32>(CHUNKING_TEST_LOG_TAG);
-    request.SetRequestHash("crc32", requestHash);
-    std::shared_ptr<Aws::IOStream> inputStream = Aws::MakeShared<Aws::StringStream>(CHUNKING_TEST_LOG_TAG, "1234567890");
-    
-    AwsChunkedStreamWrapper wrapper{&request, inputStream, 5};
-    auto chunkedStream = wrapper.GetIOStream();
-    
-    std::string output = ReadAllFromStream(std::shared_ptr<Aws::IOStream>(chunkedStream, [](Aws::IOStream*){}));
-    auto expectedStreamWithChecksum = "5\r\n12345\r\n5\r\n67890\r\n0\r\nx-amz-checksum-crc32:Jh2u5Q==\r\n\r\n";
-    EXPECT_EQ(expectedStreamWithChecksum, output);
+    withChunkedStream("1234567890", 5, [](Aws::IOStream& chunkedStream) {
+        char buffer[100];
+        
+        // First read - explicitly ask for 10 bytes (first chunk: "5\r\n12345\r\n")
+        chunkedStream.read(buffer, 10);
+        auto bytesRead = static_cast<size_t>(chunkedStream.gcount());
+        EXPECT_EQ(10ul, bytesRead);
+        std::string firstRead(buffer, bytesRead);
+        EXPECT_EQ("5\r\n12345\r\n", firstRead);
+        
+        // Second read - now we expect the rest (46 bytes: second chunk + trailer)
+        chunkedStream.read(buffer, 100);
+        bytesRead = static_cast<size_t>(chunkedStream.gcount());
+        EXPECT_EQ(46ul, bytesRead);
+        std::string secondRead(buffer, bytesRead);
+        EXPECT_EQ("5\r\n67890\r\n0\r\nx-amz-checksum-crc32:Jh2u5Q==\r\n\r\n", secondRead);
+        
+        // Subsequent reads should return 0
+        chunkedStream.read(buffer, 100);
+        bytesRead = static_cast<size_t>(chunkedStream.gcount());
+        EXPECT_EQ(0ul, bytesRead);
+    });
 }
 
 TEST_F(ChunkingInterceptorTest, ShouldWorkOnEmptyStream) {
-    StandardHttpRequest request{"www.nidia.com/juna", Http::HttpMethod::HTTP_GET};
-    auto requestHash = Aws::MakeShared<CRC32>(CHUNKING_TEST_LOG_TAG);
-    request.SetRequestHash("crc32", requestHash);
-    std::shared_ptr<Aws::IOStream> inputStream = Aws::MakeShared<Aws::StringStream>(CHUNKING_TEST_LOG_TAG, "");
-    
-    AwsChunkedStreamWrapper wrapper{&request, inputStream, 5};
-    auto chunkedStream = wrapper.GetIOStream();
-    
-    std::string output = ReadAllFromStream(std::shared_ptr<Aws::IOStream>(chunkedStream, [](Aws::IOStream*){}));
-    auto expectedStreamWithChecksum = "0\r\nx-amz-checksum-crc32:AAAAAA==\r\n\r\n";
-    EXPECT_EQ(expectedStreamWithChecksum, output);
+    withChunkedStream("", 5, [](Aws::IOStream& chunkedStream) {
+        char buffer[100];
+        chunkedStream.read(buffer, 100);
+        auto bytesRead = static_cast<size_t>(chunkedStream.gcount());
+        EXPECT_EQ(36ul, bytesRead);
+        
+        std::string output(buffer, bytesRead);
+        EXPECT_EQ("0\r\nx-amz-checksum-crc32:AAAAAA==\r\n\r\n", output);
+    });
 }