From 67df67f04b7d74e2cd655a1d47fa0094d642d154 Mon Sep 17 00:00:00 2001
From: Wiebe Cazemier <wiebe@halfgaar.net>
Date: Sun, 10 Oct 2021 21:19:40 +0200
Subject: [PATCH] Handle websocket close frames

---
 iowrapper.cpp | 34 ++++++++++++++++++++++++++++++----
 iowrapper.h   |  2 +-
 utils.cpp     | 12 ++++++++++++
 utils.h       |  2 ++
 4 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/iowrapper.cpp b/iowrapper.cpp
index cf23fa3..6f61406 100644
--- a/iowrapper.cpp
+++ b/iowrapper.cpp
@@ -406,11 +406,11 @@ ssize_t IoWrapper::readWebsocketAndOrSsl(int fd, void *buf, size_t nbytes, IoWra
             }
             else
             {
-                n = websocketBytesToReadBuffer(buf, nbytes);
+                n = websocketBytesToReadBuffer(buf, nbytes, error);
 
                 if (n > 0)
                     *error = IoWrapResult::Success;
-                else if (n == 0)
+                else if (n == 0 && *error != IoWrapResult::Disconnected)
                     *error = IoWrapResult::Wouldblock;
             }
         }
@@ -426,7 +426,7 @@ ssize_t IoWrapper::readWebsocketAndOrSsl(int fd, void *buf, size_t nbytes, IoWra
  * @param nbytes
  * @return
  */
-ssize_t IoWrapper::websocketBytesToReadBuffer(void *buf, const size_t nbytes)
+ssize_t IoWrapper::websocketBytesToReadBuffer(void *buf, const size_t nbytes, IoWrapResult *error)
 {
     const ssize_t targetBufMaxSize = nbytes;
     ssize_t nbytesRead = 0;
@@ -523,6 +523,8 @@ ssize_t IoWrapper::websocketBytesToReadBuffer(void *buf, const size_t nbytes)
 
             if (incompleteWebsocketRead.frame_bytes_left <= websocketPendingBytes.usedBytes())
             {
+                logger->logf(LOG_INFO, "Ponging websocket");
+
                 // Constructing a new temporary buffer because I need the reponse in one frame for writeAsMuchOfBufAsWebsocketFrame().
                 std::vector<char> response(incompleteWebsocketRead.frame_bytes_left);
                 websocketPendingBytes.read(response.data(), response.size());
@@ -532,11 +534,35 @@ ssize_t IoWrapper::websocketBytesToReadBuffer(void *buf, const size_t nbytes)
                 parentClient->setReadyForWriting(true);
             }
         }
+        else if (incompleteWebsocketRead.opcode == WebsocketOpcode::Close)
+        {
+            // MUST be a 2-byte unsigned integer (in network byte order) representing a status code with value /code/ defined
+            if (incompleteWebsocketRead.frame_bytes_left <= websocketPendingBytes.usedBytes()
+                && incompleteWebsocketRead.frame_bytes_left >= 2)
+            {
+                const uint8_t msb = *websocketPendingBytes.tailPtr() ^ incompleteWebsocketRead.getNextMaskingByte();
+                websocketPendingBytes.advanceTail(1);
+                const uint8_t lsb = *websocketPendingBytes.tailPtr() ^ incompleteWebsocketRead.getNextMaskingByte();
+                websocketPendingBytes.advanceTail(1);
+
+                const uint16_t code = msb << 8 | lsb;
+
+                // An actual MQTT disconnect doesn't send websocket close frames, or perhaps after the MQTT
+                // disconnect when it doesn't matter anymore. So, when users close the tab or stuff like that,
+                // we can consider it a closed transport i.e. failed connection. This means will messages
+                // will be sent.
+                parentClient->setDisconnectReason(websocketCloseCodeToString(code));
+                *error = IoWrapResult::Disconnected;
+
+                // There may be a UTF8 string with a reason in the packet still, but ignoring that for now.
+                incompleteWebsocketRead.reset();
+            }
+        }
         else
         {
             // Specs: "MQTT Control Packets MUST be sent in WebSocket binary data frames. If any other type of data frame is
             // received the recipient MUST close the Network Connection [MQTT-6.0.0-1]".
-            throw ProtocolError("Websocket frames must be 'binary' or 'ping'");
+            throw ProtocolError(formatString("Websocket frames must be 'binary' or 'ping'. Received: %d", incompleteWebsocketRead.opcode));
         }
 
         if (!incompleteWebsocketRead.sillWorkingOnFrame())
diff --git a/iowrapper.h b/iowrapper.h
index cd7fe67..00db122 100644
--- a/iowrapper.h
+++ b/iowrapper.h
@@ -115,7 +115,7 @@ class IoWrapper
 
     Logger *logger = Logger::getInstance();
 
-    ssize_t websocketBytesToReadBuffer(void *buf, const size_t nbytes);
+    ssize_t websocketBytesToReadBuffer(void *buf, const size_t nbytes, IoWrapResult *error);
     ssize_t readOrSslRead(int fd, void *buf, size_t nbytes, IoWrapResult *error);
     ssize_t writeOrSslWrite(int fd, const void *buf, size_t nbytes, IoWrapResult *error);
     ssize_t writeAsMuchOfBufAsWebsocketFrame(const void *buf, size_t nbytes, WebsocketOpcode opcode = WebsocketOpcode::Binary);
diff --git a/utils.cpp b/utils.cpp
index 50daa02..87de974 100644
--- a/utils.cpp
+++ b/utils.cpp
@@ -633,3 +633,15 @@ std::string sockaddrToString(sockaddr *addr)
 
     return "[unknown address]";
 }
+
+const std::string websocketCloseCodeToString(uint16_t code)
+{
+    switch (code) {
+    case 1000:
+        return "Normal websocket close";
+    case 1001:
+        return "Browser navigating away from page";
+    default:
+        return formatString("Websocket status code %d", code);
+    }
+}
diff --git a/utils.h b/utils.h
index 70c8352..cef7f41 100644
--- a/utils.h
+++ b/utils.h
@@ -122,5 +122,7 @@ template<typename ex> void checkWritableDir(const std::string &path)
     }
 }
 
+const std::string websocketCloseCodeToString(uint16_t code);
+
 
 #endif // UTILS_H
--
libgit2 0.21.4