From f7d9bb7ce4b9309a72954b65a21a34b384ac0b83 Mon Sep 17 00:00:00 2001
From: Badlop <badlop@process-one.net>
Date: Thu, 7 May 2015 11:41:59 +0200
Subject: [PATCH] Don't announce register feature once user is authenticated
 (#41)

---
 src/ejabberd_c2s.erl | 14 +++++++-------
 src/mod_register.erl | 13 +++++++++----
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl
index f2d16e878..5cdb8e1c9 100644
--- a/src/ejabberd_c2s.erl
+++ b/src/ejabberd_c2s.erl
@@ -457,14 +457,14 @@ wait_for_stream({xmlstreamstart, _Name, Attrs}, StateData) ->
 					    false ->
 						[]
 					end,
+				    StreamFeatures1 = TLSFeature ++ CompressFeature ++ Mechs,
+				    StreamFeatures = ejabberd_hooks:run_fold(c2s_stream_features,
+							Server, StreamFeatures1, [Server]),
 				    send_element(StateData,
 					    #xmlel{name = <<"stream:features">>,
 						   attrs = [],
 						   children =
-						    TLSFeature ++ CompressFeature ++ Mechs
-						    ++
-						    ejabberd_hooks:run_fold(c2s_stream_features,
-							Server, [], [Server])}),
+						    StreamFeatures}),
 				    fsm_next_state(wait_for_feature_request,
 					       StateData#state{
 						 server = Server,
@@ -489,7 +489,7 @@ wait_for_stream({xmlstreamstart, _Name, Attrs}, StateData) ->
 					      false ->
 						  []
 					    end,
-					StreamFeatures = [#xmlel{name = <<"bind">>,
+					StreamFeatures1 = [#xmlel{name = <<"bind">>,
 								attrs = [{<<"xmlns">>, ?NS_BIND}],
 								children = []},
 							    #xmlel{name = <<"session">>,
@@ -499,9 +499,9 @@ wait_for_stream({xmlstreamstart, _Name, Attrs}, StateData) ->
 							    RosterVersioningFeature ++
 							    StreamManagementFeature ++
 							    ejabberd_hooks:run_fold(c2s_post_auth_features,
-								Server, [], [Server]) ++
-							    ejabberd_hooks:run_fold(c2s_stream_features,
 								Server, [], [Server]),
+					StreamFeatures = ejabberd_hooks:run_fold(c2s_stream_features,
+								Server, StreamFeatures1, [Server]),
 					send_element(StateData,
 						    #xmlel{name = <<"stream:features">>,
 							    attrs = [],
diff --git a/src/mod_register.erl b/src/mod_register.erl
index cd68af936..1c03fb3e9 100644
--- a/src/mod_register.erl
+++ b/src/mod_register.erl
@@ -68,10 +68,15 @@ stop(Host) ->
 				     ?NS_REGISTER).
 
 stream_feature_register(Acc, _Host) ->
-    [#xmlel{name = <<"register">>,
-	    attrs = [{<<"xmlns">>, ?NS_FEATURE_IQREGISTER}],
-	    children = []}
-     | Acc].
+    case lists:keymember(<<"mechanisms">>, 2, Acc) of
+	true ->
+	    [#xmlel{name = <<"register">>,
+		    attrs = [{<<"xmlns">>, ?NS_FEATURE_IQREGISTER}],
+		    children = []}
+	     | Acc];
+	false ->
+	    Acc
+    end.
 
 unauthenticated_iq_register(_Acc, Server,
 			    #iq{xmlns = ?NS_REGISTER} = IQ, IP) ->
-- 
2.49.0