dovecot-2.0: istream-seekable: If we're immediately at EOF after...

[dovecot at dovecot.org]

details:   http://hg.dovecot.org/dovecot-2.0/rev/5d7ee047667f
changeset: 10523:5d7ee047667f
user:      Timo Sirainen <tss at iki.fi>
date:      Tue Dec 22 14:19:40 2009 -0500
description:
istream-seekable: If we're immediately at EOF after copying buffer to file, don't corrupt the buffer.

diffstat:

2 files changed, 55 insertions(+), 2 deletions(-)
src/lib/istream-seekable.c      |   24 ++++++++++++++++++++++--
src/lib/test-istream-seekable.c |   33 +++++++++++++++++++++++++++++++++

diffs (100 lines):

diff -r 19787c343ef1 -r 5d7ee047667f src/lib/istream-seekable.c
--- a/src/lib/istream-seekable.c	Tue Dec 22 14:19:15 2009 -0500
+++ b/src/lib/istream-seekable.c	Tue Dec 22 14:19:40 2009 -0500
@@ -72,7 +72,10 @@ i_stream_seekable_set_max_buffer_size(st
 
 static int copy_to_temp_file(struct seekable_istream *sstream)
 {
+	struct istream_private *stream = &sstream->istream;
 	const char *path;
+	const unsigned char *buffer;
+	size_t size;
 	int fd;
 
 	fd = sstream->fd_callback(&path, sstream->context);
@@ -89,11 +92,27 @@ static int copy_to_temp_file(struct seek
 	sstream->temp_path = i_strdup(path);
 	sstream->write_peak = sstream->buffer->used;
 
-	buffer_free(&sstream->buffer);
-
 	sstream->fd = fd;
 	sstream->fd_input =
 		i_stream_create_fd(fd, sstream->istream.max_buffer_size, TRUE);
+
+	/* read back the data we just had in our buffer */
+	i_stream_seek(sstream->fd_input, stream->istream.v_offset);
+	for (;;) {
+		buffer = i_stream_get_data(sstream->fd_input, &size);
+		if (size >= stream->pos)
+			break;
+
+		if (i_stream_read(sstream->fd_input) <= 0) {
+			i_error("istream-seekable: Couldn't read back "
+				"in-memory input");
+			i_stream_destroy(&sstream->fd_input);
+			return -1;
+		}
+	}
+	stream->buffer = buffer;
+	stream->pos = size;
+	buffer_free(&sstream->buffer);
 	return 0;
 }
 
@@ -341,6 +360,7 @@ i_stream_create_seekable(struct istream 
 	sstream->context = context;
 	sstream->buffer = buffer_create_dynamic(default_pool, BUF_INITIAL_SIZE);
         sstream->istream.max_buffer_size = max_buffer_size;
+	sstream->fd = -1;
 
 	sstream->input = i_new(struct istream *, count + 1);
 	memcpy(sstream->input, input, sizeof(*input) * count);
diff -r 19787c343ef1 -r 5d7ee047667f src/lib/test-istream-seekable.c
--- a/src/lib/test-istream-seekable.c	Tue Dec 22 14:19:15 2009 -0500
+++ b/src/lib/test-istream-seekable.c	Tue Dec 22 14:19:40 2009 -0500
@@ -121,6 +121,37 @@ static void test_istream_seekable_random
 	i_stream_unref(&input);
 }
 
+static void test_istream_seekable_eof(void)
+{
+	static const char *in_str = "foo";
+	unsigned int in_str_len = strlen(in_str);
+	struct istream *streams[2], *input;
+	const unsigned char *data;
+	size_t size;
+
+	test_begin("istream seekable eof");
+
+	streams[0] = i_stream_create_from_data(in_str, in_str_len);
+	streams[0]->seekable = FALSE;
+	streams[1] = NULL;
+
+	input = i_stream_create_seekable(streams, in_str_len, fd_callback, NULL);
+	i_stream_unref(&streams[0]);
+
+	test_assert(i_stream_read(input) == (ssize_t)in_str_len);
+	data = i_stream_get_data(input, &size);
+	test_assert(size == in_str_len);
+	test_assert(memcmp(data, in_str, in_str_len) == 0);
+
+	test_assert(i_stream_read(input) == -1);
+	data = i_stream_get_data(input, &size);
+	test_assert(size == in_str_len);
+	test_assert(memcmp(data, in_str, in_str_len) == 0);
+
+	i_stream_unref(&input);
+	test_end();
+}
+
 void test_istream_seekable(void)
 {
 	unsigned int i;
@@ -135,4 +166,6 @@ void test_istream_seekable(void)
 		test_istream_seekable_random();
 	} T_END;
 	test_end();
+
+	test_istream_seekable_eof();
 }