From 4e91dc4f2d6eba2e66dd5b59445abcc778117c60 Mon Sep 17 00:00:00 2001
From: Mike Buland <mbuland@penny-arcade.com>
Date: Mon, 4 Oct 2021 11:42:36 -0700
Subject: Added handler to check/modify http headers.

This is an extra call that happens during the http handshake that allows
the implementation to set additional headers, modify existing headers,
and cancel the whole connection if desired.
---
 src/unstable/protocolwebsocket.cpp | 29 +++++++++++++++++++++++------
 src/unstable/protocolwebsocket.h   |  6 ++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/unstable/protocolwebsocket.cpp b/src/unstable/protocolwebsocket.cpp
index 8d4a124..08ac176 100644
--- a/src/unstable/protocolwebsocket.cpp
+++ b/src/unstable/protocolwebsocket.cpp
@@ -66,6 +66,11 @@ void Bu::ProtocolWebSocket::onNewData( Bu::Client * /*pClient*/ )
     }
 }
 
+bool onProcessHeaders( Bu::StringList & /*lHeadersOut*/ )
+{
+    return true;
+}
+
 void Bu::ProtocolWebSocket::writeMessage( const Bu::String &sData,
         Bu::ProtocolWebSocket::Operation eOp )
 {
@@ -237,12 +242,24 @@ bool Bu::ProtocolWebSocket::processHeaders()
 
     DEBUG( Bu::println("accept: %1").arg( mbOut.getString() ) );
 
-    pClient->write("HTTP/1.1 101 Switching Protocols\r\n"
-        "Upgrade: websocket\r\n"
-        "Connection: Upgrade\r\n"
-        "Sec-WebSocket-Accept: " + mbOut.getString() + "\r\n"
-        "\r\n"
-        );
+    Bu::StringList lHeadersOut;
+    lHeadersOut.append("Upgrade: websocket");
+    lHeadersOut.append("Connection: Upgrade");
+    lHeadersOut.append("Sec-WebSocket-Accept: " + mbOut.getString());
+
+    if( !onProcessHeaders( lHeadersOut ) )
+    {
+        pClient->disconnect();
+        return false;
+    }
+
+    Bu::String sHeaderBlock("HTTP/1.1 101 Switching Protocols\r\n");
+    for( Bu::StringList::iterator i = lHeadersOut.begin(); i; i++ )
+    {
+        sHeaderBlock += (*i) + "\r\n";
+    }
+    sHeaderBlock += "\r\n";
+    pClient->write( sHeaderBlock );
 
     DEBUG( Bu::println("websocket: Switching protocols.") );
 
diff --git a/src/unstable/protocolwebsocket.h b/src/unstable/protocolwebsocket.h
index 6778ea5..e5ff804 100644
--- a/src/unstable/protocolwebsocket.h
+++ b/src/unstable/protocolwebsocket.h
@@ -35,6 +35,12 @@ namespace Bu
         virtual void onNewConnection( Bu::Client *pClient );
         virtual void onNewData( Bu::Client *pClient );
 
+        /**
+         * Return true to continue the connection, return false to disconnect.
+         * Feel free to modify the headers out list to add or remove http
+         * headers to go to the client.
+         */
+        virtual bool onProcessHeaders( Bu::StringList &lHeadersOut );
         virtual void onHandshakeComplete()=0;
         virtual void onNewMessage( const Bu::String &sData, Operation eOp )=0;
         void writeMessage( const Bu::String &sData, Operation eOp=Text );
-- 
cgit v1.2.3