From f9b886e52493740f297a7091387f2e492d8f50f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Viktor=20S=C3=B6derqvist?= <viktor.soderqvist@est.tech>
Date: Wed, 23 Nov 2022 11:42:00 +0100
Subject: Add tests for ws subprotocol negotiation

---
 test/handlers/ws_subprotocol_h.erl | 31 +++++++++++++++++++++++++++++++
 test/ws_SUITE.erl                  | 38 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 test/handlers/ws_subprotocol_h.erl

diff --git a/test/handlers/ws_subprotocol_h.erl b/test/handlers/ws_subprotocol_h.erl
new file mode 100644
index 0000000..509d74d
--- /dev/null
+++ b/test/handlers/ws_subprotocol_h.erl
@@ -0,0 +1,31 @@
+%% Feel free to use, reuse and abuse the code in this file.
+
+-module(ws_subprotocol_h).
+
+-export([init/2]).
+-export([websocket_handle/2]).
+-export([websocket_info/2]).
+
+init(Req, State) ->
+	Protos = cowboy_req:parse_header(<<"sec-websocket-protocol">>, Req),
+	init_protos(Req, State, Protos).
+
+init_protos(Req, State, undefined) ->
+	{ok, cowboy_req:reply(400, #{}, <<"undefined">>, Req), State};
+init_protos(Req, State, []) ->
+	{ok, cowboy_req:reply(400, #{}, <<"nomatch">>, Req), State};
+init_protos(Req0, State, [<<"echo">> | _]) ->
+	Req = cowboy_req:set_resp_header(<<"sec-websocket-protocol">>, <<"echo">>, Req0),
+	{cowboy_websocket, Req, State};
+init_protos(Req, State, [_ | Protos]) ->
+	init_protos(Req, State, Protos).
+
+websocket_handle({text, Data}, State) ->
+	{[{text, Data}], State};
+websocket_handle({binary, Data}, State) ->
+	{[{binary, Data}], State};
+websocket_handle(_Frame, State) ->
+	{[], State}.
+
+websocket_info(_Info, State) ->
+	{[], State}.
diff --git a/test/ws_SUITE.erl b/test/ws_SUITE.erl
index 5b7c528..709eb89 100644
--- a/test/ws_SUITE.erl
+++ b/test/ws_SUITE.erl
@@ -40,7 +40,8 @@ groups() ->
 init_per_suite(Config) ->
 	Routes = [
 		{"/", ws_echo_h, []},
-		{"/reject", ws_reject_h, []}
+		{"/reject", ws_reject_h, []},
+		{"/subprotocol", ws_subprotocol_h, []}
 	],
 	{ok, _} = cowboy:start_clear(ws, [], #{
 		enable_connect_protocol => true,
@@ -219,6 +220,41 @@ send_many_close(Config) ->
 	{ws, close} = gun:await(ConnPid, StreamRef),
 	gun:close(ConnPid).
 
+subprotocol_match(Config) ->
+	doc("Websocket subprotocol successfully negotiated."),
+	Protocols = [{P, gun_ws_h} || P <- [<<"dummy">>, <<"echo">>, <<"junk">>]],
+	{ok, ConnPid} = gun:open("localhost", config(port, Config)),
+	{ok, _} = gun:await_up(ConnPid),
+	StreamRef = gun:ws_upgrade(ConnPid, "/subprotocol", [], #{
+		protocols => Protocols
+	}),
+	{upgrade, [<<"websocket">>], _} = gun:await(ConnPid, StreamRef),
+	Frame = {text, <<"Hello!">>},
+	gun:ws_send(ConnPid, StreamRef, Frame),
+	{ws, Frame} = gun:await(ConnPid, StreamRef),
+	gun:close(ConnPid).
+
+subprotocol_nomatch(Config) ->
+	doc("Websocket subprotocol negotiation failure."),
+	Protocols = [{P, gun_ws_h} || P <- [<<"dummy">>, <<"junk">>]],
+	{ok, ConnPid} = gun:open("localhost", config(port, Config)),
+	{ok, _} = gun:await_up(ConnPid),
+	StreamRef = gun:ws_upgrade(ConnPid, "/subprotocol", [], #{
+		protocols => Protocols
+	}),
+	{response, nofin, 400, _} = gun:await(ConnPid, StreamRef),
+	{ok, <<"nomatch">>} = gun:await_body(ConnPid, StreamRef),
+	gun:close(ConnPid).
+
+subprotocol_required_but_missing(Config) ->
+	doc("Websocket subprotocol not negotiated but required by the server."),
+	{ok, ConnPid} = gun:open("localhost", config(port, Config)),
+	{ok, _} = gun:await_up(ConnPid),
+	StreamRef = gun:ws_upgrade(ConnPid, "/subprotocol", []),
+	{response, nofin, 400, _} = gun:await(ConnPid, StreamRef),
+	{ok, <<"undefined">>} = gun:await_body(ConnPid, StreamRef),
+	gun:close(ConnPid).
+
 %% Internal.
 
 do_await_enable_connect_protocol(http, _) ->
-- 
cgit v1.2.3