Make the buf api safe to continue to use even if an error has occured.

1 files changed, 17 insertions, 8 deletions

diff --git a/misc/iobuf.c b/misc/iobuf.c
index a7304c6..cf5bc73 100644
--- a/misc/iobuf.c
+++ b/misc/iobuf.c
@@ -14,20 +14,25 @@ buf_init(struct io_buffer *iob, size_t size)
 {
     iob->buf_off = 0;
     iob->buf_len = size;
+    iob->error = 0;
     iob->buf = malloc(size);
-    if (iob->buf == NULL)
+    if (iob->buf == NULL) {
+        iob->error = ENOMEM;
         return ENOMEM;
-    else
+    } else
         return 0;
 }
 
 int
 buf_grow(struct io_buffer *iob, size_t addlen)
 {
+    if (iob->error)
+        return iob->error;
     char *nbuf = realloc(iob->buf, iob->buf_len + addlen);
-    if (nbuf == NULL)
+    if (nbuf == NULL) {
+        iob->error = ENOMEM;
         return ENOMEM;
-    else {
+    } else {
         iob->buf = nbuf;
         iob->buf_len += addlen;
         return 0;
@@ -37,13 +42,15 @@ buf_grow(struct io_buffer *iob, size_t addlen)
 int
 buf_print(struct io_buffer *iob, const char *fmt, ...)
 {
+    if (iob->error)
+        return iob->error;
     int np;
     va_list ap;
     va_start(ap, fmt);
     np = vsnprintf(NULL, 0, fmt, ap);
     va_end(ap);
-    while (np + 1 > iob->buf_len - iob->buf_off)
-        if (buf_grow(iob, GROWLEN) != 0)
+    if (np + 1 > iob->buf_len - iob->buf_off)
+        if (buf_grow(iob, (1 + (np + 1) / GROWLEN) * GROWLEN) != 0)
             return ENOMEM;
     va_start(ap, fmt);
     vsnprintf(iob->buf + iob->buf_off, np + 1, fmt, ap);
@@ -55,8 +62,10 @@ buf_print(struct io_buffer *iob, const char *fmt, ...)
 int
 buf_write(struct io_buffer *iob, const void *buf, size_t len)
 {
-    while (iob->buf_len - iob->buf_off < len)
-        if (buf_grow(iob, GROWLEN) != 0)
+    if (iob->error)
+        return iob->error;
+    if (len > iob->buf_len - iob->buf_off)
+        if (buf_grow(iob, (1 + len / GROWLEN) * GROWLEN) != 0)
             return ENOMEM;
     bcopy(buf, iob->buf + iob->buf_off, len);
     iob->buf_off += len;