More strict handling of connection: close

Author: uNetworkingAB

src/AsyncSocket.h

@@ -50,6 +50,7 @@ struct AsyncSocket {
     template <bool> friend struct TemplatedApp;
     template <bool, typename> friend struct WebSocketContextData;
     template <typename, typename> friend struct TopicTree;
+    template <bool> friend struct HttpResponse;
 
 private:
     /* Helper, do not use directly (todo: move to uSockets or de-crazify) */

src/HttpResponse.h

@@ -132,6 +132,21 @@ struct HttpResponse : public AsyncSocket<SSL> {
 
             httpResponseData->markDone();
 
+            /* We need to check if we should close this socket here now */
+            if (!Super::isCorked()) {
+                if (httpResponseData->state & HttpResponseData<SSL>::HTTP_CONNECTION_CLOSE) {
+                    if ((httpResponseData->state & HttpResponseData<SSL>::HTTP_RESPONSE_PENDING) == 0) {
+                        if (((AsyncSocket<SSL> *) this)->getBufferedAmount() == 0) {
+                            ((AsyncSocket<SSL> *) this)->shutdown();
+                            /* We need to force close after sending FIN since we want to hinder
+                                * clients from keeping to send their huge data */
+                            ((AsyncSocket<SSL> *) 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 @@ struct HttpResponse : public AsyncSocket<SSL> {
             /* 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<SSL>::HTTP_CONNECTION_CLOSE) {
+                        if ((httpResponseData->state & HttpResponseData<SSL>::HTTP_RESPONSE_PENDING) == 0) {
+                            if (((AsyncSocket<SSL> *) this)->getBufferedAmount() == 0) {
+                                ((AsyncSocket<SSL> *) this)->shutdown();
+                                /* We need to force close after sending FIN since we want to hinder
+                                * clients from keeping to send their huge data */
+                                ((AsyncSocket<SSL> *) this)->close();
+                            }
+                        }
+                    }
+                }
             }
 
             return success;
@@ -472,6 +501,19 @@ struct HttpResponse : public AsyncSocket<SSL> {
                 /* 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<SSL> *httpResponseData = getHttpResponseData();
+            if (httpResponseData->state & HttpResponseData<SSL>::HTTP_CONNECTION_CLOSE) {
+                if ((httpResponseData->state & HttpResponseData<SSL>::HTTP_RESPONSE_PENDING) == 0) {
+                    if (((AsyncSocket<SSL> *) this)->getBufferedAmount() == 0) {
+                        ((AsyncSocket<SSL> *) this)->shutdown();
+                        /* We need to force close after sending FIN since we want to hinder
+                        * clients from keeping to send their huge data */
+                        ((AsyncSocket<SSL> *) this)->close();
+                    }
+                }
+            }
         } else {
             /* We are already corked, or can't cork so let's just call the handler */
             handler();