Merge pull request #32 from mirage/with-read-line-limit

#25 rebased

src/mirage_channel.ml

@@ -28,7 +28,7 @@ module type S = sig
   val read_char: t -> (char Mirage_flow.or_eof, error) result Lwt.t
   val read_some: ?len:int -> t -> (Cstruct.t Mirage_flow.or_eof, error) result Lwt.t
   val read_exactly: len:int -> t -> (Cstruct.t list Mirage_flow.or_eof, error) result Lwt.t
-  val read_line: t -> (Cstruct.t list Mirage_flow.or_eof, error) result Lwt.t
+  val read_line: ?len:int -> t -> (Cstruct.t list Mirage_flow.or_eof, error) result Lwt.t
   val write_char: t -> char -> unit
   val write_string: t -> string -> int -> int -> unit
   val write_buffer: t -> Cstruct.t -> unit
@@ -47,14 +47,17 @@ module Make(Flow: Mirage_flow.S) = struct
 
   type flow = Flow.flow
 
-  type error = [`Read_zero | `Flow of Flow.error]
+  type error = [`Line_too_long|`Read_zero | `Flow of Flow.error]
   type write_error = Flow.write_error
 
   let pp_error ppf = function
   | `Flow e    -> Flow.pp_error ppf e
   | `Read_zero ->
       Fmt.string ppf
         "FLOW.read returned 0 bytes in violation of the specification"
+  | `Line_too_long ->
+      Fmt.string ppf
+        "Unable to read a line because it is too long"
 
   let pp_write_error = Flow.pp_write_error
 
@@ -137,42 +140,51 @@ module Make(Flow: Mirage_flow.S) = struct
     loop [] len
 
   (* Read until a character is found *)
-  let read_until t ch =
+  let read_until ?len t ch =
     get_ibuf t >>=~ fun buf ->
-    let len = Cstruct.len buf in
+    (* Scan up to the length of the buffer or the supplied limit, whichever
+       is smaller. *)
+    let scan_len = 
+      let len' = Cstruct.len buf in
+      match len with None -> len' | Some x -> min x len' in
     let rec scan off =
-      if off = len then None
+      if off = scan_len then None
       else if Cstruct.get_char buf off = ch then Some off else scan (off+1)
     in
     match scan 0 with
     | None -> (* not found, return what we have until EOF *)
-       t.ibuf <- None; (* basically guaranteeing that next read is EOF *)
-       Lwt.return (Ok (`Not_found buf))
+      t.ibuf <- Some (Cstruct.shift buf scan_len);
+      Lwt.return (Ok (`Not_found (Cstruct.sub buf 0 scan_len)))
     | Some off -> (* found, so split the buffer *)
-       let hd = Cstruct.sub buf 0 off in
-       t.ibuf <- Some (Cstruct.shift buf (off+1));
-       Lwt.return (Ok (`Found hd))
+      let hd = Cstruct.sub buf 0 off in
+      t.ibuf <- Some (Cstruct.shift buf (off+1));
+      Lwt.return (Ok (`Found hd))
 
   (* This reads a line of input, which is terminated either by a CRLF
      sequence, or the end of the channel (which counts as a line).
      @return Returns a stream of views that terminates at EOF. *)
-  let read_line t =
-    let rec get acc =
-      read_until t '\n' >>= function
-      | Error e -> Lwt.return (Error e)
-      | Ok `Eof -> Lwt.return (Ok (`Data acc))
-      | Ok (`Not_found buf) when Cstruct.len buf = 0 -> Lwt.return (Ok (`Data acc))
-      | Ok (`Not_found buf) -> get (buf::acc)
-      | Ok (`Found buf) ->
-          (* chop the CR if present *)
-          let buflen = Cstruct.len buf in
-          let buf =
-            if buflen > 0 && (Cstruct.get_char buf (buflen-1) = '\r') then
-              Cstruct.sub buf 0 (buflen-1) else buf
-          in
-          Lwt.return (Ok (`Data (buf :: acc)))
+  let read_line ?len t =
+    let rec get ?len acc =
+      match len with
+      | Some 0 -> Lwt.return (Error `Line_too_long)
+      | _ ->
+        read_until ?len t '\n' >>= function
+        | Error e -> Lwt.return (Error e)
+        | Ok `Eof -> Lwt.return (Ok (`Data acc))
+        | Ok (`Not_found buf) when Cstruct.len buf = 0 -> Lwt.return (Ok (`Data acc))
+        | Ok (`Not_found buf) ->
+          let len = match len with None -> None | Some l -> Some (l - (Cstruct.len buf)) in
+          get ?len (buf::acc)
+        | Ok (`Found buf) ->
+            (* chop the CR if present *)
+            let buflen = Cstruct.len buf in
+            let buf =
+              if buflen > 0 && (Cstruct.get_char buf (buflen-1) = '\r') then
+                Cstruct.sub buf 0 (buflen-1) else buf
+            in
+            Lwt.return (Ok (`Data (buf :: acc)))
     in
-    get [] >>=~ fun bits -> Lwt.return (Ok (`Data (List.rev bits)))
+    get ?len [] >>=~ fun bits -> Lwt.return (Ok (`Data (List.rev bits)))
 
   (* Output functions *)
 

src/mirage_channel.mli

@@ -65,10 +65,19 @@ module type S = sig
   (** [read_exactly len t] reads [len] bytes from the channel [t] or fails
       with [Eof]. *)
 
-  val read_line: t -> (Cstruct.t list Mirage_flow.or_eof, error) result Lwt.t
+  val read_line: ?len:int -> t -> (Cstruct.t list Mirage_flow.or_eof, error) result Lwt.t
   (** [read_line t] reads a line of input, which is terminated
       either by a CRLF sequence, or the end of the channel (which
       counts as a line).
+
+      If [?len] is provided then the maximum length of the line returned will be
+      [len] bytes. If the line is longer than [len] then an error will be
+      returned. With [len = 0], [read_line] always returns an error.
+
+      If the input data is untrusted then care should be taken to ensure [len]
+      is set to an application-specific small value to bound the amount of
+      memory allocated by [read_line].
+
       @return Returns a list of views that terminates at EOF. *)
 
   val write_char: t -> char -> unit
@@ -100,4 +109,6 @@ module type S = sig
 end
 
 (** Functor to create a CHANNEL from a flow implementation *)
-module Make(F: Mirage_flow.S): S with type flow = F.flow
+module Make(F: Mirage_flow.S)
+  : S with type flow = F.flow
+       and type error = private [> `Read_zero | `Flow of F.error | `Line_too_long ]

test/test_channel.ml

@@ -33,6 +33,109 @@ let test_read_line () =
   | Ok `Eof -> fail "eof"
   | Error e -> fail "error: %a" Channel.pp_error e
 
+(* The line is longer than the limit *)
+let test_read_line_len () =
+  let input = "I am the very model of a modern major general" in
+  let f = F.make ~input:(F.input_string input) () in
+  let c = Channel.create f in
+  Channel.read_line ~len:5 c >|= function
+  | Ok (`Data _) -> fail "read a line which was too big"
+  | Ok `Eof -> fail "eof"
+  | Error _ -> ()
+
+(* The line is shorter than the limit and bounded by \r\n *)
+let test_read_line_len2 () =
+  let input = "I\r\n am the very model of a modern major general" in
+  let f = F.make ~input:(F.input_string input) () in
+  let c = Channel.create f in
+  Channel.read_line ~len:5 c >|= function
+  | Ok (`Data buf) -> Alcotest.(check string) "read line" "I" (Cstruct.copyv buf)
+  | Ok `Eof -> fail "eof"
+  | Error e -> fail "error: %a" Channel.pp_error e
+
+(* The line is shorter than the limit and bounded by EOF *)
+let test_read_line_len3 () =
+  let input = "I am the very model of a modern major general" in
+  let f = F.make ~input:(F.input_string input) () in
+  let c = Channel.create f in
+  Channel.read_line ~len:50 c >|= function
+  | Ok (`Data buf) -> Alcotest.(check string) "read line" input (Cstruct.copyv buf)
+  | Ok `Eof -> fail "eof"
+  | Error e -> fail "error: %a" Channel.pp_error e
+
+type channel = V : (module Mirage_channel.S with type t = 'a and type error = [> `Line_too_long ]) * 'a -> channel
+
+let channel_from_raw_string s =
+  let consumed = ref false in
+  let module Flow = struct
+    type flow = unit
+    type error = |
+    type write_error = Mirage_flow.write_error
+
+    let pp_error : error Fmt.t = fun _ -> function _ -> .
+    let pp_write_error : Mirage_flow.write_error Fmt.t =
+      fun ppf `Closed -> Fmt.string ppf "Flow closed"
+
+    let read () =
+      if not !consumed
+      then ( consumed := true
+           ; Lwt.return_ok (`Data (Cstruct.of_string s)) )
+      else Lwt.return_ok `Eof
+    let write _ _ = assert false
+    let writev _ _ = assert false
+    let close _ = Lwt.return ()
+  end in
+  let module Channel = Mirage_channel.Make(Flow) in
+  V ((module Channel), Channel.create ())
+
+let test_read_line_len4 () =
+  let V ((module Channel), c) = channel_from_raw_string "foo" in
+  Channel.read_line ~len:3 c >|= function
+  | Ok (`Data bufs) -> fail "Unexpected data: %S"
+                         Cstruct.(to_string (concat bufs))
+  | Ok `Eof -> fail "eof"
+  | Error e -> match e with
+    | `Line_too_long -> ()
+    | e -> fail "Unexpected error: %a" Channel.pp_error e
+
+let test_read_line_len5 () =
+  let V ((module Channel), c) = channel_from_raw_string "foo\r" in
+  Channel.read_line ~len:3 c >|= function
+  | Ok (`Data bufs) -> fail "Unexpected data: %S"
+                         Cstruct.(to_string (concat bufs))
+  | Ok `Eof -> fail "eof"
+  | Error e -> match e with
+    | `Line_too_long -> ()
+    | e -> fail "Unexpected error: %a" Channel.pp_error e
+
+let test_read_line_len6 () =
+  let V ((module Channel), c) = channel_from_raw_string "foo\r\n" in
+  Channel.read_line ~len:3 c >|= function
+  | Ok (`Data bufs) -> fail "Unexpected data: %S"
+                         Cstruct.(to_string (concat bufs))
+  | Ok `Eof -> fail "eof"
+  | Error e -> match e with
+    | `Line_too_long -> ()
+    | e -> fail "Unexpected error: %a" Channel.pp_error e
+
+let test_read_line_len7 () =
+  let V ((module Channel), c) = channel_from_raw_string "foo\r\n" in
+  Channel.read_line ~len:4 c >|= function
+  | Ok (`Data bufs) -> fail "Unexpected data: %S"
+                         Cstruct.(to_string (concat bufs))
+  | Ok `Eof -> fail "eof"
+  | Error e -> match e with
+    | `Line_too_long -> ()
+    | e -> fail "Unexpected error: %a" Channel.pp_error e
+
+let test_read_line_len8 () =
+  let V ((module Channel), c) = channel_from_raw_string "foo\r\n" in
+  Channel.read_line ~len:5 c >|= function
+  | Ok (`Data bufs) ->
+    Alcotest.(check string) "read line" "foo" Cstruct.(to_string (concat bufs))
+  | Ok `Eof -> fail "eof"
+  | Error e -> fail "Unexpected error: %a" Channel.pp_error e
+
 let test_read_exactly () =
   let input = "I am the very model of a modern major general" in
   let f = F.make ~input:(F.input_string input) () in
@@ -69,4 +172,12 @@ let suite = [
   "read_line"       , `Quick, test_read_line;
   "read_exactly"    , `Quick, test_read_exactly;
   "write after read EOF", `Quick, test_read_until_eof_then_write;
+  "read_line_len"   , `Quick, test_read_line_len;
+  "read_line_len2"  , `Quick, test_read_line_len2;
+  "read_line_len3"  , `Quick, test_read_line_len3;
+  "read_line_len4"  , `Quick, test_read_line_len4;
+  "read_line_len5"  , `Quick, test_read_line_len5;
+  "read_line_len6"  , `Quick, test_read_line_len6;
+  "read_line_len7"  , `Quick, test_read_line_len7;
+  "read_line_len8"  , `Quick, test_read_line_len8;
 ]