Fix onWritable calling markDone resulting in potential heap-use-after-free

src/HttpContext.h

@@ -306,7 +306,7 @@ private:
 
                 /* We expect the developer to return whether or not write was successful (true).
                  * If write was never called, the developer should still return true so that we may drain. */
-                bool success = httpResponseData->onWritable(httpResponseData->offset);
+                bool success = httpResponseData->callOnWritable(httpResponseData->offset);
 
                 /* The developer indicated that their onWritable failed. */
                 if (!success) {

src/HttpResponse.h

@@ -71,16 +71,6 @@ private:
         Super::write(buf, length);
     }
 
-    /* When we are done with a response we mark it like so */
-    void markDone(HttpResponseData<SSL> *httpResponseData) {
-        httpResponseData->onAborted = nullptr;
-        /* Also remove onWritable so that we do not emit when draining behind the scenes. */
-        httpResponseData->onWritable = nullptr;
-
-        /* We are done with this request */
-        httpResponseData->state &= ~HttpResponseData<SSL>::HTTP_RESPONSE_PENDING;
-    }
-
     /* Called only once per request */
     void writeMark() {
         /* You can disable this altogether */
@@ -137,7 +127,7 @@ private:
             /* Terminating 0 chunk */
             Super::write("\r\n0\r\n\r\n", 7);
 
-            markDone(httpResponseData);
+            httpResponseData->markDone();
 
             /* tryEnd can never fail when in chunked mode, since we do not have tryWrite (yet), only write */
             Super::timeout(HTTP_TIMEOUT_S);
@@ -188,7 +178,7 @@ private:
 
             /* Remove onAborted function if we reach the end */
             if (httpResponseData->offset == totalSize) {
-                markDone(httpResponseData);
+                httpResponseData->markDone();
             }
 
             return success;

src/HttpResponseData.h

@@ -32,6 +32,36 @@ template <bool SSL>
 struct HttpResponseData : AsyncSocketData<SSL>, HttpParser {
     template <bool> friend struct HttpResponse;
     template <bool> friend struct HttpContext;
+
+    /* When we are done with a response we mark it like so */
+    void markDone() {
+        onAborted = nullptr;
+        /* Also remove onWritable so that we do not emit when draining behind the scenes. */
+        onWritable = nullptr;
+
+        /* We are done with this request */
+        state &= ~HttpResponseData<SSL>::HTTP_RESPONSE_PENDING;
+    }
+
+    /* Caller of onWritable. It is possible onWritable calls markDone so we need to borrow it. */
+    bool callOnWritable(uintmax_t offset) {
+        /* Borrow real onWritable */
+        MoveOnlyFunction<bool(uintmax_t)> borrowedOnWritable = std::move(onWritable);
+
+        /* Set onWritable to placeholder */
+        onWritable = [](uintmax_t) {return true;};
+
+        /* Run borrowed onWritable */
+        bool ret = borrowedOnWritable(offset);
+
+        /* If we still have onWritable (the placeholder) then move back the real one */
+        if (onWritable) {
+            /* We haven't reset onWritable, so give it back */
+            onWritable = std::move(borrowedOnWritable);
+        }
+
+        return ret;
+    }
 private:
     /* Bits of status */
     enum {