From 73747dd444d8720de9463e5e85b6c2d608bbcc5c Mon Sep 17 00:00:00 2001
From: lganzzzo <lganzzzo@gmail.com>
Date: Fri, 17 Dec 2021 03:42:46 +0200
Subject: [PATCH] HttpProcessor: Fix connection upgrade sequence.

---
 src/oatpp/web/server/HttpProcessor.cpp | 58 ++++++++++++++++----------
 1 file changed, 35 insertions(+), 23 deletions(-)

diff --git a/src/oatpp/web/server/HttpProcessor.cpp b/src/oatpp/web/server/HttpProcessor.cpp
index 94aa59e9..50e6703f 100644
--- a/src/oatpp/web/server/HttpProcessor.cpp
+++ b/src/oatpp/web/server/HttpProcessor.cpp
@@ -190,17 +190,8 @@ HttpProcessor::ConnectionState HttpProcessor::processNextRequest(ProcessingResou
         response->putHeaderIfNotExists(protocol::http::Header::CONNECTION, protocol::http::Header::Value::CONNECTION_CLOSE);
         break;
 
-      case ConnectionState::DELEGATED: {
-        auto handler = response->getConnectionUpgradeHandler();
-        if(handler) {
-          handler->handleConnection(resources.connection, response->getConnectionUpgradeParameters());
-          connectionState = ConnectionState::DELEGATED;
-        } else {
-          OATPP_LOGW("[oatpp::web::server::HttpProcessor::processNextRequest()]", "Warning. ConnectionUpgradeHandler not set!");
-          connectionState = ConnectionState::CLOSING;
-        }
+      default:
         break;
-      }
 
     }
 
@@ -211,6 +202,18 @@ HttpProcessor::ConnectionState HttpProcessor::processNextRequest(ProcessingResou
 
   response->send(resources.connection.object.get(), &resources.headersOutBuffer, contentEncoderProvider.get());
 
+  /* Delegate connection handling to another handler only after the response is sent to the client */
+  if(connectionState == ConnectionState::DELEGATED) {
+    auto handler = response->getConnectionUpgradeHandler();
+    if(handler) {
+      handler->handleConnection(resources.connection, response->getConnectionUpgradeParameters());
+      connectionState = ConnectionState::DELEGATED;
+    } else {
+      OATPP_LOGW("[oatpp::web::server::HttpProcessor::processNextRequest()]", "Warning. ConnectionUpgradeHandler not set!");
+      connectionState = ConnectionState::CLOSING;
+    }
+  }
+
   return connectionState;
 
 }
@@ -370,17 +373,8 @@ HttpProcessor::Coroutine::Action HttpProcessor::Coroutine::onResponseFormed() {
       m_currentResponse->putHeaderIfNotExists(protocol::http::Header::CONNECTION, protocol::http::Header::Value::CONNECTION_CLOSE);
       break;
 
-    case ConnectionState::DELEGATED: {
-      auto handler = m_currentResponse->getConnectionUpgradeHandler();
-      if(handler) {
-        handler->handleConnection(m_connection, m_currentResponse->getConnectionUpgradeParameters());
-        m_connectionState = ConnectionState::DELEGATED;
-      } else {
-        OATPP_LOGW("[oatpp::web::server::HttpProcessor::Coroutine::onResponseFormed()]", "Warning. ConnectionUpgradeHandler not set!");
-        m_connectionState = ConnectionState::CLOSING;
-      }
+    default:
       break;
-    }
 
   }
 
@@ -393,9 +387,27 @@ HttpProcessor::Coroutine::Action HttpProcessor::Coroutine::onResponseFormed() {
 }
   
 HttpProcessor::Coroutine::Action HttpProcessor::Coroutine::onRequestDone() {
-  
-  if(m_connectionState == ConnectionState::ALIVE) {
-    return yieldTo(&HttpProcessor::Coroutine::parseHeaders);
+
+  switch (m_connectionState) {
+    case ConnectionState::ALIVE:
+      return yieldTo(&HttpProcessor::Coroutine::parseHeaders);
+
+    /* Delegate connection handling to another handler only after the response is sent to the client */
+    case ConnectionState::DELEGATED: {
+      auto handler = m_currentResponse->getConnectionUpgradeHandler();
+      if(handler) {
+        handler->handleConnection(m_connection, m_currentResponse->getConnectionUpgradeParameters());
+        m_connectionState = ConnectionState::DELEGATED;
+      } else {
+        OATPP_LOGW("[oatpp::web::server::HttpProcessor::Coroutine::onResponseFormed()]", "Warning. ConnectionUpgradeHandler not set!");
+        m_connectionState = ConnectionState::CLOSING;
+      }
+      break;
+    }
+
+    default:
+      break;
+
   }
   
   return finish();