Add CborReaderProxy to support process CBOR stream in chunks

Author: muhammad-othman

extensions/src/AWSSDK.Extensions.CborProtocol/Internal/CborReaderProxy.cs

@@ -0,0 +1,274 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * 
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ * 
+ *  http://aws.amazon.com/apache2.0
+ * 
+ * or in the "license" file accompanying this file. This file 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.
+ */
+using Amazon.Runtime.Internal.Util;
+using System;
+using System.Buffers;
+using System.Collections.Generic;
+using System.Formats.Cbor;
+using System.IO;
+
+namespace Amazon.Extensions.CborProtocol.Internal
+{
+    public class CborReaderProxy : IDisposable
+    {
+        /// <summary>
+        /// Enum to track the type of CBOR container (map or array)
+        /// for state management within the CborReaderProxy.
+        /// </summary>
+        private enum CborContainerType
+        {
+            Map,
+            Array
+        }
+
+        private static readonly ILogger _logger = Logger.GetLogger(typeof(CborReaderProxy));
+        private readonly Stack<CborContainerType> _nestingStack = new Stack<CborContainerType>();
+        private readonly Stream _stream;
+        private byte[] _buffer;
+        private CborReader _internalCborReader;
+        private int _currentChunkSize;
+
+        public CborReaderProxy(Stream stream)
+        {
+            _stream = stream ?? throw new ArgumentNullException(nameof(stream));
+            _buffer = ArrayPool<byte>.Shared.Rent(AWSConfigs.CborReaderInitialBufferSize);
+
+            _currentChunkSize = _stream.Read(_buffer, 0, AWSConfigs.CborReaderInitialBufferSize);
+            var memorySlice = new ReadOnlyMemory<byte>(_buffer, 0, _currentChunkSize);
+
+            // We must allow multiple root values because when refilling the new chunk is just a fragment of the whole stream.
+            _internalCborReader = new CborReader(memorySlice, allowMultipleRootLevelValues: true);
+        }
+
+        /// <summary>
+        /// This method is called when a read operation fails because it needs more
+        /// data than is currently available in the buffer. It handles stitching leftover
+        /// data with a new chunk from the stream and, if necessary, resizing the buffer.
+        /// </summary>
+        /// <param name="bytesToSkip">Number of bytes to skip before reading new data (e.g., 1 to skip CBOR break byte 0xFF)</param>
+        private void RefillBuffer(int bytesToSkip = 0)
+        {
+            int leftoverBytesCount = _internalCborReader.BytesRemaining;
+
+            // Determine where the leftover bytes start
+            int leftoverStartIndex = _currentChunkSize - leftoverBytesCount;
+
+            // If we are skipping bytes, we need to move the start forward
+            leftoverStartIndex += bytesToSkip;
+            leftoverBytesCount = leftoverBytesCount - bytesToSkip;
+
+            // If the leftover data completely fills the buffer, grow it
+            if (leftoverBytesCount >= _buffer.Length)
+            {
+                int newSize = _buffer.Length * 2;
+                var newBuffer = ArrayPool<byte>.Shared.Rent(newSize);
+
+                Buffer.BlockCopy(_buffer, leftoverStartIndex, newBuffer, 0, leftoverBytesCount);
+                ArrayPool<byte>.Shared.Return(_buffer);
+                _buffer = newBuffer;
+            }
+            else if (leftoverBytesCount > 0)  // Shift leftovers (after skipping) to the beginning of the buffer
+            {
+                Buffer.BlockCopy(_buffer, leftoverStartIndex, _buffer, 0, leftoverBytesCount);
+            }
+
+            // Read from stream into buffer after leftovers
+            int bytesReadFromStream = _stream.Read(_buffer, leftoverBytesCount, _buffer.Length - leftoverBytesCount);
+
+            // Update the total size of valid data in our buffer.
+            _currentChunkSize = leftoverBytesCount + bytesReadFromStream;
+
+            // Check for a malformed stream: if we have leftovers but the stream is empty,
+            // it means the CBOR data was truncated.
+            if (bytesReadFromStream == 0 && leftoverBytesCount > 0)
+            {
+                throw new CborContentException("Stream ended unexpectedly with an incomplete CBOR data item.");
+            }
+
+            var newMemorySlice = new ReadOnlyMemory<byte>(_buffer, 0, _currentChunkSize);
+            _internalCborReader.Reset(newMemorySlice);
+
+            _logger.InfoFormat("Buffer refilled: read {0} byte(s), total in buffer now: {1}.", bytesReadFromStream, _currentChunkSize);
+        }
+
+        /// <summary>
+        /// Executes a CBOR read operation, refilling the buffer and retrying if a CborContentException is thrown.
+        /// </summary>
+        /// <typeparam name="T">The return type of the CBOR read operation.</typeparam>
+        /// <param name="readOperation">A delegate representing the read operation to execute.</param>
+        /// <returns>The result of the read operation.</returns>
+        /// <exception cref="CborContentException">
+        /// Thrown if too many retries are attempted or if the stream ends unexpectedly.
+        /// </exception>
+        private T ExecuteRead<T>(Func<CborReader, T> readOperation)
+        {
+            int maxRetries = 64;
+            int retryCount = 0;
+
+            while (true)
+            {
+                try
+                {
+                    return readOperation(_internalCborReader);
+                }
+                catch (CborContentException ex)
+                {
+                    if (_currentChunkSize == 0 && _internalCborReader.BytesRemaining == 0)
+                    {
+                        // Fail fast if we’ve already consumed all input and nothing remains to refill.
+                        throw;
+                    }
+
+                    if (++retryCount > maxRetries)
+                    {
+                        throw new CborContentException("Too many retries during CBOR stream parsing. Possible malformed or infinite data.");
+                    }
+
+                    _logger.Debug(ex, $"CborContentException caught (attempt #{retryCount}), attempting to refill buffer.");
+                    // Attempt to refill and retry the operation.
+                    RefillBuffer();
+                }
+            }
+        }
+
+
+        private bool IsNextByteEndOfContainer()
+        {
+            int unreadOffset = _currentChunkSize - _internalCborReader.BytesRemaining;
+            if (unreadOffset < _currentChunkSize)
+                return _buffer[unreadOffset] == 0xFF; // 0xFF indicates "break" in indefinite-length map/array
+
+            return false;
+        }
+
+        private void ReadEndContainer(CborContainerType expectedType, CborReaderState expectedEndState, Action<CborReader> readEndAction)
+        {
+            ExecuteRead(r =>
+            {
+                if (_nestingStack.Count == 0 || _nestingStack.Peek() != expectedType)
+                    throw new CborContentException($"Unexpected end of {expectedType.ToString().ToLowerInvariant()}.");
+
+                var state = CborReaderState.Finished;
+                try
+                {
+                    state = r.PeekState();
+                }
+                catch (CborContentException)
+                {
+                    _logger.InfoFormat($"CborContentException caught during PeekState while expecting end of {expectedType.ToString().ToLowerInvariant()}");
+                }
+
+                if (state == expectedEndState)
+                {
+                    readEndAction(r);
+                    _nestingStack.Pop();
+                    return true;
+                }
+
+                if (state == CborReaderState.Finished)
+                {
+                    if (IsNextByteEndOfContainer())
+                    {
+                        RefillBuffer(1); // Skip 0xFF
+                    }
+                    else
+                    {
+                        RefillBuffer(0); // This means we are in a definite-length map/array which doesn't end with 0xFF.
+                    }
+                    _nestingStack.Pop();
+                    return true;
+                }
+
+                throw new CborContentException($"Expected end of {expectedType.ToString().ToLowerInvariant()} but could not parse it.");
+            });
+        }
+
+
+        public void ReadEndMap()
+        {
+            ReadEndContainer(CborContainerType.Map, CborReaderState.EndMap, (reader) => reader.ReadEndMap());
+        }
+
+        public void ReadEndArray()
+        {
+            ReadEndContainer(CborContainerType.Array, CborReaderState.EndArray, (r) => r.ReadEndArray());
+        }
+
+
+        public int? ReadStartMap() => ExecuteRead(reader =>
+        {
+            var result = reader.ReadStartMap();
+            _nestingStack.Push(CborContainerType.Map);
+            return result;
+        });
+
+        public int? ReadStartArray() => ExecuteRead(reader =>
+        {
+            var result = reader.ReadStartArray();
+            _nestingStack.Push(CborContainerType.Array);
+            return result;
+        });
+
+
+        public CborReaderState PeekState()
+        {
+            return ExecuteRead(r =>
+            {
+                try
+                {
+                    return r.PeekState();
+                }
+                catch (CborContentException ex)
+                {
+                    // Translate a Break code to the appropriate container end state
+                    // based on our own nesting stack.
+                    if (_nestingStack.Count > 0)
+                    {
+                        var inferredState = _nestingStack.Peek() == CborContainerType.Map
+                            ? CborReaderState.EndMap
+                            : CborReaderState.EndArray;
+
+                        _logger.Debug(ex, $"CborContentException during PeekState interpreted as {inferredState} due to nesting stack.");
+                        return inferredState;
+                    }
+                    // If our stack is empty, it's a genuine error.
+                    throw;
+                }
+            });
+        }
+
+        public string ReadTextString() => ExecuteRead(r => r.ReadTextString());
+        public int ReadInt32() => ExecuteRead(r => r.ReadInt32());
+        public long ReadInt64() => ExecuteRead(r => r.ReadInt64());
+        public decimal ReadDecimal() => ExecuteRead(r => r.ReadDecimal());
+        public double ReadDouble() => ExecuteRead(r => r.ReadDouble());
+        public bool ReadBoolean() => ExecuteRead(r => r.ReadBoolean());
+        public float ReadSingle() => ExecuteRead(r => r.ReadSingle());
+        public CborTag ReadTag() => ExecuteRead(r => r.ReadTag());
+        public byte[] ReadByteString() => ExecuteRead(r => r.ReadByteString());
+        public void ReadNull() => ExecuteRead(r => { r.ReadNull(); return true; });
+        public void SkipValue() => ExecuteRead(r => { r.SkipValue(); return true; });
+        public int CurrentDepth => _internalCborReader.CurrentDepth;
+
+        public void Dispose()
+        {
+            if (_buffer != null)
+            {
+                ArrayPool<byte>.Shared.Return(_buffer);
+                _buffer = null;
+            }
+        }
+    }
+}

extensions/src/AWSSDK.Extensions.CborProtocol/Internal/Transform/CborUnmarshallerContext.cs

@@ -29,7 +29,6 @@ namespace Amazon.Extensions.CborProtocol.Internal.Transform
     public class CborUnmarshallerContext : UnmarshallerContext
     {
         private bool disposed = false;
-        private byte[] rentedBuffer = null;
         private readonly Stack<string> _pathStack = new Stack<string>();
 
         // There isn't a direct way to check if the reader is at the start of the document,
@@ -39,7 +38,8 @@ public class CborUnmarshallerContext : UnmarshallerContext
         public override bool IsEndElement => Reader.PeekState() == CborReaderState.Finished;
         public override bool IsStartElement => Reader.PeekState() == CborReaderState.StartMap;
 
-        public CborReader Reader { get; set; }
+        public Stream Stream { get; set; }
+        public CborReaderProxy Reader { get; set; }
         public override string CurrentPath => string.Join("/", _pathStack.Reverse());
         public override int CurrentDepth => Reader.CurrentDepth;
 
@@ -76,8 +76,6 @@ IRequestContext requestContext
                 responseStream = WrappingStream;
             }
 
-            long? streamSize = null;
-
             if (responseData != null)
             {
                 bool parsedContentLengthHeader = long.TryParse(
@@ -105,79 +103,14 @@ out long contentLength
                         requestContext
                     );
                 }
-                if (parsedContentLengthHeader && responseData.ContentLength == contentLength)
-                {
-                    streamSize = contentLength;
-                }
             }
 
             this.WebResponseData = responseData;
             this.MaintainResponseBody = maintainResponseBody;
             this.IsException = isException;
 
-            Reader = CreateCborReaderFromStream(
-                FlexibleChecksumStream ?? CrcStream ?? responseStream,
-                streamSize
-            );
-        }
-
-        private CborReader CreateCborReaderFromStream(Stream stream, long? streamSize = null)
-        {
-            int totalRead = 0;
-            if (streamSize.HasValue)
-            {
-                // TODO: update this method to read the stream in chunks incase the stream is larger than int.MaxValue
-                // If we know the size, we can read directly into a buffer of exact size
-                rentedBuffer = ArrayPool<byte>.Shared.Rent((int)streamSize.Value);
-
-                while (totalRead < streamSize.Value)
-                {
-                    int bytesRead = stream.Read(rentedBuffer, totalRead, (int)(streamSize.Value - totalRead));
-                    if (bytesRead == 0)
-                    {
-                        // If no bytes are read, it means we've reached the end of the stream before reading the whole streamSize.
-                        throw new EndOfStreamException($"Expected {streamSize.Value} bytes but only read {totalRead}.");
-                    }
-
-                    totalRead += bytesRead;
-                }
-            }
-            else
-            {
-                const int InitialBufferSize = 1024 * 8; // 8kb
-                rentedBuffer = ArrayPool<byte>.Shared.Rent(InitialBufferSize);
-
-                while (true)
-                {
-                    int read = stream.Read(rentedBuffer, totalRead, rentedBuffer.Length - totalRead);
-                    if (read == 0)
-                        break;
-
-                    totalRead += read;
-
-                    if (totalRead == rentedBuffer.Length)
-                    {
-                        // Expand the buffer size by doubling it
-                        var newBuffer = ArrayPool<byte>.Shared.Rent(rentedBuffer.Length * 2);
-                        try
-                        {
-                            Buffer.BlockCopy(rentedBuffer, 0, newBuffer, 0, totalRead);
-                        }
-                        catch
-                        {
-                            ArrayPool<byte>.Shared.Return(newBuffer);
-                            ArrayPool<byte>.Shared.Return(rentedBuffer);
-                            rentedBuffer = null;
-                            throw;
-                        }
-                        ArrayPool<byte>.Shared.Return(rentedBuffer);
-                        rentedBuffer = newBuffer;
-                    }
-                }
-            }
-
-            var actualBytes = new ReadOnlyMemory<byte>(rentedBuffer, 0, totalRead);
-            return new CborReader(actualBytes);
+            Stream = FlexibleChecksumStream ?? CrcStream ?? responseStream;
+            Reader = new CborReaderProxy(Stream);
         }
 
         /// <summary>
@@ -204,9 +137,10 @@ protected override void Dispose(bool disposing)
         {
             if (!this.disposed)
             {
-                if (disposing && rentedBuffer != null)
+                if (disposing && Reader != null)
                 {
-                    ArrayPool<byte>.Shared.Return(rentedBuffer);
+                    Reader.Dispose();
+                    Reader = null;
                 }
                 disposed = true;
             }