From f7f40cf9a6b2599a8246a2282a0e49df4c7ebeca Mon Sep 17 00:00:00 2001
From: Holger Weiss <holger@zedat.fu-berlin.de>
Date: Sun, 8 May 2016 15:36:51 +0200
Subject: [PATCH] Let client retry HTTP upload on file size mismatch

Let the main mod_http_upload process look at the size of an HTTP upload
rather than performing this check in the ejabberd_http handler.  This
way, the upload slot won't be invalidated if the size of the uploaded
file doesn't match the size requested for the slot.  The PUT request is
still rejected, but the client now has a chance to retry the upload.
---
 src/mod_http_upload.erl | 35 ++++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/src/mod_http_upload.erl b/src/mod_http_upload.erl
index b01ca9d2d..0bb383700 100644
--- a/src/mod_http_upload.erl
+++ b/src/mod_http_upload.erl
@@ -321,22 +321,24 @@ init({ServerHost, Opts}) ->
       -> {reply, {ok, pos_integer(), binary(),
 		      pos_integer() | undefined,
 		      pos_integer() | undefined}, state()} |
-	 {reply, {error, binary()}, state()} | {noreply, state()}.
+	 {reply, {error, atom()}, state()} | {noreply, state()}.
 
-handle_call({use_slot, Slot}, _From, #state{file_mode = FileMode,
-					    dir_mode = DirMode,
-					    get_url = GetPrefix,
-					    thumbnail = Thumbnail,
-					    docroot = DocRoot} = State) ->
+handle_call({use_slot, Slot, Size}, _From, #state{file_mode = FileMode,
+						  dir_mode = DirMode,
+						  get_url = GetPrefix,
+						  thumbnail = Thumbnail,
+						  docroot = DocRoot} = State) ->
     case get_slot(Slot, State) of
 	{ok, {Size, Timer}} ->
 	    timer:cancel(Timer),
 	    NewState = del_slot(Slot, State),
 	    Path = str:join([DocRoot | Slot], <<$/>>),
-	    {reply, {ok, Size, Path, FileMode, DirMode, GetPrefix, Thumbnail},
+	    {reply, {ok, Path, FileMode, DirMode, GetPrefix, Thumbnail},
 	     NewState};
+	{ok, {_WrongSize, _Timer}} ->
+	    {reply, {error, size_mismatch}, State};
 	error ->
-	    {reply, {error, <<"Invalid slot">>}, State}
+	    {reply, {error, invalid_slot}, State}
     end;
 handle_call(get_docroot, _From, #state{docroot = DocRoot} = State) ->
     {reply, {ok, DocRoot}, State};
@@ -406,9 +408,8 @@ process(LocalPath, #request{method = Method, host = Host, ip = IP})
 process(_LocalPath, #request{method = 'PUT', host = Host, ip = IP,
 			     data = Data} = Request) ->
     {Proc, Slot} = parse_http_request(Request),
-    case catch gen_server:call(Proc, {use_slot, Slot}) of
-	{ok, Size, Path, FileMode, DirMode, GetPrefix, Thumbnail}
-	    when byte_size(Data) == Size ->
+    case catch gen_server:call(Proc, {use_slot, Slot, byte_size(Data)}) of
+	{ok, Path, FileMode, DirMode, GetPrefix, Thumbnail} ->
 	    ?DEBUG("Storing file from ~s for ~s: ~s",
 		   [?ADDR_TO_STR(IP), Host, Path]),
 	    case store_file(Path, Data, FileMode, DirMode,
@@ -422,13 +423,13 @@ process(_LocalPath, #request{method = 'PUT', host = Host, ip = IP,
 			       [Path, ?ADDR_TO_STR(IP), Host, ?FORMAT(Error)]),
 		    http_response(Host, 500)
 	    end;
-	{ok, Size, Path, _FileMode, _DirMode, _GetPrefix, _Thumbnail} ->
-	    ?INFO_MSG("Rejecting file ~s from ~s for ~s: Size is ~B, not ~B",
-		      [Path, ?ADDR_TO_STR(IP), Host, byte_size(Data), Size]),
+	{error, size_mismatch} ->
+	    ?INFO_MSG("Rejecting file from ~s for ~s: Unexpected size (~B)",
+		      [?ADDR_TO_STR(IP), Host, byte_size(Data)]),
 	    http_response(Host, 413);
-	{error, Error} ->
-	    ?INFO_MSG("Rejecting file from ~s for ~s: ~p",
-		      [?ADDR_TO_STR(IP), Host, Error]),
+	{error, invalid_slot} ->
+	    ?INFO_MSG("Rejecting file from ~s for ~s: Invalid slot",
+		      [?ADDR_TO_STR(IP), Host]),
 	    http_response(Host, 403);
 	Error ->
 	    ?ERROR_MSG("Cannot handle PUT request from ~s for ~s: ~p",