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.

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 );