Skip to content

Commit bc68c0b

Browse files
hannesmhannesm
committed
rework reassembly
1 parent cad4ac5 commit bc68c0b

File tree

2 files changed

+115
-19
lines changed

2 files changed

+115
-19
lines changed

‎src/state.ml‎

Lines changed: 55 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ module Reassembly_queue = struct
9191

9292
(* insert segment, potentially coalescing existing ones *)
9393
let insert_seg t (seq, fin, data) =
94-
(* they may overlap, the oldest seg wins *)
94+
(* they may overlap, the newest seg wins *)
9595
(* (1) figure out the place whereafter to insert the seg *)
9696
(* (2) peek whether the next seg can be already coalesced *)
9797
let inserted, segq =
@@ -123,39 +123,75 @@ module Reassembly_queue = struct
123123
(* there are three cases:
124124
- (a) the new seq is before the existing e.seq -> prepend
125125
(and figure out whether to merge with e)
126-
seq < e.seq
126+
seq <= e.seq
127127
- (b) the new seq is within e.seq + len e -> append (partially)
128128
seq <= e.seq + len
129129
- (c) the new seq is way behind e.seq + len e -> move along
130130
seq > e.seq + len
131131
*)
132-
if Sequence.less seq e.seq then
132+
if Sequence.less_equal seq e.seq then
133133
(* case (a) *)
134134
let seq_e = Sequence.addi seq (Cstruct.length data) in
135-
if Sequence.less_equal e.seq seq_e then
136-
(* we've to merge e into seq *)
137-
let skip_data = Sequence.sub seq_e e.seq in
138-
if Cstruct.length data >= skip_data then
139-
let data = Cstruct.shift data skip_data in
135+
(* case (1): a new segment that is way before the existing one:
136+
seq <= e.seq && seq_e <= e.seq -> e must be retained
137+
case (2): a new segment that is partially before the existing:
138+
seq <= e.seq && seq_e > e.seq -> e may be partially retained:
139+
(i) seq_e >= e.seq_e -> drop e
140+
(ii) seq_e < e.seq_e -> retain the last bytes of e
141+
*)
142+
if Sequence.less_equal seq_e e.seq then
143+
if Sequence.equal seq_e e.seq then
140144
let e = { seq ; fin = fin || e.fin ; data = e.data @ [ data ] } in
141145
Some (e, Sequence.addi seq (Cstruct.lenv e.data)), e :: acc
142146
else
143-
None, e :: acc
147+
let e' = { seq ; fin ; data = [ data ] } in
148+
Some (e', Sequence.addi seq (Cstruct.length data)), e :: e' :: acc
144149
else
145-
let e' = { seq ; fin ; data = [ data ] } in
146-
Some (e', seq_e), e :: e' :: acc
150+
let e_seq_e = Sequence.addi e.seq (Cstruct.lenv e.data) in
151+
if Sequence.greater_equal seq_e e_seq_e then
152+
let e' = { seq ; fin ; data = [ data ] } in
153+
Some (e', seq_e), e' :: acc
154+
else
155+
(* we've to retain some parts of seq *)
156+
let post =
157+
let retain_data = Sequence.sub e_seq_e seq_e in
158+
let skip_data = Cstruct.lenv e.data - retain_data in
159+
Cstruct.shiftv (List.rev e.data) skip_data
160+
in
161+
let e = { seq ; fin = fin || e.fin ; data = List.rev (data :: post) } in
162+
Some (e, Sequence.addi seq (Cstruct.lenv e.data)), e :: acc
147163
else
148-
let e_end = Sequence.addi e.seq (Cstruct.lenv e.data) in
149-
if Sequence.less_equal seq e_end then
164+
let e_seq_e = Sequence.addi e.seq (Cstruct.lenv e.data) in
165+
if Sequence.less_equal seq e_seq_e then
150166
(* case (b) we append to the thing *)
151-
let skip_data = Sequence.sub e_end seq in
152-
if Cstruct.length data >= skip_data then
153-
let data = Cstruct.shift data skip_data in
167+
if Sequence.equal seq e_seq_e then
154168
let e = { e with fin = fin || e.fin ; data = data :: e.data } in
155-
Some (e, Sequence.addi e_end (Cstruct.length data)), e :: acc
169+
Some (e, Sequence.addi e_seq_e (Cstruct.length data)), e :: acc
156170
else
157-
(* we just throw it away *)
158-
Some (e, e_end), e :: acc
171+
let overlap = Sequence.sub e_seq_e seq in
172+
let pre =
173+
let rec cut_some amount = function
174+
| [] -> []
175+
| hd :: tl ->
176+
if Cstruct.length hd < amount then
177+
cut_some (amount - Cstruct.length hd) tl
178+
else
179+
Cstruct.sub hd amount (Cstruct.length hd - amount) :: tl
180+
in
181+
cut_some overlap e.data
182+
in
183+
let seq_e = Sequence.addi seq (Cstruct.length data) in
184+
let end_ = Sequence.max e_seq_e seq_e in
185+
let post =
186+
if Sequence.greater e_seq_e seq_e then
187+
let retain_data = Sequence.sub e_seq_e seq_e in
188+
let skip_data = Cstruct.lenv e.data - retain_data in
189+
Cstruct.shiftv (List.rev e.data) skip_data
190+
else
191+
[]
192+
in
193+
let e = { e with fin = fin || e.fin ; data = post @ data :: pre } in
194+
Some (e, end_), e :: acc
159195
else
160196
(None, e :: acc))
161197
(None, []) t

‎test/reassembly.ml‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,61 @@ let take_works_taking_before_2 () =
212212
| None -> ()
213213
| Some _ -> Alcotest.fail "there shouldn't be anything"
214214

215+
let overlap_1 () =
216+
let r = insert_seg empty (Sequence.of_int32 16l, false, Cstruct.of_string "AAAAAAAA") in
217+
let r = insert_seg r (Sequence.of_int32 8l, false, Cstruct.of_string "BBBBBBBBBB") in
218+
let r = insert_seg r (Sequence.of_int32 0l, false, Cstruct.of_string "CCCCCCCCCC") in
219+
let r', s = maybe_take r Sequence.zero in
220+
Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
221+
match s with
222+
| None -> Alcotest.fail "should be some data"
223+
| Some (s, _) -> Alcotest.(check string "data is good" "CCCCCCCCCCBBBBBBBBAAAAAA"
224+
(Cstruct.to_string s))
225+
226+
let overlap_2 () =
227+
let r = insert_seg empty (Sequence.of_int32 8l, false, Cstruct.of_string "BBBBBBBBBB") in
228+
let r = insert_seg r (Sequence.of_int32 16l, false, Cstruct.of_string "AAAAAAAA") in
229+
let r = insert_seg r (Sequence.of_int32 0l, false, Cstruct.of_string "CCCCCCCCCC") in
230+
let r', s = maybe_take r Sequence.zero in
231+
Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
232+
match s with
233+
| None -> Alcotest.fail "should be some data"
234+
| Some (s, _) -> Alcotest.(check string "data is good" "CCCCCCCCCCBBBBBBAAAAAAAA"
235+
(Cstruct.to_string s))
236+
237+
let overlap_3 () =
238+
let r = insert_seg empty (Sequence.of_int32 16l, false, Cstruct.of_string "AAAAAAAA") in
239+
let r = insert_seg r (Sequence.of_int32 8l, false, Cstruct.of_string "BBBBBBBB") in
240+
let r = insert_seg r (Sequence.of_int32 0l, false, Cstruct.of_string "CCCCCCCCCC") in
241+
let r', s = maybe_take r Sequence.zero in
242+
Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
243+
match s with
244+
| None -> Alcotest.fail "should be some data"
245+
| Some (s, _) -> Alcotest.(check string "data is good" "CCCCCCCCCCBBBBBBAAAAAAAA"
246+
(Cstruct.to_string s))
247+
248+
let regr_187 () =
249+
let r = insert_seg empty (Sequence.of_int32 16l, false, Cstruct.of_string "000002om") in
250+
let r = insert_seg r (Sequence.of_int32 8l, false, Cstruct.of_string "001001nn001002nm001003nl") in
251+
let r = insert_seg r (Sequence.of_int32 0l, false, Cstruct.of_string "002000mo") in
252+
let r', s = maybe_take r Sequence.zero in
253+
Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
254+
match s with
255+
| None -> Alcotest.fail "should be some data"
256+
| Some (s, _) -> Alcotest.(check string "data is good" "002000mo001001nn001002nm001003nl"
257+
(Cstruct.to_string s))
258+
259+
let regr_188 () =
260+
let r = insert_seg empty (Sequence.of_int32 24l, false, Cstruct.of_string "000003ol") in
261+
let r = insert_seg r (Sequence.of_int32 8l, false, Cstruct.of_string "001001nn001002nm001003nl001004nk") in
262+
let r = insert_seg r (Sequence.of_int32 0l, false, Cstruct.of_string "002000mo002001mn") in
263+
let r', s = maybe_take r Sequence.zero in
264+
Alcotest.(check int "reassembly queue is now empty" 0 (length r'));
265+
match s with
266+
| None -> Alcotest.fail "should be some data"
267+
| Some (s, _) -> Alcotest.(check string "data is good" "002000mo002001mn001002nm001003nl001004nk"
268+
(Cstruct.to_string s))
269+
215270
let tests = [
216271
"empty reassembly queue", `Quick, empty_is_empty ;
217272
"non-empty reassembly queue", `Quick, added_is_nonempty ;
@@ -228,4 +283,9 @@ let tests = [
228283
"take works", `Quick, take_works ;
229284
"take works taking before", `Quick, take_works_taking_before ;
230285
"take works taking before 2", `Quick, take_works_taking_before_2 ;
286+
"overlap 1", `Quick, overlap_1 ;
287+
"overlap 2", `Quick, overlap_2 ;
288+
"overlap 3", `Quick, overlap_3 ;
289+
"regression 187", `Quick, regr_187 ;
290+
"regression 188", `Quick, regr_188 ;
231291
]

0 commit comments

Comments
 (0)