lib: istream - Add i_stream_read_limited().

diff --git a/src/lib/istream-private.h b/src/lib/istream-private.h
index c1180a57469bff6b0babd2a4d4ed63ef22f7ad80..e8542b264eff7ec4691d30c684e4f2e800a58523 100644 (file)
--- a/src/lib/istream-private.h
+++ b/src/lib/istream-private.h
@@ -37,7 +37,7 @@ struct istream_private {
        const unsigned char *buffer;
        unsigned char *w_buffer; /* may be NULL */
 
-       size_t buffer_size, max_buffer_size, init_buffer_size;
+       size_t buffer_size, max_buffer_size, init_buffer_size, data_limit;
        size_t skip, pos, try_alloc_limit;
        /* If seeking backwards within the buffer, the next read() will
           return again pos..high_pos */

diff --git a/src/lib/istream.c b/src/lib/istream.c
index 1f2728ff28c1ea4dafb69efde60d026e535b7746..bf539c7e3a272493481c3bbfe7c91c2be73c2137 100644 (file)
--- a/src/lib/istream.c
+++ b/src/lib/istream.c
@@ -794,6 +794,27 @@ int i_stream_read_data(struct istream *stream, const unsigned char **data_r,
        return -1;
 }
 
+int i_stream_read_limited(struct istream *stream, const unsigned char **data_r,
+                         size_t *size_r, size_t limit)
+{
+       struct istream_private *_stream = stream->real_stream;
+       int ret;
+
+       *data_r = i_stream_get_data(stream, size_r);
+       if (*size_r >= limit) {
+               *size_r = limit;
+               return 1;
+       }
+
+       _stream->data_limit = limit;
+       ret = i_stream_read_more(stream, data_r, size_r);
+       _stream->data_limit = 0;
+
+       if (*size_r >= limit)
+               *size_r = limit;
+       return ret;
+}
+
 void i_stream_compress(struct istream_private *stream)
 {
        i_assert(stream->memarea == NULL ||
@@ -887,10 +908,20 @@ bool i_stream_try_alloc(struct istream_private *stream,
                }
        }
 
-       *size_r = stream->buffer_size - stream->pos;
-       if (stream->try_alloc_limit > 0 &&
-           *size_r > stream->try_alloc_limit)
-               *size_r = stream->try_alloc_limit;
+       if (stream->data_limit == 0 ||
+           (stream->buffer_size - stream->skip) < stream->data_limit) {
+               *size_r = stream->buffer_size - stream->pos;
+               if (stream->try_alloc_limit > 0 &&
+                   *size_r > stream->try_alloc_limit)
+                       *size_r = stream->try_alloc_limit;
+       } else {
+               size_t buffered = (stream->pos - stream->skip);
+
+               if (buffered >= stream->data_limit)
+                       *size_r = 0;
+               else
+                       *size_r = stream->data_limit - buffered;
+       }
        return *size_r > 0;
 }
 

diff --git a/src/lib/istream.h b/src/lib/istream.h
index ed3d793cd095c01302a9b1cbc7b4080bf8f119cf..671d3ffdbd79568400b5a2042eb1ce5eb8576f16 100644 (file)
--- a/src/lib/istream.h
+++ b/src/lib/istream.h
@@ -228,6 +228,14 @@ i_stream_read_more(struct istream *stream, const unsigned char **data_r,
        i_assert(ret != -2); /* stream must have space for at least 1 byte */
        return ret;
 }
+/* Like i_stream_read_more(), but tries to avoid buffering more than the
+   indicated limit. Use this function to prevent growing the stream buffer
+   beyond what the application is willing to read immediately. Since this
+   function doesn't fully prevent buffering beyond the limit, the amount of data
+   actually buffered can exceed the limit. However, *size_r will allways be <=
+   limit to avoid confusion. */
+int i_stream_read_limited(struct istream *stream, const unsigned char **data_r,
+                         size_t *size_r, size_t limit);
 /* Return the timestamp when istream last successfully read something.
    The timestamp is 0 if nothing has ever been read. */
 void i_stream_get_last_read_time(struct istream *stream, struct timeval *tv_r);