review.fuel-infra Code Review - packages/trusty/rabbitmq-server.git/blob

   1 -module(sockjs_cowboy_handler).
   2 -behaviour(cowboy_http_handler).
   3 -behaviour(cowboy_http_websocket_handler).
   4
   5 %% Cowboy http callbacks
   6 -export([init/3, handle/2, terminate/2]).
   7
   8 %% Cowboy ws callbacks
   9 -export([websocket_init/3, websocket_handle/3,
  10          websocket_info/3, websocket_terminate/3]).
  11
  12 -include("sockjs_internal.hrl").
  13
  14 %% --------------------------------------------------------------------------
  15
  16 init({_Any, http}, Req, Service) ->
  17     case sockjs_handler:is_valid_ws(Service, {cowboy, Req}) of
  18         {true, {cowboy, _Req1}, _Reason} ->
  19             {upgrade, protocol, cowboy_http_websocket};
  20         {false, {cowboy, Req1}, _Reason} ->
  21             {ok, Req1, Service}
  22     end.
  23
  24 handle(Req, Service) ->
  25     {cowboy, Req3} = sockjs_handler:handle_req(Service, {cowboy, Req}),
  26     {ok, Req3, Service}.
  27
  28 terminate(_Req, _Service) ->
  29     ok.
  30
  31 %% --------------------------------------------------------------------------
  32
  33 websocket_init(_TransportName, Req,
  34                Service = #service{logger        = Logger,
  35                                   subproto_pref = SubProtocolPref}) ->
  36     Req3 = case cowboy_http_req:header(<<"Sec-Websocket-Protocol">>, Req) of
  37                {undefined, Req1} ->
  38                    Req1;
  39                {SubProtocols, Req1} ->
  40                    SelectedSubProtocol =
  41                      choose_subprotocol_bin(SubProtocols, SubProtocolPref),
  42                    {ok, Req2} = cowboy_http_req:set_resp_header(
  43                                   <<"Sec-Websocket-Protocol">>,
  44                                   SelectedSubProtocol, Req1),
  45                    Req2
  46            end,
  47
  48     Req4 = Logger(Service, {cowboy, Req3}, websocket),
  49
  50     Service1 = Service#service{disconnect_delay = 5*60*1000},
  51
  52     {Info, Req5} = sockjs_handler:extract_info(Req4),
  53     SessionPid = sockjs_session:maybe_create(undefined, Service1, Info),
  54     {RawWebsocket, {cowboy, Req7}} =
  55         case sockjs_handler:get_action(Service, Req5) of
  56             {{match, WS}, Req6} when WS =:= websocket orelse
  57                                      WS =:= rawwebsocket ->
  58                 {WS, Req6}
  59         end,
  60     self() ! go,
  61     {ok, Req7, {RawWebsocket, SessionPid}}.
  62
  63 websocket_handle({text, Data}, Req, {RawWebsocket, SessionPid} = S) ->
  64     case sockjs_ws_handler:received(RawWebsocket, SessionPid, Data) of
  65         ok       -> {ok, Req, S};
  66         shutdown -> {shutdown, Req, S}
  67     end;
  68 websocket_handle(_Unknown, Req, S) ->
  69     {shutdown, Req, S}.
  70
  71 websocket_info(go, Req, {RawWebsocket, SessionPid} = S) ->
  72     case sockjs_ws_handler:reply(RawWebsocket, SessionPid) of
  73         wait          -> {ok, Req, S};
  74         {ok, Data}    -> self() ! go,
  75                          {reply, {text, Data}, Req, S};
  76         {close, <<>>} -> {shutdown, Req, S};
  77         {close, Data} -> self() ! shutdown,
  78                          {reply, {text, Data}, Req, S}
  79     end;
  80 websocket_info(shutdown, Req, S) ->
  81     {shutdown, Req, S}.
  82
  83 websocket_terminate(_Reason, _Req, {RawWebsocket, SessionPid}) ->
  84     sockjs_ws_handler:close(RawWebsocket, SessionPid),
  85     ok.
  86
  87 %% --------------------------------------------------------------------------
  88
  89 choose_subprotocol_bin(SubProtocols, Pref) ->
  90     choose_subprotocol(re:split(SubProtocols, ", *"), Pref).
  91 choose_subprotocol(SubProtocols, undefined) ->
  92     erlang:hd(lists:reverse(lists:sort(SubProtocols)));
  93 choose_subprotocol(SubProtocols, Pref) ->
  94     case lists:filter(fun (E) -> lists:member(E, SubProtocols) end, Pref) of
  95         [Hd | _] -> Hd;
  96         []       -> choose_subprotocol(SubProtocols, undefined)
  97     end.