* src/mod_register.erl: Restrict registration frequency per IP or

user
* src/ejabberd_c2s.erl: Pass IP to the c2s_unauthenticated_iq hook
* src/ejabberd_config.erl: Added registration_timeout option
* src/treap.erl: Treaps implementation

SVN Revision: 1299

5 files changed, 275 insertions, 21 deletions

diff --git a/ChangeLog b/ChangeLog
index 23bf00dcc..1c8cdf018 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,12 @@
+2008-04-23  Alexey Shchepin  <alexey@process-one.net>
+
+	* src/mod_register.erl: Restrict registration frequency per IP or
+	user
+	* src/ejabberd_c2s.erl: Pass IP to the c2s_unauthenticated_iq hook
+	* src/ejabberd_config.erl: Added registration_timeout option
+
+	* src/treap.erl: Treaps implementation
+
 2008-04-22  Badlop  <badlop@process-one.net>
 
 	* src/ejabberd_auth.erl: Improve anonymous authentication to not
diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl
index 53bb41caf..567cd6b6e 100644
--- a/src/ejabberd_c2s.erl
+++ b/src/ejabberd_c2s.erl
@@ -1908,7 +1908,8 @@ process_unauthenticated_stanza(StateData, El) ->
 	    Res = ejabberd_hooks:run_fold(c2s_unauthenticated_iq,
 					  StateData#state.server,
 					  empty,
-					  [StateData#state.server, IQ]),
+					  [StateData#state.server, IQ,
+					   StateData#state.ip]),
 	    case Res of
 		empty ->
 		    % The only reasonable IQ's here are auth and register IQ's
diff --git a/src/ejabberd_config.erl b/src/ejabberd_config.erl
index 6d110ed13..7a9cade1f 100644
--- a/src/ejabberd_config.erl
+++ b/src/ejabberd_config.erl
@@ -337,6 +337,8 @@ process_term(Term, State) ->
 	    add_option({domain_balancing_component_number, Domain}, N, State);
 	{watchdog_admins, Admins} ->
 	    add_option(watchdog_admins, Admins, State);
+	{registration_timeout, Timeout} ->
+	    add_option(registration_timeout, Timeout, State);
 	{loglevel, Loglevel} ->
 	    ejabberd_loglevel:set(Loglevel),
 	    State;
diff --git a/src/mod_register.erl b/src/mod_register.erl
index fc5dbc17e..b8f99df3f 100644
--- a/src/mod_register.erl
+++ b/src/mod_register.erl
@@ -32,7 +32,7 @@
 -export([start/2,
 	 stop/1,
 	 stream_feature_register/1,
-	 unauthenticated_iq_register/3,
+	 unauthenticated_iq_register/4,
 	 process_iq/3]).
 
 -include("ejabberd.hrl").
@@ -48,6 +48,10 @@ start(Host, Opts) ->
  		       ?MODULE, stream_feature_register, 50),
     ejabberd_hooks:add(c2s_unauthenticated_iq, Host,
  		       ?MODULE, unauthenticated_iq_register, 50),
+    mnesia:create_table(mod_register_ip,
+			[{ram_copies, [node()]},
+			 {local_content, true},
+			 {attributes, [key, value]}]),
     ok.
 
 stop(Host) ->
@@ -63,20 +67,30 @@ stream_feature_register(Acc) ->
     [{xmlelement, "register",
       [{"xmlns", ?NS_FEATURE_IQREGISTER}], []} | Acc].
 
-unauthenticated_iq_register(_Acc, Server, #iq{xmlns = ?NS_REGISTER} = IQ) ->
+unauthenticated_iq_register(_Acc,
+			    Server, #iq{xmlns = ?NS_REGISTER} = IQ, IP) ->
+    Address = case IP of
+		 {A, _Port} -> A;
+		  _ -> undefined
+	      end,
     ResIQ = process_iq(jlib:make_jid("", "", ""),
  		       jlib:make_jid("", Server, ""),
- 		       IQ),
+ 		       IQ,
+		       Address),
     Res1 = jlib:replace_from_to(jlib:make_jid("", Server, ""),
  				jlib:make_jid("", "", ""),
  				jlib:iq_to_xml(ResIQ)),
     jlib:remove_attr("to", Res1);
 
-unauthenticated_iq_register(Acc, _Server, _IQ) ->
+unauthenticated_iq_register(Acc, _Server, _IQ, _IP) ->
     Acc.
 
+process_iq(From, To, IQ) ->
+    process_iq(From, To, IQ, jlib:jid_tolower(jlib:jid_remove_resource(From))).
+
 process_iq(From, To,
-	   #iq{type = Type, lang = Lang, sub_el = SubEl, id = ID} = IQ) ->
+	   #iq{type = Type, lang = Lang, sub_el = SubEl, id = ID} = IQ,
+	   Source) ->
     case Type of
 	set ->
 	    UTag = xml:get_subtag(SubEl, "username"),
@@ -151,7 +165,8 @@ process_iq(From, To,
 			    ejabberd_auth:set_password(User, Server, Password),
 			    IQ#iq{type = result, sub_el = [SubEl]};
 			_ ->
-			    case try_register(User, Server, Password) of
+			    case try_register(User, Server, Password,
+					      Source) of
 				ok ->
 				    IQ#iq{type = result, sub_el = [SubEl]};
 				{error, Error} ->
@@ -179,7 +194,7 @@ process_iq(From, To,
     end.
 
 
-try_register(User, Server, Password) ->
+try_register(User, Server, Password, Source) ->
     case jlib:is_nodename(User) of
 	false ->
 	    {error, ?ERR_BAD_REQUEST};
@@ -190,19 +205,24 @@ try_register(User, Server, Password) ->
 		deny ->
 		    {error, ?ERR_CONFLICT};
 		allow ->
-		    case ejabberd_auth:try_register(User, Server, Password) of
-			{atomic, ok} ->
-			    send_welcome_message(JID),
-			    send_registration_notifications(JID),
-			    ok;
-			{atomic, exists} ->
-			    {error, ?ERR_CONFLICT};
-			{error, invalid_jid} ->
-			    {error, ?ERR_JID_MALFORMED};
-			{error, not_allowed} ->
-			    {error, ?ERR_NOT_ALLOWED};
-			{error, _Reason} ->
-			    {error, ?ERR_INTERNAL_SERVER_ERROR}
+		    case check_timeout(Source) of
+			true ->
+			    case ejabberd_auth:try_register(User, Server, Password) of
+				{atomic, ok} ->
+				    send_welcome_message(JID),
+				    send_registration_notifications(JID),
+				    ok;
+				{atomic, exists} ->
+				    {error, ?ERR_CONFLICT};
+				{error, invalid_jid} ->
+				    {error, ?ERR_JID_MALFORMED};
+				{error, not_allowed} ->
+				    {error, ?ERR_NOT_ALLOWED};
+				{error, _Reason} ->
+				    {error, ?ERR_INTERNAL_SERVER_ERROR}
+			    end;
+			false ->
+			    {error, ?ERR_RESOURCE_CONSTRAINT}
 		    end
 	    end
     end.
@@ -249,3 +269,61 @@ send_registration_notifications(UJID) ->
 	_ ->
 	    ok
     end.
+
+
+check_timeout(undefined) ->
+    true;
+check_timeout(Source) ->
+    Timeout = case ejabberd_config:get_local_option(registration_timeout) of
+		  undefined -> 600;
+		  TO -> TO
+	      end,
+    if
+	is_integer(Timeout) ->
+	    {MSec, Sec, _USec} = now(),
+	    Priority = -(MSec * 1000000 + Sec),
+	    CleanPriority = Priority + Timeout,
+	    F = fun() ->
+			Treap = case mnesia:read(mod_register_ip, treap,
+						 write) of
+				    [] ->
+					treap:empty();
+				    [{mod_register_ip, treap, T}] -> T
+				end,
+			Treap1 = clean_treap(Treap, CleanPriority),
+			case treap:lookup(Source, Treap1) of
+			    error ->
+				Treap2 = treap:insert(Source, Priority, [],
+						      Treap1),
+				mnesia:write({mod_register_ip, treap, Treap2}),
+				true;
+			    {ok, _, _} ->
+				mnesia:write({mod_register_ip, treap, Treap1}),
+				false
+			end
+		end,
+	    case mnesia:transaction(F) of
+		{atomic, Res} ->
+		    Res;
+		{aborted, Reason} ->
+		    ?ERROR_MSG("mod_register: timeout check error: ~p~n",
+			       [Reason]),
+		    true
+	    end;
+	true ->
+	    true
+    end.
+
+clean_treap(Treap, CleanPriority) ->
+    case treap:is_empty(Treap) of
+	true ->
+	    Treap;
+	false ->
+	    {_Key, Priority, _Value} = treap:get_root(Treap),
+	    if
+		Priority > CleanPriority ->
+		    clean_treap(treap:delete_root(Treap), CleanPriority);
+		true ->
+		    Treap
+	    end
+    end.
diff --git a/src/treap.erl b/src/treap.erl
new file mode 100644
index 000000000..d7b070b9e
--- /dev/null
+++ b/src/treap.erl
@@ -0,0 +1,164 @@
+%%%----------------------------------------------------------------------
+%%% File    : treap.erl
+%%% Author  : Alexey Shchepin <alexey@process-one.net>
+%%% Purpose : Treaps implementation
+%%% Created : 22 Apr 2008 by Alexey Shchepin <alexey@process-one.net>
+%%%
+%%%
+%%% ejabberd, Copyright (C) 2002-2008   Process-one
+%%%
+%%% This program is free software; you can redistribute it and/or
+%%% modify it under the terms of the GNU General Public License as
+%%% published by the Free Software Foundation; either version 2 of the
+%%% License, or (at your option) any later version.
+%%%
+%%% This program is distributed in the hope that it will be useful,
+%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
+%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+%%% General Public License for more details.
+%%%
+%%% You should have received a copy of the GNU General Public License
+%%% along with this program; if not, write to the Free Software
+%%% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
+%%% 02111-1307 USA
+%%%
+%%%----------------------------------------------------------------------
+
+-module(treap).
+
+-export([empty/0,
+	 insert/4,
+	 delete/2,
+	 delete_root/1,
+	 get_root/1,
+	 lookup/2,
+	 is_empty/1]).
+
+empty() ->
+    nil.
+
+insert(Key, Priority, Value, Tree) ->
+    HashKey = {erlang:phash2(Key), Key},
+    insert1(Tree, HashKey, Priority, Value).
+
+insert1(nil, HashKey, Priority, Value) ->
+    {HashKey, Priority, Value, nil, nil};
+insert1({HashKey1, Priority1, Value1, Left, Right},
+	HashKey, Priority, Value) ->
+    if
+	HashKey < HashKey1 ->
+	    heapify({HashKey1, Priority1, Value1,
+		     insert1(Left, HashKey, Priority, Value),
+		     Right});
+	HashKey > HashKey1 ->
+	    heapify({HashKey1, Priority1, Value1,
+		     Left,
+		     insert1(Right, HashKey, Priority, Value)});
+	true ->
+	    erlang:error(key_exists)
+    end.
+
+heapify(nil) ->
+    nil;
+heapify({_HashKey, _Priority, _Value, nil, nil} = Tree) ->
+    Tree;
+heapify({HashKey, Priority, Value,
+	 nil = Left,
+	 {HashKeyR, PriorityR, ValueR, LeftR, RightR}} = Tree) ->
+    if
+	PriorityR > Priority ->
+	    {HashKeyR, PriorityR, ValueR,
+	     {HashKey, Priority, Value, Left, LeftR},
+	     RightR};
+	true ->
+	    Tree
+    end;
+heapify({HashKey, Priority, Value,
+	 {HashKeyL, PriorityL, ValueL, LeftL, RightL},
+	 nil = Right} = Tree) ->
+    if
+	PriorityL > Priority ->
+	    {HashKeyL, PriorityL, ValueL,
+	     LeftL,
+	     {HashKey, Priority, Value, RightL, Right}};
+	true ->
+	    Tree
+    end;
+heapify({HashKey, Priority, Value,
+	 {HashKeyL, PriorityL, ValueL, LeftL, RightL} = Left,
+	 {HashKeyR, PriorityR, ValueR, LeftR, RightR} = Right} = Tree) ->
+    if
+	PriorityR > Priority ->
+	    {HashKeyR, PriorityR, ValueR,
+	     {HashKey, Priority, Value, Left, LeftR},
+	     RightR};
+	PriorityL > Priority ->
+	    {HashKeyL, PriorityL, ValueL,
+	     LeftL,
+	     {HashKey, Priority, Value, RightL, Right}};
+	true ->
+	    Tree
+    end.
+
+
+delete(Key, Tree) ->
+    HashKey = {erlang:phash2(Key), Key},
+    delete1(Tree, HashKey).
+
+delete1(HashKey, {HashKey1, Priority1, Value1, Left, Right} = Tree) ->
+    if
+	HashKey < HashKey1 ->
+	    {HashKey1, Priority1, Value1, delete1(HashKey, Left), Right};
+	HashKey > HashKey1 ->
+	    {HashKey1, Priority1, Value1, Left, delete1(HashKey, Right)};
+	true ->
+	    delete_root(Tree)
+    end.
+
+delete_root({HashKey, Priority, Value, Left, Right}) ->
+    case {Left, Right} of
+	{nil, nil} ->
+	    nil;
+	{_, nil} ->
+	    Left;
+	{nil, _} ->
+	    Right;
+	{{HashKeyL, PriorityL, ValueL, LeftL, RightL},
+	 {HashKeyR, PriorityR, ValueR, LeftR, RightR}} ->
+	    if
+		PriorityL > PriorityR ->
+		    {HashKeyL, PriorityL, ValueL,
+		     LeftL,
+		     delete_root({HashKey, Priority, Value, RightL, Right})};
+		true ->
+		    {HashKeyR, PriorityR, ValueR,
+		     delete_root({HashKey, Priority, Value, Left, LeftR}),
+		     RightR}
+	    end
+    end.
+
+is_empty(nil) ->
+    true;
+is_empty({_HashKey, _Priority, _Value, _Left, _Right}) ->
+    false.
+
+get_root({{_Hash, Key}, Priority, Value, _Left, _Right}) ->
+    {Key, Priority, Value}.
+
+
+lookup(Key, Tree) ->
+    HashKey = {erlang:phash2(Key), Key},
+    lookup1(Tree, HashKey).
+
+lookup1(nil, _HashKey) ->
+    error;
+lookup1({HashKey1, Priority1, Value1, Left, Right}, HashKey) ->
+    if
+	HashKey < HashKey1 ->
+	    lookup1(Left, HashKey);
+	HashKey > HashKey1 ->
+	    lookup1(Right, HashKey);
+	true ->
+	    {ok, Priority1, Value1}
+    end.
+