provide ciphers with {de,en}crypt_into functionality (#231)

* Mirage_crypto.Block.ECB with {de,en}crypt_into Also provide unsafe_{en,de}crypt_into for further performance. * Mirage_crypto.Block.CBC now has {de,en}crypt_into functionality This may avoid buffer allocations. There are as well unsafe functions for those feeling bounds checks are unnecessary. * counters: add an offset parameter * Mirage_crypto.Block.CTR with {de,en}crypt_into * GCM and ChaCha have {de,en}crypt_into now * CCM16 with {de,en}crypt_into * minor adjustments to speed * Apply suggestions from code review Co-authored-by: Reynir Björnsson <[email protected]> * revise bounds checks (cc @reynir @palainp), also check off >= 0 * revise block_size check * update documentation, esp off < 0 * poly1305: mac_into appropriate bounds checks, also unsafe_mac_into * ccm: remove maclen argument, and ensure tag_size = block_size * add tailcall annotations, remove an argument from ccm's loop --------- Co-authored-by: Reynir Björnsson <[email protected]>
mirage · Jun 29, 2024 · ba299d8 · ba299d8
1 parent 98f01b1
commit ba299d8
Show file tree

Hide file tree

Showing 16 changed files with 787 additions and 253 deletions.
diff --git a/bench/speed.ml b/bench/speed.ml
@@ -45,6 +45,15 @@ let throughput title f =
     Printf.printf "    % 5d:  %04f MB/s  (%d iters in %.03f s)\n%!"
       size (bw /. mb) iters time
 
+let throughput_into ?(add = 0) title f =
+  Printf.printf "\n* [%s]\n%!" title ;
+  sizes |> List.iter @@ fun size ->
+    Gc.full_major () ;
+    let dst = Bytes.create (size + add) in
+    let (iters, time, bw) = burn (f dst) size in
+    Printf.printf "    % 5d:  %04f MB/s  (%d iters in %.03f s)\n%!"
+      size (bw /. mb) iters time
+
 let count_period = 10.
 
 let count f n =
@@ -347,55 +356,128 @@ let benchmarks = [
         fst ecdh_shares);
 
   bm "chacha20-poly1305" (fun name ->
-      let key = Mirage_crypto.Chacha20.of_secret (Mirage_crypto_rng.generate 32)
+      let key = Chacha20.of_secret (Mirage_crypto_rng.generate 32)
       and nonce = Mirage_crypto_rng.generate 8 in
-      throughput name (Mirage_crypto.Chacha20.authenticate_encrypt ~key ~nonce)) ;
+      throughput_into ~add:Chacha20.tag_size name
+        (fun dst cs -> Chacha20.authenticate_encrypt_into ~key ~nonce cs ~src_off:0 dst ~dst_off:0 ~tag_off:(String.length cs) (String.length cs))) ;
+
+  bm "chacha20-poly1305-unsafe" (fun name ->
+      let key = Chacha20.of_secret (Mirage_crypto_rng.generate 32)
+      and nonce = Mirage_crypto_rng.generate 8 in
+      throughput_into ~add:Chacha20.tag_size name
+        (fun dst cs -> Chacha20.unsafe_authenticate_encrypt_into ~key ~nonce cs ~src_off:0 dst ~dst_off:0 ~tag_off:(String.length cs) (String.length cs))) ;
 
   bm "aes-128-ecb" (fun name ->
     let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 16) in
-    throughput name (fun cs -> AES.ECB.encrypt ~key cs)) ;
+    throughput_into name
+      (fun dst cs -> AES.ECB.encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-192-ecb" (fun name ->
+    let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 24) in
+    throughput_into name (fun dst cs -> AES.ECB.encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-192-ecb-unsafe" (fun name ->
+    let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 24) in
+    throughput_into name (fun dst cs -> AES.ECB.unsafe_encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-256-ecb" (fun name ->
+    let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 32) in
+    throughput_into name (fun dst cs -> AES.ECB.encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-256-ecb-unsafe" (fun name ->
+    let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 32) in
+    throughput_into name (fun dst cs -> AES.ECB.unsafe_encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-128-ecb-unsafe" (fun name ->
+    let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 16) in
+    throughput_into name
+      (fun dst cs -> AES.ECB.unsafe_encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
 
   bm "aes-128-cbc-e" (fun name ->
     let key = AES.CBC.of_secret (Mirage_crypto_rng.generate 16)
     and iv  = Mirage_crypto_rng.generate 16 in
-    throughput name (fun cs -> AES.CBC.encrypt ~key ~iv cs)) ;
+    throughput_into name
+      (fun dst cs -> AES.CBC.encrypt_into ~key ~iv cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-128-cbc-e-unsafe" (fun name ->
+    let key = AES.CBC.of_secret (Mirage_crypto_rng.generate 16)
+    and iv  = Mirage_crypto_rng.generate 16 in
+    throughput_into name
+      (fun dst cs -> AES.CBC.unsafe_encrypt_into ~key ~iv cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-128-cbc-e-unsafe-inplace" (fun name ->
+    let key = AES.CBC.of_secret (Mirage_crypto_rng.generate 16)
+    and iv  = Mirage_crypto_rng.generate 16 in
+    throughput name
+      (fun cs ->
+         let b = Bytes.unsafe_of_string cs in
+         AES.CBC.unsafe_encrypt_into_inplace ~key ~iv b ~dst_off:0 (String.length cs))) ;
 
   bm "aes-128-cbc-d" (fun name ->
     let key = AES.CBC.of_secret (Mirage_crypto_rng.generate 16)
     and iv  = Mirage_crypto_rng.generate 16 in
-    throughput name (fun cs -> AES.CBC.decrypt ~key ~iv cs)) ;
+    throughput_into name
+      (fun dst cs -> AES.CBC.decrypt_into ~key ~iv cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-128-cbc-d-unsafe" (fun name ->
+    let key = AES.CBC.of_secret (Mirage_crypto_rng.generate 16)
+    and iv  = Mirage_crypto_rng.generate 16 in
+    throughput_into name
+      (fun dst cs -> AES.CBC.unsafe_decrypt_into ~key ~iv cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
 
   bm "aes-128-ctr" (fun name ->
     let key = Mirage_crypto_rng.generate 16 |> AES.CTR.of_secret
     and ctr = Mirage_crypto_rng.generate 16 |> AES.CTR.ctr_of_octets in
-    throughput name (fun cs -> AES.CTR.encrypt ~key ~ctr cs)) ;
+    throughput_into name (fun dst cs -> AES.CTR.encrypt_into ~key ~ctr cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "aes-128-ctr-unsafe" (fun name ->
+    let key = Mirage_crypto_rng.generate 16 |> AES.CTR.of_secret
+    and ctr = Mirage_crypto_rng.generate 16 |> AES.CTR.ctr_of_octets in
+    throughput_into name (fun dst cs -> AES.CTR.unsafe_encrypt_into ~key ~ctr cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
 
   bm "aes-128-gcm" (fun name ->
     let key = AES.GCM.of_secret (Mirage_crypto_rng.generate 16)
     and nonce = Mirage_crypto_rng.generate 12 in
-    throughput name (fun cs -> AES.GCM.authenticate_encrypt ~key ~nonce cs));
+    throughput_into ~add:AES.GCM.tag_size name
+      (fun dst cs -> AES.GCM.authenticate_encrypt_into ~key ~nonce cs ~src_off:0 dst ~dst_off:0 ~tag_off:(String.length cs) (String.length cs)));
+
+  bm "aes-128-gcm-unsafe" (fun name ->
+    let key = AES.GCM.of_secret (Mirage_crypto_rng.generate 16)
+    and nonce = Mirage_crypto_rng.generate 12 in
+    throughput_into ~add:AES.GCM.tag_size name
+      (fun dst cs -> AES.GCM.unsafe_authenticate_encrypt_into ~key ~nonce cs ~src_off:0 dst ~dst_off:0 ~tag_off:(String.length cs) (String.length cs)));
 
   bm "aes-128-ghash" (fun name ->
     let key = AES.GCM.of_secret (Mirage_crypto_rng.generate 16)
     and nonce = Mirage_crypto_rng.generate 12 in
-    throughput name (fun cs -> AES.GCM.authenticate_encrypt ~key ~nonce ~adata:cs ""));
+    throughput_into ~add:AES.GCM.tag_size name
+      (fun dst cs -> AES.GCM.authenticate_encrypt_into ~key ~nonce ~adata:cs "" ~src_off:0 dst ~dst_off:0 ~tag_off:0 0));
+
+  bm "aes-128-ghash-unsafe" (fun name ->
+    let key = AES.GCM.of_secret (Mirage_crypto_rng.generate 16)
+    and nonce = Mirage_crypto_rng.generate 12 in
+    throughput_into ~add:AES.GCM.tag_size name
+      (fun dst cs -> AES.GCM.unsafe_authenticate_encrypt_into ~key ~nonce ~adata:cs "" ~src_off:0 dst ~dst_off:0 ~tag_off:0 0));
 
   bm "aes-128-ccm" (fun name ->
     let key   = AES.CCM16.of_secret (Mirage_crypto_rng.generate 16)
     and nonce = Mirage_crypto_rng.generate 10 in
-    throughput name (fun cs -> AES.CCM16.authenticate_encrypt ~key ~nonce cs));
-
-  bm "aes-192-ecb" (fun name ->
-    let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 24) in
-    throughput name (fun cs -> AES.ECB.encrypt ~key cs)) ;
+    throughput_into ~add:AES.CCM16.tag_size name
+      (fun dst cs -> AES.CCM16.authenticate_encrypt_into ~key ~nonce cs ~src_off:0 dst ~dst_off:0 ~tag_off:(String.length cs) (String.length cs)));
 
-  bm "aes-256-ecb" (fun name ->
-    let key = AES.ECB.of_secret (Mirage_crypto_rng.generate 32) in
-    throughput name (fun cs -> AES.ECB.encrypt ~key cs)) ;
+  bm "aes-128-ccm-unsafe" (fun name ->
+    let key   = AES.CCM16.of_secret (Mirage_crypto_rng.generate 16)
+    and nonce = Mirage_crypto_rng.generate 10 in
+    throughput_into ~add:AES.CCM16.tag_size name
+      (fun dst cs -> AES.CCM16.unsafe_authenticate_encrypt_into ~key ~nonce cs ~src_off:0 dst ~dst_off:0 ~tag_off:(String.length cs) (String.length cs)));
 
   bm "d3des-ecb" (fun name ->
     let key = DES.ECB.of_secret (Mirage_crypto_rng.generate 24) in
-    throughput name (fun cs -> DES.ECB.encrypt ~key cs)) ;
+    throughput_into name (fun dst cs -> DES.ECB.encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
+
+  bm "d3des-ecb-unsafe" (fun name ->
+    let key = DES.ECB.of_secret (Mirage_crypto_rng.generate 24) in
+    throughput_into name (fun dst cs -> DES.ECB.unsafe_encrypt_into ~key cs ~src_off:0 dst ~dst_off:0 (String.length cs))) ;
 
   bm "fortuna" (fun name ->
     let open Mirage_crypto_rng.Fortuna in

diff --git a/src/aead.ml b/src/aead.ml
@@ -10,4 +10,16 @@ module type AEAD = sig
     string -> string * string
   val authenticate_decrypt_tag : key:key -> nonce:string -> ?adata:string ->
     tag:string -> string -> string option
+  val authenticate_encrypt_into : key:key -> nonce:string ->
+    ?adata:string -> string -> src_off:int -> bytes -> dst_off:int ->
+    tag_off:int -> int -> unit
+  val authenticate_decrypt_into : key:key -> nonce:string ->
+    ?adata:string -> string -> src_off:int -> tag_off:int -> bytes ->
+    dst_off:int -> int -> bool
+  val unsafe_authenticate_encrypt_into : key:key -> nonce:string ->
+    ?adata:string -> string -> src_off:int -> bytes -> dst_off:int ->
+    tag_off:int -> int -> unit
+  val unsafe_authenticate_decrypt_into : key:key -> nonce:string ->
+    ?adata:string -> string -> src_off:int -> tag_off:int -> bytes ->
+    dst_off:int -> int -> bool
 end
diff --git a/src/ccm.ml b/src/ccm.ml
@@ -10,7 +10,7 @@ let encode_len buf ~off size value =
     | 0 -> Bytes.set_uint8 buf off num
     | m ->
       Bytes.set_uint8 buf (off + m) (num land 0xff);
-      ass (num lsr 8) (pred m)
+      (ass [@tailcall]) (num lsr 8) (pred m)
   in
   ass value (pred size)
 
@@ -74,10 +74,8 @@ let prepare_header nonce adata plen tlen =
 
 type mode = Encrypt | Decrypt
 
-let crypto_core ~cipher ~mode ~key ~nonce ~maclen ~adata data =
-  let datalen = String.length data in
-  let cbcheader = prepare_header nonce adata datalen maclen in
-  let dst = Bytes.create datalen in
+let crypto_core_into ~cipher ~mode ~key ~nonce ~adata src ~src_off dst ~dst_off len =
+  let cbcheader = prepare_header nonce adata len block_size in
 
   let small_q = 15 - String.length nonce in
   let ctr_flag_val = flags 0 0 (small_q - 1) in
@@ -93,65 +91,62 @@ let crypto_core ~cipher ~mode ~key ~nonce ~maclen ~adata data =
     cipher ~key (Bytes.unsafe_to_string block) ~src_off:dst_off block ~dst_off
   in
 
-  let cbcprep =
+  let iv =
     let rec doit iv iv_off block block_off =
       match Bytes.length block - block_off with
       | 0 -> Bytes.sub iv iv_off block_size
       | _ ->
          cbc (Bytes.unsafe_to_string iv) iv_off block block_off;
-         doit block block_off block (block_off + block_size)
+         (doit [@tailcall]) block block_off block (block_off + block_size)
     in
     doit (Bytes.make block_size '\x00') 0 cbcheader 0
   in
 
-  let rec loop iv ctr src src_off dst dst_off=
+  let rec loop ctr src src_off dst dst_off len =
     let cbcblock, cbc_off =
       match mode with
       | Encrypt -> src, src_off
       | Decrypt -> Bytes.unsafe_to_string dst, dst_off
     in
-    match String.length src - src_off with
-    | 0 -> iv
-    | x when x < block_size ->
+    if len = 0 then
+      ()
+    else if len < block_size then begin
       let buf = Bytes.make block_size '\x00' in
-      Bytes.unsafe_blit dst dst_off buf 0 x;
+      Bytes.unsafe_blit dst dst_off buf 0 len ;
       ctrblock ctr buf ;
-      Bytes.unsafe_blit buf 0 dst dst_off x ;
-      unsafe_xor_into src ~src_off dst ~dst_off x ;
-      Bytes.unsafe_blit_string cbcblock cbc_off buf 0 x;
-      Bytes.unsafe_fill buf x (block_size - x) '\x00';
-      cbc (Bytes.unsafe_to_string buf) cbc_off iv 0 ;
-      iv
-    | _ ->
+      Bytes.unsafe_blit buf 0 dst dst_off len ;
+      unsafe_xor_into src ~src_off dst ~dst_off len ;
+      Bytes.unsafe_blit_string cbcblock cbc_off buf 0 len ;
+      Bytes.unsafe_fill buf len (block_size - len) '\x00';
+      cbc (Bytes.unsafe_to_string buf) cbc_off iv 0
+    end else begin
       ctrblock ctr dst ;
       unsafe_xor_into src ~src_off dst ~dst_off block_size ;
       cbc cbcblock cbc_off iv 0 ;
-      loop iv (succ ctr) src (src_off + block_size) dst (dst_off + block_size)
+      (loop [@tailcall]) (succ ctr) src (src_off + block_size) dst (dst_off + block_size) (len - block_size)
+    end
   in
-  let last = loop cbcprep 1 data 0 dst 0 in
-  let t = Bytes.sub last 0 maclen in
-  (dst, t)
+  loop 1 src src_off dst dst_off len;
+  iv
+
+let crypto_core ~cipher ~mode ~key ~nonce ~adata data =
+  let datalen = String.length data in
+  let dst = Bytes.create datalen in
+  let t = crypto_core_into ~cipher ~mode ~key ~nonce ~adata data ~src_off:0 dst ~dst_off:0 datalen in
+  dst, t
 
 let crypto_t t nonce cipher key =
   let ctr = gen_ctr nonce 0 in
   cipher ~key (Bytes.unsafe_to_string ctr) ~src_off:0 ctr ~dst_off:0 ;
   unsafe_xor_into (Bytes.unsafe_to_string ctr) ~src_off:0 t ~dst_off:0 (Bytes.length t)
 
-let valid_nonce nonce =
-  let nsize = String.length nonce in
-  if nsize < 7 || nsize > 13 then
-    invalid_arg "CCM: nonce length not between 7 and 13: %u" nsize
-
-let generation_encryption ~cipher ~key ~nonce ~maclen ~adata data =
-  valid_nonce nonce;
-  let cdata, t = crypto_core ~cipher ~mode:Encrypt ~key ~nonce ~maclen ~adata data in
+let unsafe_generation_encryption_into ~cipher ~key ~nonce ~adata src ~src_off dst ~dst_off ~tag_off len =
+  let t = crypto_core_into ~cipher ~mode:Encrypt ~key ~nonce ~adata src ~src_off dst ~dst_off len in
   crypto_t t nonce cipher key ;
-  Bytes.unsafe_to_string cdata, Bytes.unsafe_to_string t
+  Bytes.unsafe_blit t 0 dst tag_off block_size
 
-let decryption_verification ~cipher ~key ~nonce ~maclen ~adata ~tag data =
-  valid_nonce nonce;
-  let cdata, t = crypto_core ~cipher ~mode:Decrypt ~key ~nonce ~maclen ~adata data in
+let unsafe_decryption_verification_into ~cipher ~key ~nonce ~adata src ~src_off ~tag_off dst ~dst_off len =
+  let tag = String.sub src tag_off block_size in
+  let t = crypto_core_into ~cipher ~mode:Decrypt ~key ~nonce ~adata src ~src_off dst ~dst_off len in
   crypto_t t nonce cipher key ;
-  match Eqaf.equal tag (Bytes.unsafe_to_string t) with
-  | true  -> Some (Bytes.unsafe_to_string cdata)
-  | false -> None
+  Eqaf.equal tag (Bytes.unsafe_to_string t)