5 files changed, 78 insertions, 6 deletions
diff --git a/src/cowboy.erl b/src/cowboy.erl
index 79dbb71..257172d 100644
--- a/src/cowboy.erl
+++ b/src/cowboy.erl
@@ -18,6 +18,7 @@
 -export([start_http/4]).
 -export([start_https/4]).
 -export([stop_listener/1]).
+-export([set_env/3]).
 
 %% @doc Start an HTTP listener.
 -spec start_http(any(), non_neg_integer(), any(), any()) -> {ok, pid()}.
@@ -37,3 +38,16 @@ start_https(Ref, NbAcceptors, TransOpts, ProtoOpts)
 -spec stop_listener(any()) -> ok.
 stop_listener(Ref) ->
 	ranch:stop_listener(Ref).
+
+%% @doc Convenience function for setting an environment value.
+%%
+%% Allows you to update live an environment value used by middlewares.
+%% This function is primarily intended to simplify updating the dispatch
+%% list used for routing.
+-spec set_env(any(), atom(), any()) -> ok.
+set_env(Ref, Name, Value) ->
+	Opts = ranch:get_protocol_options(Ref),
+	{_, Env} = lists:keyfind(env, 1, Opts),
+	Env2 = [{Name, Value}|lists:keydelete(Name, 1, Env)],
+	Opts2 = lists:keyreplace(env, 1, Opts, {env, Env2}),
+	ok = ranch:set_protocol_options(Ref, Opts2).
diff --git a/src/cowboy_handler.erl b/src/cowboy_handler.erl
index 65e22b7..fd2e085 100644
--- a/src/cowboy_handler.erl
+++ b/src/cowboy_handler.erl
@@ -62,7 +62,7 @@ execute(Req, Env) ->
 
 -spec handler_init(Req, #state{}, module(), any())
 	-> {ok, Req, cowboy_middleware:env()}
-	| {error, 500, Req} | {suspend, module(), function(), [any()]}
+	| {error, 500, Req} | {suspend, module(), atom(), [any()]}
 	when Req::cowboy_req:req().
 handler_init(Req, State, Handler, HandlerOpts) ->
 	Transport = cowboy_req:get(transport, Req),
@@ -137,7 +137,7 @@ handler_handle(Req, State, Handler, HandlerState) ->
 %% to receive data and buffer it indefinitely.
 -spec handler_before_loop(Req, #state{}, module(), any())
 	-> {ok, Req, cowboy_middleware:env()}
-	| {error, 500, Req} | {suspend, module(), function(), [any()]}
+	| {error, 500, Req} | {suspend, module(), atom(), [any()]}
 	when Req::cowboy_req:req().
 handler_before_loop(Req, State=#state{hibernate=true}, Handler, HandlerState) ->
 	[Socket, Transport] = cowboy_req:get([socket, transport], Req),
@@ -165,7 +165,7 @@ handler_loop_timeout(State=#state{loop_timeout=Timeout,
 %% @private
 -spec handler_loop(Req, #state{}, module(), any())
 	-> {ok, Req, cowboy_middleware:env()}
-	| {error, 500, Req} | {suspend, module(), function(), [any()]}
+	| {error, 500, Req} | {suspend, module(), atom(), [any()]}
 	when Req::cowboy_req:req().
 handler_loop(Req, State=#state{loop_buffer_size=NbBytes,
 		loop_max_buffer=Threshold, loop_timeout_ref=TRef},
@@ -195,7 +195,7 @@ handler_loop(Req, State=#state{loop_buffer_size=NbBytes,
 			handler_before_loop(Req, State#state{resp_sent=true},
 				Handler, HandlerState);
 		{timeout, TRef, ?MODULE} ->
-			terminate_request(Req, State, Handler, HandlerState,
+			handler_after_loop(Req, State, Handler, HandlerState,
 				{normal, timeout});
 		{timeout, OlderTRef, ?MODULE} when is_reference(OlderTRef) ->
 			handler_before_loop(Req, State, Handler, HandlerState);
@@ -205,12 +205,12 @@ handler_loop(Req, State=#state{loop_buffer_size=NbBytes,
 
 -spec handler_call(Req, #state{}, module(), any(), any())
 	-> {ok, Req, cowboy_middleware:env()}
-	| {error, 500, Req} | {suspend, module(), function(), [any()]}
+	| {error, 500, Req} | {suspend, module(), atom(), [any()]}
 	when Req::cowboy_req:req().
 handler_call(Req, State, Handler, HandlerState, Message) ->
 	try Handler:info(Message, Req, HandlerState) of
 		{ok, Req2, HandlerState2} ->
-			terminate_request(Req2, State, Handler, HandlerState2,
+			handler_after_loop(Req2, State, Handler, HandlerState2,
 				{normal, shutdown});
 		{loop, Req2, HandlerState2} ->
 			handler_before_loop(Req2, State, Handler, HandlerState2);
@@ -230,6 +230,25 @@ handler_call(Req, State, Handler, HandlerState, Message) ->
 		error_terminate(Req, State)
 	end.
 
+%% It is sometimes important to make a socket passive as it was initially
+%% and as it is expected to be by cowboy_protocol, right after we're done
+%% with loop handling. The browser may freely pipeline a bunch of requests
+%% if previous one was, say, a JSONP long-polling request.
+-spec handler_after_loop(Req, #state{}, module(), any(),
+	{normal, timeout | shutdown} | {error, atom()}) ->
+	{ok, Req, cowboy_middleware:env()} when Req::cowboy_req:req().
+handler_after_loop(Req, State, Handler, HandlerState, Reason) ->
+	[Socket, Transport] = cowboy_req:get([socket, transport], Req),
+	Transport:setopts(Socket, [{active, false}]),
+	{OK, _Closed, _Error} = Transport:messages(),
+	Req2 = receive
+		{OK, Socket, Data} ->
+			cowboy_req:append_buffer(Data, Req)
+	after 0 ->
+		Req
+	end,
+	terminate_request(Req2, State, Handler, HandlerState, Reason).
+
 -spec terminate_request(Req, #state{}, module(), any(),
 	{normal, timeout | shutdown} | {error, atom()}) ->
 	{ok, Req, cowboy_middleware:env()} when Req::cowboy_req:req().
diff --git a/src/cowboy_rest.erl b/src/cowboy_rest.erl
index d3a91d3..a40e1c6 100644
--- a/src/cowboy_rest.erl
+++ b/src/cowboy_rest.erl
@@ -19,6 +19,7 @@
 %% documentation available at http://wiki.basho.com/Webmachine.html
 %% at the time of writing.
 -module(cowboy_rest).
+-behaviour(cowboy_sub_protocol).
 
 -export([upgrade/4]).
 
diff --git a/src/cowboy_sub_protocol.erl b/src/cowboy_sub_protocol.erl
new file mode 100644
index 0000000..0b231d3
--- /dev/null
+++ b/src/cowboy_sub_protocol.erl
@@ -0,0 +1,37 @@
+%% Copyright (c) 2013, James Fish <[email protected]>
+%%
+%% Permission to use, copy, modify, and/or distribute this software for any
+%% purpose with or without fee is hereby granted, provided that the above
+%% copyright notice and this permission notice appear in all copies.
+%%
+%% THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+%% WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+%% MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+%% ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+%% WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+%% ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+%% OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+%% @doc Behaviour for sub protocols.
+%%
+%% Only one function needs to be implemented, <em>upgrade/4</em>.
+%% It receives the Req, the environment, the handler that the request has been
+%% routed to and the handler's options. It acts exactly the same as a
+%% middleware, so returns the same values a middleware's execute/2.
+%%
+%% Once the sub protocol has processed the request it should add the result
+%% to the environment. This is done by adding the tuple {result, Value} to the
+%% environment list. To continue handling requests on the current connection the
+%% Value should be the atom ok. Any other value will prevent the processing of
+%% subsequent requests.
+%%
+%% <em>upgrade/4</em> will be called when a handler's init/3 returns
+%% {upgrade, protocol, Module}, where Module is the module of the sub protocol.
+-module(cowboy_sub_protocol).
+
+-callback upgrade(Req, Env, module(), any())
+	-> {ok, Req, Env}
+	| {suspend, module(), atom(), any()}
+	| {halt, Req}
+	| {error, cowboy_http:status(), Req}
+	when Req::cowboy_req:req(), Env::cowboy_middleware:env().
diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index fa74223..7bb5e74 100644
--- a/src/cowboy_websocket.erl
+++ b/src/cowboy_websocket.erl
@@ -17,6 +17,7 @@
 %% Cowboy supports versions 7 through 17 of the Websocket drafts.
 %% It also supports RFC6455, the proposed standard for Websocket.
 -module(cowboy_websocket).
+-behaviour(cowboy_sub_protocol).
 
 %% API.
 -export([upgrade/4]).