diff --git a/src/AsyncSocket.h b/src/AsyncSocket.h index 0581c04..6077650 100644 --- a/src/AsyncSocket.h +++ b/src/AsyncSocket.h @@ -50,6 +50,7 @@ struct AsyncSocket { template friend struct TemplatedApp; template friend struct WebSocketContextData; template friend struct TopicTree; + template friend struct HttpResponse; private: /* Helper, do not use directly (todo: move to uSockets or de-crazify) */ diff --git a/src/HttpResponse.h b/src/HttpResponse.h index 9b3d505..2b2d4e9 100644 --- a/src/HttpResponse.h +++ b/src/HttpResponse.h @@ -132,6 +132,21 @@ private: httpResponseData->markDone(); + /* We need to check if we should close this socket here now */ + if (!Super::isCorked()) { + if (httpResponseData->state & HttpResponseData::HTTP_CONNECTION_CLOSE) { + if ((httpResponseData->state & HttpResponseData::HTTP_RESPONSE_PENDING) == 0) { + if (((AsyncSocket *) this)->getBufferedAmount() == 0) { + ((AsyncSocket *) this)->shutdown(); + /* We need to force close after sending FIN since we want to hinder + * clients from keeping to send their huge data */ + ((AsyncSocket *) this)->close(); + return true; + } + } + } + } + /* tryEnd can never fail when in chunked mode, since we do not have tryWrite (yet), only write */ Super::timeout(HTTP_TIMEOUT_S); return true; @@ -182,6 +197,20 @@ private: /* Remove onAborted function if we reach the end */ if (httpResponseData->offset == totalSize) { httpResponseData->markDone(); + + /* We need to check if we should close this socket here now */ + if (!Super::isCorked()) { + if (httpResponseData->state & HttpResponseData::HTTP_CONNECTION_CLOSE) { + if ((httpResponseData->state & HttpResponseData::HTTP_RESPONSE_PENDING) == 0) { + if (((AsyncSocket *) this)->getBufferedAmount() == 0) { + ((AsyncSocket *) this)->shutdown(); + /* We need to force close after sending FIN since we want to hinder + * clients from keeping to send their huge data */ + ((AsyncSocket *) this)->close(); + } + } + } + } } return success; @@ -472,6 +501,19 @@ public: /* This behavior should equal the behavior in HttpContext when uncorking fails */ Super::timeout(HTTP_TIMEOUT_S); } + + /* If we have no backbuffer and we are connection close and we responded fully then close */ + HttpResponseData *httpResponseData = getHttpResponseData(); + if (httpResponseData->state & HttpResponseData::HTTP_CONNECTION_CLOSE) { + if ((httpResponseData->state & HttpResponseData::HTTP_RESPONSE_PENDING) == 0) { + if (((AsyncSocket *) this)->getBufferedAmount() == 0) { + ((AsyncSocket *) this)->shutdown(); + /* We need to force close after sending FIN since we want to hinder + * clients from keeping to send their huge data */ + ((AsyncSocket *) this)->close(); + } + } + } } else { /* We are already corked, or can't cork so let's just call the handler */ handler();