lib: istream-concat - Fix snapshot handling when combining two istreams

Snapshotting wasn't handled correctly when two (or more) istreams' contents
were combined into the same buffer.

diff --git a/src/lib/istream-concat.c b/src/lib/istream-concat.c
index 874b03670c8f32ef6d01b16b72384c48ed1a8f24..d034be4ab7e2491f7431927e166627d504203b1a 100644 (file)
--- a/src/lib/istream-concat.c
+++ b/src/lib/istream-concat.c
@@ -2,6 +2,7 @@
 
 #include "lib.h"
 #include "buffer.h"
+#include "memarea.h"
 #include "istream-private.h"
 #include "istream-concat.h"
 
@@ -95,8 +96,12 @@ static void i_stream_concat_read_next(struct concat_istream *cstream)
 
        /* we already verified that the data size is less than the
           maximum buffer size */
+       cstream->istream.skip = 0;
        cstream->istream.pos = 0;
        if (data_size > 0) {
+               if (cstream->istream.memarea != NULL &&
+                   memarea_get_refcount(cstream->istream.memarea) > 1)
+                       i_stream_memarea_detach(&cstream->istream);
                if (!i_stream_try_alloc(&cstream->istream, data_size, &size))
                        i_unreached();
                i_assert(size >= data_size);
@@ -379,6 +384,5 @@ struct istream *i_stream_create_concat(struct istream *input[])
        cstream->istream.istream.readable_fd = FALSE;
        cstream->istream.istream.blocking = blocking;
        cstream->istream.istream.seekable = seekable;
-       return i_stream_create(&cstream->istream, NULL, -1,
-                              ISTREAM_CREATE_FLAG_NOOP_SNAPSHOT);
+       return i_stream_create(&cstream->istream, NULL, -1, 0);
 }

diff --git a/src/lib/test-istream-concat.c b/src/lib/test-istream-concat.c
index fd3887ddaf20a551ad3d4bcdf49083e333bf7542..297536439545af8982232e205c2beb85f92449c9 100644 (file)
--- a/src/lib/test-istream-concat.c
+++ b/src/lib/test-istream-concat.c
@@ -44,6 +44,7 @@ static void test_istream_concat_one(unsigned int buffer_size)
                for (j = 0; j < size; j++) {
                        test_assert((char)data[j] == input_string[(input->v_offset + j) % STREAM_BYTES]);
                }
+               test_assert(i_stream_read(input) <= 0);
        }
        test_assert(i_stream_read(input) == -1);
        i_stream_skip(input, i_stream_get_data_size(input));
@@ -161,6 +162,72 @@ static void test_istream_concat_early_end(void)
        test_end();
 }
 
+static void test_istream_concat_snapshot(void)
+{
+       struct istream *input;
+       const unsigned char *data;
+       size_t size;
+
+       test_begin("istream concat snapshot");
+
+       struct istream *test_istreams[] = {
+               test_istream_create("abcdefghijklmnopqrst"),
+               test_istream_create("ABCDEFGHIJKLMNOPQRSTUVWXY"),
+               test_istream_create("!\"#$%&'()*+,-./01234567890:;<="),
+               NULL
+       };
+
+       input = i_stream_create_concat(test_istreams);
+       test_istream_set_size(test_istreams[0], 20);
+       test_istream_set_size(test_istreams[1], 0);
+       test_istream_set_size(test_istreams[2], 0);
+
+       /* first stream */
+       test_istream_set_allow_eof(test_istreams[0], FALSE);
+       test_assert(i_stream_read_data(input, &data, &size, 0) == 1);
+       test_assert(size == 20);
+       test_assert(memcmp(data, "abcdefghijklmnopqrst", 20) == 0);
+
+       /* partially skip */
+       i_stream_skip(input, 12);
+
+       /* second stream */
+       test_assert(i_stream_read_data(input, &data, &size, 10) == 0);
+       test_assert(size == 8);
+       test_istream_set_allow_eof(test_istreams[0], TRUE);
+       test_istream_set_size(test_istreams[0], 0);
+       test_assert(i_stream_read_data(input, &data, &size, 10) == 0);
+       test_assert(size == 8);
+       test_istream_set_size(test_istreams[1], 10);
+       test_assert(i_stream_read_data(input, &data, &size, 10) == 1);
+       test_assert(size == 18);
+       test_istream_set_allow_eof(test_istreams[1], FALSE);
+       test_assert(i_stream_read(input) == 0);
+       test_istream_set_size(test_istreams[1], 25);
+       test_istream_set_allow_eof(test_istreams[1], TRUE);
+       test_assert(i_stream_read_data(input, &data, &size, 30) == 1);
+       test_assert(size == 33);
+       test_assert(memcmp(data, "mnopqrst"
+               "ABCDEFGHIJKLMNOPQRSTUVWXY", 33) == 0);
+
+       /* partially skip */
+       i_stream_skip(input, 12);
+
+       /* third stream */
+       test_istream_set_size(test_istreams[2], 0);
+       test_assert(i_stream_read(input) == 0);
+       test_istream_set_size(test_istreams[2], 30);
+       test_assert(i_stream_read_data(input, &data, &size, 25) == 1);
+       test_assert(size == 51);
+       test_assert(memcmp(data, "EFGHIJKLMNOPQRSTUVWXY"
+               "!\"#$%&'()*+,-./01234567890:;<=", 51) == 0);
+
+       i_stream_unref(&input);
+       for (unsigned int i = 0; test_istreams[i] != NULL; i++)
+               i_stream_unref(&test_istreams[i]);
+       test_end();
+}
+
 void test_istream_concat(void)
 {
        unsigned int i;
@@ -180,4 +247,5 @@ void test_istream_concat(void)
 
        test_istream_concat_seek_end();
        test_istream_concat_early_end();
+       test_istream_concat_snapshot();
 }