Give HttpResponse the ability to send chunked responses

src/httpresponse.cpp

@@ -48,8 +48,9 @@ class HttpResponseExtendedWrite : public ExtendedWrite {
 public:
   HttpResponseExtendedWrite(std::shared_ptr<HttpResponse> pParent,
                             uv_stream_t* pHandle,
-                            std::shared_ptr<DataSource> pDataSource) :
-      ExtendedWrite(pHandle, pDataSource), _pParent(pParent) {}
+                            std::shared_ptr<DataSource> pDataSource,
+                            bool chunked) :
+      ExtendedWrite(pHandle, pDataSource, chunked), _pParent(pParent) {}
 
   void onWriteComplete(int status) {
     delete this;
@@ -62,23 +63,32 @@ void HttpResponse::writeResponse() {
   // TODO: Optimize
   std::ostringstream response(std::ios_base::binary);
   response << "HTTP/1.1 " << _statusCode << " " << _status << "\r\n";
-  bool hasContentLengthHeader = false;
+  std::string contentLength;
   for (ResponseHeaders::const_iterator it = _headers.begin();
      it != _headers.end();
      it++) {
-    response << it->first << ": " << it->second << "\r\n";
-
     if (strcasecmp(it->first.c_str(), "Content-Length") == 0) {
-      hasContentLengthHeader = true;
+      contentLength = it->second;
+    } else {
+      response << it->first << ": " << it->second << "\r\n";
     }
   }
 
-  // Some valid responses (such as HTTP 204 and 304) must not set this header,
-  // since they can't have a body.
-  //
-  // See: https://tools.ietf.org/html/rfc7230#section-3.3.2
-  if (_pBody != nullptr && !hasContentLengthHeader) {
+  if (_statusCode == 101) {
+    // HTTP 101 must not set this header, even if there *is* body data (which is
+    // actually not a true HTTP body, but instead, just the first bytes for the
+    // switched-to protocol)
+  } else if (_chunked) {
+    response << "Transfer-Encoding: chunked\r\n";
+  } else if (!contentLength.empty()) {
+    response << "Content-Length: " << contentLength << "\r\n";
+  } else if (_pBody != nullptr) {
     response << "Content-Length: " << _pBody->size() << "\r\n";
+  } else {
+    // Some valid responses (such as HTTP 204 and 304) must not set this header,
+    // since they can't have a body.
+    //
+    // See: https://tools.ietf.org/html/rfc7230#section-3.3.2
   }
 
   response << "\r\n";
@@ -127,7 +137,7 @@ void HttpResponse::onResponseWritten(int status) {
 
   if (_pBody != NULL) {
     HttpResponseExtendedWrite* pResponseWrite = new HttpResponseExtendedWrite(
-      shared_from_this(), _pRequest->handle(), _pBody);
+      shared_from_this(), _pRequest->handle(), _pBody, this->_chunked);
     pResponseWrite->begin();
   }
 }

src/httpresponse.h

@@ -18,6 +18,7 @@ class HttpResponse : public std::enable_shared_from_this<HttpResponse>  {
   std::vector<char> _responseHeader;
   std::shared_ptr<DataSource> _pBody;
   bool _closeAfterWritten;
+  bool _chunked;
 
 public:
   HttpResponse(std::shared_ptr<HttpRequest> pRequest,
@@ -28,7 +29,8 @@ class HttpResponse : public std::enable_shared_from_this<HttpResponse>  {
       _statusCode(statusCode),
       _status(status),
       _pBody(pBody),
-      _closeAfterWritten(false)
+      _closeAfterWritten(false),
+      _chunked(false)
   {
     _headers.push_back(std::make_pair("Date", http_date_string(time(NULL))));
   }

src/uvutil.cpp

@@ -4,17 +4,46 @@
 
 
 class WriteOp {
-public:
-  uv_write_t handle;
+private:
   ExtendedWrite* pParent;
   uv_buf_t buffer;
+  // Bytes to write before writing the buffer
+  std::vector<char> prefix;
+  // Bytes to write after writing the buffer
+  std::vector<char> suffix;
+public:
+  uv_write_t handle;
 
   WriteOp(ExtendedWrite* parent, uv_buf_t data)
         : pParent(parent), buffer(data) {
     memset(&handle, 0, sizeof(uv_write_t));
     handle.data = this;
   }
 
+  void setPrefix(const char* data, size_t len) {
+    prefix.clear();
+    std::copy(data, data + len, std::back_inserter(prefix));
+  }
+
+  void setSuffix(const char* data, size_t len) {
+    suffix.clear();
+    std::copy(data, data + len, std::back_inserter(suffix));
+  }
+
+  std::vector<uv_buf_t> bufs() {
+    std::vector<uv_buf_t> res;
+    if (prefix.size() > 0) {
+      uv_buf_t buf_prefix = {&prefix[0], prefix.size()};
+      res.push_back(buf_prefix);
+    }
+    res.push_back(buffer);
+    if (suffix.size() > 0) {
+      uv_buf_t buf_suffix = {&suffix[0], suffix.size()};
+      res.push_back(buf_suffix);
+    }
+    return res;
+  }
+
   void end() {
     ASSERT_BACKGROUND_THREAD()
     pParent->_pDataSource->freeData(buffer);
@@ -71,6 +100,9 @@ void ExtendedWrite::begin() {
   next();
 }
 
+const std::string CRLF = "\r\n";
+const std::string TRAILER = "0\r\n\r\n";
+
 void ExtendedWrite::next() {
   ASSERT_BACKGROUND_THREAD()
   if (_errored) {
@@ -80,6 +112,13 @@ void ExtendedWrite::next() {
     }
     return;
   }
+  if (_completed) {
+    if (_activeWrites == 0) {
+      _pDataSource->close();
+      onWriteComplete(0);
+    }
+    return;
+  }
 
   uv_buf_t buf;
   try {
@@ -93,15 +132,41 @@ void ExtendedWrite::next() {
     return;
   }
   if (buf.len == 0) {
-    _pDataSource->freeData(buf);
-    if (_activeWrites == 0) {
-      _pDataSource->close();
-      onWriteComplete(0);
-    }
-    return;
+    // No more data is going to come.
+    // Ensure future calls to next() results in disposal (assuming that all
+    // outstanding writes are done).
+    _completed = true;
   }
+
   WriteOp* pWriteOp = new WriteOp(this, buf);
-  uv_write(&pWriteOp->handle, _pHandle, &pWriteOp->buffer, 1, &writecb);
+  if (this->_chunked) {
+    if (buf.len == 0) {
+      // In chunked mode, the last chunk must be followed by one more "\r\n".
+      pWriteOp->setSuffix(TRAILER.c_str(), TRAILER.length());
+    } else {
+      // In chunked mode, data chunks must be preceded by 1) the number of bytes
+      // in the chunk, as a hexadecimal string; and 2) "\r\n"; and succeeded by
+      // another "\r\n"
+      char prefix[16];
+      memset(prefix, 0, sizeof(prefix));
+      int len = snprintf(prefix, sizeof(prefix), "%lX\r\n", buf.len);
+      pWriteOp->setPrefix(prefix, len);
+      pWriteOp->setSuffix(CRLF.c_str(), CRLF.length());
+    }
+  } else {
+    // Non-chunked mode
+    if (buf.len == 0) {
+      // We've reached the end of the response body. Technically there's nothing
+      // to uv_write; we're passing a uv_buf_t, but it's actually empty. The
+      // reason we're doing it anyway is to avoid having two cleanup paths, one
+      // for chunked and one for non-chunked; it's hard enough to reason about
+      // as it is.
+    } else {
+      // This is the simple/common case; we're about to write some data to the
+      // socket, then we'll come back and see if there's more to write.
+    }
+  }
   _activeWrites++;
-  //uv_write(pReq)
+  auto op_bufs = pWriteOp->bufs();
+  uv_write(&pWriteOp->handle, _pHandle, &op_bufs[0], op_bufs.size(), &writecb);
 }

src/uvutil.h

@@ -67,14 +67,16 @@ class InMemoryDataSource : public DataSource {
 // not to buffer too much data in memory (happens when you try
 // to write too much data to a slow uv_stream_t).
 class ExtendedWrite {
+  bool _chunked;
   int _activeWrites;
   bool _errored;
+  bool _completed;
   uv_stream_t* _pHandle;
   std::shared_ptr<DataSource> _pDataSource;
 
 public:
-  ExtendedWrite(uv_stream_t* pHandle, std::shared_ptr<DataSource> pDataSource)
-      : _activeWrites(0), _errored(false), _pHandle(pHandle),
+  ExtendedWrite(uv_stream_t* pHandle, std::shared_ptr<DataSource> pDataSource, bool chunked)
+      : _chunked(chunked), _activeWrites(0), _errored(false), _completed(false), _pHandle(pHandle),
         _pDataSource(pDataSource) {}
   virtual ~ExtendedWrite() {}