fix: http proxy for upgrade protocol

ithewei · Oct 11, 2023 · 9f4c75e · 9f4c75e
1 parent 3e75b65
commit 9f4c75e
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/http/HttpMessage.cpp b/http/HttpMessage.cpp
@@ -443,6 +443,10 @@ bool HttpMessage::IsKeepAlive() {
     return keepalive;
 }
 
+bool HttpMessage::IsUpgrade() {
+    auto iter = headers.find("upgrade");
+    return iter != headers.end();
+}
 
 // headers
 void HttpMessage::SetHeader(const char* key, const std::string& value) {

diff --git a/http/HttpMessage.h b/http/HttpMessage.h
@@ -269,6 +269,7 @@ class HV_EXPORT HttpMessage {
 
     bool IsChunked();
     bool IsKeepAlive();
+    bool IsUpgrade();
 
     // headers
     void SetHeader(const char* key, const std::string& value);

diff --git a/http/server/HttpHandler.cpp b/http/server/HttpHandler.cpp
@@ -302,12 +302,12 @@ void HttpHandler::onMessageComplete() {
     addResponseHeaders();
 
     // upgrade ? handleUpgrade : HandleHttpRequest
-    upgrade = 0;
-    auto iter_upgrade = req->headers.find("upgrade");
-    if (iter_upgrade != req->headers.end()) {
-        upgrade = 1;
-        handleUpgrade(iter_upgrade->second.c_str());
-        status_code = resp->status_code;
+    if (upgrade) {
+        auto iter_upgrade = req->headers.find("upgrade");
+        if (iter_upgrade != req->headers.end()) {
+            handleUpgrade(iter_upgrade->second.c_str());
+            status_code = resp->status_code;
+        }
     } else {
         status_code = HandleHttpRequest();
         if (status_code != HTTP_STATUS_NEXT) {
@@ -342,6 +342,9 @@ void HttpHandler::handleRequestHeaders() {
     // keepalive
     keepalive = pReq->IsKeepAlive();
 
+    // upgrade
+    upgrade = pReq->IsUpgrade();
+
     // proxy
     proxy = forward_proxy = reverse_proxy = 0;
     if (hv::startswith(pReq->url, "http")) {
@@ -1107,6 +1110,7 @@ void HttpHandler::onProxyConnect(hio_t* upstream_io) {
     }
 
     // NOTE: start recv request continue then upstream
+    if (handler->upgrade) hio_setcb_read(io, hio_write_upstream);
     hio_read_start(io);
     // NOTE: start recv response then upstream
     hio_setcb_read(upstream_io, hio_write_upstream);