Properly freeing libpq results on errors

2 files changed, 47 insertions, 1 deletions

diff --git a/src/c/urweb.c b/src/c/urweb.c
index d4fd1844..039ba119 100644
--- a/src/c/urweb.c
+++ b/src/c/urweb.c
@@ -15,6 +15,11 @@ typedef struct regions {
   struct regions *next;
 } regions;
 
+typedef struct {
+  void (*func)(void*);
+  void *arg;
+} cleanup;
+
 struct uw_context {
   char *page, *page_front, *page_back;
   char *heap, *heap_front, *heap_back;
@@ -26,6 +31,8 @@ struct uw_context {
 
   regions *regions;
 
+  cleanup *cleanup, *cleanup_front, *cleanup_back;
+
   char error_message[ERROR_BUF_LEN];
 };
 
@@ -46,6 +53,8 @@ uw_context uw_init(size_t page_len, size_t heap_len) {
 
   ctx->regions = NULL;
 
+  ctx->cleanup_front = ctx->cleanup_back = ctx->cleanup = malloc(0);
+
   ctx->error_message[0] = 0;
 
   return ctx;
@@ -63,6 +72,7 @@ void uw_free(uw_context ctx) {
   free(ctx->page);
   free(ctx->heap);
   free(ctx->inputs);
+  free(ctx->cleanup);
   free(ctx);
 }
 
@@ -70,6 +80,7 @@ void uw_reset_keep_request(uw_context ctx) {
   ctx->page_front = ctx->page;
   ctx->heap_front = ctx->heap;
   ctx->regions = NULL;
+  ctx->cleanup_front = ctx->cleanup;
 
   ctx->error_message[0] = 0;
 }
@@ -78,6 +89,7 @@ void uw_reset_keep_error_message(uw_context ctx) {
   ctx->page_front = ctx->page;
   ctx->heap_front = ctx->heap;
   ctx->regions = NULL;
+  ctx->cleanup_front = ctx->cleanup;
 }
 
 void uw_reset(uw_context ctx) {
@@ -107,14 +119,46 @@ failure_kind uw_begin(uw_context ctx, char *path) {
 }
 
 __attribute__((noreturn)) void uw_error(uw_context ctx, failure_kind fk, const char *fmt, ...) {
+  cleanup *cl;
+
   va_list ap;
   va_start(ap, fmt);
 
   vsnprintf(ctx->error_message, ERROR_BUF_LEN, fmt, ap);
 
+  for (cl = ctx->cleanup; cl < ctx->cleanup_front; ++cl)
+    cl->func(cl->arg);
+
+  ctx->cleanup_front = ctx->cleanup;
+
   longjmp(ctx->jmp_buf, fk);
 }
 
+void uw_push_cleanup(uw_context ctx, void (*func)(void *), void *arg) {
+  if (ctx->cleanup_front >= ctx->cleanup_back) {
+    int len = ctx->cleanup_back - ctx->cleanup, newLen;
+    if (len == 0)
+      newLen = 1;
+    else
+      newLen *= 2;
+    ctx->cleanup = realloc(ctx->cleanup, newLen);
+    ctx->cleanup_front = ctx->cleanup + len;
+    ctx->cleanup_back = ctx->cleanup + newLen;
+  }
+
+  ctx->cleanup_front->func = func;
+  ctx->cleanup_front->arg = arg;
+  ++ctx->cleanup_front;
+}
+
+void uw_pop_cleanup(uw_context ctx) {
+  if (ctx->cleanup_front == ctx->cleanup)
+    uw_error(ctx, FATAL, "Attempt to pop from empty cleanup action stack");
+
+  --ctx->cleanup_front;
+  ctx->cleanup_front->func(ctx->cleanup_front->arg);
+}
+
 char *uw_error_message(uw_context ctx) {
   return ctx->error_message;
 }
diff --git a/src/cjr_print.sml b/src/cjr_print.sml
index 7d74376e..26f6149e 100644
--- a/src/cjr_print.sml
+++ b/src/cjr_print.sml
@@ -850,6 +850,8 @@ fun p_exp' par env (e, loc) =
 
                  string "uw_end_region(ctx);",
                  newline,
+                 string "uw_push_cleanup(ctx, (void (*)(void *))PQclear, res);",
+                 newline,
                  string "n = PQntuples(res);",
                  newline,
                  string "for (i = 0; i < n; ++i) {",
@@ -906,7 +908,7 @@ fun p_exp' par env (e, loc) =
                  string "}",
                  newline,
                  newline,
-                 string "PQclear(res);",
+                 string "uw_pop_cleanup(ctx);",
                  newline,
                  if wontLeakAnything then
                      box [string "uw_end_region(ctx);",