/src/support/z_parse_multipart.erl

  1%% @author Bob Ippolito <bob@mochimedia.com>
  2%% @copyright 2007 Mochi Media, Inc.
  3%%
  4%% @author Marc Worrell <marc@worrell.nl>
  5%% Date: 2009-05-13
  6%%
  7%% @doc Parse multipart/form-data request bodies. Uses a callback function to receive the next parts, can call 
  8%% a progress function to report back the progress on receiving the data.
  9%%
 10%% Adapted from mochiweb_multipart.erl, integrated with webmachine and zotonic
 11
 12%% This is the MIT license.
 13%% 
 14%% Copyright (c) 2007 Mochi Media, Inc.
 15%% 
 16%% Permission is hereby granted, free of charge, to any person obtaining a copy 
 17%% of this software and associated documentation files (the "Software"), to deal 
 18%% in the Software without restriction, including without limitation the rights 
 19%% to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
 20%% of the Software, and to permit persons to whom the Software is furnished to do 
 21%% so, subject to the following conditions:
 22%% 
 23%% The above copyright notice and this permission notice shall be included in all 
 24%% copies or substantial portions of the Software.
 25%% 
 26%% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, 
 27%% INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR 
 28%% PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 
 29%% LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, 
 30%% TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE 
 31%% OR OTHER DEALINGS IN THE SOFTWARE.
 32
 33
 34-module(z_parse_multipart).
 35-author("Marc Worrell <marc@worrell.nl").
 36
 37%% interface functions
 38-export([
 39   recv_parse/1,
 40   recv_parse/2,
 41   
 42   find_boundary/2
 43]).
 44
 45-include("zotonic.hrl").
 46
 47-define(CHUNKSIZE, 4096).
 48
 49-record(mp, {state, boundary, content_length, length, percentage=0, 
 50			 buffer, next_chunk, callback, progress, context}).
 51
 52
 53%% @doc Receive and parse the form data in the request body.  
 54%% The progress function should accept the parameters [Percentage, Context]
 55%% @spec recv_parse(Context) -> {form(), NewContext}
 56recv_parse(Context) ->
 57    recv_parse(fun(_Filename, _ContentType, _Size) -> ok end, Context).
 58
 59%% @spec recv_parse(UploadCheckFun, Context) -> {form(), NewContext}
 60recv_parse(UploadCheckFun, Context) ->
 61    Callback = fun(N) -> callback(N, #multipart_form{}, UploadCheckFun) end,
 62    {_LengthRemaining, _RestData, Form, ContextParsed} = parse_multipart_request(fun progress/4, Callback, Context),
 63    if Form#multipart_form.file =/= undefined ->
 64        % Premature end
 65        file:close(Form#multipart_form.file);
 66    true ->
 67        nop
 68    end,
 69    {Form, ContextParsed}.
 70
 71
 72%% @doc Report progress back to the page.
 73progress(Percentage, ContentLength, _ReceivedLength, Context) when ContentLength > ?CHUNKSIZE*5 ->
 74	case {	z_convert:to_bool(z_context:get_q("z_comet", Context)),
 75			z_context:get_q("z_pageid", Context), 
 76			z_context:get_q("z_trigger_id", Context)} of
 77		{true, PageId, TriggerId} when PageId /= undefined; TriggerId /= undefined ->
 78			ContextEnsured = z_context:ensure_all(Context),
 79			z_session_page:add_script("z_progress('"
 80						++z_utils:js_escape(TriggerId)++"',"
 81						++integer_to_list(Percentage)++");", ContextEnsured);
 82		_ -> nop
 83	end;
 84progress(_, _, _, _) ->
 85    nop.
 86	
 87
 88%% @doc Callback function collecting all data found in the multipart/form-data body
 89%% @spec callback(Next, function(), form()) -> function() | form()
 90callback(Next, Form, UploadCheckFun) ->
 91    case Next of
 92        {headers, Headers} ->
 93            % Find out if it is a file
 94            ContentDisposition = proplists:get_value("content-disposition", Headers),
 95            case ContentDisposition of
 96                {"form-data", [{"name", Name}, {"filename",Filename}]} ->
 97                    ContentLength = case proplists:get_value("content-length", Headers) of
 98                                        undefined -> undefined;
 99                                        {CL,_} -> z_convert:to_integer(CL)
100                                    end,
101                    ContentType = case proplists:get_value("content-type", Headers) of
102                                        undefined -> undefined;
103                                        {Mime,_} -> Mime
104                                  end,
105                    case UploadCheckFun(Filename, ContentType, ContentLength) of
106                        ok ->
107                            NF = Form#multipart_form{name=Name,
108                                                     filename=Filename, 
109                                                     content_length=ContentLength, 
110                                                     content_type=ContentType,
111                                                     tmpfile=z_tempfile:new()},
112                            fun(N) -> callback(N, NF, UploadCheckFun) end;
113                        {error, _Reason} = Error ->
114                            throw(Error)
115                    end;
116                {"form-data",[{"name",Name}]} ->
117                    NF = Form#multipart_form{name=Name, data=[]},
118                    fun(N) -> callback(N, NF, UploadCheckFun) end;
119                _ ->
120                    fun(N) -> callback(N, Form, UploadCheckFun) end
121            end;
122
123        {body, Data} ->
124            if  Form#multipart_form.filename =/= undefined ->
125                if Form#multipart_form.file =/= undefined ->
126                    file:write(Form#multipart_form.file, Data),
127                    NewForm = Form;
128                true ->
129                    case file:open(Form#multipart_form.tmpfile, [raw,write]) of
130                        {ok, File} ->
131                            file:write(File, Data),
132                            NewForm = Form#multipart_form{file=File};
133                        {error, Error} ->
134                            ?ERROR("Couldn't open ~p for writing, error: ~p~n", [Form#multipart_form.tmpfile, Error]),
135                            NewForm = Form,
136                            exit(could_not_open_file_for_writing)
137                    end
138                end;
139            true ->
140                NewForm = Form#multipart_form{data=[binary_to_list(Data), Form#multipart_form.data]}
141            end,
142            fun(N) -> callback(N, NewForm, UploadCheckFun) end;
143
144         body_end ->
145            NewForm = if Form#multipart_form.file =/= undefined ->
146                            file:close(Form#multipart_form.file),
147                            Form#multipart_form{
148                                name=undefined,
149                                data=undefined,
150                                file=undefined,
151                                tmpfile=undefined,
152                                filename=undefined,
153                                content_type=undefined,
154                                content_length=undefined,
155                                files=[{Form#multipart_form.name, Form#multipart_form.filename, Form#multipart_form.tmpfile}|Form#multipart_form.files]
156                            };
157                        Form#multipart_form.name =/= undefined ->
158                            Data = lists:flatten(Form#multipart_form.data),
159                            Form#multipart_form{
160                                name=undefined,
161                                data=undefined,
162                                args=[{Form#multipart_form.name, Data} | Form#multipart_form.args]
163                            };
164                        true ->
165                            Form
166                        end,
167            fun(N) -> callback(N, NewForm, UploadCheckFun) end;
168        
169        eof ->
170            Form
171    end.
172
173
174%% @doc Parse the multipart request
175parse_multipart_request(ProgressFunction, Callback, Context) ->
176    ReqData  = z_context:get_reqdata(Context),
177    Length   = list_to_integer(wrq:get_req_header_lc("content-length", ReqData)),
178    Boundary = iolist_to_binary(get_boundary(wrq:get_req_header_lc("content-type", ReqData))),
179    Prefix = <<"\r\n--", Boundary/binary>>,
180    BS = size(Boundary),
181    {{Chunk, Next}, ReqData1} = wrq:stream_req_body(ReqData, ?CHUNKSIZE),
182    Context1 = z_context:set_reqdata(ReqData1, Context),
183    <<"--", Boundary:BS/binary, "\r\n", Rest/binary>> = Chunk,
184    feed_mp(headers, #mp{boundary=Prefix,
185                         length=size(Chunk),
186                         content_length=Length,
187                         buffer=Rest,
188                         callback=Callback,
189                         progress=ProgressFunction,
190                         next_chunk=Next,
191                         context=Context1}).
192
193
194feed_mp(headers, State=#mp{buffer=Buffer, callback=Callback}) ->
195    {State1, P} = case find_in_binary(<<"\r\n\r\n">>, Buffer) of
196        {exact, N} ->
197            {State, N};
198        _ ->
199           S1 = read_more(State),
200           %% Assume headers must be less than ?CHUNKSIZE
201           {exact, N} = find_in_binary(<<"\r\n\r\n">>, S1#mp.buffer),
202           {S1, N}
203    end,
204    <<Headers:P/binary, "\r\n\r\n", Rest/binary>> = State1#mp.buffer,
205    NextCallback = Callback({headers, parse_headers(Headers)}),
206    feed_mp(body, State1#mp{buffer=Rest, callback=NextCallback});
207
208feed_mp(body, State=#mp{boundary=Prefix, buffer=Buffer, callback=Callback}) ->
209    case find_boundary(Prefix, Buffer) of
210        {end_boundary, Start, Skip} ->
211             <<Data:Start/binary, _:Skip/binary, Rest/binary>> = Buffer,
212             C1 = Callback({body, Data}),
213             C2 = C1(body_end),
214             {State#mp.length, Rest, C2(eof), State#mp.context};
215        {next_boundary, Start, Skip} ->
216             <<Data:Start/binary, _:Skip/binary, Rest/binary>> = Buffer,
217             C1 = Callback({body, Data}),
218             feed_mp(headers, State#mp{callback=C1(body_end), buffer=Rest});
219        {maybe, 0} ->
220            % Found a boundary, without an ending newline
221            case read_more(State) of
222                State -> throw({error, incomplete_end_boundary});
223                S1 -> feed_mp(body, S1)
224            end;
225        {maybe, Start} ->
226            <<Data:Start/binary, Rest/binary>> = Buffer,
227            feed_mp(body, read_more(State#mp{callback=Callback({body, Data}), buffer=Rest}));
228        not_found ->
229            {Data, Rest} = {Buffer, <<>>},
230            feed_mp(body, read_more(State#mp{callback=Callback({body, Data}), buffer=Rest}))
231    end.
232
233
234
235%% @doc Read more data for the feed_mp functions.
236%% @spec read_more(mp()) -> mp()
237read_more(State=#mp{next_chunk=done, content_length=ContentLength, length=Length} = State) when ContentLength =:= Length ->
238    State;
239read_more(State=#mp{next_chunk=done} = State) ->
240    throw({error, wrong_content_length});
241read_more(State=#mp{length=Length, content_length=ContentLength, 
242				percentage=Percentage,
243				buffer=Buffer, next_chunk=Next, context=Context,
244				progress=ProgressFunction}) ->
245    {Data, Next1} = Next(),
246    Buffer1 = <<Buffer/binary, Data/binary>>,
247	Length1 = Length + size(Data),
248	NewPercentage = case ContentLength of
249						0 -> 100;
250						_ -> (Length1 * 100) div ContentLength
251					end,
252	case NewPercentage > Percentage of
253		true ->
254			case ProgressFunction of
255				undefined -> nop;
256				F -> F(NewPercentage, ContentLength, Length1, Context)
257			end;
258		_ ->
259			nop
260	end,
261    State#mp{length=Length1, buffer=Buffer1, next_chunk=Next1, percentage=NewPercentage}.
262
263
264%% @doc Parse the headers of a part in the form data
265parse_headers(<<>>) ->
266    [];
267parse_headers(Binary) ->
268    parse_headers(Binary, []).
269
270parse_headers(Binary, Acc) ->
271    case find_in_binary(<<"\r\n">>, Binary) of
272        {exact, N} ->
273            <<Line:N/binary, "\r\n", Rest/binary>> = Binary,
274            parse_headers(Rest, [split_header(Line) | Acc]);
275        not_found ->
276            lists:reverse([split_header(Binary) | Acc])
277    end.
278
279split_header(Line) ->
280    {Name, [$: | Value]} = lists:splitwith(fun (C) -> C =/= $: end, binary_to_list(Line)),
281    {string:to_lower(string:strip(Name)), mochiweb_util:parse_header(Value)}.
282
283
284%% @doc Get the request boundary separating the parts in the request body
285get_boundary(ContentType) ->
286    {"multipart/form-data", Opts} = mochiweb_util:parse_header(ContentType),
287    case proplists:get_value("boundary", Opts) of
288        S when is_list(S) ->
289            S
290    end.
291
292%% @doc Find the next boundary in the data
293find_boundary(Prefix, Data) ->
294    case find_in_binary(Prefix, Data) of
295        {exact, Skip} ->
296            PrefixSkip = Skip + size(Prefix),
297            case Data of
298                <<_:PrefixSkip/binary, "\r\n", _/binary>> ->
299                    {next_boundary, Skip, size(Prefix) + 2};
300                <<_:PrefixSkip/binary, "--\r\n", _/binary>> ->
301                    {end_boundary, Skip, size(Prefix) + 4};
302                % POSTs by Adobe Flash don't have the ending newline
303                <<_:PrefixSkip/binary, "--", _/binary>> ->
304                    {end_boundary, Skip, size(Prefix) + 2};
305                _ when size(Data) < PrefixSkip + 4 ->
306                    %% Underflow
307                    {maybe, Skip};
308                _ ->
309                    %% False positive
310                    not_found
311            end;
312        {partial, Skip, Length} when (Skip + Length) =:= size(Data) ->
313            %% Underflow
314            {maybe, Skip};
315        _ ->
316            not_found
317    end.
318
319
320
321find_in_binary(B, Data) when size(B) > 0 ->
322    case size(Data) - size(B) of
323        Last when Last < 0 ->
324            partial_find(B, Data, 0, size(Data));
325        Last ->
326            find_in_binary(B, size(B), Data, 0, Last)
327    end.
328
329find_in_binary(B, BS, D, N, Last) when N =< Last->
330    case D of
331        <<_:N/binary, B:BS/binary, _/binary>> ->
332            {exact, N};
333        _ ->
334            find_in_binary(B, BS, D, 1 + N, Last)
335    end;
336find_in_binary(B, BS, D, N, Last) when N =:= 1 + Last ->
337    partial_find(B, D, N, BS - 1).
338
339partial_find(_B, _D, _N, 0) ->
340    not_found;
341partial_find(B, D, N, K) ->
342    <<B1:K/binary, _/binary>> = B,
343    case D of
344        <<_Skip:N/binary, B1:K/binary>> ->
345            {partial, N, K};
346        _ ->
347            partial_find(B, D, 1 + N, K - 1)
348    end.
349