%%%------------------------------------------------------------------- %%% File : ejabberd_sm_redis.erl %%% Author : Evgeny Khramtsov %%% Created : 11 Mar 2015 by Evgeny Khramtsov %%% %%% %%% ejabberd, Copyright (C) 2002-2016 ProcessOne %%% %%% 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., %%% 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. %%% %%%---------------------------------------------------------------------- -module(ejabberd_sm_redis). -behaviour(ejabberd_config). -behaviour(ejabberd_sm). -export([init/0, set_session/1, delete_session/4, get_sessions/0, get_sessions/1, get_sessions/2, get_sessions/3, opt_type/1]). -include("ejabberd.hrl"). -include("ejabberd_sm.hrl"). -include("logger.hrl"). %%%=================================================================== %%% API %%%=================================================================== -spec init() -> ok | {error, any()}. init() -> clean_table(). -spec set_session(#session{}) -> ok. set_session(Session) -> T = term_to_binary(Session), USKey = us_to_key(Session#session.us), SIDKey = sid_to_key(Session#session.sid), ServKey = server_to_key(element(2, Session#session.us)), USSIDKey = us_sid_to_key(Session#session.us, Session#session.sid), case ejabberd_redis:qp([["HSET", USKey, SIDKey, T], ["HSET", ServKey, USSIDKey, T]]) of [{ok, _}, {ok, _}] -> ok; Err -> ?ERROR_MSG("failed to set session for redis: ~p", [Err]) end. -spec delete_session(binary(), binary(), binary(), sid()) -> {ok, #session{}} | {error, notfound}. delete_session(LUser, LServer, _LResource, SID) -> USKey = us_to_key({LUser, LServer}), case ejabberd_redis:q(["HGETALL", USKey]) of {ok, Vals} -> Ss = decode_session_list(Vals), case lists:keyfind(SID, #session.sid, Ss) of false -> {error, notfound}; Session -> SIDKey = sid_to_key(SID), ServKey = server_to_key(element(2, Session#session.us)), USSIDKey = us_sid_to_key(Session#session.us, SID), ejabberd_redis:qp([["HDEL", USKey, SIDKey], ["HDEL", ServKey, USSIDKey]]), {ok, Session} end; Err -> ?ERROR_MSG("failed to delete session from redis: ~p", [Err]), {error, notfound} end. -spec get_sessions() -> [#session{}]. get_sessions() -> lists:flatmap( fun(LServer) -> get_sessions(LServer) end, ejabberd_sm:get_vh_by_backend(?MODULE)). -spec get_sessions(binary()) -> [#session{}]. get_sessions(LServer) -> ServKey = server_to_key(LServer), case ejabberd_redis:q(["HGETALL", ServKey]) of {ok, Vals} -> decode_session_list(Vals); Err -> ?ERROR_MSG("failed to get sessions from redis: ~p", [Err]), [] end. -spec get_sessions(binary(), binary()) -> [#session{}]. get_sessions(LUser, LServer) -> USKey = us_to_key({LUser, LServer}), case ejabberd_redis:q(["HGETALL", USKey]) of {ok, Vals} when is_list(Vals) -> decode_session_list(Vals); Err -> ?ERROR_MSG("failed to get sessions from redis: ~p", [Err]), [] end. -spec get_sessions(binary(), binary(), binary()) -> [#session{}]. get_sessions(LUser, LServer, LResource) -> USKey = us_to_key({LUser, LServer}), case ejabberd_redis:q(["HGETALL", USKey]) of {ok, Vals} when is_list(Vals) -> [S || S <- decode_session_list(Vals), element(3, S#session.usr) == LResource]; Err -> ?ERROR_MSG("failed to get sessions from redis: ~p", [Err]), [] end. %%%=================================================================== %%% Internal functions %%%=================================================================== iolist_to_list(IOList) -> binary_to_list(iolist_to_binary(IOList)). us_to_key({LUser, LServer}) -> <<"ejabberd:sm:", LUser/binary, "@", LServer/binary>>. server_to_key(LServer) -> <<"ejabberd:sm:", LServer/binary>>. us_sid_to_key(US, SID) -> term_to_binary({US, SID}). sid_to_key(SID) -> term_to_binary(SID). decode_session_list([_, Val|T]) -> [binary_to_term(Val)|decode_session_list(T)]; decode_session_list([]) -> []. clean_table() -> ?INFO_MSG("Cleaning Redis SM table...", []), lists:foreach( fun(LServer) -> ServKey = server_to_key(LServer), case ejabberd_redis:q(["HKEYS", ServKey]) of {ok, []} -> ok; {ok, Vals} -> Vals1 = lists:filter( fun(USSIDKey) -> {_, SID} = binary_to_term(USSIDKey), node(element(2, SID)) == node() end, Vals), Q1 = case Vals1 of [] -> []; _ -> ["HDEL", ServKey | Vals1] end, Q2 = lists:map( fun(USSIDKey) -> {US, SID} = binary_to_term(USSIDKey), USKey = us_to_key(US), SIDKey = sid_to_key(SID), ["HDEL", USKey, SIDKey] end, Vals1), Res = ejabberd_redis:qp(lists:delete([], [Q1|Q2])), case lists:filter( fun({ok, _}) -> false; (_) -> true end, Res) of [] -> ok; Errs -> ?ERROR_MSG("failed to clean redis table for " "server ~s: ~p", [LServer, Errs]) end; Err -> ?ERROR_MSG("failed to clean redis table for " "server ~s: ~p", [LServer, Err]) end end, ejabberd_sm:get_vh_by_backend(?MODULE)). opt_type(redis_connect_timeout) -> fun (I) when is_integer(I), I > 0 -> I end; opt_type(redis_db) -> fun (I) when is_integer(I), I >= 0 -> I end; opt_type(redis_password) -> fun iolist_to_list/1; opt_type(redis_port) -> fun (P) when is_integer(P), P > 0, P < 65536 -> P end; opt_type(redis_reconnect_timeout) -> fun (I) when is_integer(I), I > 0 -> I end; opt_type(redis_server) -> fun iolist_to_list/1; opt_type(_) -> [redis_connect_timeout, redis_db, redis_password, redis_port, redis_reconnect_timeout, redis_server].