more tests and fixes

hannesm · hannesm · commit 52f2068c86f8 · 2025-05-29T15:45:12.000+02:00
diff --git a/src/state.ml b/src/state.ml
@@ -101,20 +101,17 @@ module Reassembly_queue = struct
             (* 2 - the current "e" may be merged into the head of acc *)
             let acc' = match acc with [] -> [] | _hd :: tl -> tl in
             if Sequence.less_equal e.seq seq_end then
-              let to_cut = Sequence.sub seq_end e.seq in
-              if to_cut = 0 then
+              let overlap = Sequence.sub seq_end e.seq in
+              if overlap = 0 then
                 (* to_cut = 0, we can just merge them *)
                 let elt = { elt with fin = e.fin || elt.fin ; data = e.data @ elt.data } in
                 Some (elt, Sequence.addi elt.seq (Cstruct.lenv elt.data)), elt :: acc'
               else
-                (* we need to cut some bytes from the current hd *)
-                match elt.data with
-                | head :: tl ->
-                  let hd = Cstruct.sub head 0 (Cstruct.length head - to_cut) in
-                  let data = e.data @ hd :: tl in
-                  let elt = { elt with fin = e.fin || elt.fin ; data } in
-                  Some (elt, Sequence.addi elt.seq (Cstruct.lenv data)), elt :: acc'
-                | [] -> (inserted, e :: acc)
+                (* we need to cut some bytes from e *)
+                let data = List.rev (Cstruct.shiftv (List.rev e.data) overlap) in
+                let data = data @ elt.data in
+                let elt = { elt with fin = e.fin || elt.fin ; data } in
+                Some (elt, Sequence.addi elt.seq (Cstruct.lenv data)), elt :: acc'
             else
               (* there's still a hole, nothing to merge *)
               (inserted, e :: acc)
diff --git a/test/reassembly.ml b/test/reassembly.ml
@@ -245,6 +245,42 @@ let overlap_3 () =
   | Some (s, _) -> Alcotest.(check string "data is good" "CCCCCCCCCCBBBBBBAAAAAAAA"
                                (Cstruct.to_string s))
 
+let overlap_4 () =
+  let r = insert_seg empty (Sequence.of_int32 16l, false, Cstruct.of_string "AAAAAAAA") in
+  let r = insert_seg r (Sequence.of_int32 8l, false, Cstruct.of_string "BBBBBB") in
+  let r = insert_seg r (Sequence.of_int32 6l, false, Cstruct.of_string "CCCCCCCCCCCCCCCC") in
+  let r = insert_seg r (Sequence.of_int32 0l, false, Cstruct.of_string "DDDDDDDD") in
+  let r', s = maybe_take r Sequence.zero in
+  Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
+  match s with
+  | None -> Alcotest.fail "should be some data"
+  | Some (s, _) -> Alcotest.(check string "data is good" "DDDDDDDDCCCCCCCCCCCCCCAA"
+                               (Cstruct.to_string s))
+
+let overlap_5 () =
+  let r = insert_seg empty (Sequence.of_int32 0l, false, Cstruct.of_string "AAAAAAAAAAAAAAAA") in
+  let r = insert_seg r (Sequence.of_int32 4l, false, Cstruct.of_string "BBBBBB") in
+  let r = insert_seg r (Sequence.of_int32 10l, false, Cstruct.of_string "CCCCCC") in
+  let r = insert_seg r (Sequence.of_int32 12l, false, Cstruct.of_string "DDDDDD") in
+  let r = insert_seg r (Sequence.of_int32 14l, false, Cstruct.of_string "EE") in
+  let r', s = maybe_take r Sequence.zero in
+  Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
+  match s with
+  | None -> Alcotest.fail "should be some data"
+  | Some (s, _) -> Alcotest.(check string "data is good" "AAAABBBBBBCCDDEEDD"
+                               (Cstruct.to_string s))
+
+let overlap_6 () =
+  let r = insert_seg empty (Sequence.of_int32 0l, false, Cstruct.of_string "AAAAAAAA") in
+  let r = insert_seg r (Sequence.of_int32 10l, false, Cstruct.of_string "BBBBBB") in
+  let r = insert_seg r (Sequence.of_int32 8l, false, Cstruct.of_string "DDDD") in
+  let r', s = maybe_take r Sequence.zero in
+  Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
+  match s with
+  | None -> Alcotest.fail "should be some data"
+  | Some (s, _) -> Alcotest.(check string "data is good" "AAAAAAAADDDDBBBB"
+                               (Cstruct.to_string s))
+
 let regr_187 () =
   let r = insert_seg empty (Sequence.of_int32 16l, false, Cstruct.of_string "000002om") in
   let r = insert_seg r (Sequence.of_int32 8l, false, Cstruct.of_string "001001nn001002nm001003nl") in
@@ -286,6 +322,9 @@ let tests = [
   "overlap 1", `Quick, overlap_1 ;
   "overlap 2", `Quick, overlap_2 ;
   "overlap 3", `Quick, overlap_3 ;
+  "overlap 4", `Quick, overlap_4 ;
+  "overlap 5", `Quick, overlap_5 ;
+  "overlap 6", `Quick, overlap_6 ;
   "regression 187", `Quick, regr_187 ;
   "regression 188", `Quick, regr_188 ;
 ]
