From 2ff50630e96664c70be5421c4fa0e5cf57f7b5c6 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Fri, 26 Jun 2020 06:08:41 +0200
Subject: [PATCH] Send TCP FIN from HttpContext if WebSocket FIN sent

---
 src/HttpContext.h      | 12 +++++++++++-
 src/WebSocketContext.h |  4 ++++
 src/WebSocketData.h    |  2 ++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/HttpContext.h b/src/HttpContext.h
index 18c7049..80bda91 100644
--- a/src/HttpContext.h
+++ b/src/HttpContext.h
@@ -24,6 +24,7 @@
 #include "HttpContextData.h"
 #include "HttpResponseData.h"
 #include "AsyncSocket.h"
+#include "WebSocketData.h"
 
 #include <string_view>
 #include <iostream>
@@ -250,7 +251,16 @@ private:
                 AsyncSocket<SSL> *asyncSocket = (AsyncSocket<SSL> *) httpContextData->upgradedWebSocket;
 
                 /* Uncork here as well (note: what if we failed to uncork and we then pub/sub before we even upgraded?) */
-                /*auto [written, failed] = */asyncSocket->uncork();
+                auto [written, failed] = asyncSocket->uncork();
+
+                /* If we succeeded in uncorking, check if we have sent WebSocket FIN */
+                if (!failed) {
+                    WebSocketData *webSocketData = (WebSocketData *) asyncSocket->getAsyncSocketData();
+                    if (webSocketData->isShuttingDown) {
+                        /* In that case, also send TCP FIN (this is similar to what we have in ws drain handler) */
+                        asyncSocket->shutdown();
+                    }
+                }
 
                 /* Reset upgradedWebSocket before we return */
                 httpContextData->upgradedWebSocket = nullptr;
diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h
index 7a740e4..a578af1 100644
--- a/src/WebSocketContext.h
+++ b/src/WebSocketContext.h
@@ -302,6 +302,10 @@ private:
         /* Handle HTTP write out (note: SSL_read may trigger this spuriously, the app need to handle spurious calls) */
         us_socket_context_on_writable(SSL, getSocketContext(), [](auto *s) {
 
+            /* NOTE: Are we called here corked? If so, the below write code is broken, since
+             * we will have 0 as getBufferedAmount due to writing to cork buffer, then sending TCP FIN before
+             * we actually uncorked and sent off things */
+
             /* It makes sense to check for us_is_shut_down here and return if so, to avoid shutting down twice */
             if (us_socket_is_shut_down(SSL, (us_socket_t *) s)) {
                 return s;
diff --git a/src/WebSocketData.h b/src/WebSocketData.h
index 7b29ed4..a26e953 100644
--- a/src/WebSocketData.h
+++ b/src/WebSocketData.h
@@ -28,9 +28,11 @@
 namespace uWS {
 
 struct WebSocketData : AsyncSocketData<false>, WebSocketState<true> {
+    /* This guy has a lot of friends - why? */
     template <bool, bool> friend struct WebSocketContext;
     template <bool> friend struct WebSocketContextData;
     template <bool, bool> friend struct WebSocket;
+    template <bool> friend struct HttpContext;
 private:
     std::string fragmentBuffer;
     int controlTipLength = 0;