Add headers to mqtt_client

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,7 +54,7 @@
                       keepalive,
                       max_clientid_len = ?MAX_CLIENTID_LEN,
                       client_pid,
-                      ws_cookie, %% for websocket client
+                      ws_initial_headers, %% Headers from first HTTP request for websocket client
                       connected_at}).
 
 -type proto_state() :: #proto_state{}.
@@ -66,22 +66,23 @@
 
 init(Peername, SendFun, Opts) ->
     MaxLen = proplists:get_value(max_clientid_len, Opts, ?MAX_CLIENTID_LEN),
-    WsCookie = proplists:get_value(ws_cookie, Opts),
-	#proto_state{peername         = Peername,
-                 sendfun          = SendFun,
-                 max_clientid_len = MaxLen,
-                 client_pid       = self(),
-                 ws_cookie        = WsCookie}.
-
-info(#proto_state{client_id	   = ClientId,
-                  username     = Username,
-                  peername     = Peername,
-                  proto_ver    = ProtoVer,
-                  proto_name   = ProtoName,
-                  keepalive    = KeepAlive,
-				  clean_sess   = CleanSess,
-				  will_msg	   = WillMsg,
-                  connected_at = ConnectedAt}) ->
+    WsInitialHeaders = proplists:get_value(ws_initial_headers, Opts),
+	#proto_state{peername           = Peername,
+                 sendfun            = SendFun,
+                 max_clientid_len   = MaxLen,
+                 client_pid         = self(),
+                 ws_initial_headers = WsInitialHeaders}.
+
+info(#proto_state{client_id	         = ClientId,
+                  username           = Username,
+                  peername           = Peername,
+                  proto_ver          = ProtoVer,
+                  proto_name         = ProtoName,
+                  keepalive          = KeepAlive,
+				  clean_sess         = CleanSess,
+                  ws_initial_headers = WsInitialHeaders,
+				  will_msg	         = WillMsg,
+                  connected_at       = ConnectedAt}) ->
     [{client_id, ClientId},
      {username, Username},
      {peername, Peername},
@@ -89,36 +90,37 @@ info(#proto_state{client_id	   = ClientId,
      {proto_name, ProtoName},
      {keepalive, KeepAlive},
      {clean_sess, CleanSess},
+     {ws_initial_headers, WsInitialHeaders},
      {will_msg,	WillMsg},
      {connected_at, ConnectedAt}].
 
 clientid(#proto_state{client_id = ClientId}) ->
     ClientId.
 
-client(#proto_state{client_id  = ClientId,
-                    peername   = Peername,
-                    username   = Username,
-                    clean_sess = CleanSess,
-                    proto_ver  = ProtoVer,
-                    keepalive  = Keepalive,
-                    will_msg   = WillMsg,
-                    client_pid = Pid,
-                    ws_cookie  = WsCookie,
-                    connected_at = Time}) ->
+client(#proto_state{client_id          = ClientId,
+                    peername           = Peername,
+                    username           = Username,
+                    clean_sess         = CleanSess,
+                    proto_ver          = ProtoVer,
+                    keepalive          = Keepalive,
+                    will_msg           = WillMsg,
+                    client_pid         = Pid,
+                    ws_initial_headers = WsInitialHeaders,
+                    connected_at       = Time}) ->
     WillTopic = if
                     WillMsg =:= undefined -> undefined;
                     true -> WillMsg#mqtt_message.topic
                 end,
-    #mqtt_client{client_id  = ClientId,
-                 client_pid = Pid,
-                 username   = Username,
-                 peername   = Peername,
-                 clean_sess = CleanSess,
-                 proto_ver  = ProtoVer,
-                 keepalive  = Keepalive,
-                 will_topic = WillTopic,
-                 ws_cookie  = WsCookie,
-                 connected_at = Time}.
+    #mqtt_client{client_id          = ClientId,
+                 client_pid         = Pid,
+                 username           = Username,
+                 peername           = Peername,
+                 clean_sess         = CleanSess,
+                 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,8 +104,9 @@ init([WsPid, Req, ReplyChannel, PktOpts]) ->
     process_flag(trap_exit, true),
     {ok, Peername} = Req:get(peername),
     SendFun = fun(Payload) -> ReplyChannel({binary, Payload}) end,
-    Cookie = Req:parse_cookie(),
-    ProtoState = emqttd_protocol:init(Peername, SendFun, [{ws_cookie, Cookie}|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) ->