diff --git a/src/HttpContext.h b/src/HttpContext.h index e977bb4..a44e42c 100644 --- a/src/HttpContext.h +++ b/src/HttpContext.h @@ -252,6 +252,9 @@ private: if ((httpResponseData->state & HttpResponseData::HTTP_RESPONSE_PENDING) == 0) { if (((AsyncSocket *) s)->getBufferedAmount() == 0) { ((AsyncSocket *) s)->shutdown(); + /* We need to force close after sending FIN since we want to hinder + * clients from keeping to send their huge data */ + ((AsyncSocket *) s)->close(); } } } @@ -324,6 +327,9 @@ private: if ((httpResponseData->state & HttpResponseData::HTTP_RESPONSE_PENDING) == 0) { if (asyncSocket->getBufferedAmount() == 0) { asyncSocket->shutdown(); + /* We need to force close after sending FIN since we want to hinder + * clients from keeping to send their huge data */ + asyncSocket->close(); } } } diff --git a/src/HttpResponse.h b/src/HttpResponse.h index 555c78b..cf5732d 100644 --- a/src/HttpResponse.h +++ b/src/HttpResponse.h @@ -93,7 +93,7 @@ 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, size_t totalSize, bool optional, bool allowContentLength = true) { + bool internalEnd(std::string_view data, size_t totalSize, bool optional, bool allowContentLength = true, bool closeConnection = false) { /* Write status if not already done */ writeStatus(HTTP_200_OK); @@ -103,6 +103,12 @@ private: } HttpResponseData *httpResponseData = getHttpResponseData(); + + /* In some cases, such as when refusing huge data we want to close the connection when drained */ + if (closeConnection) { + httpResponseData->state |= HttpResponseData::HTTP_CONNECTION_CLOSE; + } + if (httpResponseData->state & HttpResponseData::HTTP_WRITE_CALLED) { /* We do not have tryWrite-like functionalities, so ignore optional in this path */ @@ -353,8 +359,8 @@ public: } /* End the response with an optional data chunk. Always starts a timeout. */ - void end(std::string_view data = {}) { - internalEnd(data, data.length(), false); + void end(std::string_view data = {}, bool closeConnection = false) { + internalEnd(data, data.length(), false, true, closeConnection); } /* Try and end the response. Returns [true, true] on success.