fixup! WIPrewrite

2 files changed, 225 insertions, 10 deletions

diff --git a/src/cowboy_rewrite_h.erl b/src/cowboy_rewrite_h.erl
index 2b40600..80c2ca4 100644
--- a/src/cowboy_rewrite_h.erl
+++ b/src/cowboy_rewrite_h.erl
@@ -27,6 +27,7 @@
 
 %% Predefined rewrites.
 -export([index/2]).
+-export([regex/2]).
 -export([slash/2]).
 
 %% The outcome determines what happens after a rewrite
@@ -106,8 +107,6 @@ record_rewrite(Event, Req) ->
 
 %% Rewrite logic.
 
-%% @todo Mark the Req as being modified so we can detect it was.
-
 rewrite([], Req) ->
 	{ok, Req};
 rewrite([{F, Opts}|Tail], Req0) when is_atom(F) ->
@@ -155,7 +154,48 @@ index(Req=#{path := Path0}, Opts) ->
 			{continue, Req}
 	end.
 
-%% @todo regex
+%% Apply a regular expression to the path using re:replace/4.
+%%
+%% The regex can be provided as text or pre-compiled.
+%% Some options may not apply; the user must carefully choose
+%% which option(s) to use if they decide to use any.
+%%
+%% The resulting replacement may including a query string.
+%% In that case the new query string segment is pre-pended
+%% to the one in Req. The result must either be a /path
+%% or a /path?query_string. Users are responsible for the
+%% generated string, this function only checks that the
+%% result starts with a /.
+
+-spec regex(cowboy_req:req(), #{
+		regex := re:mp() | iodata() | unicode:charlist(),
+		replacement := iodata() | unicode:charlist() | re:replace_fun(),
+		options => re:options(),
+		outcome => outcome()})
+	-> {outcome(), cowboy_req:req()}.
+
+regex(Req=#{path := Path0}, Opts=#{regex := Regex, replacement := Replacement}) ->
+	Options = maps:get(options, Opts, []),
+	Result = iolist_to_binary(re:replace(Path0, Regex, Replacement, Options)),
+	<<"/",_/bits>> = Result, %% Result must be a path.
+	case string:split(Result, <<"?">>) of
+		[Path0] ->
+			{continue, Req};
+		[Path] ->
+			{maps:get(outcome, Opts, continue),
+				record_rewrite({path, Path0, Path},
+					Req#{path => Path})};
+		[Path, ReQs] ->
+			#{qs := Qs0} = Req,
+			Qs = case Qs0 of
+				<<>> -> ReQs;
+				_ -> iolist_to_binary([ReQs, $&, Qs0])
+			end,
+			{maps:get(outcome, Opts, continue),
+				record_rewrite({path, Path0, Path},
+					record_rewrite({qs, Qs0, Qs},
+						Req#{path => Path, qs => Qs}))}
+	end.
 
 %% Ensure there is a slash at the end of the path.
 
diff --git a/test/rewrite_SUITE.erl b/test/rewrite_SUITE.erl
index 8d81e0f..6c215bb 100644
--- a/test/rewrite_SUITE.erl
+++ b/test/rewrite_SUITE.erl
@@ -34,7 +34,7 @@ groups() ->
 	[{rewrite, [parallel], ct_helper:all(?MODULE)}].
 
 init_opts() ->
-	Dispatch = cowboy_router:compile([{"localhost", [
+	Dispatch = cowboy_router:compile([{'_', [
 		{"/[...]", send_message_h, #{}}
 	]}]),
 	#{
@@ -72,16 +72,100 @@ do_get(Path, Config) ->
 		gun:close(ConnPid)
 	end.
 
+do_check_location(Path, Headers) ->
+	{_, Location} = lists:keyfind(<<"location">>, 1, Headers),
+	#{path := Path} = uri_string:parse(Location),
+	ok.
+
 %% Tests.
 
 %% @todo custom_fun(Config) ->
 
-%% @todo custom_mf(Config) ->
+custom_fun(Config0) ->
+	Config = init_http(?FUNCTION_NAME, init_opts(), Config0),
+	try
+		Fun = fun
+			(Req, #{clause := 1}) ->
+				{continue, Req#{path => <<"/custom/path">>}};
+			(Req, #{clause := 2}) ->
+				{continue, Req#{host => <<"ninenines.eu">>}};
+			(Req, #{clause := 3}) ->
+				{{redirect, 302}, Req#{host => <<"ninenines.eu">>, path => <<"/">>}};
+			(Req, #{clause := 4}) ->
+				{stop, Req#{path => <<"/custom/path">>}}
+		end,
 
-do_check_location(Path, Headers) ->
-	{_, Location} = lists:keyfind(<<"location">>, 1, Headers),
-	#{path := Path} = uri_string:parse(Location),
-	ok.
+		%% Custom path.
+		?REWRITE([{Fun, #{clause => 1}}]),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/", Config),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/path/to", Config),
+
+		%% Custom host.
+		?REWRITE([{Fun, #{clause => 2}}]),
+		{ok, #{host := <<"ninenines.eu">>}} = do_get("/", Config),
+		{ok, #{host := <<"ninenines.eu">>}} = do_get("/path/to", Config),
+
+		%% Custom outcome: external redirect with custom host and path.
+		?REWRITE([{Fun, #{clause => 3}}]),
+		{redirect, 302, Headers1} = do_get("/", Config),
+		{_, Location1} = lists:keyfind(<<"location">>, 1, Headers1),
+		#{host := <<"ninenines.eu">>, path := <<"/">>} = uri_string:parse(Location1),
+		{redirect, 302, Headers2} = do_get("/path/to", Config),
+		{_, Location2} = lists:keyfind(<<"location">>, 1, Headers2),
+		#{host := <<"ninenines.eu">>, path := <<"/">>} = uri_string:parse(Location2),
+
+		%% Custom outcome: stop processing.
+		?REWRITE([
+			{Fun, #{clause => 4}},
+			{fun(Req,_) -> {continue, Req#{path => <<"/the_end">>}} end, #{}}
+		]),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/", Config),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/path/to", Config)
+	after
+		cowboy:stop_listener(?FUNCTION_NAME)
+	end.
+
+custom_mf(Config0) ->
+	Config = init_http(?FUNCTION_NAME, init_opts(), Config0),
+	try
+		%% Custom path.
+		?REWRITE([{?MODULE, do_custom_mf, #{clause => 1}}]),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/", Config),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/path/to", Config),
+
+		%% Custom host.
+		?REWRITE([{?MODULE, do_custom_mf, #{clause => 2}}]),
+		{ok, #{host := <<"ninenines.eu">>}} = do_get("/", Config),
+		{ok, #{host := <<"ninenines.eu">>}} = do_get("/path/to", Config),
+
+		%% Custom outcome: external redirect with custom host and path.
+		?REWRITE([{?MODULE, do_custom_mf, #{clause => 3}}]),
+		{redirect, 302, Headers1} = do_get("/", Config),
+		{_, Location1} = lists:keyfind(<<"location">>, 1, Headers1),
+		#{host := <<"ninenines.eu">>, path := <<"/">>} = uri_string:parse(Location1),
+		{redirect, 302, Headers2} = do_get("/path/to", Config),
+		{_, Location2} = lists:keyfind(<<"location">>, 1, Headers2),
+		#{host := <<"ninenines.eu">>, path := <<"/">>} = uri_string:parse(Location2),
+
+		%% Custom outcome: stop processing.
+		?REWRITE([
+			{?MODULE, do_custom_mf, #{clause => 4}},
+			{fun(Req,_) -> {continue, Req#{path => <<"/the_end">>}} end, #{}}
+		]),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/", Config),
+		{ok, #{path := <<"/custom/path">>}} = do_get("/path/to", Config)
+	after
+		cowboy:stop_listener(?FUNCTION_NAME)
+	end.
+
+do_custom_mf(Req, #{clause := 1}) ->
+	{continue, Req#{path => <<"/custom/path">>}};
+do_custom_mf(Req, #{clause := 2}) ->
+	{continue, Req#{host => <<"ninenines.eu">>}};
+do_custom_mf(Req, #{clause := 3}) ->
+	{{redirect, 302}, Req#{host => <<"ninenines.eu">>, path => <<"/">>}};
+do_custom_mf(Req, #{clause := 4}) ->
+	{stop, Req#{path => <<"/custom/path">>}}.
 
 index(Config0) ->
 	Config = init_http(?FUNCTION_NAME, init_opts(), Config0),
@@ -118,5 +202,96 @@ index(Config0) ->
 		cowboy:stop_listener(?FUNCTION_NAME)
 	end.
 
-%% @todo slash(Config) ->
+regex(Config0) ->
+	Config = init_http(?FUNCTION_NAME, init_opts(), Config0),
+	try
+		%% String regex.
+		?REWRITE([{regex, #{
+			regex => <<"^/path/([0-9a-z/.]+)$">>,
+			replacement => <<"/post.ext?where=\\1">>
+		}}]),
+		{ok, #{path := <<"/">>}} = do_get("/", Config),
+		{ok, #{path := <<"/post.ext">>, qs := <<"where=to">>}} = do_get("/path/to", Config),
+		{ok, #{path := <<"/post.ext">>, qs := <<"where=to/">>}} = do_get("/path/to/", Config),
+		{ok, #{path := <<"/post.ext">>, qs := <<"where=to/index.html">>}} = do_get("/path/to/index.html", Config),
+		{ok, #{path := <<"/path/to_underscore">>, qs := <<>>}} = do_get("/path/to_underscore", Config),
+
+		%% Compiled regex.
+		{ok, CompiledRegex} = re:compile(<<"^/path/([0-9a-z/.]+)$">>),
+		?REWRITE([{regex, #{
+			regex => CompiledRegex,
+			replacement => <<"/post.ext?where=\\1">>
+		}}]),
+		{ok, #{path := <<"/">>}} = do_get("/", Config),
+		{ok, #{path := <<"/post.ext">>, qs := <<"where=to">>}} = do_get("/path/to", Config),
+		{ok, #{path := <<"/post.ext">>, qs := <<"where=to/">>}} = do_get("/path/to/", Config),
+		{ok, #{path := <<"/post.ext">>, qs := <<"where=to/index.html">>}} = do_get("/path/to/index.html", Config),
+		{ok, #{path := <<"/path/to_underscore">>, qs := <<>>}} = do_get("/path/to_underscore", Config),
+
+		%% Options.
+		%%
+		%% We use the 'offset' option because it makes testing simpler,
+		%% but it is clearly not an option that's going to work well
+		%% because it causes crashes when the path is shorter than
+		%% the given offset.
+		?REWRITE([{regex, #{
+			regex => <<"second">>,
+			replacement => <<"third">>,
+			options => [{offset, 5}]
+		}}]),
+		{ok, #{path := <<"/first/third">>}} = do_get("/first/second", Config),
+		{ok, #{path := <<"/second/first">>}} = do_get("/second/first", Config),
+
+		%% Custom outcome: external redirect.
+		?REWRITE([{regex, #{
+			regex => <<"^/path/([0-9a-z/.]+)$">>,
+			replacement => <<"/post.ext?where=\\1">>,
+			outcome => {redirect, 302}
+		}}]),
+		{ok, #{path := <<"/">>}} = do_get("/", Config),
+		{redirect, 302, Headers1} = do_get("/path/to", Config),
+		{_, Location1} = lists:keyfind(<<"location">>, 1, Headers1),
+		#{path := <<"/post.ext">>, query := <<"where=to">>} = uri_string:parse(Location1),
+
+		%% Custom outcome: stop processing.
+		?REWRITE([
+			{regex, #{
+				regex => <<"^/path/([0-9a-z/.]+)$">>,
+				replacement => <<"/post.ext?where=\\1">>,
+				outcome => stop
+			}},
+			{fun(Req,_) -> {continue, Req#{path => <<"/the_end">>}} end, #{}}
+		]),
+		{ok, #{path := <<"/the_end">>}} = do_get("/", Config),
+		{ok, #{path := <<"/post.ext">>, qs := <<"where=to">>}} = do_get("/path/to", Config)
+	after
+		cowboy:stop_listener(?FUNCTION_NAME)
+	end.
+
+slash(Config0) ->
+	Config = init_http(?FUNCTION_NAME, init_opts(), Config0),
+	try
+		%% Default options.
+		?REWRITE([{slash, #{}}]),
+		{ok, #{path := <<"/">>}} = do_get("/", Config),
+		{ok, #{path := <<"/path/to/">>}} = do_get("/path/to", Config),
+		{ok, #{path := <<"/path/to/">>}} = do_get("/path/to/", Config),
+
+		%% Custom outcome: external redirect.
+		?REWRITE([{slash, #{outcome => {redirect, 302}}}]),
+		{ok, #{path := <<"/">>}} = do_get("/", Config),
+		{redirect, 302, Headers1} = do_get("/path/to", Config),
+		do_check_location(<<"/path/to/">>, Headers1),
+		{ok, #{path := <<"/path/to/">>}} = do_get("/path/to/", Config),
 
+		%% Custom outcome: stop processing.
+		?REWRITE([
+			{slash, #{outcome => stop}},
+			{fun(Req,_) -> {continue, Req#{path => <<"/the_end">>}} end, #{}}
+		]),
+		{ok, #{path := <<"/the_end">>}} = do_get("/", Config),
+		{ok, #{path := <<"/path/to/">>}} = do_get("/path/to", Config),
+		{ok, #{path := <<"/the_end">>}} = do_get("/path/to/", Config)
+	after
+		cowboy:stop_listener(?FUNCTION_NAME)
+	end.