Created a websocket api so you can build site/module specific ws handlers

doc/ref/controllers/controller_websocket.rst

@@ -1,4 +1,56 @@
 
 .. include:: meta-websocket.rst
 
-.. todo:: Not yet documented.
+Provides persistent websocket connections between the client and the server. The default implementation 
+is used by mod_base when the ```{% stream %}``` tag is placed on a page. See :ref:`tag-stream` for more 
+information.
+
+Defining Custom Websocket Behaviour
+-----------------------------------
+
+You can provide your own websocket_start similar too controller websocket by setting a ws_handler 
+containing the name of a websocket handler module in the zotonic context.
+
+Example::
+
+    websocket_start(ReqData, Context) ->
+        Context1 = z_context:set(ws_handler, ?MODULE, Context),
+        controller_websocket:websocket_start(ReqData, Context).
+
+When passing a custom handler module, the default handler websocket will not be used, but the specified 
+one. Controller websocket contains the code for the default zotonic handler. It attaches itself as 
+websocket handler to the page session.
+
+It is also possible to configure a custom ```ws_handler``` by specifying it in a dispatch rule.::
+
+    {customws, ["socket", "custom"], controller_websocket, [{ws_handler, my_ws_handler}]}
+
+WebSocket Handler API
+---------------------
+
+In order to implement your own websocket handler you have to implement four callback functions. 
+When you want to sent a message to the client you call ``controller_websocket:send_data/2``.
+
+Example::
+
+    -module(my_ws_handler).
+    -export([websocket_init/1, websocket_message/2, websocket_info/2, websocket_terminate/2])
+
+    % Called when the websocket is initialized.
+    websocket_init(_Context) ->
+        erlang:start_timer(1000, self(), <<"Hello!">>).
+
+    % Called when a message arrives on the websocket.
+    websocket_message(Msg, Context) ->
+        controller_websocket:websocket_send_data(self(), ["You said: ", Msg]).
+
+    % Called when another type of message arrives.
+    websocket_info(Msg, _Context) ->
+        controller_websocket:websocket_send_data(self(), Msg),
+        erlang:start_timer(5000, self(), <<"Hello again!">>).
+
+    % Called when the websocket terminates.
+    websocket_terminate(_Reason, _Context) ->
+        ok.
+
+

modules/mod_base/controllers/controller_websocket.erl

@@ -26,13 +26,20 @@
     charsets_provided/2,
     content_types_provided/2,
     provide_content/2,
-    websocket_start/2,
-
-    handle_message/2
+    websocket_start/2,    
+    websocket_send_data/2
+]).
+
+% websocket handler exports.
+-export([
+    websocket_init/1,
+    websocket_message/2,
+    websocket_info/2,
+    websocket_terminate/2
 ]).
 
 -include_lib("webmachine_controller.hrl").
--include_lib("include/zotonic.hrl").
+-include_lib("zotonic.hrl").
 
 init(_Args) -> {ok, []}.
 
@@ -67,29 +74,40 @@ provide_content(ReqData, Context) ->
 websocket_start(ReqData, Context) ->
     ContextReq = ?WM_REQ(ReqData, Context),
     Context1 = z_context:ensure_all(ContextReq),
-    case z_context:get_req_header("sec-websocket-version", Context1) of
+    Context2 = case z_context:get(ws_handler, Context1) of
+        undefined ->
+            z_context:set(ws_handler, ?MODULE, Context1);
+        _Hdlr -> Context1
+    end,
+    case z_context:get_req_header("sec-websocket-version", Context2) of
         undefined ->
-            case z_context:get_req_header("sec-websocket-key1", Context1) of
+            case z_context:get_req_header("sec-websocket-key1", Context2) of
                 undefined ->
-                    z_websocket_hixie75:start(ReqData, Context1);
+                    z_websocket_hixie75:start(ReqData, Context2);
                 WsKey1 ->
-                    z_websocket_hybi00:start(WsKey1, ReqData, Context1)
+                    z_websocket_hybi00:start(WsKey1, ReqData, Context2)
             end;
         "7" ->
             % http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-07
-            z_websocket_hybi17:start(ReqData, Context1);
+            z_websocket_hybi17:start(ReqData, Context2);
         "8" ->
             % http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-10
-            z_websocket_hybi17:start(ReqData, Context1);
+            z_websocket_hybi17:start(ReqData, Context2);
         "13" ->
             % http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-17
-            z_websocket_hybi17:start(ReqData, Context1)
+            z_websocket_hybi17:start(ReqData, Context2)
     end.
 
+%% @doc Send Data over websocket Pid to the client.
+websocket_send_data(Pid, Data) ->
+    Pid ! {send_data, Data}.
 
+%% Called during initialization of the websocket.
+websocket_init(Context) ->
+    z_session_page:websocket_attach(self(), Context).
 
 %% Handle a message from the browser, should contain an url encoded request. Sends result script back to browser.
-handle_message(Msg, Context) ->
+websocket_message(Msg, Context) ->
     Qs = mochiweb_util:parse_qs(Msg),
     Context1 = z_context:set('q', Qs, Context),
 
@@ -111,3 +129,9 @@ handle_message(Msg, Context) ->
     z_utils:erase_process_dict(),
     ok.
 
+websocket_info(_Msg, _Context) -> 
+    ok.
+
+websocket_terminate(_Reason, _Context) -> 
+    ok.
+

modules/mod_base/support/z_websocket_hixie75.erl

@@ -19,6 +19,8 @@
 -module(z_websocket_hixie75).
 -author("Marc Worrell <marc@worrell.nl>").
 
+-include_lib("zotonic.hrl").
+
 -export([
     start/2,
 
@@ -34,7 +36,13 @@
 % First draft protocol version, this code should be removed in due time.
 start(ReqData, Context1) ->
     Hostname = m_site:get(hostname, Context1),
-    WebSocketPath = iolist_to_binary(["/websocket?z_pageid=", mochiweb_util:quote_plus(z_context:get_q("z_pageid", Context1))]),
+
+    Qs = [[K, $=, mochiweb_util:quote_plus(V)] || {K, V} <- wrq:req_qs(ReqData)],
+    WebSocketPath = case Qs of
+        [] -> iolist_to_binary(wrq:path(ReqData));
+        _ -> iolist_to_binary([wrq:path(ReqData), $?, Qs])
+    end,
+
     Protocol = case wrq:is_ssl(ReqData) of true -> "https"; _ -> "http" end,
     Socket = webmachine_request:socket(ReqData),
     Data = ["HTTP/1.1 101 Web Socket Protocol Handshake", 13, 10,
@@ -79,14 +87,14 @@ handle_data(<<>>, nolength, <<255,_T/binary>>, _Socket, _Context) ->
     % A packet of <<0,255>> signifies that the ua wants to close the connection
     ua_close_request;
 handle_data(Msg, nolength, <<255,T/binary>>, Socket, Context) ->
-    controller_websocket:handle_message(Msg, Context),
+    handle_message(Msg, Context),
     handle_data(none, nolength, T, Socket, Context);
 handle_data(Msg, nolength, <<H,T/binary>>, Socket, Context) ->
     handle_data(<<Msg/binary, H>>, nolength, T, Socket, Context);
 
 %% Extract frame with length bytes
 handle_data(Msg, 0, T, Socket, Context) ->
-    controller_websocket:handle_message(Msg, Context),
+    handle_message(Msg, Context),
     handle_data(none, nolength, T, Socket, Context);
 handle_data(Msg, Length, <<H,T/binary>>, Socket, Context) when is_integer(Length) and Length > 0 ->
     handle_data(<<Msg/binary, H>>, Length-1, T, Socket, Context);
@@ -96,6 +104,25 @@ handle_data(Msg, Length, <<>>, Socket, Context) ->
     z_websocket_hixie75:receive_loop(Msg, Length, Socket, Context).
 
 
+% Call the handler. We are initializing.
+handle_init(Context) ->
+    H = z_context:get(ws_handler, Context),
+    H:websocket_init(Context).
+
+% Call the handler, new message arrived.
+handle_message(Msg, Context) ->
+    H = z_context:get(ws_handler, Context),
+    H:websocket_message(Msg, Context).
+
+handle_info(Msg, Context) ->
+    H = z_context:get(ws_handler, Context),
+    H:websocket_info(Msg, Context).
+
+handle_terminate(Reason, Context) ->
+    H = z_context:get(ws_handler, Context),
+    H:websocket_terminate(Reason, Context).
+
+
 %% @doc Unpack the length bytes
 %% author: Davide Marquês (From yaws_websockets.erl)
 unpack_length(Binary) ->
@@ -135,21 +162,31 @@ unpack_length(Binary, LenBytes, Length) ->
 start_send_loop(Socket, Context) ->
     % We want to receive any exit signal (including 'normal') from the socket's process.
     process_flag(trap_exit, true),
-    z_session_page:websocket_attach(self(), Context),
+    handle_init(Context),
     send_loop(Socket, Context).
 
 send_loop(Socket, Context) ->
     receive
         {send_data, Data} ->
             case send(Socket, [0, Data, 255]) of
-                ok -> z_websocket_hixie75:send_loop(Socket, Context);
-                closed -> closed
+                ok -> 
+                    send_loop(Socket, Context);
+                {error, closed} -> 
+                    handle_terminate({error, closed}, Context),
+                    closed;
+                _ ->
+                    handle_terminate(normal, Context),
+                    normal
             end;
-        {'EXIT', _FromPid, _Reason} ->
-            % Exit of the socket's process, stop sending data.
+        {'EXIT', _FromPid, normal} ->
+            handle_terminate(normal, Context),
+            normal;
+        {'EXIT', _FromPid, Reason} ->
+            handle_terminate({error, {exit, Reason}}, Context),
             exit;
-        _ ->
-            z_websocket_hixie75:send_loop(Socket, Context)
+        Msg ->
+            handle_info(Msg, Context),
+            send_loop(Socket, Context)
     end.
 
 
@@ -158,9 +195,5 @@ send_loop(Socket, Context) ->
 send(undefined, _Data) ->
     ok;
 send(Socket, Data) ->
-    case mochiweb_socket:send(Socket, iolist_to_binary(Data)) of
-        ok -> ok;
-        {error, closed} -> closed;
-        _ -> exit(normal)
-    end.
+    mochiweb_socket:send(Socket, iolist_to_binary(Data)).
 

modules/mod_base/support/z_websocket_hybi00.erl

@@ -19,14 +19,22 @@
 -module(z_websocket_hybi00).
 -author("Marc Worrell <marc@worrell.nl>").
 
+-include_lib("zotonic.hrl").
+
 -export([
     start/3
 ]).
 
 
 start(WsKey1, ReqData, Context1) ->
     Hostname = m_site:get(hostname, Context1),
-    WebSocketPath = iolist_to_binary(["/websocket?z_pageid=", mochiweb_util:quote_plus(z_context:get_q("z_pageid", Context1))]),
+
+    Qs = [[K, $=, mochiweb_util:quote_plus(V)] || {K, V} <- wrq:req_qs(ReqData)],
+    WebSocketPath = case Qs of
+        [] -> iolist_to_binary(wrq:path(ReqData));
+        _ -> iolist_to_binary([wrq:path(ReqData), $?, Qs])
+    end,
+
     Protocol = case wrq:is_ssl(ReqData) of true -> "https"; _ -> "http" end,
     Socket = webmachine_request:socket(ReqData),
 

modules/mod_base/support/z_websocket_hybi17.erl

@@ -123,11 +123,11 @@ unmask_data(Opcode, <<O:8>>, MaskKey, RemainingData, Socket, Context, Acc) ->
 
 % Text frame
 handle_frame(RemainingData, 1, Message, Socket, Context) ->
-    controller_websocket:handle_message(Message, Context),
+    handle_message(Message, Context),
     handle_data(RemainingData, Socket, Context);
 % Binary frame
 handle_frame(RemainingData, 2, Message, Socket, Context) ->
-    controller_websocket:handle_message(Message, Context),
+    handle_message(Message, Context),
     handle_data(RemainingData, Socket, Context);
 % Close control frame
 handle_frame(_RemainingData, 8, _Message, Socket, Context) ->
@@ -142,6 +142,23 @@ handle_frame(RemainingData, 9, Message, Socket, Context) ->
 handle_frame(RemainingData, 10, _Message, Socket, Context) ->
     handle_data(RemainingData, Socket, Context).
 
+% Call the handler about the initialization
+handle_init(Context) ->
+    W = z_context:get(ws_handler,  Context),
+    W:websocket_init(Context).
+
+handle_message(Message, Context) ->
+    H = z_context:get(ws_handler, Context),
+    H:websocket_message(Message, Context).
+
+handle_info(Message, Context) ->
+    H = z_context:get(ws_handler, Context),
+    H:websocket_info(Message, Context).
+
+handle_terminate(Reason, Context) ->
+    H = z_context:get(ws_handler, Context),
+    H:websocket_terminate(Reason, Context).
+
 
 %% @TODO: log any errors
 close(_Reason, Socket, _Context) ->
@@ -157,20 +174,32 @@ close(_Reason, Socket, _Context) ->
 start_send_loop(Socket, Context) ->
     % We want to receive any exit signal (including 'normal') from the socket's process.
     process_flag(trap_exit, true),
-    z_session_page:websocket_attach(self(), Context),
+    handle_init(Context),
     send_loop(Socket, Context).
 
 
 send_loop(Socket, Context) ->
     receive
         {send_data, Data} ->
-            send_frame(Socket, Data),
-            z_websocket_hybi17:send_loop(Socket, Context);
-        {'EXIT', _FromPid, _Reason} ->
-            % Exit of the socket's process, stop sending data.
+            case send_frame(Socket, Data) of
+                ok -> 
+                    send_loop(Socket, Context);
+                {error, closed} ->
+                    handle_terminate({error, closed}, Context),
+                    closed;
+                _ ->
+                    handle_terminate(normal, Context),
+                    normal
+            end;
+        {'EXIT', _FromPid, normal} ->
+            handle_terminate(normal, Context),
+            normal;
+        {'EXIT', _FromPid, Msg} ->
+            handle_terminate({error, {exit, Msg}}, Context),
             exit;
-        _ ->
-            z_websocket_hybi17:send_loop(Socket, Context)
+        Msg ->
+            handle_info(Msg, Context),
+            send_loop(Socket, Context)
     after ?PING_TIMEOUT ->
         send_ping(Socket),
         send_loop(Socket, Context)
@@ -193,11 +222,7 @@ send_ping(Socket) ->
 send(undefined, _Data) ->
     ok;
 send(Socket, Data) ->
-    case mochiweb_socket:send(Socket, iolist_to_binary(Data)) of
-        ok -> ok;
-        {error, closed} -> exit(closed);
-        _ -> exit(normal)
-    end.
+    mochiweb_socket:send(Socket, iolist_to_binary(Data)). 
 
 
 

src/support/z_context.erl

@@ -406,6 +406,8 @@ output1([C|Rest], Context, Acc) ->
         case proplists:get_value(format, Args, "html") of
             "html" ->
                 [ <<"\n\n<script type='text/javascript'>\n$(function() {\n">>, Script, <<"\n});\n</script>\n">> ];
+            "js" -> 
+                [ $\n, Script, $\n ];
             "escapejs" ->
                 z_utils:js_escape(Script)
         end.

src/support/z_session_page.erl

@@ -338,7 +338,7 @@ ping_comet_ws(#page_state{comet_pid=undefined, websocket_pid=undefined} = State)
     State;
 ping_comet_ws(#page_state{websocket_pid=WsPid} = State) when is_pid(WsPid) ->
     try
-        State#page_state.websocket_pid ! {send_data, lists:reverse(State#page_state.script_queue)},
+        controller_websocket:websocket_send_data(WsPid, lists:reverse(State#page_state.script_queue)),
         State#page_state{script_queue=[]}
     catch _M : _E ->
         State#page_state{websocket_pid=undefined}