mod_msg_filter allows the filtering of "message" stanzas across an HTTP service.

mod_msg_filter.erl

%%
%% mod_msg_filter allows the filtering of "message"
%% stanzas across an HTTP service. The URL of the
%% service must be passed as part of the module's
%% configuration. Both JIDs and their resources are
%% passed as part of the query string and the result
%% is expected to be one of:
%%
%% <status value="denied">
%%  <stanza1><error/></stanza1>
%%  <stanza2><error/></stanza2>
%% </status>
%%
%% or:
%%
%% <status value="allowed">
%%  <stanza1><noop/></stanza1>
%%  <stanza2><noop/></stanza2>
%% </status>
%%
%% The values of the <error> tags or <noop> tags will
%% be cached in mnesia using 2 keys that look like:
%%
%% {bare_jid1, bare_jid2, resource1, resource2}
%% {bare_jid2, bare_jid1, resource2, resource1}
%%
%% The <error> tags will then be sent over to both JIDs
%% if the <status> has a "value" of "denied", otherwise
%% the original message is let through.
%%
%% The mnesia cache can be flushed if the ejabberd
%% server is hit on a request handler that maps to this
%% module (for example: /mod_msg_manager/). A stanza must
%% be POSTed that looks like:
%%
%% <flush jid="user@domain.com"/>
%%
%% Note that no resource is included at all. This can be
%% used if ejabberd is part of a system that has billing
%% restrictions on chatting but allows presence to go
%% through all the time.
%%
-module(mod_msg_filter).
-author('hisham.mardambey@gmail.com').

-behaviour(gen_server).
-export([start_link/3]).
-export([init/1, handle_call/3, 
		 handle_cast/2, handle_info/2, 
		 terminate/2, code_change/3]).

-behaviour(gen_mod).
-export([start/2, stop/1, 
		 filter_packet/1, process/2]).

-include("ejabberd.hrl").
-include("jlib.hrl").
-include("mod_roster.hrl").
-include("web/ejabberd_http.hrl").

%% Url the HTTP service uses to fetch permissions
%%
-define(DEFAULT_PERMS_URL, "http://localhost/").

%% This is the response the HTTP service get back 
%% after its checked for someone's permissions
%%
-define(DENIED, "denied").

%% This is the response the HTTP service receives 
%% after its checked a member's permissions and 
%% they are allowed to chat.
%%
-define(ALLOWED, "allowed").

%% Global configuration
%%
-record(config, {perms_url}).

%% Represents a permission pair that represents 
%% whether 2 JIDs can chat or not.
%% id = {bare_jid1, bare_jid2, resource1, resource2}
%% status = allowed | denied
%% jid1Stanza = <error> to send to jid1
%% jid2Stanza = <error> to send to jid2
%%
-record(perms, {id, status, jid1Stanza, jid2Stanza}).

%% Starts the module and reads in the configuration.
%% Starts the supervisor as well.
%%
start(Host, Opts) ->	
	?INFO_MSG("Loading module 'mod_msg_filter'", []),
	PermsUrl = gen_mod:get_opt(perms_url, Opts, ?DEFAULT_PERMS_URL),
	Proc = gen_mod:get_module_proc(Host, ?MODULE),
	ChildSpec = {Proc,
         {?MODULE, start_link, [Host, Opts, #config{perms_url = PermsUrl}]},
         permanent,
         1000,
         worker,
         [?MODULE]},
	supervisor:start_child(ejabberd_sup, ChildSpec).	

%% Shuts down the module.
%%
stop(Host) ->
	?INFO_MSG("Unloading module 'mod_msg_filter'", []),	
	Proc = gen_mod:get_module_proc(Host, ?MODULE),
    supervisor:terminate_child(ejabberd_sup, Proc),
    supervisor:delete_child(ejabberd_sup, Proc),	
	ok.

start_link(_Host, _Opts, Config) ->
	?DEBUG("start_link: ~p", [Config]),
    gen_server:start_link({local, ?MODULE}, ?MODULE, [Config], []).

init([Config]) ->
	inets:start(),	
	db_init(),
	ejabberd_hooks:add(filter_packet, global, ?MODULE, filter_packet, 100),
    {ok, Config}.

%% Handles all gen_server calls.
%%
handle_call(Request, _From, Config) ->
	?DEBUG("handle_call: ~p ~p" , [Request, Config]),

	Reply = case Request of
				%% being asked to get perms over http
				{permissions, {From, To, R1, R2}} -> perms_get_http(From, To, R1, R2, Config);
				_ -> ok
			end,
	
    {reply, Reply, Config}.

handle_cast(_Msg, State) ->
    {noreply, State}.

handle_info(_Info, State) ->
    {noreply, State}.

terminate(_Reason, _State) ->
	ejabberd_hooks:delete(filter_packet, global, ?MODULE, filter_packet, 100),	
    ok.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%% Return drop to drop the packet, or the original input to let it through.
%% From and To are jid records.
%%
filter_packet(drop) -> drop;
filter_packet({From, To, _Packet} = Input) ->
	?DEBUG("filter_packet(~p)", [Input]),	
	%% It probably doesn't make any sense to block packets to oneself.
	%% If this is not out packet or ourselves we'll check_stanza and 
	%% either return drop, allow, or send another stanza instead
	R = if From#jid.luser == To#jid.luser, From#jid.lserver == To#jid.lserver -> Input;
		From#jid.luser == "admin" -> Input; %% allow stanzas from admin
		true -> check_stanza(Input)
		end,

	case R of
		{drop, _} -> drop;
		{drop, _, _} -> drop;
		_ -> R
	end.

%% If this is anything but a "message" stanza move on. 
%% If this is a "message" then we'll inspect its From 
%% and To and figure out if we should let it through 
%% or spoof error stanzas to both sides and drop the 
%% original message.
%%
check_stanza({From, To, {xmlelement, "message", _Attrs, _Els}} = Input) ->
	
	?DEBUG("check_stanza(message): Got message from ~p to ~p", [jlib:jid_to_string(From), jlib:jid_to_string(To)]),
	case can_chat(From, To) of 
		allowed -> Input;                      %% allowed to chat, let packet through
		{denied, FromStanza, ToStanza} ->      %% not allowed, send respective stanzas
			J = #jid{ 
					 user="admin", 
					 server=From#jid.server, 
					 resource="",
					 luser="admin",
					 lserver=From#jid.lserver, 
					 lresource=""
			},
			spoof_error(J, From, FromStanza), %% send error to From, with admin as sender
			spoof_error(J, To, ToStanza),     %% send error to To, with admin as sender
			drop;                             %% drop the packet
		_ -> drop                             %% drop on an unrecognized reponse
	end;
check_stanza(Input) -> 
	?DEBUG("check_stanza: letting packet through: ~p", [Input]),
	Input.

%% Creates and routes a "message" stanza of type 
%% error that has an embedded tag given as the 
%% Error parameter (usually an <error>).
%%
spoof_error(From, To, Error) -> 
	ejabberd_router:route(From, To,								 
	{
		xmlelement, 
		"message", 
		[{"to", jlib:jid_to_string(To)}, {"type","error"}], 
		[Error]
	}).

%% Creates a composite key using {From, To} 
%% (bare JIDs) and attempts look it up in the 
%% database. If not found then we attempt to 
%% get and cache in the db the permissions using 
%% the HTTP service. Once we have the permissions 
%%  
can_chat(From, To) ->	
	Key = {
		   jlib:jid_to_string(jlib:jid_remove_resource(From)),
		   jlib:jid_to_string(jlib:jid_remove_resource(To)),
		   From#jid.resource,
		   To#jid.resource
	},
	perms_parse(Key, db_read(perms, Key)).

%% Parse a permissions entry.
%% 
perms_parse(Key, %% no permissions, load and cache using http service
	[]) ->
		?DEBUG("perms_parse(~p): not found in db, getting over http.", [Key]),
		perms_get_http_and_cache(Key);

perms_parse(_,  %% allowed, lets the message through
	[#perms{id=Id , status=allowed, jid1Stanza=_S1, jid2Stanza=_S2}]) ->
		?DEBUG("perms_parse(~p): allowed, letting message through.", [Id]),
		allowed;

perms_parse(_,  %% denied, return stanzas to they can be sent to both sides
	[#perms{id=Id , status=denied, jid1Stanza=S1, jid2Stanza=S2}]) ->
		?DEBUG("perms_parse(~p): denied, sending respective stanzas.", [Id]),
		{denied, S1, S2};

perms_parse(Key, %% denied, unknown response
	Response) ->
		?DEBUG("perms_parse(~p): denied, unrecognized response: ~p", [Key, Response]),
		denied.

%% Gets permissions for given Key over HTTP 
%% service and caches them in the database.
%%
perms_get_http_and_cache(Key) ->
	?DEBUG("perms_get_http_and_cache(~p): trying to get and cache perms.", [Key]),
	case gen_server:call(?MODULE, {permissions, Key}) of	
		{?ALLOWED, _Stanza1, _Stanza2} -> 
			?DEBUG("perms_get_http_and_cache(~p): allowed, letting message through.", [Key]),
			perms_cache(Key, allowed),
			allowed;
		{?DENIED, Stanza1, Stanza2} ->
			?DEBUG("perms_get_http_and_cache(~p): denied, sending respective stanzas.", [Key]),
			perms_cache(Key, denied, Stanza1, Stanza2),
			{denied, Stanza1, Stanza2};
		Unknown ->
			?DEBUG("perms_get_http_and_cache(~p): unknown response ~p", [Key, Unknown]),
			denied %% send denied but do not cache it as this was error
	end.

%% Cache permissions for Key in the database.
%%
perms_cache({Jid1, Jid2, R1, R2} = _Key, allowed) ->                           %% cache allowed with noop as stanzas
	?DEBUG("perms_cache(~p ~p): allowed", [Jid1, Jid2]),
	perms_cache(Jid1, Jid2, R1, R2, allowed, noop, noop).
perms_cache({Jid1, Jid2, R1, R2} = _Key, denied, Stanza1, Stanza2) ->          %% cache denied with given stanzas 
	?DEBUG("perms_cache(~p ~p): denied", [Jid1, Jid2]),
	perms_cache(Jid1, Jid2, R1, R2, denied, Stanza1, Stanza2).
perms_cache(Jid1, Jid2, R1, R2, Status, Stanza1, Stanza2) -> %% cache Status for {Jid1,Jid2} and {Jid2, Jid1}
	?DEBUG("perms_cache: ~p ~p ~p", [Status,
		#perms{id={Jid1, Jid2, R1, R2}, status=Status, jid1Stanza=Stanza1, jid2Stanza=Stanza2},
		#perms{id={Jid2, Jid1, R2, R1}, status=Status, jid1Stanza=Stanza2, jid2Stanza=Stanza1}
	]),
	
	db_transaction_write([
		#perms{id={Jid1, Jid2, R1, R2}, status=Status, jid1Stanza=Stanza1, jid2Stanza=Stanza2},
		#perms{id={Jid2, Jid1, R2, R1}, status=Status, jid1Stanza=Stanza2, jid2Stanza=Stanza1}
	]).

%% Flush persmissions for the given JID. 
%% This works by matching all objects in 
%% the cache that look like:
%% {perms, {Jid, '_', '_', '_'}, '_', '_', '_'}
%% {perms, {'_', Jid, '_', '_'}, '_', '_', '_'}
%%
perms_flush(Jid) ->
	Fun = fun() ->
		L1 = mnesia:match_object(perms, {perms, {Jid, '_', '_', '_'}, '_', '_', '_'}, read),
		L2 = mnesia:match_object(perms, {perms, {'_', Jid, '_', '_'}, '_', '_', '_'}, read),
		perms_delete(L1),
		perms_delete(L2)
	end,
	mnesia:transaction(Fun).

perms_delete([]) -> ok;
perms_delete([Item|Tail]) -> 
	Key = Item#perms.id,
	?DEBUG("perms_delete(~p)", [Key]),
	mnesia:delete({perms, Key}),
	perms_delete(Tail).

%% Writes the given list into the 
%% database using a transaction.
%%
db_transaction_write(List) ->
	Fun = fun() ->
				  db_write(List)
		  end,
	mnesia:transaction(Fun).

%% Writes the given list into the database.
%% Meant to be used within a transaction.
%%
db_write([]) -> ok;
db_write([Item|Tail]) ->
	?DEBUG("db_write(~p)", [Item]),
	mnesia:write(Item),
	db_write(Tail).

%% Given Table and Key, tries to fetch Key from 
%% Table using a transactional read.
%%
db_read(Table, Key) -> 
	Fun = fun() ->
		mnesia:read(Table, Key)
	end,
	
	case mnesia:transaction(Fun) of
		{atomic, []} ->
			?DEBUG("db_read(~p): []", [Key]),
			[];
		{atomic, List} ->
			?DEBUG("db_read(~p): ~p", [Key, List]),
			List;		
		Unknown ->
			?DEBUG("db_read(~p): ~p", [Key, Unknown]),
			[]
	end.

%% Remove all white-space from the given input.
%%
trim_whitespace(Input) -> re:replace(Input, "\\s+", " ", [global]).

%% This loop is spawned off and is responsible for making HTTP 
%% requests to fetch "permissions" and "blocklist" for Jids.
%%
perms_get_http(From, To, R1, R2, Config) ->					
	% construct url			
	Url = io_lib:format(Config#config.perms_url, [From, To, R1, R2]),
							
	% make http request and parse out result
	{ok, {{_Version, 200, _ReasonPhrase}, _Headers, Body}} = http:request(Url),
	
	Result = trim_whitespace(Body),
	?DEBUG("Got permission response: ~p", [Result]),
	
	% send the result back
	case xml_stream:parse_element(Result) of

		{xmlelement, "status", [{"value", Status}], [Stanza1, Stanza2]} -> %% valid response structure
			?DEBUG("perms_get_http: got valid perms response: ~p ~p ~p", [Status, Stanza1, Stanza2]),
			{Status, Stanza1, Stanza2};

		Unknown -> %% invalid response, deny request
			?DEBUG("perms_get_http: got invalid perms response: ~p", [Unknown]),
			{denied}  
	end.
	

%% Initializes mnesia and creates needed schema 
%% and tables.
%%
db_init() ->
	%% TODO: should we catch the error from create_schema?
	mnesia:create_schema([node()]),
	db_create_table(perms, [{attributes, record_info(fields, perms)}]),
	mnesia:start().

%% Creates a table if it does not exist.
%% Returns one of:
%% "already_exists": table already exists 
%% "ok": table created
%% "{error, reason}": unknown error with the mnesia error 
db_create_table(Name, ArgList) ->
	case mnesia:create_table(Name, ArgList) of 
		{aborted, {already_exists, _Table}} -> already_exists;
		{aborted, Reason} -> {error, Reason};
		{atomic, ok} -> ok;
		Err -> {error, Err}
	end.

%% Given a Jid delete all entries that have 
%% it in its key from the cache.
%%
process([], #request{method = 'POST',
             data = Data,
             host = Host,
             ip = ClientIp
            }) ->
    {200, [], process_http_request(Data, Host, ClientIp)};
    
process(Path, Request) ->
    ?DEBUG("Got request to ~p: ~p", [Path, Request]),
    {200, [], "Try POSTing a stanza."}.

%% If the first character of Data is "<", it is considered a stanza to process.
%% Otherwise, an error is returned.
%%
process_http_request([$< | _ ] = Data, _Host, _ClientIp) ->
    Stanza = xml_stream:parse_element(Data),
	case Stanza of 
		{xmlelement, "flush", [{"jid", Jid}], []} ->
			F = perms_flush(Jid),
			?DEBUG("perms_flush: ~p", [F]),
			successful_http_request();
		Unknown -> 
			?DEBUG("process_http_request: unknown request ~p", [Unknown]),
			unknown_http_request()
	end;

process_http_request(_Data, _Host, _ClientIp) -> unknown_http_request().

successful_http_request() ->
       {xmlelement, "status", [{"value", "success"}], []}.

unknown_http_request() ->
       {xmlelement, "status", [{"value", "error"}], []}.