[PATCH 03/19] httpd: rework httpd to use auto-free for memstream

[Eric Wong <e@80x24.org>]

This helps us avoid memory leaks after OOM.
---
 mwrap_httpd.h | 154 +++++++++++++++++++++++++-------------------------
 1 file 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;
 }