Merge pull request #241 from emqtt/issue#231

Issue#231

include/emqttd.hrl

@@ -82,6 +82,10 @@
 %%------------------------------------------------------------------------------
 %% MQTT Client
 %%------------------------------------------------------------------------------
+
+-type header_key() :: atom() | binary() | string().
+-type header_val() :: atom() | binary() | string() | integer().
+
 -record(mqtt_client, {
     client_id     :: binary() | undefined,
     client_pid    :: pid(),
@@ -91,6 +95,7 @@
     proto_ver     :: 3 | 4,
     keepalive = 0,
     will_topic    :: undefined | binary(),
+    ws_initial_headers :: list({header_key(), header_val()}),
     connected_at  :: erlang:timestamp()
 }).
 

src/emqttd_protocol.erl

@@ -54,6 +54,7 @@
                       keepalive,
                       max_clientid_len = ?MAX_CLIENTID_LEN,
                       client_pid,
+                      ws_initial_headers, %% Headers from first HTTP request for websocket client
                       connected_at}).
 
 -type proto_state() :: #proto_state{}.
@@ -65,10 +66,12 @@
 
 init(Peername, SendFun, Opts) ->
     MaxLen = proplists:get_value(max_clientid_len, Opts, ?MAX_CLIENTID_LEN),
+    WsInitialHeaders = proplists:get_value(ws_initial_headers, Opts),
 	#proto_state{peername           = Peername,
                  sendfun            = SendFun,
                  max_clientid_len   = MaxLen,
-                 client_pid       = self()}.
+                 client_pid         = self(),
+                 ws_initial_headers = WsInitialHeaders}.
 
 info(#proto_state{client_id	         = ClientId,
                   username           = Username,
@@ -77,6 +80,7 @@ info(#proto_state{client_id	   = ClientId,
                   proto_name         = ProtoName,
                   keepalive          = KeepAlive,
 				  clean_sess         = CleanSess,
+                  ws_initial_headers = WsInitialHeaders,
 				  will_msg	         = WillMsg,
                   connected_at       = ConnectedAt}) ->
     [{client_id, ClientId},
@@ -86,6 +90,7 @@ info(#proto_state{client_id	   = ClientId,
      {proto_name, ProtoName},
      {keepalive, KeepAlive},
      {clean_sess, CleanSess},
+     {ws_initial_headers, WsInitialHeaders},
      {will_msg,	WillMsg},
      {connected_at, ConnectedAt}].
 
@@ -100,6 +105,7 @@ client(#proto_state{client_id  = ClientId,
                     keepalive          = Keepalive,
                     will_msg           = WillMsg,
                     client_pid         = Pid,
+                    ws_initial_headers = WsInitialHeaders,
                     connected_at       = Time}) ->
     WillTopic = if
                     WillMsg =:= undefined -> undefined;
@@ -113,6 +119,7 @@ client(#proto_state{client_id  = ClientId,
                  proto_ver          = ProtoVer,
                  keepalive          = Keepalive,
                  will_topic         = WillTopic,
+                 ws_initial_headers = WsInitialHeaders,
                  connected_at       = Time}.
 
 %% CONNECT – Client requests a connection to a Server

src/emqttd_ws_client.erl

@@ -104,7 +104,9 @@ init([WsPid, Req, ReplyChannel, PktOpts]) ->
     process_flag(trap_exit, true),
     {ok, Peername} = Req:get(peername),
     SendFun = fun(Payload) -> ReplyChannel({binary, Payload}) end,
-    ProtoState = emqttd_protocol:init(Peername, SendFun, PktOpts),
+    Headers = mochiweb_request:get(headers, Req),
+    HeadersList = mochiweb_headers:to_list(Headers),
+    ProtoState = emqttd_protocol:init(Peername, SendFun, [{ws_initial_headers, HeadersList}|PktOpts]),
     {ok, #client_state{ws_pid = WsPid, request = Req, proto_state = ProtoState}}.
 
 handle_call(_Req, _From, State) ->