feat: require \r\n before all boundaries but the first

c-cube · Dec 2, 2024 · 21c0f7f · 21c0f7f
1 parent 099777b
commit 21c0f7f
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 8 deletions.
diff --git a/src/multipart_form/tiny_httpd_multipart_form_data.ml b/src/multipart_form/tiny_httpd_multipart_form_data.ml
@@ -27,6 +27,7 @@ type st = {
   boundary: string;
   ic: Iostream.In.t;
   buf: buf;  (** Used to split on the boundary *)
+  mutable first: bool;  (** Are we parsing the first boundary? *)
   mutable eof_split: bool;
   buf_out: buf;  (** Used to return output slices *)
   mutable st_out: out_state;
@@ -37,6 +38,7 @@ let create ?(buf_size = 64 * 1024) ?(out_buf_size = 8 * 1024) ~boundary ic : st
   let ic = (ic : #Iostream.In.t :> Iostream.In.t) in
   {
     boundary;
+    first = true;
     ic;
     buf = { bs = Bytes.create buf_size; len = 0 };
     eof_split = false;
@@ -46,7 +48,14 @@ let create ?(buf_size = 64 * 1024) ?(out_buf_size = 8 * 1024) ~boundary ic : st
 
 type chunk = Delim | Eof | Read of int
 
-let[@inline] min_len_ (self : st) : int = 2 + String.length self.boundary
+let[@inline] prefix_size_ (self : st) : int =
+  if self.first then
+    2
+  else
+    4
+
+let[@inline] min_len_ (self : st) : int =
+  prefix_size_ self + String.length self.boundary
 
 exception Found_boundary of int
 
@@ -74,15 +83,27 @@ let rec read_chunk_ (self : st) buf i_buf len : chunk =
   ) else (
     try
       let i = ref 0 in
-      let end_pos = min len self.buf.len - 2 - String.length self.boundary in
+      let end_pos =
+        min len self.buf.len - prefix_size_ self - String.length self.boundary
+      in
       while !i <= end_pos do
         if
-          Bytes.unsafe_get self.buf.bs !i = '-'
+          self.first
+          && Bytes.unsafe_get self.buf.bs !i = '-'
           && Bytes.unsafe_get self.buf.bs (!i + 1) = '-'
           && Utils_.string_eq
                ~a:(Bytes.unsafe_to_string self.buf.bs)
                ~a_start:(!i + 2) ~b:self.boundary
                ~len:(String.length self.boundary)
+          || (not self.first)
+             && Bytes.unsafe_get self.buf.bs !i = '\r'
+             && Bytes.unsafe_get self.buf.bs (!i + 1) = '\n'
+             && Bytes.unsafe_get self.buf.bs (!i + 2) = '-'
+             && Bytes.unsafe_get self.buf.bs (!i + 3) = '-'
+             && Utils_.string_eq
+                  ~a:(Bytes.unsafe_to_string self.buf.bs)
+                  ~a_start:(!i + 4) ~b:self.boundary
+                  ~len:(String.length self.boundary)
         then
           raise_notrace (Found_boundary !i);
         incr i
@@ -93,7 +114,8 @@ let rec read_chunk_ (self : st) buf i_buf len : chunk =
       Read n_read
     with
     | Found_boundary 0 ->
-      shift_left_ self.buf (2 + String.length self.boundary);
+      shift_left_ self.buf (prefix_size_ self + String.length self.boundary);
+      self.first <- false;
       Delim
     | Found_boundary n ->
       let n_read = min n len in
@@ -189,6 +211,7 @@ and parse_headers_rec (self : st) acc : Headers.t =
       )
     | i ->
       let line = Bytes.sub_string self.buf_out.bs 0 i in
+      Printf.eprintf "parse header line %S\n%!" line;
       shift_left_ self.buf_out (i + 2);
       if line = "" then
         List.rev acc

diff --git a/tests/multipart_form/t_chunk.ml b/tests/multipart_form/t_chunk.ml
@@ -40,11 +40,14 @@ let () =
   test
     "hello--YOLO\n\
     \  world\n\
-    \    what is the meaning of--YOLOthis??--YOLOok ok ok--YOLO";
+    \    what is the meaning of\r\n\
+     --YOLOthis??\r\n\
+     --YOLOok ok ok\r\n\
+     --YOLO";
   pf "T2\n";
-  test "--YOLO--YOLOah bon--YOLOaight--YOLO--YOLO";
+  test "--YOLO\r\n--YOLOah bon\r\n--YOLOaight\r\n--YOLO\r\n--YOLO";
   pf "T3\n";
   test
-    (spf "--YOLO%s--YOLO--YOLO%s--YOLO%s" (String.make 400 'a')
+    (spf "--YOLO%s\r\n--YOLO\r\n--YOLO%s\r\n--YOLO%s" (String.make 400 'a')
        (String.make 512 'b') (String.make 400 'c'));
   ()
diff --git a/tests/multipart_form/t_parse.ml b/tests/multipart_form/t_parse.ml
@@ -46,10 +46,12 @@ let () =
      ohlook: here\r\n\
      \r\n\
      and now for the b-o-d-y 👏\n\
+     \r\n\
      --YOLO\r\n\
      more: headers\r\n\
      \r\n\
      and another body\r\n\
+     \r\n\
      --YOLO--";
   pf "T1\n";
   test
@@ -60,8 +62,11 @@ let () =
         \r\n\
         and now for the bigger body:\n\
         %s\n\
+        \r\n\
         --YOLO\r\n\
         more: headers\r\n\
         \r\n\
-        and another body--YOLO--" (String.make 500 'a'));
+        and another body\r\n\
+        --YOLO--"
+       (String.make 500 'a'));
   ()