From 950f0e734f84f25bb954ec07c083ba608ef0052b Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Mon, 26 Feb 2024 22:50:30 -0500 Subject: [PATCH] fix bugs --- src/core/IO.ml | 49 ++++++++++++++++++++------------------------ src/core/request.ml | 30 +++++++++++++++------------ src/core/request.mli | 3 ++- src/unix/dir.ml | 5 +++-- 4 files changed, 44 insertions(+), 43 deletions(-) diff --git a/src/core/IO.ml b/src/core/IO.ml index 525b0472..249da955 100644 --- a/src/core/IO.ml +++ b/src/core/IO.ml @@ -235,7 +235,7 @@ module Input = struct Buf.clear buf; let continue = ref true in while !continue do - let slice = self#fill_buf () in + let slice = fill_buf self in if slice.len = 0 then continue := false else ( @@ -246,7 +246,7 @@ module Input = struct done; Buf.contents_and_clear buf - (** put [n] bytes from the input into bytes *) + (** Read [n] bytes from the input into [bytes]. *) let read_exactly_ ~too_short (self : #t) (bytes : bytes) (n : int) : unit = assert (Bytes.length bytes >= n); let offset = ref 0 in @@ -295,49 +295,44 @@ module Input = struct | () -> Some (Buf.contents_and_clear buf) | exception End_of_file -> None - let reading_exactly_ ~skip_on_close ~close_rec ~size (arg : t) : t = + (* helper for making a new input stream that either contains at most [size] + bytes, or contains exactly [size] bytes. *) + let reading_exactly_ ~skip_on_close ~close_rec ~size ~bytes (arg : t) : t = let remaining_size = ref size in object + inherit t_from_refill ~bytes () + method close () = if !remaining_size > 0 && skip_on_close then skip arg !remaining_size; if close_rec then close arg - method fill_buf () = - if !remaining_size > 0 then - fill_buf arg - else - Slice.empty - - method input bs i len = + method private refill (slice : Slice.t) = + slice.off <- 0; + slice.len <- 0; if !remaining_size > 0 then ( - let slice = fill_buf arg in - let n = min len (min slice.len !remaining_size) in - Bytes.blit slice.bytes slice.off bs i n; + let sub = fill_buf arg in + let n = + min !remaining_size (min sub.len (Bytes.length slice.bytes)) + in + Bytes.blit sub.bytes sub.off slice.bytes 0 n; + Slice.consume sub n; remaining_size := !remaining_size - n; - Slice.consume slice n; - n - ) else - 0 - - method consume n = - if n > !remaining_size then - invalid_arg "reading_exactly: consuming too much"; - remaining_size := !remaining_size - n; - consume arg n + slice.len <- n + ) end (** new stream with maximum size [max_size]. @param close_rec if true, closing this will also close the input stream *) - let limit_size_to ~close_rec ~max_size (arg : t) : t = - reading_exactly_ ~size:max_size ~skip_on_close:false ~close_rec arg + let limit_size_to ~close_rec ~max_size ~bytes (arg : t) : t = + reading_exactly_ ~size:max_size ~skip_on_close:false ~bytes ~close_rec arg (** New stream that consumes exactly [size] bytes from the input. If fewer bytes are read before [close] is called, we read and discard the remaining quota of bytes before [close] returns. @param close_rec if true, closing this will also close the input stream *) - let reading_exactly ~close_rec ~size (arg : t) : t = - reading_exactly_ ~size ~close_rec ~skip_on_close:true arg + let reading_exactly ~close_rec ~size ~bytes (arg : t) : t = + reading_exactly_ ~size ~close_rec ~skip_on_close:true ~bytes arg let read_chunked ~(bytes : bytes) ~fail (ic : #t) : t = let first = ref true in diff --git a/src/core/request.ml b/src/core/request.ml index c5404d85..de982367 100644 --- a/src/core/request.ml +++ b/src/core/request.ml @@ -71,17 +71,17 @@ let read_stream_chunked_ ~bytes (bs : #IO.Input.t) : IO.Input.t = Log.debug (fun k -> k "body: start reading chunked stream..."); IO.Input.read_chunked ~bytes ~fail:(fun s -> Bad_req (400, s)) bs -let limit_body_size_ ~max_size (bs : #IO.Input.t) : IO.Input.t = +let limit_body_size_ ~max_size ~bytes (bs : #IO.Input.t) : IO.Input.t = Log.debug (fun k -> k "limit size of body to max-size=%d" max_size); - IO.Input.limit_size_to ~max_size ~close_rec:false bs + IO.Input.limit_size_to ~max_size ~close_rec:false ~bytes bs -let limit_body_size ~max_size (req : IO.Input.t t) : IO.Input.t t = - { req with body = limit_body_size_ ~max_size req.body } +let limit_body_size ~max_size ~bytes (req : IO.Input.t t) : IO.Input.t t = + { req with body = limit_body_size_ ~max_size ~bytes req.body } (** read exactly [size] bytes from the stream *) -let read_exactly ~size (bs : #IO.Input.t) : IO.Input.t = +let read_exactly ~size ~bytes (bs : #IO.Input.t) : IO.Input.t = Log.debug (fun k -> k "body: must read exactly %d bytes" size); - IO.Input.reading_exactly bs ~close_rec:false ~size + IO.Input.reading_exactly bs ~close_rec:false ~bytes ~size (* parse request, but not body (yet) *) let parse_req_start ~client_addr ~get_time_s ~buf (bs : IO.Input.t) : @@ -151,23 +151,27 @@ let parse_req_start ~client_addr ~get_time_s ~buf (bs : IO.Input.t) : let parse_body_ ~tr_stream ~bytes (req : IO.Input.t t) : IO.Input.t t resp_result = try - let size = + let size, has_size = match Headers.get_exn "Content-Length" req.headers |> int_of_string with - | n -> n (* body of fixed size *) - | exception Not_found -> 0 + | n -> n, true (* body of fixed size *) + | exception Not_found -> 0, false | exception _ -> bad_reqf 400 "invalid content-length" in let body = match get_header ~f:String.trim req "Transfer-Encoding" with - | None -> read_exactly ~size @@ tr_stream req.body + | None -> read_exactly ~size ~bytes @@ tr_stream req.body + | Some "chunked" when has_size -> + bad_reqf 400 "specifying both transfer-encoding and content-length" | Some "chunked" -> (* body sent by chunks *) let bs : IO.Input.t = read_stream_chunked_ ~bytes @@ tr_stream req.body in - if size > 0 then - limit_body_size_ ~max_size:size bs - else + if size > 0 then ( + (* TODO: ensure we recycle [bytes] when the new input is closed *) + let bytes = Bytes.create 4096 in + limit_body_size_ ~max_size:size ~bytes bs + ) else bs | Some s -> bad_reqf 500 "cannot handle transfer encoding: %s" s in diff --git a/src/core/request.mli b/src/core/request.mli index 25981a66..70ae741d 100644 --- a/src/core/request.mli +++ b/src/core/request.mli @@ -107,7 +107,8 @@ val start_time : _ t -> float (** time stamp (from {!Unix.gettimeofday}) after parsing the first line of the request @since 0.11 *) -val limit_body_size : max_size:int -> IO.Input.t t -> IO.Input.t t +val limit_body_size : + max_size:int -> bytes:bytes -> IO.Input.t t -> IO.Input.t t (** Limit the body size to [max_size] bytes, or return a [413] error. @since 0.3 diff --git a/src/unix/dir.ml b/src/unix/dir.ml index d596a78f..f29e9db5 100644 --- a/src/unix/dir.ml +++ b/src/unix/dir.ml @@ -253,9 +253,10 @@ let add_vfs_ ~on_fs ~top ~config ~vfs:((module VFS : VFS) as vfs) ~prefix server (Printexc.to_string e) in let req = - Request.limit_body_size ~max_size:config.max_upload_size req + Request.limit_body_size ~bytes:(Bytes.create 4096) + ~max_size:config.max_upload_size req in - IO.Input.iter write req.Request.body; + IO.Input.iter write req.body; close (); Log.debug (fun k -> k "dir: done uploading file to %S" path); Response.make_raw ~code:201 "upload successful")