Fixes a bug that caused an infinite loop when writing to a fixed-size…

… buffer that is too small (#227) Adds a test to demonstrate how to use a user handler to grow the buffer when necessary.
amazon-ion · Feb 9, 2021 · e44782e · e44782e
1 parent 06d78fc
commit e44782e
Show file tree

Hide file tree

Showing 2 changed files with 130 additions and 0 deletions.
diff --git a/ionc/ion_stream.c b/ionc/ion_stream.c
@@ -777,6 +777,10 @@ iERR ion_stream_write(ION_STREAM *stream, BYTE *buf, SIZE len, SIZE *p_bytes_wri
     dst = stream->_curr;
     if (to_write > src_remaining) {
       to_write = src_remaining;
+    } else if (to_write < 1 && _ion_stream_is_fully_buffered(stream)) {
+      // There is no space available for writing. Since this is not a file or a user stream, space can never become
+      // available.
+      FAILWITH(IERR_BUFFER_TOO_SMALL);
     }
     memcpy(dst, src, to_write);
     if (stream->_dirty_start == NULL) {

diff --git a/test/test_ion_stream.cpp b/test/test_ion_stream.cpp
@@ -22,6 +22,7 @@
 #include <ionc/ion.h>
 #include "ion_helpers.h"
 #include "ion_test_util.h"
+#include "ion_assert.h"
 
 struct _test_in_memory_paged_stream_context {
     BYTE *data;
@@ -103,3 +104,128 @@ TEST(IonStream, ContinuesOverPageBoundary) {
     ION_ASSERT_OK(ion_timestamp_equals(&expected, &ts, &is_equal, &g_IonEventDecimalContext));
     ASSERT_TRUE(is_equal);
 }
+
+TEST(IonStream, BufferTooSmall) {
+    iENTER;
+    hWRITER writer = NULL;
+    uint8_t buf[2]; // This buffer is too small to hold the output.
+    ION_WRITER_OPTIONS options = { 0 };
+    options.output_as_binary = TRUE;
+    ION_ASSERT_OK(ion_writer_open_buffer(&writer, buf, sizeof(buf), &options));
+    ION_ASSERT_OK(ion_writer_write_int32(writer, 1));
+
+    SIZE len;
+    ION_ASSERT_FAIL(ion_writer_finish(writer, &len));
+    ION_ASSERT_FAIL(ion_writer_close(writer));
+}
+
+iERR grow_buffer_if_necessary(_ion_user_stream *stream) {
+    iENTER;
+    _test_in_memory_paged_stream_context *context = (_test_in_memory_paged_stream_context *)stream->handler_state;
+    context->data_len = stream->curr - context->data;
+    if (stream->curr == stream->limit) {
+        // The buffer is full and must grow.
+        context->page_size *= 2;
+        stream->curr = (BYTE *)malloc(context->page_size);
+        stream->limit = stream->curr + context->page_size;
+        memcpy(stream->curr, context->data, context->data_len);
+        free(context->data);
+        context->data = stream->curr;
+        stream->curr += context->data_len;
+    }
+    iRETURN;
+}
+
+TEST(IonStream, WriteToUserStream) {
+    // Write output to a user-controlled stream backed by a buffer. The user handler will grow the buffer
+    // when necessary.
+
+    hWRITER writer = NULL;
+    SIZE image_length;
+    SIZE initial_buffer_length = 1;
+    ION_STREAM *stream = NULL;
+
+    ION_WRITER_OPTIONS options;
+    memset(&options, 0, sizeof(ION_WRITER_OPTIONS));
+    options.output_as_binary = TRUE;
+
+    _test_in_memory_paged_stream_context context;
+    memset(&context, 0, sizeof(_test_in_memory_paged_stream_context));
+    context.data = (BYTE *)(malloc(initial_buffer_length));
+    context.page_size = initial_buffer_length;
+
+    ION_ASSERT_OK(ion_stream_open_handler_out(&grow_buffer_if_necessary, &context, &stream));
+    ION_ASSERT_OK(ion_writer_open(&writer, stream, &options));
+
+    ION_STREAM_USER_PAGED *paged_stream = (ION_STREAM_USER_PAGED *)stream;
+    paged_stream->_user_stream.curr = context.data;
+    paged_stream->_user_stream.limit = context.data + initial_buffer_length;
+
+    ION_ASSERT_OK(ion_writer_start_container(writer, tid_STRUCT));
+    ION_STRING fieldNameString;
+    ion_string_assign_cstr(&fieldNameString, "str_col1", strlen("str_col1"));
+    ION_ASSERT_OK(ion_writer_write_field_name(writer, &fieldNameString));
+
+    ION_STRING value_string;
+    ion_string_assign_cstr(&value_string, "str_val1", strlen("str_val1"));
+    ION_ASSERT_OK(ion_writer_write_string(writer, &value_string));
+
+    ion_string_assign_cstr(&fieldNameString, "str_col2", strlen("str_col2"));
+    ION_ASSERT_OK(ion_writer_write_field_name(writer, &fieldNameString));
+
+    ion_string_assign_cstr(&value_string, "str_val1", strlen("str_val1"));
+    ION_ASSERT_OK(ion_writer_write_string(writer, &value_string));
+
+    ion_string_assign_cstr(&fieldNameString, "str_col3", strlen("str_col3"));
+    ION_ASSERT_OK(ion_writer_write_field_name(writer, &fieldNameString));
+
+    ion_string_assign_cstr(&value_string, "str_val1", strlen("str_val1"));
+    ION_ASSERT_OK(ion_writer_write_string(writer, &value_string));
+
+    ION_ASSERT_OK(ion_writer_finish_container(writer));
+    ION_ASSERT_OK(ion_writer_flush(writer, &image_length));
+    ION_ASSERT_OK(ion_writer_close(writer));
+    ION_ASSERT_OK(ion_stream_close(stream));
+
+    ASSERT_EQ(image_length, context.data_len); // 72 bytes in this case.
+    ASSERT_EQ(128, context.page_size); // The next power of 2 above 72 (we doubled the size on each grow).
+
+    // Verify the output.
+
+    hREADER reader = NULL;
+    ION_TYPE type;
+    ION_STRING stringValue;
+
+    ION_ASSERT_OK(ion_reader_open_buffer(&reader, context.data, context.data_len, NULL));
+    ION_ASSERT_OK(ion_reader_next(reader, &type));
+    ASSERT_EQ(tid_STRUCT, type);
+    ION_ASSERT_OK(ion_reader_step_in(reader));
+
+    ION_ASSERT_OK(ion_reader_next(reader, &type));
+    ASSERT_EQ(tid_STRING, type);
+    ION_ASSERT_OK(ion_reader_get_field_name(reader, &fieldNameString));
+    assertStringsEqual("str_col1", (char *)fieldNameString.value, fieldNameString.length);
+    ION_ASSERT_OK(ion_reader_read_string(reader, &stringValue));
+    assertStringsEqual("str_val1", (char *)stringValue.value, stringValue.length);
+
+    ION_ASSERT_OK(ion_reader_next(reader, &type));
+    ASSERT_EQ(tid_STRING, type);
+    ION_ASSERT_OK(ion_reader_get_field_name(reader, &fieldNameString));
+    assertStringsEqual("str_col2", (char *)fieldNameString.value, fieldNameString.length);
+    ION_ASSERT_OK(ion_reader_read_string(reader, &stringValue));
+    assertStringsEqual("str_val1", (char *)stringValue.value, stringValue.length);
+
+    ION_ASSERT_OK(ion_reader_next(reader, &type));
+    ASSERT_EQ(tid_STRING, type);
+    ION_ASSERT_OK(ion_reader_get_field_name(reader, &fieldNameString));
+    assertStringsEqual("str_col3", (char *)fieldNameString.value, fieldNameString.length);
+    ION_ASSERT_OK(ion_reader_read_string(reader, &stringValue));
+    assertStringsEqual("str_val1", (char *)stringValue.value, stringValue.length);
+
+    ION_ASSERT_OK(ion_reader_step_out(reader));
+    ION_ASSERT_OK(ion_reader_next(reader, &type));
+    ASSERT_EQ(tid_EOF, type);
+    ION_ASSERT_OK(ion_reader_close(reader));
+
+    free(context.data);
+}