Handle SStream overflow by truncating appends (#2921)

Co-authored-by: pancake <[email protected]>

SStream.c

@@ -18,6 +18,47 @@
 #include "cs_priv.h"
 #include "utils.h"
 
+static size_t SStream_remaining(SStream *ss)
+{
+	assert(ss);
+	if (ss->index >= SSTREAM_BUF_LEN) {
+		ss->index = SSTREAM_BUF_LEN - 1;
+		ss->buffer[ss->index] = '\0';
+		return 0;
+	}
+	return SSTREAM_BUF_LEN - ss->index - 1;
+}
+
+static void SStream_concat_raw(SStream *ss, const char *s, size_t len)
+{
+	size_t remaining = SStream_remaining(ss);
+	if (len > remaining) {
+		len = remaining;
+	}
+	if (len > 0) {
+		memcpy(ss->buffer + ss->index, s, len);
+		ss->index += len;
+	}
+	ss->buffer[ss->index] = '\0';
+}
+
+static void SStream_concat1_raw(SStream *ss, char c)
+{
+	if (SStream_remaining(ss) == 0) {
+		return;
+	}
+	ss->buffer[ss->index] = c;
+	ss->index++;
+	ss->buffer[ss->index] = '\0';
+}
+
+static void SStream_close_markup(SStream *ss)
+{
+	if (ss->markup_stream && ss->prefixed_by_markup) {
+		SStream_concat1_raw(ss, '>');
+	}
+}
+
 void SStream_Init(SStream *ss)
 {
 	assert(ss);
@@ -177,19 +218,9 @@ void SStream_concat0(SStream *ss, const char *s)
 	SSTREAM_RETURN_IF_CLOSED(ss);
 	if (s[0] == '\0')
 		return;
-	unsigned int len = (unsigned int)strlen(s);
 
-	SSTREAM_OVERFLOW_CHECK(ss, len);
-
-	memcpy(ss->buffer + ss->index, s, len);
-	ss->index += len;
-	ss->buffer[ss->index] = '\0';
-	if (ss->markup_stream && ss->prefixed_by_markup) {
-		SSTREAM_OVERFLOW_CHECK(ss, 1);
-		ss->buffer[ss->index] = '>';
-		ss->index += 1;
-		ss->buffer[ss->index] = '\0';
-	}
+	SStream_concat_raw(ss, s, strlen(s));
+	SStream_close_markup(ss);
 #else
 	ss->buffer[ss->index] = '\0';
 #endif
@@ -206,16 +237,8 @@ void SStream_concat1(SStream *ss, const char c)
 	if (c == '\0')
 		return;
 
-	SSTREAM_OVERFLOW_CHECK(ss, 1);
-
-	ss->buffer[ss->index] = c;
-	ss->index++;
-	ss->buffer[ss->index] = '\0';
-	if (ss->markup_stream && ss->prefixed_by_markup) {
-		SSTREAM_OVERFLOW_CHECK(ss, 1);
-		ss->buffer[ss->index] = '>';
-		ss->index++;
-	}
+	SStream_concat1_raw(ss, c);
+	SStream_close_markup(ss);
 #else
 	ss->buffer[ss->index] = '\0';
 #endif
@@ -231,21 +254,24 @@ void SStream_concat(SStream *ss, const char *fmt, ...)
 	SSTREAM_RETURN_IF_CLOSED(ss);
 	va_list ap;
 	int ret;
-
-	va_start(ap, fmt);
-	ret = cs_vsnprintf(ss->buffer + ss->index,
-			   sizeof(ss->buffer) - (ss->index + 1), fmt, ap);
-	va_end(ap);
-	if (ret < 0) {
+	size_t remaining = SStream_remaining(ss);
+	if (remaining == 0) {
 		return;
 	}
-	SSTREAM_OVERFLOW_CHECK(ss, ret);
-	ss->index += ret;
-	if (ss->markup_stream && ss->prefixed_by_markup) {
-		SSTREAM_OVERFLOW_CHECK(ss, 1);
-		ss->buffer[ss->index] = '>';
-		ss->index += 1;
+
+	va_start(ap, fmt);
+	ret = cs_vsnprintf(ss->buffer + ss->index, remaining + 1, fmt, ap);
+	va_end(ap);
+	if (ret < 0) {
+		ss->buffer[ss->index] = '\0';
+		return;
 	}
+	if ((size_t)ret > remaining) {
+		ss->index += remaining;
+	} else {
+		ss->index += (size_t)ret;
+	}
+	SStream_close_markup(ss);
 #else
 	ss->buffer[ss->index] = '\0';
 #endif

SStream.h

@@ -25,14 +25,6 @@ typedef struct SStream {
 	bool unsigned_num; ///< Print all numbers as unsigned. Set with CS_OPT_UNSIGNED.
 } SStream;
 
-#define SSTREAM_OVERFLOW_CHECK(OS, len) \
-	do { \
-		if (OS->index + len + 1 > SSTREAM_BUF_LEN) { \
-			fprintf(stderr, "Buffer overflow caught!\n"); \
-			return; \
-		} \
-	} while (0)
-
 #define SSTREAM_RETURN_IF_CLOSED(OS) \
 	do { \
 		if (OS->is_closed) \

tests/unit/sstream.c

@@ -59,6 +59,43 @@ static bool test_overflow_check()
 	return true;
 }
 
+static bool test_truncating_appends()
+{
+	printf("Test test_truncating_appends\n");
+
+	char too_long[SSTREAM_BUF_LEN + 1] = { 0 };
+	memset(too_long, 'A', SSTREAM_BUF_LEN);
+	char full[SSTREAM_BUF_LEN] = { 0 };
+	memset(full, 'A', sizeof(full) - 1);
+
+	SStream OS = { 0 };
+	SStream_Init(&OS);
+	SStream_concat0(&OS, too_long);
+	CHECK_OS_EQUAL_RET_FALSE(OS, full);
+	CHECK_INT_EQUAL_RET_FALSE(OS.index, SSTREAM_BUF_LEN - 1);
+	SStream_concat0(&OS, "B");
+	CHECK_OS_EQUAL_RET_FALSE(OS, full);
+	CHECK_INT_EQUAL_RET_FALSE(OS.index, SSTREAM_BUF_LEN - 1);
+
+	SStream_Flush(&OS, NULL);
+	SStream_concat(&OS, "%s", too_long);
+	CHECK_OS_EQUAL_RET_FALSE(OS, full);
+	CHECK_INT_EQUAL_RET_FALSE(OS.index, SSTREAM_BUF_LEN - 1);
+
+	char almost_full[SSTREAM_BUF_LEN] = { 0 };
+	memset(almost_full, 'C', SSTREAM_BUF_LEN - 2);
+	SStream_Flush(&OS, NULL);
+	SStream_concat0(&OS, almost_full);
+	SStream_concat1(&OS, 'D');
+	almost_full[SSTREAM_BUF_LEN - 2] = 'D';
+	CHECK_OS_EQUAL_RET_FALSE(OS, almost_full);
+	CHECK_INT_EQUAL_RET_FALSE(OS.index, SSTREAM_BUF_LEN - 1);
+	SStream_concat1(&OS, 'E');
+	CHECK_OS_EQUAL_RET_FALSE(OS, almost_full);
+	CHECK_INT_EQUAL_RET_FALSE(OS.index, SSTREAM_BUF_LEN - 1);
+	return true;
+}
+
 static bool test_markup_os()
 {
 	printf("Test test_markup_os\n");
@@ -670,6 +707,7 @@ int main()
 	bool result = true;
 	result &= test_markup_os();
 	result &= test_overflow_check();
+	result &= test_truncating_appends();
 	result &= test_printint8();
 	result &= test_printint16();
 	result &= test_printint32();