Merge pull request #3059 from esl/prepared-queries-auth-token

Use prepared queries in mod_auth_token

big_tests/tests/oauth_SUITE.erl

@@ -76,7 +76,7 @@ suite() ->
 %%--------------------------------------------------------------------
 
 init_per_suite(Config0) ->
-    case is_pgsql_available(Config0) of
+    case mongoose_helper:is_rdbms_enabled(domain()) of
         true ->
             Config = dynamic_modules:stop_running(mod_last, Config0),
             Host = ct:get_config({hosts, mim, domain}),
@@ -95,7 +95,7 @@ init_per_suite(Config0) ->
             dynamic_modules:start(Host, mod_auth_token, AuthOpts),
             escalus:init_per_suite([{auth_opts, AuthOpts} | Config]);
         false ->
-            {skip, "PostgreSQL not available - check env configuration"}
+            {skip, "RDBMS not available"}
     end.
 
 end_per_suite(Config) ->
@@ -311,8 +311,8 @@ token_removed_on_user_removal(Config) ->
                 escalus:assert(is_iq_result, [IQ], escalus:wait_for_stanza(Bob))
         end,
     escalus:story(Config, [{bob, 1}], S),
-    %% then token database doesn't contain user's tokens
-    {selected, []} = get_users_token(Config, bob).
+    %% then token database doesn't contain user's tokens (cleanup is done after IQ result)
+    mongoose_helper:wait_until(fun() -> get_users_token(Config, bob) end, {selected, []}).
 
 provision_token_login(Config) ->
     %% given
@@ -464,16 +464,5 @@ serialize(ServerSideToken) ->
 to_lower(B) when is_binary(B) ->
     list_to_binary(string:to_lower(binary_to_list(B))).
 
-is_pgsql_available(_) ->
-    Q = [<<"SELECT version();">>],
-    %% TODO: hardcoded RDBMSHost
-    RDBMSHost = domain(),
-    case rpc(mim(), mongoose_rdbms, sql_query, [RDBMSHost, Q]) of
-        {selected, [{<<"PostgreSQL", _/binary>>}]} ->
-            true;
-        _ ->
-            false
-    end.
-
 domain() ->
     ct:get_config({hosts, mim, domain}).

doc/modules/mod_auth_token.md

@@ -36,7 +36,7 @@ Three token types are supported:
   Access tokens can't be revoked.
   An access token is valid only until its expiry date is reached.
 
-- _refresh tokens_: These are longer lived tokens which are tracked by the server and therefore require persistent storage (as of now only PostgreSQL is supported).
+- _refresh tokens_: These are longer lived tokens which are tracked by the server and therefore require persistent storage in a relational database.
   Refresh tokens can be used as a payload for the X-OAUTH authentication mechanism and to grant access to the system.
   Also they can result in a new set of tokens being returned upon successful authentication.
   They can be revoked - if a refresh token hasn't been revoked, it is valid until it has expired.

priv/mssql2012.sql

@@ -165,6 +165,11 @@ CREATE TABLE [dbo].[offline_message](
 )WITH (PAD_INDEX = OFF, STATISTICS_NORECOMPUTE = OFF, IGNORE_DUP_KEY = OFF, ALLOW_ROW_LOCKS = ON, ALLOW_PAGE_LOCKS = ON) ON [PRIMARY]
 ) ON [PRIMARY] TEXTIMAGE_ON [PRIMARY]
 
+CREATE TABLE auth_token (
+  owner NVARCHAR(250) NOT NULL PRIMARY KEY,
+  seq_no BIGINT NOT NULL,
+);
+
 GO
 SET ANSI_PADDING OFF
 GO

priv/mysql.sql

@@ -306,6 +306,12 @@ CREATE TABLE offline_message(
   ROW_FORMAT=DYNAMIC;
 CREATE INDEX i_offline_message USING BTREE ON offline_message(server, username, id);
 
+CREATE TABLE auth_token(
+    owner varchar(250)      NOT NULL PRIMARY KEY,
+    seq_no BIGINT UNSIGNED NOT NULL
+) CHARACTER SET utf8mb4
+  ROW_FORMAT=DYNAMIC;
+
 CREATE TABLE muc_light_rooms(
     id BIGINT UNSIGNED      NOT NULL AUTO_INCREMENT,
     luser VARCHAR(250)      NOT NULL,

src/mod_auth_token.erl

@@ -60,6 +60,9 @@
                            | {ok, module(), jid:user(), binary()}
                            | error().
 
+-callback start(LServer) -> ok when
+      LServer :: jid:lserver().
+
 -callback revoke(Owner) -> ok | not_found when
       Owner :: jid:jid().
 
@@ -81,6 +84,7 @@
 -spec start(jid:server(), list()) -> ok.
 start(Domain, Opts) ->
     gen_mod:start_backend_module(?MODULE, default_opts(Opts)),
+    mod_auth_token_backend:start(Domain),
     mod_disco:register_feature(Domain, ?NS_ESL_TOKEN_AUTH),
     IQDisc = gen_mod:get_opt(iqdisc, Opts, no_queue),
     [ ejabberd_hooks:add(Hook, Domain, ?MODULE, Handler, Priority)

src/mod_auth_token_rdbms.erl

@@ -1,67 +1,69 @@
 -module(mod_auth_token_rdbms).
 -behaviour(mod_auth_token).
 
--export([get_valid_sequence_number/1,
+-export([start/1,
+         get_valid_sequence_number/1,
          revoke/1,
          clean_tokens/1]).
 
 -include("jlib.hrl").
 -include("mongoose.hrl").
 
+-spec start(jid:lserver()) -> ok.
+start(LServer) ->
+    prepare_queries(LServer).
+
 %% Assumption: all sequence numbers less than the current valid one
 %%             are not valid.
 -spec get_valid_sequence_number(JID) -> integer() when
       JID :: jid:jid().
 get_valid_sequence_number(#jid{lserver = LServer} = JID) ->
     BBareJID = jid:to_binary(jid:to_bare(JID)),
-    EBareJID = mongoose_rdbms:escape_string(BBareJID),
-    Q = valid_sequence_number_query(EBareJID),
-    {atomic, [{updated, _}, {selected, [{BSeqNo}]}]} = mongoose_rdbms:sql_transaction(LServer, Q),
-    mongoose_rdbms:result_to_integer(BSeqNo).
-
-valid_sequence_number_query(EOwner) ->
-    SqlOwner = iolist_to_binary(mongoose_rdbms:use_escaped_string(EOwner)),
-    [<<"WITH existing AS (SELECT at.seq_no "
-       "                  FROM auth_token at "
-       "                  WHERE at.owner = ", SqlOwner/bytes, ") "
-       "INSERT INTO auth_token "
-         "SELECT ", SqlOwner/bytes, ", 1 "
-       "WHERE NOT EXISTS (SELECT * FROM existing); ">>,
-     <<"SELECT seq_no "
-       "FROM auth_token "
-       "WHERE owner = ", SqlOwner/bytes, "; ">>].
+    {atomic, Selected} =
+        mongoose_rdbms:sql_transaction(
+          LServer, fun() -> get_sequence_number_t(LServer, BBareJID) end),
+    mongoose_rdbms:selected_to_integer(Selected).
 
 -spec revoke(JID) -> ok | not_found when
       JID :: jid:jid().
 revoke(#jid{lserver = LServer} = JID) ->
     BBareJID = jid:to_binary(jid:to_bare(JID)),
-    EBareJID = mongoose_rdbms:escape_string(BBareJID),
-    RevokeQuery = revoke_query(EBareJID),
-    QueryResult = mongoose_rdbms:sql_query(LServer, RevokeQuery),
-    ?LOG_DEBUG(#{what => auth_token_revoke, sql_query => RevokeQuery,
-                 sql_result => QueryResult}),
+    QueryResult = execute_revoke_token(LServer, BBareJID),
+    ?LOG_DEBUG(#{what => auth_token_revoke, owner => BBareJID, sql_result => QueryResult}),
     case QueryResult of
         {updated, 1} -> ok;
         {updated, 0} -> not_found
     end.
 
--spec revoke_query(mongoose_rdbms:escaped_string()) ->
-    mongoose_rdbms:sql_query_part().
-revoke_query(EOwner) ->
-    [<<"UPDATE auth_token SET seq_no=seq_no+1 WHERE owner = ">>,
-     mongoose_rdbms:use_escaped_string(EOwner)].
-
 -spec clean_tokens(Owner) -> ok when
       Owner :: jid:jid().
 clean_tokens(#jid{lserver = LServer} = Owner) ->
     BBareJID = jid:to_binary(jid:to_bare(Owner)),
-    EBareJID = mongoose_rdbms:escape_string(BBareJID),
-    Q = clean_tokens_query(EBareJID),
-    {updated, _} = mongoose_rdbms:sql_query(LServer, Q),
+    execute_delete_token(LServer, BBareJID),
     ok.
 
--spec clean_tokens_query(mongoose_rdbms:escaped_string()) ->
-    mongoose_rdbms:sql_query_part().
-clean_tokens_query(EOwner) ->
-    [<<"DELETE FROM auth_token WHERE owner = ">>,
-     mongoose_rdbms:use_escaped_string(EOwner)].
+-spec prepare_queries(jid:lserver()) -> ok.
+prepare_queries(LServer) ->
+    mongoose_rdbms:prepare(auth_token_select, auth_token, [owner],
+                           <<"SELECT seq_no FROM auth_token WHERE owner = ?">>),
+    mongoose_rdbms:prepare(auth_token_revoke, auth_token, [owner],
+                           <<"UPDATE auth_token SET seq_no=seq_no+1 WHERE owner = ?">>),
+    mongoose_rdbms:prepare(auth_token_delete, auth_token, [owner],
+                           <<"DELETE from auth_token WHERE owner = ?">>),
+    rdbms_queries:prepare_upsert(LServer, auth_token_upsert, auth_token,
+                                 [<<"owner">>, <<"seq_no">>], [], [<<"owner">>]),
+    ok.
+
+-spec execute_revoke_token(jid:lserver(), jid:literal_jid()) -> mongoose_rdbms:query_result().
+execute_revoke_token(LServer, Owner) ->
+    mongoose_rdbms:execute_successfully(LServer, auth_token_revoke, [Owner]).
+
+-spec execute_delete_token(jid:lserver(), jid:literal_jid()) -> mongoose_rdbms:query_result().
+execute_delete_token(LServer, Owner) ->
+    mongoose_rdbms:execute_successfully(LServer, auth_token_delete, [Owner]).
+
+-spec get_sequence_number_t(jid:lserver(), jid:literal_jid()) -> mongoose_rdbms:query_result().
+get_sequence_number_t(LServer, Owner) ->
+    {updated, _} =
+        rdbms_queries:execute_upsert(LServer, auth_token_upsert, [Owner, 1], [], [Owner]),
+    mongoose_rdbms:execute_successfully(LServer, auth_token_select, [Owner]).

src/rdbms/rdbms_queries.erl

@@ -129,20 +129,26 @@ mysql_and_pgsql_insert(Table, Columns) ->
      join(Placeholders, ", "),
      ")"].
 
-upsert_mysql_query(Table, InsertFields, UpdateFields, _) ->
+upsert_mysql_query(Table, InsertFields, UpdateFields, [Key | _]) ->
     Insert = mysql_and_pgsql_insert(Table, InsertFields),
-    OnConflict = mysql_on_conflict(UpdateFields),
+    OnConflict = mysql_on_conflict(UpdateFields, Key),
     [Insert, OnConflict].
 
 upsert_pgsql_query(Table, InsertFields, UpdateFields, UniqueKeyFields) ->
     Insert = mysql_and_pgsql_insert(Table, InsertFields),
     OnConflict = pgsql_on_conflict(UpdateFields, UniqueKeyFields),
     [Insert, OnConflict].
 
-mysql_on_conflict(UpdateFields) ->
+mysql_on_conflict([], Key) ->
+    %% Update field to itself (no-op), there is no 'DO NOTHING' in MySQL
+    [" ON DUPLICATE KEY UPDATE ", Key, " = ", Key];
+mysql_on_conflict(UpdateFields, _) ->
     [" ON DUPLICATE KEY UPDATE ",
      update_fields_on_conflict(UpdateFields)].
 
+pgsql_on_conflict([], UniqueKeyFields) ->
+    JoinedKeys = join(UniqueKeyFields, ", "),
+    [" ON CONFLICT (", JoinedKeys, ") DO NOTHING"];
 pgsql_on_conflict(UpdateFields, UniqueKeyFields) ->
     JoinedKeys = join(UniqueKeyFields, ", "),
     [" ON CONFLICT (", JoinedKeys, ")"
@@ -164,9 +170,11 @@ upsert_mssql_query(Table, InsertFields, UpdateFields, UniqueKeyFields) ->
         " ON (", join(UniqueConstraint, " AND "), ")"
      " WHEN NOT MATCHED THEN INSERT"
        " (", JoinedInsertFields, ")"
-         " VALUES (", join(Placeholders, ", "), ")"
-     " WHEN MATCHED THEN UPDATE"
-       " SET ", update_fields_on_conflict(UpdateFields),";"].
+         " VALUES (", join(Placeholders, ", "), ")" | mssql_on_conflict(UpdateFields)].
+
+mssql_on_conflict([]) -> ";";
+mssql_on_conflict(UpdateFields) ->
+     [" WHEN MATCHED THEN UPDATE SET ", update_fields_on_conflict(UpdateFields), ";"].
 
 %% F can be either a fun or a list of queries
 %% TODO: We should probably move the list of queries transaction

test/auth_tokens_SUITE.erl

@@ -59,6 +59,7 @@ init_per_testcase(Test, Config)
     Config1;
 
 init_per_testcase(validity_period_test, Config) ->
+    mock_rdbms_backend(),
     mock_mongoose_metrics(),
     mock_gen_iq_handler(),
     mock_ejabberd_commands(),
@@ -86,6 +87,7 @@ end_per_testcase(Test, C)
     C;
 
 end_per_testcase(validity_period_test, C) ->
+    meck:unload(mod_auth_token_rdbms),
     meck:unload(mongoose_metrics),
     meck:unload(gen_iq_handler),
     meck:unload(ejabberd_commands),
@@ -239,10 +241,12 @@ mock_mongoose_metrics() ->
     ok.
 
 mock_rdbms_backend() ->
-    gen_mod:start_backend_module(?TESTED, [{backend, rdbms}]),
     meck:new(mod_auth_token_rdbms, []),
+    meck:expect(mod_auth_token_rdbms, start, fun(_) -> ok end),
     meck:expect(mod_auth_token_rdbms, get_valid_sequence_number,
-                fun (_) -> valid_seq_no_threshold() end).
+                fun (_) -> valid_seq_no_threshold() end),
+    gen_mod:start_backend_module(?TESTED, [{backend, rdbms}]),
+    ok.
 
 mock_keystore() ->
     ejabberd_hooks:add(get_key, <<"localhost">>, ?MODULE, mod_keystore_get_key, 50).