httpd: rework httpd to use auto-free for memstream

This helps us avoid memory leaks after OOM.

1 files changed, 77 insertions, 77 deletions

diff --git a/mwrap_httpd.h b/mwrap_httpd.h
index 1f706ea..63dfcfa 100644
--- a/mwrap_httpd.h
+++ b/mwrap_httpd.h
@@ -36,9 +36,10 @@ enum mw_qev {
         MW_QEV_WR = POLLOUT
 };
 
-struct mw_membuf { /* for open_memstream */
+struct mw_fbuf {
         char *ptr;
         size_t len;
+        FILE *fp;
 };
 
 struct mw_wbuf { /* for response headers + bodies */
@@ -78,8 +79,7 @@ struct mw_h1d { /* the daemon + listener, a singleton */
         unsigned alive;
         struct cds_list_head conn; /* <=> mw_h1.nd */
         /* use open_memstream + fwrite to implement a growing pollfd array */
-        FILE *pfp; /* see poll_add, poll_detach */
-        struct mw_membuf pbuf; /* pollfd vector */
+        struct mw_fbuf pb; /* pollfd vector */
         pthread_t tid;
         size_t pid_len;
         char pid_str[10];
@@ -202,60 +202,70 @@ static enum mw_qev h1_send_flush(struct mw_h1 *h1)
         return h1->persist ? MW_QEV_RD : h1_close(h1);
 }
 
-static FILE *memstream_new(struct mw_membuf *mb)
+static FILE *fbuf_init(struct mw_fbuf *fb)
 {
-        FILE *fp = open_memstream(&mb->ptr, &mb->len);
-        if (!fp) fprintf(stderr, "open_memstream: %m\n");
-        return fp;
+        fb->ptr = NULL;
+        fb->fp = open_memstream(&fb->ptr, &fb->len);
+        if (!fb->fp) fprintf(stderr, "open_memstream: %m\n");
+        return fb->fp;
 }
 
-static FILE *wbuf_new(struct mw_membuf *mb)
+static FILE *wbuf_init(struct mw_fbuf *fb)
 {
         static const struct mw_wbuf pad;
-        FILE *fp = memstream_new(mb);
-        if (fp) /* pad space is populated before h1_send_flush */
-                fwrite(&pad, 1, sizeof(pad), fp);
-        return fp;
+        if (fbuf_init(fb)) /* pad space is populated before h1_send_flush */
+                fwrite(&pad, 1, sizeof(pad), fb->fp);
+        return fb->fp;
 }
 
-static int err_close(FILE *fp)
+static int fbuf_close(struct mw_fbuf *fb)
 {
-        int e = ferror(fp) | fclose(fp);
+        int e = ferror(fb->fp) | fclose(fb->fp);
+        fb->fp = NULL;
         if (e) fprintf(stderr, "ferror|fclose: %m\n");
         return e;
 }
 
+/* supported by modern gcc + clang */
+#define AUTO_CLOFREE __attribute__((__cleanup__(cleanup_clofree)))
+static void cleanup_clofree(void *ptr)
+{
+        struct mw_fbuf *fb = ptr;
+        if (fb->fp) fclose(fb->fp);
+        free(fb->ptr);
+}
+
 static enum mw_qev h1_res_oneshot(struct mw_h1 *h1, const char *buf, size_t len)
 {
-        struct mw_membuf mb;
-        FILE *fp = wbuf_new(&mb);
+        struct mw_fbuf fb;
 
-        if (!fp)
+        if (!wbuf_init(&fb))
                 return h1_close(h1);
-        fwrite(buf, 1, len, fp);
-        if (err_close(fp))
+
+        fwrite(buf, 1, len, fb.fp);
+        if (fbuf_close(&fb))
                 return h1_close(h1);
 
-        /* fill in the zero padding we added at wbuf_new */
+        /* fill in the zero padding we added at wbuf_init */
         mwrap_assert(!h1->wbuf);
-        struct mw_wbuf *wbuf = h1->wbuf = (struct mw_wbuf *)mb.ptr;
+        struct mw_wbuf *wbuf = h1->wbuf = (struct mw_wbuf *)fb.ptr;
         wbuf->iov_nr = 1;
-        wbuf->iov[0].iov_len = mb.len - sizeof(*wbuf);
+        wbuf->iov[0].iov_len = fb.len - sizeof(*wbuf);
         wbuf->iov[0].iov_base = wbuf->bytes;
         return h1_send_flush(h1);
 }
 
 #define FPUTS(STR, fp) fwrite(STR, sizeof(STR) - 1, 1, fp)
-static enum mw_qev h1_200(struct mw_h1 *h1, FILE *fp, struct mw_membuf *mb)
+static enum mw_qev h1_200(struct mw_h1 *h1, struct mw_fbuf *fb)
 {
         /*
          * the HTTP header goes at the END of the body buffer,
          * we'll rely on iovecs via sendmsg(2) to reorder and clamp it
          */
-        off_t clen = ftello(fp);
+        off_t clen = ftello(fb->fp);
         if (clen < 0) {
                 fprintf(stderr, "ftello: %m\n");
-                fclose(fp);
+                fbuf_close(fb);
                 return h1_close(h1);
         }
         clen -= sizeof(struct mw_wbuf);
@@ -265,18 +275,18 @@ static enum mw_qev h1_200(struct mw_h1 *h1, FILE *fp, struct mw_membuf *mb)
                 "Pragma: no-cache\r\n"
                 "Cache-Control: no-cache, max-age=0, must-revalidate\r\n"
                 "Content-Type: text/html; charset=UTF-8\r\n"
-                "Content-Length: ", fp);
-        fprintf(fp, "%zu", (size_t)clen);
-        FPUTS("\r\n\r\n", fp);
+                "Content-Length: ", fb->fp);
+        fprintf(fb->fp, "%zu", (size_t)clen);
+        FPUTS("\r\n\r\n", fb->fp);
 
-        if (err_close(fp))
+        if (fbuf_close(fb))
                 return h1_close(h1);
 
-        /* fill in the zero-padding we added at wbuf_new */
+        /* fill in the zero-padding we added at wbuf_init */
         mwrap_assert(!h1->wbuf);
-        struct mw_wbuf *wbuf = h1->wbuf = (struct mw_wbuf *)mb->ptr;
+        struct mw_wbuf *wbuf = h1->wbuf = (struct mw_wbuf *)fb->ptr;
         wbuf->iov_nr = 2;
-        wbuf->iov[0].iov_len = mb->len - ((size_t)clen + sizeof(*wbuf));
+        wbuf->iov[0].iov_len = fb->len - ((size_t)clen + sizeof(*wbuf));
         wbuf->iov[0].iov_base = wbuf->bytes + (size_t)clen;
         wbuf->iov[1].iov_len = clen;
         wbuf->iov[1].iov_base = wbuf->bytes;
@@ -447,9 +457,8 @@ static off_t write_loc_name(FILE *fp, const struct src_loc *l)
 
 static struct h1_src_loc *accumulate(unsigned long min, size_t *hslc, FILE *lp)
 {
-        struct mw_membuf mb;
-        FILE *fp = memstream_new(&mb);
-        if (!fp) return NULL;
+        struct mw_fbuf fb;
+        if (!fbuf_init(&fb)) return NULL;
         rcu_read_lock();
         struct cds_lfht *t = CMM_LOAD_SHARED(totals);
         struct cds_lfht_iter iter;
@@ -471,17 +480,17 @@ static struct h1_src_loc *accumulate(unsigned long min, size_t *hslc, FILE *lp)
                 hsl.max_life = uatomic_read(&l->max_lifespan);
                 hsl.sl = l;
                 hsl.lname_len = write_loc_name(lp, l);
-                fwrite(&hsl, sizeof(hsl), 1, fp);
+                fwrite(&hsl, sizeof(hsl), 1, fb.fp);
         }
         rcu_read_unlock();
 
         struct h1_src_loc *hslv;
-        if (err_close(fp)) {
+        if (fbuf_close(&fb)) {
                 hslv = NULL;
         } else {
-                *hslc = mb.len / sizeof(*hslv);
-                mwrap_assert((mb.len % sizeof(*hslv)) == 0);
-                hslv = (struct h1_src_loc *)mb.ptr;
+                *hslc = fb.len / sizeof(*hslv);
+                mwrap_assert((fb.len % sizeof(*hslv)) == 0);
+                hslv = (struct h1_src_loc *)fb.ptr;
         }
         return hslv;
 }
@@ -507,23 +516,21 @@ static enum mw_qev each_at(struct mw_h1 *h1, struct mw_h1req *h1r)
 
         if (!l) return h1_404(h1);
 
-        struct mw_membuf lname;
-        FILE *lp = memstream_new(&lname);
-        if (!lp) return h1_close(h1);
-        if (write_loc_name(lp, l) < 0) return h1_close(h1);
-        if (err_close(lp))
+        AUTO_CLOFREE struct mw_fbuf lb;
+        if (!fbuf_init(&lb)) return h1_close(h1);
+        if (write_loc_name(lb.fp, l) < 0) return h1_close(h1);
+        if (fbuf_close(&lb))
                 return h1_close(h1);
 
-        struct mw_membuf mb;
-        FILE *fp = wbuf_new(&mb);
+        struct mw_fbuf html;
+        FILE *fp = wbuf_init(&html);
         if (!fp) return h1_close(h1);
         FPUTS("<html><head><title>", fp);
-        write_html(fp, lname.ptr, lname.len);
+        write_html(fp, lb.ptr, lb.len);
         FPUTS("</title></head><body><p>live allocations at", fp);
         if (bt_req_depth) FPUTS("<br/>", fp);
         else fputc('\n', fp);
-        write_html(fp, lname.ptr, lname.len);
-        free(lname.ptr);
+        write_html(fp, lb.ptr, lb.len);
 
         show_age(fp);
         FPUTS("<table><tr><th>size</th><th>generation</th>"
@@ -539,7 +546,7 @@ static enum mw_qev each_at(struct mw_h1 *h1, struct mw_h1req *h1r)
         }
         rcu_read_unlock();
         FPUTS("</table></body></html>", fp);
-        return h1_200(h1, fp, &mb);
+        return h1_200(h1, &html);
 }
 
 /* /$PID/each/$MIN endpoint */
@@ -556,28 +563,26 @@ static enum mw_qev each_gt(struct mw_h1 *h1, struct mw_h1req *h1r,
         }
 
         size_t hslc;
-        struct mw_membuf ln;
-        FILE *lp = memstream_new(&ln);
-        if (!lp) return h1_close(h1);
-        AUTO_FREE struct h1_src_loc *hslv = accumulate(min, &hslc, lp);
+        AUTO_CLOFREE struct mw_fbuf lb;
+        if (!fbuf_init(&lb)) return h1_close(h1);
+        AUTO_FREE struct h1_src_loc *hslv = accumulate(min, &hslc, lb.fp);
         if (!hslv)
                 return h1_close(h1);
 
-        if (err_close(lp))
+        if (fbuf_close(&lb))
                 return h1_close(h1);
 
-        char *n = ln.ptr;
+        char *n = lb.ptr;
         for (size_t i = 0; i < hslc; ++i) {
                 hslv[i].loc_name = n;
                 n += hslv[i].lname_len;
-                if (hslv[i].lname_len < 0) {
-                        free(ln.ptr);
+                if (hslv[i].lname_len < 0)
                         return h1_close(h1);
-                }
         }
 
-        struct mw_membuf mb;
-        FILE *fp = wbuf_new(&mb);
+        struct mw_fbuf html;
+        FILE *fp = wbuf_init(&html);
+        if (!fp) return h1_close(h1);
         fprintf(fp, "<html><head><title>mwrap each &gt;%lu"
                 "</title></head><body><p>mwrap each &gt;%lu "
                 "(change `%lu' in URL to adjust filtering)", min, min, min);
@@ -624,9 +629,8 @@ static enum mw_qev each_gt(struct mw_h1 *h1, struct mw_h1req *h1r,
                 write_html(fp, hsl->loc_name, hsl->lname_len);
                 FPUTS("</a></td></tr>", fp);
         }
-        free(ln.ptr);
         FPUTS("</table></body></html>", fp);
-        return h1_200(h1, fp, &mb);
+        return h1_200(h1, &html);
 }
 
 static enum mw_qev h1_dispatch(struct mw_h1 *h1, struct mw_h1req *h1r)
@@ -864,13 +868,11 @@ static int poll_add(struct mw_h1d *h1d, int fd, short events)
 {
         struct pollfd pfd;
 
-        if (!h1d->pfp) {
-                h1d->pfp = memstream_new(&h1d->pbuf);
-                if (!h1d->pfp) return -1;
-        }
+        if (!h1d->pb.fp && !fbuf_init(&h1d->pb))
+                return -1;
         pfd.fd = fd;
         pfd.events = events;
-        fwrite(&pfd, 1, sizeof(pfd), h1d->pfp);
+        fwrite(&pfd, 1, sizeof(pfd), h1d->pb.fp);
         return 0; /* success */
 }
 
@@ -879,20 +881,18 @@ static struct pollfd *poll_detach(struct mw_h1d *h1d, nfds_t *nfds)
         struct pollfd *pfd = NULL; /* our return value */
 
         /* not sure how to best recover from ENOMEM errors in stdio */
-        if (h1d->pfp) {
-                if (err_close(h1d->pfp)) {
+        if (h1d->pb.fp) {
+                if (fbuf_close(&h1d->pb)) {
                         exit(EXIT_FAILURE);
                 } else {
-                        mwrap_assert(h1d->pbuf.len % sizeof(*pfd) == 0);
-                        pfd = (struct pollfd *)h1d->pbuf.ptr;
-                        *nfds = h1d->pbuf.len / sizeof(*pfd);
+                        mwrap_assert(h1d->pb.len % sizeof(*pfd) == 0);
+                        pfd = (struct pollfd *)h1d->pb.ptr;
+                        *nfds = h1d->pb.len / sizeof(*pfd);
                 }
         }
 
         /* prepare a new poll buffer the next loop */
-        h1d->pbuf.len = 0;
-        h1d->pbuf.ptr = NULL;
-        h1d->pfp = NULL;
+        memset(&h1d->pb, 0, sizeof(h1d->pb));
 
         return pfd;
 }