From: Stephan Bosch <stephan.bosch@open-xchange.com>
Date: Fri, 16 Apr 2021 10:25:43 +0000 (+0200)
Subject: lib: istream - Add i_stream_read_limited().
X-Git-Tag: 2.3.18~288
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=377f1cbe883eb65aca96e5e244100cfe1ffb7236;p=thirdparty%2Fdovecot%2Fcore.git

lib: istream - Add i_stream_read_limited().
---

diff --git a/src/lib/istream-private.h b/src/lib/istream-private.h
index c1180a5746..e8542b264e 100644
--- 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 1f2728ff28..bf539c7e3a 100644
--- 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 ed3d793cd0..671d3ffdbd 100644
--- 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);