Source file fragments.ml
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
let src = Logs.Src.create "ipv4-fragments" ~doc:"IPv4 fragmentation"
module Log = (val Logs.src_log src : Logs.LOG)
module V = struct
  type t = int64 * Cstruct.t * bool * int * (int * Cstruct.t) list
  let weight (_, _, _, _, v) = Cstruct.lenv (List.map snd v)
end
module K = struct
  type t = Ipaddr.V4.t * Ipaddr.V4.t * int * int
  let compare (src, dst, proto, id) (src', dst', proto', id') =
    let (&&&) a b = match a with 0 -> b | x -> x in
    let int_cmp : int -> int -> int = compare in
    Ipaddr.V4.compare src src' &&&
    Ipaddr.V4.compare dst dst' &&&
    int_cmp proto proto' &&&
    int_cmp id id'
end
module Cache = Lru.F.Make(K)(V)
let rec insert_sorted ((frag_start, _) as frag) = function
  | [] -> [ frag ]
  | ((frag'_start, _) as frag')::tl ->
    if frag'_start <= frag_start
    then frag::frag'::tl
    else frag'::insert_sorted frag tl
type r = Bad | Hole
let attempt_reassemble fragments =
  Log.debug (fun m -> m "reassemble %a"
                Fmt.(list ~sep:(any "; ") (pair ~sep:(any ", len ") int int))
                (List.map (fun (off, data) -> off, Cstruct.length data) fragments)) ;
  
  
  let len =
    
    let off, data = List.hd fragments in
    off + Cstruct.length data
  in
  let rec check until = function
    | [] -> if until = 0 then Ok () else Error Hole
    | (start, d)::tl ->
      let until' = start + (Cstruct.length d) in
      if until = until'
      then check start tl
      else if until' > until
      then Error Bad
      else Error Hole
  in
  Result.bind
    (check len fragments)
    (fun () ->
       let buf = Cstruct.create_unsafe len in
       List.iter (fun (off, data) ->
           Cstruct.blit data 0 buf off (Cstruct.length data))
         fragments ;
       Ok buf)
let max_number_of_fragments = 16
let max_duration = Duration.of_sec 10
let process cache ts (packet : Ipv4_packet.t) payload =
  let add_trim key value cache =
    let cache' = Cache.add key value cache in
    Cache.trim cache'
  in
  if packet.off land 0x3FFF = 0 then 
    
    cache, Some (packet, payload)
  else
    let offset, more =
      (packet.off land 0x1FFF) lsl 3, 
      packet.off land 0x2000 = 0x2000
    and key = (packet.src, packet.dst, packet.proto, packet.id)
    in
    let v = (ts, packet.options, not more, 1, [(offset, payload)]) in
    match Cache.find key cache with
    | None ->
      Log.debug (fun m -> m "%a none found, inserting into cache" Ipv4_packet.pp packet) ;
      add_trim key v cache, None
    | Some (ts', options, finished, cnt, frags) ->
      if Int64.sub ts ts' >= max_duration then begin
        Log.warn (fun m -> m "%a found some, but timestamp exceeded duration %a, dropping old segments and inserting new segment into cache" Ipv4_packet.pp packet Duration.pp max_duration) ;
        add_trim key v cache, None
      end else
        let cache' = Cache.promote key cache in
        let all_frags = insert_sorted (offset, payload) frags
        and try_reassemble = finished || not more
        and options' = if offset = 0 then packet.options else options
        in
        Log.debug (fun m -> m "%d found, finished %b more %b try_reassemble %b"
                      cnt finished more try_reassemble) ;
        let maybe_add_to_cache c =
          if cnt < max_number_of_fragments then
            add_trim key (ts', options', try_reassemble, succ cnt, all_frags) c
          else
            (Log.warn (fun m -> m "%a dropping from cache, maximum number of fragments exceeded"
                          Ipv4_packet.pp packet) ;
             Cache.remove key c)
        in
        if try_reassemble then
          match attempt_reassemble all_frags with
          | Ok p ->
            Log.debug (fun m -> m "%a reassembled to payload %d" Ipv4_packet.pp packet (Cstruct.length p)) ;
            let packet' = { packet with options = options' ; off = 0 } in
            Cache.remove key cache', Some (packet', p)
          | Error Bad ->
            Log.warn (fun m -> m "%a dropping from cache, bad fragments (%a)"
                         Ipv4_packet.pp packet
                         Fmt.(list ~sep:(any "; ") (pair ~sep:(any ", ") int int))
                         (List.map (fun (s, d) -> (s, Cstruct.length d)) all_frags)) ;
            Log.debug (fun m -> m "full fragments: %a"
                          Fmt.(list ~sep:(any "@.") Cstruct.hexdump_pp)
                          (List.map snd all_frags)) ;
            Cache.remove key cache', None
          | Error Hole -> maybe_add_to_cache cache', None
        else
          maybe_add_to_cache cache', None
let fragment ~mtu hdr payload =
  let rec frag1 acc hdr hdr_buf offset data_size payload =
    let more = Cstruct.length payload > data_size in
    let hdr' =
      
      let off = (offset / 8) lor (if more then 0x2000 else 0) in
      { hdr with Ipv4_packet.off }
    in
    let this_payload, rest =
      if more then Cstruct.split payload data_size else payload, Cstruct.empty
    in
    let payload_len = Cstruct.length this_payload in
    Ipv4_wire.set_checksum hdr_buf 0;
    (match Ipv4_packet.Marshal.into_cstruct ~payload_len hdr' hdr_buf with
     
     | Error msg -> invalid_arg msg
     | Ok () -> ());
    let acc' = Cstruct.append hdr_buf this_payload :: acc in
    if more then
      let offset = offset + data_size in
      (frag1[@tailcall]) acc' hdr hdr_buf offset data_size rest
    else
      acc'
  in
  let hdr_size =
    
    let opt_size = (Cstruct.length hdr.Ipv4_packet.options + 3) / 4 * 4 in
    opt_size + Ipv4_wire.sizeof_ipv4
  in
  let data_size =
    let full = mtu - hdr_size in
    (full / 8) * 8
  in
  if data_size <= 0 then
    []
  else
    List.rev (frag1 [] hdr (Cstruct.create hdr_size) data_size data_size payload)