fix: AsyncAbstractResponse might loose part of send buffer

AsyncAbstractResponse::_ack could allocate temp buffer with size larger than
available sock buffer (i.e. to fit headers) and eventually lossing the remainder on transfer
due to not checking if the complete data was added to sock buff.

Refactoring code in favor of having a dedicated std::vector object acting as accumulating
buffer and more carefull control on amount of data actually copied to sockbuff

Closes #315

Author: vortigont

src/ESPAsyncWebServer.h

@@ -1336,8 +1336,10 @@ class AsyncWebServerResponse {
   bool _sendContentLength;
   bool _chunked;
   size_t _headLength;
+  // amount of data sent for content part of the response (excluding all headers)
   size_t _sentLength;
   size_t _ackedLength;
+  // amount of response bytes (including all headers) written to sockbuff for delivery
   size_t _writtenLength;
   WebResponseState _state;
 

src/WebResponseImpl.h

@@ -39,12 +39,17 @@ class AsyncAbstractResponse : public AsyncWebServerResponse {
   // in-flight queue credits
   size_t _in_flight_credit{2};
 #endif
-  String _head;
+  // buffer to accumulate all response headers
+  String _assembled_headers;
+  // amount of headers buffer sent
+  size_t _assembled_headers_sent{0};
   // Data is inserted into cache at begin().
   // This is inefficient with vector, but if we use some other container,
   // we won't be able to access it as contiguous array of bytes when reading from it,
   // so by gaining performance in one place, we'll lose it in another.
   std::vector<uint8_t> _cache;
+  // intermediate buffer to copy outbound data to
+  std::vector<uint8_t> _send_buffer;
   size_t _readDataFromCacheOrContent(uint8_t *data, const size_t len);
   size_t _fillBufferAndProcessTemplates(uint8_t *buf, size_t maxLen);
 

src/WebResponses.cpp

@@ -327,8 +327,8 @@ size_t AsyncBasicResponse::_ack(AsyncWebServerRequest *request, size_t len, uint
 
 /*
  * Abstract Response
- * */
-
+ * 
+ */
 AsyncAbstractResponse::AsyncAbstractResponse(AwsTemplateProcessor callback) : _callback(callback) {
   // In case of template processing, we're unable to determine real response size
   if (callback) {
@@ -340,7 +340,7 @@ AsyncAbstractResponse::AsyncAbstractResponse(AwsTemplateProcessor callback) : _c
 
 void AsyncAbstractResponse::_respond(AsyncWebServerRequest *request) {
   addHeader(T_Connection, T_close, false);
-  _assembleHead(_head, request->version());
+  _assembleHead(_assembled_headers, request->version());
   _state = RESPONSE_HEADERS;
   _ack(request, 0, 0);
 }
@@ -364,127 +364,139 @@ size_t AsyncAbstractResponse::_ack(AsyncWebServerRequest *request, size_t len, u
     async_ws_log_d("(chunk) out of in-flight credits");
   }
 
-  _in_flight -= (_in_flight > len) ? len : _in_flight;
-  // get the size of available sock space
+  _in_flight -= std::min(len, _in_flight);
+  // for response data we need to control the queue and in-flight fragmentation. Sending small chunks could give low latency,
+  // but flood asynctcp's queue and fragment socket buffer space for large responses.
+  // Let's ignore polled acks and acks in case when we have more in-flight data then the available socket buff space.
+  // That way we could balance on having half the buffer in-flight while another half is filling up, while minimizing events in asynctcp q
+  if (_in_flight > request->client()->space()) {
+    // async_ws_log_d("defer user call %u/%u", _in_flight, space);
+    // take the credit back since we are ignoring this ack and rely on other inflight data
+    if (len) {
+      --_in_flight_credit;
+    }
+    return 0;
+  }
 #endif
 
+  // this is not functionally needed in AsyncAbstractResponse itself, but keept for compatibility if some of the derived classes are rely on it somehow
   _ackedLength += len;
-  size_t space = request->client()->space();
 
-  size_t headLen = _head.length();
+  // send headers
   if (_state == RESPONSE_HEADERS) {
-    if (space >= headLen) {
-      _state = RESPONSE_CONTENT;
-      space -= headLen;
-    } else {
-      String out = _head.substring(0, space);
-      _head = _head.substring(space);
-      _writtenLength += request->client()->write(out.c_str(), out.length());
-#if ASYNCWEBSERVER_USE_CHUNK_INFLIGHT
-      _in_flight += out.length();
+    // copy headers buffer to sock buffer
+    size_t const added = request->client()->add(_assembled_headers.c_str() + _assembled_headers_sent, _assembled_headers.length() - _assembled_headers_sent);
+    _writtenLength += added;
+    if (added < _assembled_headers.length()){
+      // we were not able to fit all headers in current buff, send this part here and return later for the rest
+      _assembled_headers_sent += added;
+      #if ASYNCWEBSERVER_USE_CHUNK_INFLIGHT
+      _in_flight += added;
       --_in_flight_credit;  // take a credit
-#endif
-      return out.length();
+      #endif
+      if (!request->client()->send()){
+        // something is wrong
+        request->client()->close();
+      }
+      return added;
     }
+    // otherwise we've added all the (ramainder) headers in current buff
+    _state = RESPONSE_CONTENT;
+    _assembled_headers = String();  // clear
   }
 
-  if (_state == RESPONSE_CONTENT) {
+  // if there are leftovers in buffer from the previous run, let's deal with it first
+  if (_state == RESPONSE_CONTENT && _send_buffer.size()){
+    size_t const written = request->client()->add(reinterpret_cast<char*>(_send_buffer.data()), _send_buffer.size());
+    if (written != _send_buffer.size()){
+      // we were not able to fit entire buffer again?! OK, let's send partial now and come back later
+      _writtenLength += written;
 #if ASYNCWEBSERVER_USE_CHUNK_INFLIGHT
-    // for response data we need to control the queue and in-flight fragmentation. Sending small chunks could give low latency,
-    // but flood asynctcp's queue and fragment socket buffer space for large responses.
-    // Let's ignore polled acks and acks in case when we have more in-flight data then the available socket buff space.
-    // That way we could balance on having half the buffer in-flight while another half is filling up, while minimizing events in asynctcp q
-    if (_in_flight > space) {
-      // async_ws_log_d("defer user call %u/%u", _in_flight, space);
-      //  take the credit back since we are ignoring this ack and rely on other inflight data
-      if (len) {
-        --_in_flight_credit;
-      }
-      return 0;
-    }
+      _in_flight += written;
+      --_in_flight_credit;  // take a credit
 #endif
+      request->client()->send();
+      _send_buffer.erase(_send_buffer.begin(), _send_buffer.begin() + written);
+      return written;
+    } else
+      _send_buffer.clear();
+    // OK buffer depleted, we can go on for more data within same sockbuff
+  }
 
-    size_t outLen;
-    if (_chunked) {
+  // send content body
+  if (_state == RESPONSE_CONTENT) {
+    size_t const space = request->client()->space();
+    if (_chunked || !_sendContentLength) {
       if (space <= 8) {
         return 0;
       }
-
-      outLen = space;
-    } else if (!_sendContentLength) {
-      outLen = space;
+      _send_buffer.resize(space);
     } else {
-      outLen = ((_contentLength - _sentLength) > space) ? space : (_contentLength - _sentLength);
-    }
-
-    uint8_t *buf = (uint8_t *)malloc(outLen + headLen);
-    if (!buf) {
-      async_ws_log_e("Failed to allocate");
-      request->abort();
-      return 0;
-    }
-
-    if (headLen) {
-      memcpy(buf, _head.c_str(), _head.length());
+      _send_buffer.resize(std::min(space, _contentLength - _sentLength));
     }
 
-    size_t readLen = 0;
-
     if (_chunked) {
       // HTTP 1.1 allows leading zeros in chunk length. Or spaces may be added.
-      // See RFC2616 sections 2, 3.6.1.
-      readLen = _fillBufferAndProcessTemplates(buf + headLen + 6, outLen - 8);
+      // See RFC2616 sections 2, 3.6.1 https://datatracker.ietf.org/doc/html/rfc2616#section-3.6.1
+      size_t const readLen = _fillBufferAndProcessTemplates(_send_buffer.data() + 6, _send_buffer.size() - 8);  // reserve 8 bytes for chunk size data
       if (readLen == RESPONSE_TRY_AGAIN) {
-        free(buf);
-        return 0;
+        _send_buffer.clear(); // we won't send anything, do not release mem but leave it for later
+      } else {
+        sprintf(reinterpret_cast<char*>(_send_buffer.data()), "%04x\r\n", readLen);   // print chunk size in buffer
+        _send_buffer.at(readLen + 6) = '\r';
+        _send_buffer.at(readLen + 7) = '\n';
+        _send_buffer.resize(readLen + 8);     // set internal vector's size to match added data
+        if (!readLen){
+          // last chunk?
+          _state = RESPONSE_WAIT_ACK;
+        }
+        _sentLength += readLen;               // not sure if that is needed for chunked data?
       }
-      outLen = sprintf((char *)buf + headLen, "%04x", readLen) + headLen;
-      buf[outLen++] = '\r';
-      buf[outLen++] = '\n';
-      outLen += readLen;
-      buf[outLen++] = '\r';
-      buf[outLen++] = '\n';
     } else {
-      readLen = _fillBufferAndProcessTemplates(buf + headLen, outLen);
+      size_t const readLen = _fillBufferAndProcessTemplates(_send_buffer.data(), _send_buffer.size());
       if (readLen == RESPONSE_TRY_AGAIN) {
-        free(buf);
-        return 0;
+        _send_buffer.clear(); // won't release mem but leave it for later
+      } else if (readLen == 0){
+        // no more data to send
+        _state = RESPONSE_WAIT_ACK;
+        _send_buffer.clear();
+      } else {
+        _send_buffer.resize(readLen);       // set internal vector's size to match added data
+        _sentLength += readLen;
+        if (_sendContentLength && _sentLength == _contentLength){
+          // it was last piece of content
+          _state = RESPONSE_WAIT_ACK;
+        }
       }
-      outLen = readLen + headLen;
-    }
-
-    if (headLen) {
-      _head = emptyString;
     }
 
-    if (outLen) {
-      _writtenLength += request->client()->write((const char *)buf, outLen);
+    size_t written{0};
+    if (_send_buffer.size()){
+      written = request->client()->add(reinterpret_cast<char*>(_send_buffer.data()), _send_buffer.size());
+      if (written != _send_buffer.size()){
+        // we were not able to send entire buffer now somehow, leave it for later
+        // (this should not happen normally unless connection's TCP window suddenly changed or some other thread highjacked our sock buffer :) )
+        _send_buffer.erase(_send_buffer.begin(), _send_buffer.begin() + written);
+      } else
+        _send_buffer.clear();
+
+      _writtenLength += written;
 #if ASYNCWEBSERVER_USE_CHUNK_INFLIGHT
-      _in_flight += outLen;
+      _in_flight += written;
       --_in_flight_credit;  // take a credit
 #endif
     }
-
-    if (_chunked) {
-      _sentLength += readLen;
-    } else {
-      _sentLength += outLen - headLen;
-    }
-
-    free(buf);
-
-    if ((_chunked && readLen == 0) || (!_sendContentLength && outLen == 0) || (!_chunked && _sentLength == _contentLength)) {
-      _state = RESPONSE_WAIT_ACK;
-    }
-    return outLen;
-
-  } else if (_state == RESPONSE_WAIT_ACK) {
-    if (!_sendContentLength || _ackedLength >= _writtenLength) {
-      _state = RESPONSE_END;
-      if (!_chunked && !_sendContentLength) {
-        request->client()->close(true);
-      }
-    }
+    // wether or not we have a new content in buffer now we must send anyway whatever is in sockbuff (maybe empty body)
+    // might be other data in sockbuff with hdrs or previous buffer data
+    request->client()->send();
+    return written;
+  }
+  
+  if (_state == RESPONSE_WAIT_ACK) {
+    // we do not need to wait for any acks actually if we won't send any more data,
+    // connection would be closed gracefully with last piece of data
+    _state = RESPONSE_END;
+    request->client()->close();
   }
   return 0;
 }
@@ -512,8 +524,8 @@ size_t AsyncAbstractResponse::_fillBufferAndProcessTemplates(uint8_t *data, size
   // Now we've read 'len' bytes, either from cache or from file
   // Search for template placeholders
   uint8_t *pTemplateStart = data;
-  while ((pTemplateStart < &data[len]) && (pTemplateStart = (uint8_t *)memchr(pTemplateStart, TEMPLATE_PLACEHOLDER, &data[len - 1] - pTemplateStart + 1))
-  ) {  // data[0] ... data[len - 1]
+  while ((pTemplateStart < &data[len]) && (pTemplateStart = (uint8_t *)memchr(pTemplateStart, TEMPLATE_PLACEHOLDER, &data[len - 1] - pTemplateStart + 1)) ) {
+    // data[0] ... data[len - 1]
     uint8_t *pTemplateEnd =
       (pTemplateStart < &data[len - 1]) ? (uint8_t *)memchr(pTemplateStart + 1, TEMPLATE_PLACEHOLDER, &data[len - 1] - pTemplateStart) : nullptr;
     // temporary buffer to hold parameter name