From ad67d9e5980832b3d5f8a054b7b55ad78bb67c40 Mon Sep 17 00:00:00 2001
From: Nelson Vides <nelson.vides@erlang-solutions.com>
Date: Thu, 7 Sep 2023 12:23:45 +0200
Subject: [PATCH] Fix invalid username in scram authentication

---
 big_tests/tests/login_SUITE.erl | 11 +++++++-
 src/sasl/cyrsasl_scram.erl      | 45 ++++++++++++++++++++-------------
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/big_tests/tests/login_SUITE.erl b/big_tests/tests/login_SUITE.erl
index 8db4ed16bb..0b3c3098b4 100644
--- a/big_tests/tests/login_SUITE.erl
+++ b/big_tests/tests/login_SUITE.erl
@@ -92,7 +92,8 @@ all_tests() ->
     [log_one,
      log_non_existent_plain,
      log_one_scram_sha1,
-     log_non_existent_scram
+     log_non_existent_scram,
+     log_bad_user_fails
     ].
 
 access_tests() ->
@@ -378,6 +379,14 @@ log_non_existent_scram(Config) ->
     R = log_non_existent([{escalus_auth_method, <<"SCRAM-SHA-1">>} | Config]),
     {expected_challenge, _, _} = R.
 
+log_bad_user_fails(Config) ->
+    Config1 = [{escalus_auth_method, <<"SCRAM-SHA-1">>} | Config],
+    [{kate, UserSpec}] = escalus_users:get_users([kate]),
+    UserSpec1 = lists:keyreplace(username, 1, UserSpec, {username, <<" kate">>}),
+    {error, {connection_step_failed, _, R}} = escalus_client:start(Config1, UserSpec1, <<"res">>),
+    {expected_challenge, got, Xmlel} = R,
+    #xmlel{name = <<"failure">>} = Xmlel.
+
 log_non_existent(Config) ->
     [{kate, UserSpec}] = escalus_users:get_users([kate]),
     {error, {connection_step_failed, _, R}} = escalus_client:start(Config, UserSpec, <<"res">>),
diff --git a/src/sasl/cyrsasl_scram.erl b/src/sasl/cyrsasl_scram.erl
index 9d89e85ad0..ea0fad00f9 100644
--- a/src/sasl/cyrsasl_scram.erl
+++ b/src/sasl/cyrsasl_scram.erl
@@ -8,6 +8,7 @@
 -export([mech_new/3, mech_step/2]).
 
 -include("mongoose.hrl").
+-include("jlib.hrl").
 
 -behaviour(cyrsasl).
 
@@ -24,32 +25,40 @@ mech_new(LServer, Creds, #{sha := Sha,
                            auth_mech := AuthMech,
                            scram_plus := ScramPlus}) ->
     ChannelBinding = calculate_channel_binding(Socket, ScramPlus, Sha, AuthMech),
-    Fun = fun(Username, St0) ->
-                  JID = jid:make_bare(Username, LServer),
-                  HostType = mongoose_credentials:host_type(Creds),
-                  case get_scram_attributes(HostType, JID, Sha) of
-                      {AuthModule, {StoredKey, ServerKey, Salt, ItCount}} ->
-                          Creds1 = fast_scram:mech_get(creds, St0, Creds),
-                          R = [{username, Username}, {auth_module, AuthModule}],
-                          Creds2 = mongoose_credentials:extend(Creds1, R),
-                          St1 = fast_scram:mech_set(creds, Creds2, St0),
-                          ExtraConfig = #{it_count => ItCount, salt => Salt,
-                                          auth_data => #{stored_key => StoredKey,
-                                                         server_key => ServerKey}},
-                          {St1, ExtraConfig};
-                      {error, Reason, User} ->
-                          {error, {Reason, User}}
-                  end
-          end,
     {ok, St0} = fast_scram:mech_new(
                         #{entity => server,
                           hash_method => Sha,
-                          retrieve_mechanism => Fun,
+                          retrieve_mechanism => retrieve_mechanism_fun(LServer, Creds, Sha),
                           nonce_size => ?NONCE_LENGTH,
                           channel_binding => ChannelBinding}),
     St1 = fast_scram:mech_set(creds, Creds, St0),
     {ok, St1}.
 
+retrieve_mechanism_fun(LServer, Creds, Sha) ->
+    fun(Username, St0) ->
+            case jid:make_bare(Username, LServer) of
+                error -> {error, {invalid_username, Username}};
+                JID ->
+                    retrieve_mechanism_continue(JID, Creds, Sha, St0)
+            end
+    end.
+
+retrieve_mechanism_continue(#jid{luser = Username} = JID, Creds, Sha, St0) ->
+    HostType = mongoose_credentials:host_type(Creds),
+    case get_scram_attributes(HostType, JID, Sha) of
+        {AuthModule, {StoredKey, ServerKey, Salt, ItCount}} ->
+            Creds1 = fast_scram:mech_get(creds, St0, Creds),
+            R = [{username, Username}, {auth_module, AuthModule}],
+            Creds2 = mongoose_credentials:extend(Creds1, R),
+            St1 = fast_scram:mech_set(creds, Creds2, St0),
+            ExtraConfig = #{it_count => ItCount, salt => Salt,
+                            auth_data => #{stored_key => StoredKey,
+                                           server_key => ServerKey}},
+            {St1, ExtraConfig};
+        {error, Reason, User} ->
+            {error, {Reason, User}}
+    end.
+
 mech_step(State, ClientIn) ->
     case fast_scram:mech_step(State, ClientIn) of
         {continue, Msg, NewState} ->