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