Add CivetWebSocketHandler to the C++ API

Patch from Nick Bulleid
(https://groups.google.com/forum/#!topic/civetweb/-iqfDnVMw2s)

include/CivetServer.h

@@ -89,6 +89,67 @@ class CIVETWEB_API CivetHandler
 };
 
 /**
+ * Basic interface for a websocket handler.  Handlers implementations
+ * must be reentrant.
+ */
+class CIVETWEB_API CivetWebSocketHandler
+{
+  public:
+	/**
+	 * Destructor
+	 */
+	virtual ~CivetWebSocketHandler()
+	{
+	}
+
+	/**
+	 * Callback method for when the client intends to establish a websocket
+	 *connection, before websocket handshake.
+	 *
+	 * @param server - the calling server
+	 * @param conn - the connection information
+	 * @returns true to keep socket open, false to close it
+	 */
+	virtual bool handleConnection(CivetServer *server,
+	                              const struct mg_connection *conn);
+
+	/**
+	 * Callback method for when websocket handshake is successfully completed,
+	 *and connection is ready for data exchange.
+	 *
+	 * @param server - the calling server
+	 * @param conn - the connection information
+	 */
+	virtual void handleReadyState(CivetServer *server,
+	                              struct mg_connection *conn);
+
+	/**
+	 * Callback method for when a data frame has been received from the client.
+	 *
+	 * @param server - the calling server
+	 * @param conn - the connection information
+	 * @bits: first byte of the websocket frame, see websocket RFC at
+	 *http://tools.ietf.org/html/rfc6455, section 5.2
+	 * @data, data_len: payload, with mask (if any) already applied.
+	 * @returns true to keep socket open, false to close it
+	 */
+	virtual bool handleData(CivetServer *server,
+	                        struct mg_connection *conn,
+	                        int bits,
+	                        char *data,
+	                        size_t data_len);
+
+	/**
+	 * Callback method for when the connection is closed.
+	 *
+	 * @param server - the calling server
+	 * @param conn - the connection information
+	 */
+	virtual void handleClose(CivetServer *server,
+	                         const struct mg_connection *conn);
+};
+
+/**
  * CivetServer
  *
  * Basic class for embedded web server.  This has an URL mapping built-in.
@@ -153,6 +214,27 @@ class CIVETWEB_API CivetServer
 	}
 
 	/**
+	 * addWebSocketHandler
+	 *
+	 * Adds a WebSocket handler for a specific URI.  If there is existing URI
+	 *handler, it will
+	 * be replaced with this one.
+	 *
+	 * URI's are ordered and prefix (REST) URI's are supported.
+	 *
+	 *  @param uri - URI to match.
+	 *  @param handler - handler instance to use.
+	 */
+	void addWebSocketHandler(const std::string &uri,
+	                         CivetWebSocketHandler *handler);
+
+	void
+	addWebSocketHandler(const std::string &uri, CivetWebSocketHandler &handler)
+	{
+		addWebSocketHandler(uri, &handler);
+	}
+
+	/**
 	 * removeHandler(const std::string &)
 	 *
 	 * Removes a handler.
@@ -378,6 +460,16 @@ class CIVETWEB_API CivetServer
 	 */
 	static int requestHandler(struct mg_connection *conn, void *cbdata);
 
+	static int webSocketConnectionHandler(const struct mg_connection *conn,
+	                                      void *cbdata);
+	static void webSocketReadyHandler(struct mg_connection *conn, void *cbdata);
+	static int webSocketDataHandler(struct mg_connection *conn,
+	                                int bits,
+	                                char *data,
+	                                size_t data_len,
+	                                void *cbdata);
+	static void webSocketCloseHandler(const struct mg_connection *conn,
+	                                  void *cbdata);
 	/**
 	 * closeHandler(struct mg_connection *)
 	 *

src/CivetServer.cpp

@@ -55,6 +55,48 @@ CivetHandler::handleOptions(CivetServer *server, struct mg_connection *conn)
 	return false;
 }
 
+bool
+CivetWebSocketHandler::handleConnection(CivetServer *server,
+                                        const struct mg_connection *conn)
+{
+	UNUSED_PARAMETER(server);
+	UNUSED_PARAMETER(conn);
+	return true;
+}
+
+void
+CivetWebSocketHandler::handleReadyState(CivetServer *server,
+                                        struct mg_connection *conn)
+{
+	UNUSED_PARAMETER(server);
+	UNUSED_PARAMETER(conn);
+	return;
+}
+
+bool
+CivetWebSocketHandler::handleData(CivetServer *server,
+                                  struct mg_connection *conn,
+                                  int bits,
+                                  char *data,
+                                  size_t data_len)
+{
+	UNUSED_PARAMETER(server);
+	UNUSED_PARAMETER(conn);
+	UNUSED_PARAMETER(bits);
+	UNUSED_PARAMETER(data);
+	UNUSED_PARAMETER(data_len);
+	return true;
+}
+
+void
+CivetWebSocketHandler::handleClose(CivetServer *server,
+                                   const struct mg_connection *conn)
+{
+	UNUSED_PARAMETER(server);
+	UNUSED_PARAMETER(conn);
+	return;
+}
+
 int
 CivetServer::requestHandler(struct mg_connection *conn, void *cbdata)
 {
@@ -90,6 +132,92 @@ CivetServer::requestHandler(struct mg_connection *conn, void *cbdata)
 	return 0; // No handler found
 }
 
+int
+CivetServer::webSocketConnectionHandler(const struct mg_connection *conn,
+                                        void *cbdata)
+{
+	const struct mg_request_info *request_info = mg_get_request_info(conn);
+	assert(request_info != NULL);
+	CivetServer *me = (CivetServer *)(request_info->user_data);
+	assert(me != NULL);
+
+	// Happens when a request hits the server before the context is saved
+	if (me->context == NULL)
+		return 0;
+
+	CivetWebSocketHandler *handler = (CivetWebSocketHandler *)cbdata;
+
+	if (handler) {
+		return handler->handleConnection(me, conn) ? 0 : 1;
+	}
+
+	return 1; // No handler found, close connection
+}
+
+void
+CivetServer::webSocketReadyHandler(struct mg_connection *conn, void *cbdata)
+{
+	const struct mg_request_info *request_info = mg_get_request_info(conn);
+	assert(request_info != NULL);
+	CivetServer *me = (CivetServer *)(request_info->user_data);
+	assert(me != NULL);
+
+	// Happens when a request hits the server before the context is saved
+	if (me->context == NULL)
+		return;
+
+	CivetWebSocketHandler *handler = (CivetWebSocketHandler *)cbdata;
+
+	if (handler) {
+		handler->handleReadyState(me, conn);
+	}
+}
+
+int
+CivetServer::webSocketDataHandler(struct mg_connection *conn,
+                                  int bits,
+                                  char *data,
+                                  size_t data_len,
+                                  void *cbdata)
+{
+	const struct mg_request_info *request_info = mg_get_request_info(conn);
+	assert(request_info != NULL);
+	CivetServer *me = (CivetServer *)(request_info->user_data);
+	assert(me != NULL);
+
+	// Happens when a request hits the server before the context is saved
+	if (me->context == NULL)
+		return 0;
+
+	CivetWebSocketHandler *handler = (CivetWebSocketHandler *)cbdata;
+
+	if (handler) {
+		return handler->handleData(me, conn, bits, data, data_len) ? 1 : 0;
+	}
+
+	return 1; // No handler found
+}
+
+void
+CivetServer::webSocketCloseHandler(const struct mg_connection *conn,
+                                   void *cbdata)
+{
+	const struct mg_request_info *request_info = mg_get_request_info(conn);
+	assert(request_info != NULL);
+	CivetServer *me = (CivetServer *)(request_info->user_data);
+	assert(me != NULL);
+
+	// Happens when a request hits the server before the context is saved
+	if (me->context == NULL)
+		return;
+
+	CivetWebSocketHandler *handler = (CivetWebSocketHandler *)cbdata;
+
+	if (handler) {
+		handler->handleClose(me, conn);
+	}
+}
+
 CivetServer::CivetServer(const char **options,
                          const struct mg_callbacks *_callbacks)
     : context(0)
@@ -141,6 +269,19 @@ CivetServer::addHandler(const std::string &uri, CivetHandler *handler)
 }
 
 void
+CivetServer::addWebSocketHandler(const std::string &uri,
+                                 CivetWebSocketHandler *handler)
+{
+	mg_set_websocket_handler(context,
+	                         uri.c_str(),
+	                         webSocketConnectionHandler,
+	                         webSocketReadyHandler,
+	                         webSocketDataHandler,
+	                         webSocketCloseHandler,
+	                         handler);
+}
+
+void
 CivetServer::removeHandler(const std::string &uri)
 {
 	mg_set_request_handler(context, uri.c_str(), NULL, NULL);