More garbage friendly cstruct wrapper

src/cstruct_ext.ml

@@ -0,0 +1,25 @@
+(* [or_empty cs] is [cs] unless [Cstruct.is_empty cs] where it is [Cstruct.empty].
+   An empty cstruct will keep a live reference to its underlying buffer. It is
+   preferable to keep a reference live to [Cstruct.empty.buffer] than
+   [cs.buffer] if [cs] is empty. *)
+let [@ocaml.inline always] or_empty cs =
+  if Cstruct.is_empty cs then
+    Cstruct.empty
+  else cs
+
+(* [Cstruct.sub cs 0 cs.len] will result in an empty cstruct **that keeps a
+   live reference to [cs.buffer]**!! *)
+let sub cs off len = or_empty (Cstruct.sub cs off len)
+
+(* [Cstruct.shift cs cs.len] has a similar story *)
+let shift cs len = or_empty (Cstruct.shift cs len)
+
+(* If we append the empty cstruct and we don't need a fresh copy we can do
+   nothing. This has different semantics than Cstruct.append. *)
+let append_nocopy cs cs' =
+  if Cstruct.is_empty cs then
+    cs'
+  else if Cstruct.is_empty cs' then
+    cs
+  else
+    Cstruct.append cs cs'

src/engine.ml

@@ -122,7 +122,7 @@ let encrypt_and_out protocol { my; _ } key
   let buf, enc_off, enc_len = Packet.Tls_crypt.encode protocol (key, p) in
   let encrypted =
     Mirage_crypto.Cipher_block.AES.CTR.encrypt ~key:my_key ~ctr
-      (Cstruct.sub buf enc_off enc_len)
+      (Cstruct_ext.sub buf enc_off enc_len)
   in
   Cstruct.blit encrypted 0 buf enc_off enc_len;
   buf
@@ -1027,7 +1027,7 @@ let unpad block_size cs =
   let l = Cstruct.length cs in
   let amount = Cstruct.get_uint8 cs (pred l) in
   let len = l - amount in
-  if len >= 0 && amount <= block_size then Ok (Cstruct.sub cs 0 len)
+  if len >= 0 && amount <= block_size then Ok (Cstruct_ext.sub cs 0 len)
   else Error (`Msg "bad padding")
 
 let out ?add_timestamp (ctx : keys) hmac_algorithm compress rng data =
@@ -1125,7 +1125,7 @@ let data_out ?add_timestamp (ctx : keys) hmac_algorithm compress protocol rng
 let static_out ~add_timestamp ctx hmac_algorithm compress protocol rng data =
   let ctx, payload = out ~add_timestamp ctx hmac_algorithm compress rng data in
   let prefix = Packet.encode_protocol protocol (Cstruct.length payload) in
-  let out = Cstruct.append prefix payload in
+  let out = Cstruct_ext.append_nocopy prefix payload in
   Log.debug (fun m ->
       m "sending %d bytes data (enc %d) out id %lu" (Cstruct.length data)
         (Cstruct.length payload) ctx.my_replay_id);
@@ -1692,7 +1692,7 @@ let incoming ?(is_not_taken = fun _ip -> false) state control_crypto buf =
         if Cstruct.is_empty linger then Ok (state, out, payloads, act_opt)
         else multi linger (state, out, payloads, act_opt)
   in
-  let r = multi (Cstruct.append state.linger buf) (state, [], [], None) in
+  let r = multi (Cstruct_ext.append_nocopy state.linger buf) (state, [], [], None) in
   let+ s', out, payloads, act_opt = udp_ignore r in
   Log.debug (fun m -> m "out state is %a" State.pp s');
   Log.debug (fun m ->
@@ -1960,7 +1960,7 @@ let handle_static_client t s keys ev =
                   let acc = Option.fold d ~none:acc ~some:(fun p -> p :: acc) in
                   process_one acc linger
           in
-          let+ t, payloads = process_one [] (Cstruct.append t.linger cs) in
+          let+ t, payloads = process_one [] (Cstruct_ext.append_nocopy t.linger cs) in
           (t, [], List.rev payloads, None)
       | s, ev ->
           Result.error_msgf

src/packet.ml

@@ -184,7 +184,7 @@ let decode_control ~hmac_len buf =
   let* header, off = decode_header ~hmac_len buf in
   let+ () = guard (Cstruct.length buf >= off + 4) `Partial in
   let sequence_number = Cstruct.BE.get_uint32 buf off
-  and payload = Cstruct.shift buf (off + 4) in
+  and payload = Cstruct_ext.shift buf (off + 4) in
   (header, sequence_number, payload)
 
 let decode_ack_or_control op ~hmac_len buf =
@@ -217,7 +217,7 @@ let decode_protocol proto buf =
       let* () = guard (Cstruct.length buf >= 2) `Tcp_partial in
       let plen = Cstruct.BE.get_uint16 buf 0 in
       let+ () = guard (Cstruct.length buf - 2 >= plen) `Tcp_partial in
-      (Cstruct.sub buf 2 plen, Cstruct.shift buf (plen + 2))
+      (Cstruct.sub buf 2 plen, Cstruct_ext.shift buf (plen + 2))
   | `Udp -> Ok (buf, Cstruct.empty)
 
 let decode_key_op proto buf =
@@ -227,7 +227,7 @@ let decode_key_op proto buf =
   let opkey = Cstruct.get_uint8 buf 0 in
   let op, key = (opkey lsr 3, opkey land 0x07) in
   let+ op = int_to_operation op in
-  (op, key, Cstruct.shift buf 1, linger)
+  (op, key, Cstruct_ext.shift buf 1, linger)
 
 let operation = function
   | `Ack _ -> Ack
@@ -426,7 +426,7 @@ module Tls_crypt = struct
     let* hdr, off = decode_decrypted_header clear_hdr buf in
     let+ () = guard (Cstruct.length buf >= off + 4) `Partial in
     let sequence_number = Cstruct.BE.get_uint32 buf off
-    and payload = Cstruct.shift buf (off + 4) in
+    and payload = Cstruct_ext.shift buf (off + 4) in
     (hdr, sequence_number, payload)
 
   let decode_decrypted_ack_or_control clear_hdr op buf =
@@ -448,7 +448,7 @@ module Tls_crypt = struct
     and timestamp = Cstruct.BE.get_uint32 buf 12
     and hmac = Cstruct.sub buf 16 hmac_len in
     ( { local_session; replay_id; timestamp; hmac },
-      Cstruct.shift buf clear_hdr_len )
+      Cstruct_ext.shift buf clear_hdr_len )
 end
 
 type ack = [ `Ack of header ]
@@ -604,9 +604,9 @@ let decode_tls_data ?(with_premaster = false) buf =
       let+ peer_info =
         if Cstruct.length buf <= peer_info_start + 2 then Ok None
         else
-          let data = Cstruct.shift buf peer_info_start in
+          let data = Cstruct_ext.shift buf peer_info_start in
           let len = Cstruct.BE.get_uint16 data 0 in
-          let data = Cstruct.shift data 2 in
+          let data = Cstruct_ext.shift data 2 in
           let* () = guard (Cstruct.length data >= len) `Partial in
           if Cstruct.length data > len then
             Log.warn (fun m ->