Add headers to gun_ws_upgrade message

Also improves the code and documentation about this message.
It was incorrectly specified that a gun_ws_upgrade message
could be sent on error; instead a gun_response is sent.

6 files changed, 25 insertions, 32 deletions

diff --git a/doc/src/guide/websocket.asciidoc b/doc/src/guide/websocket.asciidoc
index 83a8384..35580ca 100644
--- a/doc/src/guide/websocket.asciidoc
+++ b/doc/src/guide/websocket.asciidoc
@@ -40,17 +40,25 @@ The fourth argument is those same options. This function call
 will crash if the options are incorrect, unlike when passing
 them through `gun:open/{2,3}`.
 
-The success or failure of this operation will be sent as a
-message.
+When the upgrade succeeds, a `gun_ws_upgrade` message is sent.
+If the server does not understand Websocket or refused the
+upgrade, a `gun_response` message is sent. If Gun couldn't
+perform the upgrade due to an error (for example attempting
+to upgrade to Websocket on an HTTP/1.0 connection) then a
+`gun_error` message is sent.
 
-@todo hmm we want the headers to be sent in the gun_ws_upgrade ok message too
+When the server does not understand Websocket, it may send
+a meaningful response which should be processed. In the
+following example we however ignore it:
 
 [source,erlang]
 receive
-	{gun_ws_upgrade, ConnPid, ok} ->
+	{gun_ws_upgrade, ConnPid, ok, Headers} ->
 		upgrade_success(ConnPid);
-	{gun_ws_upgrade, ConnPid, error, IsFin, Status, Headers} ->
+	{gun_response, ConnPid, _, _, Status, Headers} ->
 		exit({ws_upgrade_failed, Status, Headers});
+	{gun_error, ConnPid, StreamRef, Reason} ->
+		exit({ws_upgrade_failed, Reason});
 	%% More clauses here as needed.
 after 1000 ->
 	exit(timeout);
diff --git a/doc/src/manual/gun.asciidoc b/doc/src/manual/gun.asciidoc
index 7c12fc3..baaadcf 100644
--- a/doc/src/manual/gun.asciidoc
+++ b/doc/src/manual/gun.asciidoc
@@ -194,22 +194,12 @@ Reason = any():: Error reason.
 
 General error.
 
-=== {gun_ws_upgrade, ConnPid, ok}
+=== {gun_ws_upgrade, ConnPid, ok, Headers}
 
 ConnPid = pid():: The pid of the Gun connection process.
-
-Successful upgrade to the Websocket protocol.
-
-@todo Yeah we need the headers.
-
-=== {gun_ws_upgrade, ConnPid, error, IsFin, Status, Headers}
-
-ConnPid = pid():: The pid of the Gun connection process.
-IsFin = fin | nofin:: Whether this message terminates the response.
-Status = binary():: Status line for the response.
 Headers = [{binary(), binary()}]:: Headers sent with the response.
 
-Failed upgrade to the Websocket protocol.
+Successful upgrade to the Websocket protocol.
 
 === {gun_ws, ConnPid, Frame}
 
diff --git a/src/gun.erl b/src/gun.erl
index ae30494..9ebb34a 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -393,7 +393,7 @@ flush_pid(ServerPid) ->
 			flush_pid(ServerPid);
 		{gun_error, ServerPid, _} ->
 			flush_pid(ServerPid);
-		{gun_ws_upgrade, ServerPid, _} ->
+		{gun_ws_upgrade, ServerPid, _, _} ->
 			flush_pid(ServerPid);
 		{gun_ws, ServerPid, _} ->
 			flush_pid(ServerPid);
diff --git a/src/gun_http.erl b/src/gun_http.erl
index 6e65341..67c235f 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -452,11 +452,11 @@ ws_validate_extensions(_, _, _, _) ->
 	close.
 
 %% @todo Validate protocols.
-ws_handshake_protocols(Buffer, State, _Headers, Extensions, _GunProtocols = []) ->
+ws_handshake_protocols(Buffer, State, Headers, Extensions, _GunProtocols = []) ->
 	Protocols = [],
-	ws_handshake_end(Buffer, State, Extensions, Protocols).
+	ws_handshake_end(Buffer, State, Headers, Extensions, Protocols).
 
-ws_handshake_end(Buffer, #http_state{owner=Owner, socket=Socket, transport=Transport}, Extensions, Protocols) ->
+ws_handshake_end(Buffer, #http_state{owner=Owner, socket=Socket, transport=Transport}, Headers, Extensions, Protocols) ->
 	%% Send ourselves the remaining buffer, if any.
 	_ = case Buffer of
 		<<>> ->
@@ -465,4 +465,4 @@ ws_handshake_end(Buffer, #http_state{owner=Owner, socket=Socket, transport=Trans
 			{OK, _, _} = Transport:messages(),
 			self() ! {OK, Socket, Buffer}
 	end,
-	gun_ws:init(Owner, Socket, Transport, Extensions, Protocols).
+	gun_ws:init(Owner, Socket, Transport, Headers, Extensions, Protocols).
diff --git a/src/gun_ws.erl b/src/gun_ws.erl
index 66d0fa2..a8154d6 100644
--- a/src/gun_ws.erl
+++ b/src/gun_ws.erl
@@ -16,7 +16,7 @@
 
 -export([check_options/1]).
 -export([name/0]).
--export([init/5]).
+-export([init/6]).
 -export([handle/2]).
 -export([send/2]).
 -export([down/1]).
@@ -56,8 +56,8 @@ do_check_options([Opt|_]) ->
 name() -> ws.
 
 %% @todo Protocols
-init(Owner, Socket, Transport, Extensions, _Protocols) ->
-	Owner ! {gun_ws_upgrade, self(), ok},
+init(Owner, Socket, Transport, Headers, Extensions, _Protocols) ->
+	Owner ! {gun_ws_upgrade, self(), ok, Headers},
 	{upgrade, ?MODULE, #ws_state{owner=Owner, socket=Socket, transport=Transport, extensions=Extensions}}.
 
 %% Do not handle anything if we received a close frame.
diff --git a/test/ws_SUITE.erl b/test/ws_SUITE.erl
index 2df4329..ce2c303 100644
--- a/test/ws_SUITE.erl
+++ b/test/ws_SUITE.erl
@@ -143,16 +143,11 @@ log_output() ->
 
 connect(Path) ->
 	{ok, Pid} = gun:open("127.0.0.1", 33080, #{retry=>0}),
-	receive
-		{gun_up, Pid, http} ->
-			ok
-	after 1000 ->
-		error(open_timeout)
-	end,
+	{ok, http} = gun:await_up(Pid),
 	Ref = monitor(process, Pid),
 	gun:ws_upgrade(Pid, Path, [], #{compress => true}),
 	receive
-		{gun_ws_upgrade, Pid, ok} ->
+		{gun_ws_upgrade, Pid, ok, _} ->
 			ok;
 		Msg ->
 			ct:print("Unexpected message ~p", [Msg]),