Properly handle the Websocket protocol header

We support v10.stomp, v11.stomp and v12.stomp at the moment. The better one will be selected if a client advertises two or more.
rabbitmq · Mar 30, 2016 · d81a9b8 · d81a9b8
1 parent 71e0511
commit d81a9b8
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 5 deletions.
diff --git a/src/rabbit_ws_handler.erl b/src/rabbit_ws_handler.erl
@@ -36,7 +36,18 @@
 init(_, _Req, _Opts) ->
     {upgrade, protocol, cowboy_websocket}.
 
-websocket_init(_TransportName, Req, [{type, FrameType}]) ->
+websocket_init(_TransportName, Req0, [{type, FrameType}]) ->
+    Req = case cowboy_req:header(<<"sec-websocket-protocol">>, Req0) of
+        {undefined, _} -> Req0;
+        {ProtocolHd, _} ->
+            Protocols = parse_sec_websocket_protocol_req(ProtocolHd),
+            case filter_stomp_protocols(Protocols) of
+                [] -> Req0;
+                [StompProtocol|_] ->
+                    cowboy_req:set_resp_header(<<"sec-websocket-protocol">>,
+                        StompProtocol, Req0)
+            end
+    end,
     {Peername, _} = cowboy_req:peer(Req),
     [Socket, Transport] = cowboy_req:get([socket, transport], Req),
     {ok, Sockname} = Transport:sockname(Socket),
@@ -72,6 +83,25 @@ websocket_terminate(_Reason, _Req, #state{pid=Pid}) ->
     rabbit_ws_client:sockjs_closed(Pid),
     ok.
 
+%% When moving to Cowboy 2, this code should be replaced
+%% with a simple call to cow_http_hd:parse_sec_websocket_protocol_req/1.
+
+parse_sec_websocket_protocol_req(Bin) ->
+    Protocols = binary:split(Bin, [<<$,>>, <<$\s>>], [global]),
+    [P || P <- Protocols, P =/= <<>>].
+
+%% The protocols v10.stomp, v11.stomp and v12.stomp are registered
+%% at IANA: https://www.iana.org/assignments/websocket/websocket.xhtml
+
+filter_stomp_protocols(Protocols) ->
+    lists:reverse(lists:sort(lists:filter(
+        fun(<< "v1", C, ".stomp">>)
+            when C =:= $2; C =:= $1; C =:= $0 -> true;
+           (_) ->
+            false
+        end,
+        Protocols))).
+
 %% SockJS connection handling.
 
 %% The following functions are replicating the functionality

diff --git a/test/src/rabbit_ws_test_raw_websocket.erl b/test/src/rabbit_ws_test_raw_websocket.erl
@@ -24,6 +24,13 @@ connection_test() ->
     {close, _} = rfc6455_client:close(WS),
     ok.
 
+connection_with_protocols_test() ->
+    WS = rfc6455_client:new("ws://127.0.0.1:15674/stomp/websocket", self(),
+        undefined, ["v11.stomp", "v10.stomp", "v12.stomp"]),
+    {ok, _} = rfc6455_client:open(WS),
+    {close, _} = rfc6455_client:close(WS),
+    ok.
+
 
 raw_send(WS, Command, Headers) ->
     raw_send(WS, Command, Headers, <<>>).

diff --git a/test/src/rfc6455_client.erl b/test/src/rfc6455_client.erl
@@ -16,16 +16,19 @@
 
 -module(rfc6455_client).
 
--export([new/2, new/3, open/1, recv/1, send/2, close/1, close/2]).
+-export([new/2, new/3, new/4, open/1, recv/1, send/2, close/1, close/2]).
 
 -record(state, {host, port, addr, path, ppid, socket, data, phase}).
 
 %% --------------------------------------------------------------------------
 
 new(WsUrl, PPid) ->
-    new(WsUrl, PPid, undefined).
+    new(WsUrl, PPid, undefined, []).
 
 new(WsUrl, PPid, AuthInfo) ->
+    new(WsUrl, PPid, AuthInfo, []).
+
+new(WsUrl, PPid, AuthInfo, Protocols) ->
     crypto:start(),
     "ws://" ++ Rest = WsUrl,
     [Addr, Path] = split("/", Rest, 1),
@@ -40,7 +43,7 @@ new(WsUrl, PPid, AuthInfo) ->
                    path = "/" ++ Path,
                    ppid = PPid},
     spawn(fun () ->
-                  start_conn(State, AuthInfo)
+                  start_conn(State, AuthInfo, Protocols)
           end).
 
 open(WS) ->
@@ -82,7 +85,7 @@ close(WS, WsReason) ->
 
 %% --------------------------------------------------------------------------
 
-start_conn(State, AuthInfo) ->
+start_conn(State, AuthInfo, Protocols) ->
     {ok, Socket} = gen_tcp:connect(State#state.host, State#state.port,
                                    [binary,
                                     {packet, 0}]),
@@ -97,13 +100,19 @@ start_conn(State, AuthInfo) ->
                 ++ "\r\n"
     end,
 
+    ProtocolHd = case Protocols of
+        [] -> "";
+        _  -> "Sec-Websocket-Protocol: " ++ string:join(Protocols, ", ")
+    end,
+
     Key = base64:encode_to_string(crypto:rand_bytes(16)),
     gen_tcp:send(Socket,
         "GET " ++ State#state.path ++ " HTTP/1.1\r\n" ++
         "Host: " ++ State#state.addr ++ "\r\n" ++
         "Upgrade: websocket\r\n" ++
         "Connection: Upgrade\r\n" ++
         AuthHd ++
+        ProtocolHd ++
         "Sec-WebSocket-Key: " ++ Key ++ "\r\n" ++
         "Origin: null\r\n" ++
         "Sec-WebSocket-Version: 13\r\n\r\n"),