From b883e238bfe638ee286104594c32baddf78420ae Mon Sep 17 00:00:00 2001
From: JianBo He <heeejianbo@gmail.com>
Date: Mon, 12 Oct 2020 17:05:06 +0800
Subject: [PATCH] fix(ws): take ws_cookie to clientinfo

see: https://github.com/emqx/emqx/issues/3747#issuecomment-702268570
---
 src/emqx_channel.erl        | 13 +++++++++++--
 test/emqx_channel_SUITE.erl | 12 ++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/emqx_channel.erl b/src/emqx_channel.erl
index 289ed50f0..71b89f153 100644
--- a/src/emqx_channel.erl
+++ b/src/emqx_channel.erl
@@ -183,8 +183,9 @@ init(ConnInfo = #{peername := {PeerHost, _Port},
                      is_bridge    => false,
                      is_superuser => false
                     }, Options),
-    #channel{conninfo   = ConnInfo,
-             clientinfo = ClientInfo,
+    {NClientInfo, NConnInfo} = take_ws_cookie(ClientInfo, ConnInfo),
+    #channel{conninfo   = NConnInfo,
+             clientinfo = NClientInfo,
              topic_aliases = #{inbound => #{},
                                outbound => #{}
                               },
@@ -213,6 +214,14 @@ setting_peercert_infos(Peercert, ClientInfo, Options) ->
                end,
     ClientInfo#{username => Username, dn => DN, cn => CN}.
 
+take_ws_cookie(ClientInfo, ConnInfo) ->
+    case maps:take(ws_cookie, ConnInfo) of
+        {WsCookie, NConnInfo} ->
+            {ClientInfo#{ws_cookie => WsCookie}, NConnInfo};
+        _ ->
+            {ClientInfo, ConnInfo}
+    end.
+
 %%--------------------------------------------------------------------
 %% Handle incoming packet
 %%--------------------------------------------------------------------
diff --git a/test/emqx_channel_SUITE.erl b/test/emqx_channel_SUITE.erl
index c90c7510c..613da40ac 100644
--- a/test/emqx_channel_SUITE.erl
+++ b/test/emqx_channel_SUITE.erl
@@ -687,6 +687,18 @@ t_terminate(_) ->
     ok = emqx_channel:terminate(sock_error, channel(#{conn_state => connected})),
     ok = emqx_channel:terminate({shutdown, kicked}, channel(#{conn_state => connected})).
 
+t_ws_cookie_init(_) ->
+    WsCookie = [{<<"session_id">>, <<"xyz">>}],
+    ConnInfo = #{socktype => ws,
+                 peername => {{127,0,0,1}, 3456},
+                 sockname => {{127,0,0,1}, 1883},
+                 peercert => nossl,
+                 conn_mod => emqx_ws_connection,
+                 ws_cookie => WsCookie
+                },
+    Channel = emqx_channel:init(ConnInfo, [{zone, zone}]),
+    ?assertMatch(#{ws_cookie := WsCookie}, emqx_channel:info(clientinfo, Channel)).
+
 %%--------------------------------------------------------------------
 %% Helper functions
 %%--------------------------------------------------------------------