In the k5buf module, add a function to append formatted data to a

buffer.

ticket: 6200
status: open

git-svn-id: svn://anonsvn.mit.edu/krb5/trunk@20932 dc483132-0cff-0310-8789-dd5450dbe970

2 files changed, 114 insertions, 3 deletions

diff --git a/src/util/support/k5buf.c b/src/util/support/k5buf.c
index 747490280..23fb20365 100644
--- a/src/util/support/k5buf.c
+++ b/src/util/support/k5buf.c
@@ -115,6 +115,71 @@ void krb5int_buf_add_len(struct k5buf *buf, const char *data, size_t len)
     buf->data[buf->len] = '\0';
 }
 
+void krb5int_buf_add_fmt(struct k5buf *buf, const char *fmt, ...)
+{
+    va_list ap;
+    int r;
+    size_t remaining;
+    char *tmp;
+
+    if (buf->buftype == ERROR)
+        return;
+    remaining = buf->space - buf->len;
+
+    if (buf->buftype == FIXED) {
+        /* Format the data directly into the fixed buffer. */
+        va_start(ap, fmt);
+        r = vsnprintf(buf->data + buf->len, remaining, fmt, ap);
+        va_end(ap);
+        if (SNPRINTF_OVERFLOW(r, remaining))
+            buf->buftype = ERROR;
+        else
+            buf->len += (unsigned int) r;
+        return;
+    }
+
+    /* Optimistically format the data directly into the dynamic buffer. */
+    assert(buf->buftype == DYNAMIC);
+    va_start(ap, fmt);
+    r = vsnprintf(buf->data + buf->len, remaining, fmt, ap);
+    va_end(ap);
+    if (!SNPRINTF_OVERFLOW(r, remaining)) {
+        buf->len += (unsigned int) r;
+        return;
+    }
+
+    if (r >= 0) {
+        /* snprintf correctly told us how much space is required. */
+        if (!ensure_space(buf, r))
+            return;
+        remaining = buf->space - buf->len;
+        va_start(ap, fmt);
+        r = vsnprintf(buf->data + buf->len, remaining, fmt, ap);
+        va_end(ap);
+        if (SNPRINTF_OVERFLOW(r, remaining))  /* Shouldn't ever happen. */
+            buf->buftype = ERROR;
+        else
+            buf->len += (unsigned int) r;
+        return;
+    }
+
+    /* It's a pre-C99 snprintf implementation, or something else went
+       wrong.  Fall back to asprintf. */
+    va_start(ap, fmt);
+    r = vasprintf(&tmp, fmt, ap);
+    va_end(ap);
+    if (r < 0) {
+        buf->buftype = ERROR;
+        return;
+    }
+    if (ensure_space(buf, r)) {
+        /* Copy the temporary string into buf, including terminator. */
+        memcpy(buf->data + buf->len, tmp, r + 1);
+        buf->len += r;
+    }
+    free(tmp);
+}
+
 void krb5int_buf_truncate(struct k5buf *buf, size_t len)
 {
     if (buf->buftype == ERROR)
diff --git a/src/util/support/t_k5buf.c b/src/util/support/t_k5buf.c
index d76398a96..b7f284b46 100644
--- a/src/util/support/t_k5buf.c
+++ b/src/util/support/t_k5buf.c
@@ -81,7 +81,8 @@ static void test_realloc()
 {
     struct k5buf buf;
     char data[1024], *s;
-    ssize_t i, len;
+    ssize_t len;
+    size_t i;
 
     for (i = 0; i < sizeof(data); i++)
 	data[i] = 'a';
@@ -176,8 +177,7 @@ static void test_overflow()
 static void test_error()
 {
     struct k5buf buf;
-    char storage[1], *s;
-    ssize_t len;
+    char storage[1];
 
     /* Cause an overflow and then perform actions afterwards. */
     krb5int_buf_init_fixed(&buf, storage, sizeof(storage));
@@ -226,6 +226,51 @@ static void test_binary()
     fail_if(s[3] != 'a' || s[4] != 0 || s[5] != 'b', "binary");
 }
 
+static void test_fmt()
+{
+    struct k5buf buf;
+    char *s, storage[10], data[1024];
+    ssize_t len;
+    size_t i;
+
+    for (i = 0; i < sizeof(data) - 1; i++)
+	data[i] = 'a';
+    data[i] = '\0';
+
+    /* Format some text into a non-empty fixed buffer. */
+    krb5int_buf_init_fixed(&buf, storage, sizeof(storage));
+    krb5int_buf_add(&buf, "foo");
+    krb5int_buf_add_fmt(&buf, " %d ", 3);
+    check_buf(&buf, "fmt 1");
+    s = krb5int_buf_cstr(&buf);
+    len = krb5int_buf_len(&buf);
+    fail_if(!s || len != 6 || strcmp(s, "foo 3 ") != 0, "fmt 1");
+
+    /* Overflow the same buffer with formatted text. */
+    krb5int_buf_add_fmt(&buf, "%d%d%d%d", 1, 2, 3, 4);
+    check_buf(&buf, "fmt 2");
+    s = krb5int_buf_cstr(&buf);
+    len = krb5int_buf_len(&buf);
+    fail_if(buf.buftype != ERROR || s != NULL || len != -1, "fmt 2");
+
+    /* Format some text into a non-empty dynamic buffer. */
+    krb5int_buf_init_dynamic(&buf);
+    krb5int_buf_add(&buf, "foo");
+    krb5int_buf_add_fmt(&buf, " %d ", 3);
+    check_buf(&buf, "fmt 3");
+    s = krb5int_buf_cstr(&buf);
+    len = krb5int_buf_len(&buf);
+    fail_if(!s || len != 6 || strcmp(s, "foo 3 ") != 0, "fmt 3");
+
+    /* Format more text into the same buffer, causing a big resize. */
+    krb5int_buf_add_fmt(&buf, "%s", data);
+    check_buf(&buf, "fmt 4");
+    fail_if(buf.space != 2048, "fmt 4");
+    s = krb5int_buf_cstr(&buf);
+    len = krb5int_buf_len(&buf);
+    fail_if(!s || len != 1029 || strcmp(s + 6, data) != 0, "fmt 4");
+}
+
 int main()
 {
     test_basic();
@@ -234,5 +279,6 @@ int main()
     test_error();
     test_truncate();
     test_binary();
+    test_fmt();
     return 0;
 }