diff --git a/SStream.c b/SStream.c index 33946e2df..833ab5984 100644 --- a/SStream.c +++ b/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 diff --git a/SStream.h b/SStream.h index 5b74e195f..2fe33ec7d 100644 --- a/SStream.h +++ b/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) \ diff --git a/tests/unit/sstream.c b/tests/unit/sstream.c index 4dc8adfb1..d539fc666 100644 --- a/tests/unit/sstream.c +++ b/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();