Support user quotas for HTTP File Upload

mod_http_upload_quota implements two features:

- When a "hard quota" is exceeded during a file upload, old files are
  removed until the disk usage equals or falls below the "soft quota".

- Once a day, all uploaded files (and directories) older than a
  configurable number of days are deleted.

src/mod_http_upload_quota.erl

@@ -0,0 +1,345 @@
+%%%-------------------------------------------------------------------
+%%% File    : mod_http_upload_quota.erl
+%%% Author  : Holger Weiss <holger@zedat.fu-berlin.de>
+%%% Purpose : Quota management for HTTP File Upload (XEP-0363)
+%%% Created : 15 Oct 2015 by Holger Weiss <holger@zedat.fu-berlin.de>
+%%%-------------------------------------------------------------------
+
+-module(mod_http_upload_quota).
+-author('holger@zedat.fu-berlin.de').
+
+-define(GEN_SERVER, gen_server).
+-define(PROCNAME, ?MODULE).
+-define(TIMEOUT, timer:hours(24)).
+-define(INITIAL_TIMEOUT, timer:minutes(10)).
+-define(FORMAT(Error), file:format_error(Error)).
+
+-behaviour(?GEN_SERVER).
+-behaviour(gen_mod).
+
+%% gen_mod/supervisor callbacks.
+-export([start_link/3,
+	 start/2,
+	 stop/1,
+	 mod_opt_type/1]).
+
+%% gen_server callbacks.
+-export([init/1,
+	 handle_call/3,
+	 handle_cast/2,
+	 handle_info/2,
+	 terminate/2,
+	 code_change/3]).
+
+%% ejabberd_hooks callback.
+-export([handle_slot_request/5]).
+
+-include("jlib.hrl").
+-include("logger.hrl").
+-include_lib("kernel/include/file.hrl").
+
+-record(state,
+	{server_host                    :: binary(),
+	 access_soft_quota              :: atom(),
+	 access_hard_quota              :: atom(),
+	 max_days                       :: pos_integer() | infinity,
+	 docroot                        :: binary(),
+	 disk_usage = dict:new()        :: term(),
+	 timers                         :: [timer:tref()]}).
+
+-type state() :: #state{}.
+
+%%--------------------------------------------------------------------
+%% gen_mod/supervisor callbacks.
+%%--------------------------------------------------------------------
+
+-spec start_link(binary(), atom(), gen_mod:opts())
+      -> {ok, pid()} | ignore | {error, _}.
+
+start_link(ServerHost, Proc, Opts) ->
+    ?GEN_SERVER:start_link({local, Proc}, ?MODULE, {ServerHost, Opts}, []).
+
+-spec start(binary(), gen_mod:opts()) -> {ok, _} | {ok, _, _} | {error, _}.
+
+start(ServerHost, Opts) ->
+    Proc = mod_http_upload:get_proc_name(ServerHost, ?PROCNAME),
+    Spec = {Proc,
+	    {?MODULE, start_link, [ServerHost, Proc, Opts]},
+	    permanent,
+	    3000,
+	    worker,
+	    [?MODULE]},
+    supervisor:start_child(ejabberd_sup, Spec).
+
+-spec stop(binary()) -> ok.
+
+stop(ServerHost) ->
+    Proc = mod_http_upload:get_proc_name(ServerHost, ?PROCNAME),
+    ok = supervisor:terminate_child(ejabberd_sup, Proc),
+    ok = supervisor:delete_child(ejabberd_sup, Proc).
+
+-spec mod_opt_type(atom()) -> fun((term()) -> term()) | [atom()].
+
+mod_opt_type(access_soft_quota) ->
+    fun(A) when is_atom(A) -> A end;
+mod_opt_type(access_hard_quota) ->
+    fun(A) when is_atom(A) -> A end;
+mod_opt_type(max_days) ->
+    fun(I) when is_integer(I), I > 0 -> I;
+       (infinity) -> infinity
+    end;
+mod_opt_type(_) ->
+    [access_soft_quota, access_hard_quota, max_days].
+
+%%--------------------------------------------------------------------
+%% gen_server callbacks.
+%%--------------------------------------------------------------------
+
+-spec init({binary(), gen_mod:opts()}) -> {ok, state()}.
+
+init({ServerHost, Opts}) ->
+    process_flag(trap_exit, true),
+    AccessSoftQuota = gen_mod:get_opt(access_soft_quota, Opts,
+				      fun(A) when is_atom(A) -> A end,
+				      soft_upload_quota),
+    AccessHardQuota = gen_mod:get_opt(access_hard_quota, Opts,
+				      fun(A) when is_atom(A) -> A end,
+				      hard_upload_quota),
+    MaxDays = gen_mod:get_opt(max_days, Opts,
+			      fun(I) when is_integer(I), I > 0 -> I;
+				 (infinity) -> infinity
+			      end,
+			      infinity),
+    DocRoot1 = gen_mod:get_module_opt(ServerHost, mod_http_upload, docroot,
+				      fun iolist_to_binary/1,
+				      <<"@HOME@/upload">>),
+    DocRoot2 = mod_http_upload:expand_home(str:strip(DocRoot1, right, $/)),
+    Timers = if MaxDays == infinity -> [];
+		true ->
+		     {ok, T1} = timer:send_after(?INITIAL_TIMEOUT, sweep),
+		     {ok, T2} = timer:send_interval(?TIMEOUT, sweep),
+		     [T1, T2]
+	     end,
+    ejabberd_hooks:add(http_upload_slot_request, ServerHost, ?MODULE,
+		       handle_slot_request, 50),
+    {ok, #state{server_host = ServerHost,
+		access_soft_quota = AccessSoftQuota,
+		access_hard_quota = AccessHardQuota,
+		max_days = MaxDays,
+		docroot = DocRoot2,
+		timers = Timers}}.
+
+-spec handle_call(_, {pid(), _}, state()) -> {noreply, state()}.
+
+handle_call(Request, From, State) ->
+    ?ERROR_MSG("Got unexpected request from ~p: ~p", [From, Request]),
+    {noreply, State}.
+
+-spec handle_cast(_, state()) -> {noreply, state()}.
+
+handle_cast({handle_slot_request, #jid{user = U, server = S} = JID, Path, Size},
+	    #state{server_host = ServerHost,
+		   access_soft_quota = AccessSoftQuota,
+		   access_hard_quota = AccessHardQuota,
+		   disk_usage = DiskUsage} = State) ->
+    HardQuota = case acl:match_rule(ServerHost, AccessHardQuota, JID) of
+		  Hard when is_integer(Hard), Hard > 0 ->
+		      Hard * 1024 * 1024;
+		  _ ->
+		      0
+		end,
+    SoftQuota = case acl:match_rule(ServerHost, AccessSoftQuota, JID) of
+		  Soft when is_integer(Soft), Soft > 0 ->
+		      Soft * 1024 * 1024;
+		  _ ->
+		      0
+		end,
+    OldSize = case dict:find({U, S}, DiskUsage) of
+		{ok, Value} ->
+		    Value;
+		error ->
+		    undefined
+	      end,
+    NewSize = case {HardQuota, SoftQuota} of
+		{0, 0} ->
+		    ?DEBUG("No quota specified for ~s",
+			   [jlib:jid_to_string(JID)]),
+		    Size;
+		{0, _} ->
+		      ?WARNING_MSG("No hard quota specified for ~s",
+				   [jlib:jid_to_string(JID)]),
+		      enforce_quota(Path, Size, OldSize, SoftQuota, SoftQuota);
+		{_, 0} ->
+		      ?WARNING_MSG("No soft quota specified for ~s",
+				   [jlib:jid_to_string(JID)]),
+		      enforce_quota(Path, Size, OldSize, HardQuota, HardQuota);
+		_ when SoftQuota > HardQuota ->
+		      ?WARNING_MSG("Bad quota for ~s (soft: ~p, hard: ~p)",
+				   [jlib:jid_to_string(JID),
+				    SoftQuota, HardQuota]),
+		      enforce_quota(Path, Size, OldSize, SoftQuota, SoftQuota);
+		_ ->
+		      ?DEBUG("Enforcing quota for ~s",
+			     [jlib:jid_to_string(JID)]),
+		      enforce_quota(Path, Size, OldSize, SoftQuota, HardQuota)
+	      end,
+    {noreply, State#state{disk_usage = dict:store({U, S}, NewSize, DiskUsage)}};
+handle_cast(Request, State) ->
+    ?ERROR_MSG("Got unexpected request: ~p", [Request]),
+    {noreply, State}.
+
+-spec handle_info(_, state()) -> {noreply, state()}.
+
+handle_info(sweep, #state{server_host = ServerHost,
+			  docroot = DocRoot,
+			  max_days = MaxDays} = State)
+    when is_integer(MaxDays), MaxDays > 0 ->
+    ?DEBUG("Got 'sweep' message for ~s", [ServerHost]),
+    case file:list_dir(DocRoot) of
+      {ok, Entries} ->
+	  BackThen = secs_since_epoch() - (MaxDays * 86400),
+	  DocRootS = binary_to_list(DocRoot),
+	  PathNames = lists:map(fun(Entry) -> DocRootS ++ "/" ++ Entry end,
+				Entries),
+	  UserDirs = lists:filter(fun filelib:is_dir/1, PathNames),
+	  lists:foreach(fun(UserDir) -> delete_old_files(UserDir, BackThen) end,
+			UserDirs);
+      {error, Error} ->
+	  ?ERROR_MSG("Cannot open document root ~s: ~s",
+		     [DocRoot, ?FORMAT(Error)])
+    end,
+    {noreply, State};
+handle_info(Info, State) ->
+    ?ERROR_MSG("Got unexpected info: ~p", [Info]),
+    {noreply, State}.
+
+-spec terminate(normal | shutdown | {shutdown, _} | _, _) -> ok.
+
+terminate(Reason, #state{server_host = ServerHost, timers = Timers}) ->
+    ?DEBUG("Stopping upload quota process for ~s: ~p", [ServerHost, Reason]),
+    ejabberd_hooks:delete(http_upload_slot_request, ServerHost, ?MODULE,
+			  handle_slot_request, 50),
+    lists:foreach(fun(Timer) -> timer:cancel(Timer) end, Timers).
+
+-spec code_change({down, _} | _, state(), _) -> {ok, state()}.
+
+code_change(_OldVsn, #state{server_host = ServerHost} = State, _Extra) ->
+    ?DEBUG("Updating upload quota process for ~s", [ServerHost]),
+    {ok, State}.
+
+%%--------------------------------------------------------------------
+%% ejabberd_hooks callback.
+%%--------------------------------------------------------------------
+
+-spec handle_slot_request(term(), jid(), binary(), non_neg_integer(), binary())
+      -> term().
+
+handle_slot_request(allow, #jid{lserver = ServerHost} = JID, Path, Size,
+		    _Lang) ->
+    Proc = mod_http_upload:get_proc_name(ServerHost, ?PROCNAME),
+    ?GEN_SERVER:cast(Proc, {handle_slot_request, JID, Path, Size}),
+    allow;
+handle_slot_request(Acc, _JID, _Path, _Size, _Lang) -> Acc.
+
+%%--------------------------------------------------------------------
+%% Internal functions.
+%%--------------------------------------------------------------------
+
+-spec enforce_quota(file:filename_all(), non_neg_integer(),
+		    non_neg_integer() | undefined, non_neg_integer(),
+		    non_neg_integer())
+      -> non_neg_integer().
+
+enforce_quota(_UserDir, SlotSize, OldSize, _MinSize, MaxSize)
+    when is_integer(OldSize), OldSize + SlotSize =< MaxSize ->
+    OldSize + SlotSize;
+enforce_quota(UserDir, SlotSize, _OldSize, MinSize, MaxSize) ->
+    Files = lists:sort(fun({_PathA, _SizeA, TimeA}, {_PathB, _SizeB, TimeB}) ->
+			       TimeA > TimeB
+		       end, gather_file_info(UserDir)),
+    {DelFiles, OldSize, NewSize} =
+	lists:foldl(fun({_Path, Size, _Time}, {[], AccSize, AccSize})
+			    when AccSize + Size + SlotSize =< MinSize ->
+			    {[], AccSize + Size, AccSize + Size};
+		       ({Path, Size, _Time}, {[], AccSize, AccSize}) ->
+			    {[Path], AccSize + Size, AccSize};
+		       ({Path, Size, _Time}, {AccFiles, AccSize, NewSize}) ->
+			    {[Path | AccFiles], AccSize + Size, NewSize}
+		    end, {[], 0, 0}, Files),
+    if OldSize + SlotSize > MaxSize ->
+	    lists:foreach(fun(File) -> del_file_and_dir(File) end, DelFiles),
+	    file:del_dir(UserDir), % In case it's empty, now.
+	    NewSize + SlotSize;
+       true ->
+	    OldSize + SlotSize
+    end.
+
+-spec delete_old_files(file:filename_all(), integer()) -> ok.
+
+delete_old_files(UserDir, Timestamp) ->
+    FileInfo = gather_file_info(UserDir),
+    case [Path || {Path, _Size, Time} <- FileInfo, Time < Timestamp] of
+      [] ->
+	  ok;
+      OldFiles ->
+	  lists:foreach(fun(File) -> del_file_and_dir(File) end, OldFiles),
+	  file:del_dir(UserDir) % In case it's empty, now.
+    end.
+
+-spec gather_file_info(file:filename_all())
+      -> [{binary(), non_neg_integer(), non_neg_integer()}].
+
+gather_file_info(Dir) when is_binary(Dir) ->
+    gather_file_info(binary_to_list(Dir));
+gather_file_info(Dir) ->
+    case file:list_dir(Dir) of
+      {ok, Entries} ->
+	  lists:foldl(fun(Entry, Acc) ->
+			      Path = Dir ++ "/" ++ Entry,
+			      case file:read_file_info(Path, [{time, posix}]) of
+				{ok, #file_info{type = directory}} ->
+				    gather_file_info(Path) ++ Acc;
+				{ok, #file_info{type = regular,
+						mtime = Time,
+						size = Size}} ->
+				    [{Path, Size, Time} | Acc];
+				{ok, _Info} ->
+				    ?DEBUG("Won't stat(2) non-regular file ~s",
+					   [Path]),
+				    Acc;
+				{error, Error} ->
+				    ?ERROR_MSG("Cannot stat(2) ~s: ~s",
+					       [Path, ?FORMAT(Error)]),
+				    Acc
+			      end
+		      end, [], Entries);
+      {error, enoent} ->
+	  ?DEBUG("Directory ~s doesn't exist", [Dir]),
+	  [];
+      {error, Error} ->
+	  ?ERROR_MSG("Cannot open directory ~s: ~s", [Dir, ?FORMAT(Error)]),
+	  []
+    end.
+
+-spec del_file_and_dir(file:name_all()) -> ok.
+
+del_file_and_dir(File) ->
+    case file:delete(File) of
+      ok ->
+	  ?INFO_MSG("Removed ~s", [File]),
+	  Dir = filename:dirname(File),
+	  case file:del_dir(Dir) of
+	    ok ->
+		?DEBUG("Removed ~s", [Dir]);
+	    {error, Error} ->
+		?INFO_MSG("Cannot remove ~s: ~s", [Dir, ?FORMAT(Error)])
+	  end;
+      {error, Error} ->
+	  ?WARNING_MSG("Cannot remove ~s: ~s", [File, ?FORMAT(Error)])
+    end.
+
+-spec secs_since_epoch() -> non_neg_integer().
+
+secs_since_epoch() ->
+    {MegaSecs, Secs, _MicroSecs} = os:timestamp(),
+    MegaSecs * 1000000 + Secs.