fix bugs

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

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

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

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")