From f2198ccfac140798ee066c5a3c5235b396473205 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Wed, 29 Jul 2020 08:41:53 +0200
Subject: [PATCH] Fix >2gb streamed HTTP responses

---
 src/HttpResponse.h     | 34 ++++++++++++++++++++++++++--------
 src/HttpResponseData.h |  4 ++--
 src/Utilities.h        | 17 +++++++++++++++++
 src/WebSocketContext.h |  2 +-
 4 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/src/HttpResponse.h b/src/HttpResponse.h
index 36ae770..84a71e5 100644
--- a/src/HttpResponse.h
+++ b/src/HttpResponse.h
@@ -61,6 +61,15 @@ private:
         Super::write(buf, length);
     }
 
+    /* Write an unsigned 64-bit integer */
+    void writeUnsigned64(uint64_t value) {
+        char buf[20];
+        int length = utils::u64toa(value, buf);
+
+        /* For now we do this copy */
+        Super::write(buf, length);
+    }
+
     /* Write an unsigned 32-bit integer */
     void writeUnsigned(unsigned int value) {
         char buf[10];
@@ -90,13 +99,13 @@ private:
 
     /* Returns true on success, indicating that it might be feasible to write more data.
      * Will start timeout if stream reaches totalSize or write failure. */
-    bool internalEnd(std::string_view data, int totalSize, bool optional, bool allowContentLength = true) {
+    bool internalEnd(std::string_view data, size_t totalSize, bool optional, bool allowContentLength = true) {
         /* Write status if not already done */
         writeStatus(HTTP_200_OK);
 
         /* If no total size given then assume this chunk is everything */
         if (!totalSize) {
-            totalSize = (int) data.length();
+            totalSize = data.length();
         }
 
         HttpResponseData<SSL> *httpResponseData = getHttpResponseData();
@@ -132,7 +141,7 @@ private:
                 if (allowContentLength) {
                     /* Even zero is a valid content-length */
                     Super::write("Content-Length: ", 16);
-                    writeUnsigned(totalSize);
+                    writeUnsigned64(totalSize);
                     Super::write("\r\n\r\n", 4);
                 } else {
                     Super::write("\r\n", 2);
@@ -146,11 +155,20 @@ private:
              * if it failed to drain any prior failed header writes */
 
             /* Write as much as possible without causing backpressure */
-            auto [written, failed] = Super::write(data.data(), (int) data.length(), optional);
+            size_t written = 0;
+            bool failed = false;
+            while (written < data.length() && !failed) {
+                /* uSockets only deals with int sizes, so pass chunks of max signed int size */
+                auto writtenFailed = Super::write(data.data() + written, (int) std::min<size_t>(data.length() - written, INT_MAX), optional);
+
+                written += writtenFailed.first;
+                failed = writtenFailed.second;
+            }
+
             httpResponseData->offset += written;
 
             /* Success is when we wrote the entire thing without any failures */
-            bool success = (unsigned int) written == data.length() && !failed;
+            bool success = written == data.length() && !failed;
 
             /* If we are now at the end, start a timeout. Also start a timeout if we failed. */
             if (!success || httpResponseData->offset == totalSize) {
@@ -368,12 +386,12 @@ public:
 
     /* End the response with an optional data chunk. Always starts a timeout. */
     void end(std::string_view data = {}) {
-        internalEnd(data, (int) data.length(), false);
+        internalEnd(data, data.length(), false);
     }
 
     /* Try and end the response. Returns [true, true] on success.
      * Starts a timeout in some cases. Returns [ok, hasResponded] */
-    std::pair<bool, bool> tryEnd(std::string_view data, int totalSize = 0) {
+    std::pair<bool, bool> tryEnd(std::string_view data, size_t totalSize = 0) {
         return {internalEnd(data, totalSize, true), hasResponded()};
     }
 
@@ -446,7 +464,7 @@ public:
     }
 
     /* Attach handler for writable HTTP response */
-    HttpResponse *onWritable(fu2::unique_function<bool(int)> &&handler) {
+    HttpResponse *onWritable(fu2::unique_function<bool(size_t)> &&handler) {
         HttpResponseData<SSL> *httpResponseData = getHttpResponseData();
 
         httpResponseData->onWritable = std::move(handler);
diff --git a/src/HttpResponseData.h b/src/HttpResponseData.h
index 1e04cef..5746f4f 100644
--- a/src/HttpResponseData.h
+++ b/src/HttpResponseData.h
@@ -43,11 +43,11 @@ private:
     };
 
     /* Per socket event handlers */
-    fu2::unique_function<bool(int)> onWritable;
+    fu2::unique_function<bool(size_t)> onWritable;
     fu2::unique_function<void()> onAborted;
     fu2::unique_function<void(std::string_view, bool)> inStream; // onData
     /* Outgoing offset */
-    int offset = 0;
+    size_t offset = 0;
 
     /* Current state (content-length sent, status sent, write called, etc */
     int state = 0;
diff --git a/src/Utilities.h b/src/Utilities.h
index 8ff80ad..c9be037 100644
--- a/src/Utilities.h
+++ b/src/Utilities.h
@@ -60,6 +60,23 @@ inline int u32toa(uint32_t value, char *dst) {
     return ret;
 }
 
+inline int u64toa(uint64_t value, char *dst) {
+    char temp[20];
+    char *p = temp;
+    do {
+        *p++ = (char) ((value % 10) + '0');
+        value /= 10;
+    } while (value > 0);
+
+    int ret = (int) (p - temp);
+
+    do {
+        *dst++ = *--p;
+    } while (p != temp);
+
+    return ret;
+}
+
 }
 }
 
diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h
index f0b0a74..1b10655 100644
--- a/src/WebSocketContext.h
+++ b/src/WebSocketContext.h
@@ -316,7 +316,7 @@ private:
 
             /* We store old backpressure since it is unclear whether write drained anything,
              * however, in case of coming here with 0 backpressure we still need to emit drain event */
-            int backpressure = asyncSocket->getBufferedAmount();
+            unsigned int backpressure = asyncSocket->getBufferedAmount();
 
             /* Drain as much as possible */
             asyncSocket->write(nullptr, 0);