From 3b28f0d9bbc9e5745fb95e48e6daf9179461116b Mon Sep 17 00:00:00 2001
From: Kirilll Zaborsky <qrilka@gmail.com>
Date: Thu, 24 Sep 2015 15:19:52 +0300
Subject: inets: scheme validation fun for http_uri

http_uri:parse_scheme function should allow checking
scheme of URIs otherwise it could be easily abused to
reach limit number of atoms in the VM
---
 lib/inets/doc/src/http_uri.xml      | 13 ++++++++++++-
 lib/inets/src/http_lib/http_uri.erl | 31 ++++++++++++++++++++++++-------
 lib/inets/test/uri_SUITE.erl        | 23 ++++++++++++++++++++++-
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/lib/inets/doc/src/http_uri.xml b/lib/inets/doc/src/http_uri.xml
index 47c40da96a..64e6c7a6cc 100644
--- a/lib/inets/doc/src/http_uri.xml
+++ b/lib/inets/doc/src/http_uri.xml
@@ -117,7 +117,8 @@
         <v>Options = [Option]</v> 
         <v>Option = {ipv6_host_with_brackets, boolean()} | 
                     {scheme_defaults, scheme_defaults()} |
-                    {fragment, boolean()}]</v>
+                    {fragment, boolean()} |
+                    {schema_validation_fun, fun()}]</v>
         <v>Result = {Scheme, UserInfo, Host, Port, Path, Query} |
                     {Scheme, UserInfo, Host, Port, Path, Query, Fragment}</v>
 	<v>UserInfo = user_info()</v>
@@ -141,6 +142,16 @@
         <p>If the fragment option is <c>true</c>, the URI fragment is returned as
           part of the parsing result, otherwise it is ignored.</p>
 
+        <p>Scheme validation fun is to be defined as follows:
+
+	    <code>
+fun(SchemeStr :: string()) ->
+	valid |	{error, Reason :: term()}.
+	    </code>
+
+        It is called before scheme string gets converted into scheme atom and
+        thus possible atom leak could be prevented</p>
+
         <marker id="encode"></marker>
       </desc>
     </func>
diff --git a/lib/inets/src/http_lib/http_uri.erl b/lib/inets/src/http_lib/http_uri.erl
index 79591eec29..6fe8c1776d 100644
--- a/lib/inets/src/http_lib/http_uri.erl
+++ b/lib/inets/src/http_lib/http_uri.erl
@@ -138,16 +138,33 @@ parse_scheme(AbsURI, Opts) ->
 	{error, no_scheme} ->
 	    {error, no_scheme};
 	{SchemeStr, Rest} ->
-	    Scheme = list_to_atom(http_util:to_lower(SchemeStr)),
-	    SchemeDefaults = which_scheme_defaults(Opts), 
-	    case lists:keysearch(Scheme, 1, SchemeDefaults) of
-		{value, {Scheme, DefaultPort}} ->
-		    {Scheme, DefaultPort, Rest};
-		false ->
-		    {Scheme, no_default_port, Rest}
+	    case extract_scheme(SchemeStr, Opts) of
+		{error, Error} ->
+		    {error, Error};
+		{ok, Scheme} ->
+		    SchemeDefaults = which_scheme_defaults(Opts),
+		    case lists:keysearch(Scheme, 1, SchemeDefaults) of
+			{value, {Scheme, DefaultPort}} ->
+			    {Scheme, DefaultPort, Rest};
+			false ->
+			    {Scheme, no_default_port, Rest}
+		    end
 	    end
     end.
 
+extract_scheme(Str, Opts) ->
+    case lists:keysearch(scheme_validation_fun, 1, Opts) of
+	{value, {scheme_validation_fun, Fun}} when is_function(Fun) ->
+	    case Fun(Str) of
+		valid ->
+		    {ok, list_to_atom(http_util:to_lower(Str))};
+		{error, Error} ->
+		    {error, Error}
+	    end;
+	_ ->
+	    {ok, list_to_atom(http_util:to_lower(Str))}
+    end.
+
 parse_uri_rest(Scheme, DefaultPort, "//" ++ URIPart, Opts) ->
     {Authority, PathQueryFragment} =
         split_uri(URIPart, "[/?#]", {URIPart, ""}, 1, 0),
diff --git a/lib/inets/test/uri_SUITE.erl b/lib/inets/test/uri_SUITE.erl
index bfcd7bd339..2642b8fd4e 100644
--- a/lib/inets/test/uri_SUITE.erl
+++ b/lib/inets/test/uri_SUITE.erl
@@ -49,7 +49,8 @@ all() ->
      queries,
      fragments,
      escaped,
-     hexed_query
+     hexed_query,
+     scheme_validation
     ].
 
 %%--------------------------------------------------------------------
@@ -175,6 +176,26 @@ hexed_query(Config) when is_list(Config) ->
     verify_uri(URI2, Verify2),
     verify_uri(URI3, Verify3).
 
+scheme_validation(Config) when is_list(Config) ->
+    {ok, {http,[],"localhost",80,"/",""}} =
+	http_uri:parse("http://localhost#fragment"),
+
+    ValidationFun =
+	fun("http") -> valid;
+	   (_) -> {error, bad_scheme}
+	end,
+
+    {ok, {http,[],"localhost",80,"/",""}} =
+	http_uri:parse("http://localhost#fragment",
+		       [{scheme_validation_fun, ValidationFun}]),
+    {error, bad_scheme} =
+	http_uri:parse("https://localhost#fragment",
+		       [{scheme_validation_fun, ValidationFun}]),
+    %% non-fun scheme_validation_fun works as no option passed
+    {ok, {https,[],"localhost",443,"/",""}} =
+	http_uri:parse("https://localhost#fragment",
+		       [{scheme_validation_fun, none}]).
+
 
 %%--------------------------------------------------------------------
 %% Internal Functions ------------------------------------------------
-- 
cgit v1.2.3