lib: istream-chain - Fix snapshot handling when link istream is destroyed

diff --git a/src/lib/istream-chain.c b/src/lib/istream-chain.c
index ea141f53e25013d8b15f3be6814aa07a66d6eef1..6ee65a14aa49b3774c210592ec7d0678b66cc12a 100644 (file)
--- a/src/lib/istream-chain.c
+++ b/src/lib/istream-chain.c
@@ -304,6 +304,39 @@ static void i_stream_chain_close(struct iostream_private *stream,
        }
 }
 
+static struct istream_snapshot *
+i_stream_chain_snapshot(struct istream_private *stream,
+                       struct istream_snapshot *prev_snapshot)
+{
+       if (stream->buffer == stream->w_buffer) {
+               /* Two or more istreams have been combined. Snapshot the
+                  w_buffer's contents that contains their data. */
+               i_assert(stream->memarea != NULL);
+               return i_stream_default_snapshot(stream, prev_snapshot);
+       }
+       /* Individual istreams are being read. Snapshot the istream directly. */
+       struct chain_istream *cstream =
+               container_of(stream, struct chain_istream, istream);
+       struct istream_chain_link *link = cstream->chain.head;
+       if (link == NULL || link->stream == NULL)
+               return prev_snapshot;
+
+       struct istream_private *_link_stream = link->stream->real_stream;
+       struct istream_snapshot *snapshot = i_new(struct istream_snapshot, 1);
+       snapshot->prev_snapshot =
+               _link_stream->snapshot(_link_stream, prev_snapshot);
+       if (snapshot->prev_snapshot == prev_snapshot) {
+               /* The link stream didn't implement snapshotting in any way.
+                  This could cause trouble if the link stream is freed while
+                  it's still referred to in this snapshot. Fix this by
+                  referencing the link istream. Normally avoid doing this,
+                  since the extra references can cause unexpected problems. */
+               snapshot->istream = link->stream;
+               i_stream_ref(snapshot->istream);
+       }
+       return snapshot;
+}
+
 struct istream *i_stream_create_chain(struct istream_chain **chain_r)
 {
        struct chain_istream *cstream;
@@ -317,6 +350,7 @@ struct istream *i_stream_create_chain(struct istream_chain **chain_r)
                i_stream_chain_set_max_buffer_size;
 
        cstream->istream.read = i_stream_chain_read;
+       cstream->istream.snapshot = i_stream_chain_snapshot;
 
        cstream->istream.istream.readable_fd = FALSE;
        cstream->istream.istream.blocking = FALSE;

diff --git a/src/lib/istream-private.h b/src/lib/istream-private.h
index afabe2359943e14e91aa06b56844235fa2fdc566..b612b9c7f2a95aa28e9fbb051b88908ea50a0792 100644 (file)
--- a/src/lib/istream-private.h
+++ b/src/lib/istream-private.h
@@ -73,6 +73,7 @@ struct istream_private {
 struct istream_snapshot {
        struct istream_snapshot *prev_snapshot;
        struct memarea *old_memarea;
+       struct istream *istream;
        void (*free)(struct istream_snapshot *snapshot);
 };
 

diff --git a/src/lib/istream.c b/src/lib/istream.c
index c63fa2b5e13c6ce27651fbaf945f2f2fda4d35ae..5fc511247cd6b03060827c934d5d86266241e667 100644 (file)
--- a/src/lib/istream.c
+++ b/src/lib/istream.c
@@ -256,6 +256,7 @@ void i_stream_snapshot_free(struct istream_snapshot **_snapshot)
        else {
                if (snapshot->old_memarea != NULL)
                        memarea_unref(&snapshot->old_memarea);
+               i_stream_unref(&snapshot->istream);
                i_free(snapshot);
        }
 }

diff --git a/src/lib/test-istream-chain.c b/src/lib/test-istream-chain.c
index 162fc6522c4992301791eeb2f8ff34d451d52977..3307ddcdf3d22d81faf940ec01e428eb8ec822ca 100644 (file)
--- a/src/lib/test-istream-chain.c
+++ b/src/lib/test-istream-chain.c
@@ -82,27 +82,27 @@ static void test_istream_chain_early_end(void)
 
 static void test_istream_chain_accumulate(void)
 {
-       struct istream *input;
-       struct istream  *test_input, *test_input2, *test_input3, *test_input4,
-               *test_input5;
+       struct istream *input, *tmp_istream;
+       struct istream *test_istreams[5];
        struct istream_chain *chain;
        const unsigned char *data;
        size_t size;
 
        test_begin("istream chain accumulate");
 
-       test_input = test_istream_create("abcdefghijklmnopqrst");
-       test_input2 = test_istream_create("ABCDEFGHIJKLMNOPQRSTUVWXY");
-       test_input3 = test_istream_create("!\"#$%&'()*+,-./01234567890:;<=");
-       test_input4 = test_istream_create("z1y2x3w4v5u6t7s8r9q0p.o,n");
-       test_input5 = test_istream_create("aAbBcCdDeEfFgGhHiIjJ");
+       test_istreams[0] = test_istream_create("abcdefghijklmnopqrst");
+       test_istreams[1] = test_istream_create("ABCDEFGHIJKLMNOPQRSTUVWXY");
+       test_istreams[2] = test_istream_create("!\"#$%&'()*+,-./01234567890:;<=");
+       test_istreams[3] = test_istream_create("z1y2x3w4v5u6t7s8r9q0p.o,n");
+       test_istreams[4] = test_istream_create("aAbBcCdDeEfFgGhHiIjJ");
 
        input = i_stream_create_chain(&chain);
        /* no input */
        test_assert(i_stream_read(input) == 0);
 
        /* first stream */
-       i_stream_chain_append(chain, test_input);
+       i_stream_chain_append(chain, test_istreams[0]);
+       tmp_istream = test_istreams[0]; i_stream_unref(&tmp_istream);
        test_assert(i_stream_read_data(input, &data, &size, 0) == 1);
        test_assert(size == 20);
        test_assert(memcmp(data, "abcdefghijklmnopqrst", 20) == 0);
@@ -111,17 +111,18 @@ static void test_istream_chain_accumulate(void)
        i_stream_skip(input, 12);
 
        /* second stream */
-       i_stream_chain_append(chain, test_input2);
-       test_istream_set_size(test_input2, 0);
+       i_stream_chain_append(chain, test_istreams[1]);
+       tmp_istream = test_istreams[1]; i_stream_unref(&tmp_istream);
+       test_istream_set_size(test_istreams[1], 0);
        test_assert(i_stream_read_data(input, &data, &size, 10) == 0);
        test_assert(size == 8);
-       test_istream_set_size(test_input2, 10);
+       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_input2, FALSE);
+       test_istream_set_allow_eof(test_istreams[1], FALSE);
        test_assert(i_stream_read(input) == 0);
-       test_istream_set_size(test_input2, 25);
-       test_istream_set_allow_eof(test_input2, TRUE);
+       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"
@@ -131,10 +132,11 @@ static void test_istream_chain_accumulate(void)
        i_stream_skip(input, 12);
 
        /* third stream */
-       i_stream_chain_append(chain, test_input3);
-       test_istream_set_size(test_input3, 0);
+       i_stream_chain_append(chain, test_istreams[2]);
+       tmp_istream = test_istreams[2]; i_stream_unref(&tmp_istream);
+       test_istream_set_size(test_istreams[2], 0);
        test_assert(i_stream_read(input) == 0);
-       test_istream_set_size(test_input3, 30);
+       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"
@@ -145,7 +147,8 @@ static void test_istream_chain_accumulate(void)
        i_stream_skip(input, 12);
 
        /* forth stream */
-       i_stream_chain_append(chain, test_input4);
+       i_stream_chain_append(chain, test_istreams[3]);
+       tmp_istream = test_istreams[3]; i_stream_unref(&tmp_istream);
        test_assert(i_stream_read_data(input, &data, &size, 40) == 1);
        test_assert(size == 64);
        test_assert(memcmp(data, "QRSTUVWXY"
@@ -156,7 +159,8 @@ static void test_istream_chain_accumulate(void)
        i_stream_skip(input, 6);
 
        /* fifth stream */
-       i_stream_chain_append(chain, test_input5);
+       i_stream_chain_append(chain, test_istreams[4]);
+       tmp_istream = test_istreams[4]; i_stream_unref(&tmp_istream);
        test_assert(i_stream_read_data(input, &data, &size, 60) == 1);
        test_assert(size == 78);
        test_assert(memcmp(data, "WXY"
@@ -184,12 +188,6 @@ static void test_istream_chain_accumulate(void)
        test_assert(size == 0);
 
        i_stream_unref(&input);
-
-       i_stream_unref(&test_input);
-       i_stream_unref(&test_input2);
-       i_stream_unref(&test_input3);
-       i_stream_unref(&test_input4);
-       i_stream_unref(&test_input5);
        test_end();
 }