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

2 files changed, 18 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;
diff --git a/misc/iobuf.h b/misc/iobuf.h
index ef213d1..bdaf6eb 100644
--- a/misc/iobuf.h
+++ b/misc/iobuf.h
@@ -5,6 +5,7 @@ struct io_buffer {
     size_t buf_off;
     size_t buf_len;
     char *buf;
+    int error;
 };
 
 int buf_init(struct io_buffer *iob, size_t size);