From fa8b2a5ac5492a3110c81a3c2e5fff838d58cf41 Mon Sep 17 00:00:00 2001
From: zhouzb <zhouzb@emqx.io>
Date: Wed, 17 Mar 2021 15:32:05 +0800
Subject: [PATCH] chore(test-cswsh): add test case for cswsh

---
 test/emqx_ws_connection_SUITE.erl | 45 ++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/test/emqx_ws_connection_SUITE.erl b/test/emqx_ws_connection_SUITE.erl
index 71c6e40e5..d83d7d859 100644
--- a/test/emqx_ws_connection_SUITE.erl
+++ b/test/emqx_ws_connection_SUITE.erl
@@ -43,7 +43,9 @@ all() -> emqx_ct:all(?MODULE).
 
 init_per_testcase(TestCase, Config) when
     TestCase =/= t_ws_sub_protocols_mqtt_equivalents,
-    TestCase =/= t_ws_sub_protocols_mqtt ->
+    TestCase =/= t_ws_sub_protocols_mqtt,
+    TestCase =/= t_ws_check_origin,
+    TestCase =/= t_ws_non_check_origin ->
     %% Mock cowboy_req
     ok = meck:new(cowboy_req, [passthrough, no_history, no_link]),
     ok = meck:expect(cowboy_req, peer, fun(_) -> {{127,0,0,1}, 3456} end),
@@ -85,7 +87,9 @@ init_per_testcase(_, Config) ->
 
 end_per_testcase(TestCase, Config) when
         TestCase =/= t_ws_sub_protocols_mqtt_equivalents,
-        TestCase =/= t_ws_sub_protocols_mqtt ->
+        TestCase =/= t_ws_sub_protocols_mqtt,
+        TestCase =/= t_ws_check_origin,
+        TestCase =/= t_ws_non_check_origin ->
     lists:foreach(fun meck:unload/1,
                   [cowboy_req,
                    emqx_zone,
@@ -171,6 +175,41 @@ t_ws_sub_protocols_mqtt_equivalents(_) ->
         start_ws_client(#{protocols => [<<"not-mqtt">>]})),
     emqx_ct_helpers:stop_apps([]).
 
+t_ws_check_origin(_) ->
+    emqx_ct_helpers:start_apps([],
+        fun(emqx) ->
+            {ok, Listeners} = application:get_env(emqx, listeners),
+            NListeners = lists:map(fun(#{listen_on := 8083, opts := Opts} = Listener) ->
+                                       NOpts = proplists:delete(check_origin_enable, Opts),
+                                       Listener#{opts => [{check_origin_enable, true} | NOpts]};
+                                      (Listener) ->
+                                       Listener
+                                   end, Listeners),
+            application:set_env(emqx, listeners, NListeners),
+            ok;
+           (_) -> ok
+        end),
+    {ok, _} = application:ensure_all_started(gun),
+    ?assertMatch({gun_upgrade, _},
+        start_ws_client(#{protocols => [<<"mqtt">>],
+                          headers => [{<<"origin">>, <<"http://localhost:18083">>}]})),
+    ?assertMatch({gun_response, {_, 500, _}},
+        start_ws_client(#{protocols => [<<"mqtt">>],
+                          headers => [{<<"origin">>, <<"http://localhost:18080">>}]})),
+    emqx_ct_helpers:stop_apps([]).
+
+t_ws_non_check_origin(_) ->
+    emqx_ct_helpers:start_apps([]),
+    {ok, _} = application:ensure_all_started(gun),
+    ?assertMatch({gun_upgrade, _},
+        start_ws_client(#{protocols => [<<"mqtt">>],
+                          headers => [{<<"origin">>, <<"http://localhost:18083">>}]})),
+    ?assertMatch({gun_upgrade, _},
+        start_ws_client(#{protocols => [<<"mqtt">>],
+                          headers => [{<<"origin">>, <<"http://localhost:18080">>}]})),
+    emqx_ct_helpers:stop_apps([]).
+
+
 t_init(_) ->
     Opts = [{idle_timeout, 300000},
             {fail_if_no_subprotocol, false},
@@ -433,7 +472,7 @@ ws_client(State) ->
     receive
         {gun_up, WPID, _Proto} ->
             #{protocols := Protos} = State,
-            StreamRef = gun:ws_upgrade(WPID, "/mqtt", [],
+            StreamRef = gun:ws_upgrade(WPID, "/mqtt", maps:get(headers, State, []),
                 #{protocols => [{P, gun_ws_h} || P <- Protos]}),
             ws_client(State#{wref => StreamRef});
         {gun_down, _WPID, _, Reason, _, _} ->