From a73212b51f8f9a32a0f6c06e4f8ade547f844400 Mon Sep 17 00:00:00 2001
From: swrup <swrup@protonmail.com>
Date: Thu, 12 Feb 2026 14:21:50 +0100
Subject: [PATCH]

---
 src/crypto.ml | 58 +++++++++++++++++++--------------------------------
 1 file changed, 21 insertions(+), 37 deletions(-)

diff --git a/src/crypto.ml b/src/crypto.ml
index bf437fce..92d1dadb 100644
--- a/src/crypto.ml
+++ b/src/crypto.ml
@@ -1,9 +1,6 @@
 open Syntax
 
 module Binary_format_rsa = struct
-  (* TODO tests:
-         - need to strip leading zeros?
-         - endianess ok? *)
   (* RSA public key binary format
        https://www.gnupg.org/documentation/manuals/gcrypt/MPI-formats.html
        := { uint16_be: n size; uint16_be: e size; n; e}
@@ -12,19 +9,9 @@ module Binary_format_rsa = struct
        leading zeroes are stripped unless they are required to keep a value positive
        no 0-termination *)
 
-  let rev_string len s = String.init len (fun i -> s.[len - 1 - i])
-
-  (* we reverse bytes because Z.of_bits reads bytes in little endian *)
-  let z_of_bits_be src pos len =
-    String.sub src pos len |> rev_string len |> Z.of_bits
-
-  let z_to_bits_be z =
-    let bits = Z.to_bits z in
-    rev_string (String.length bits) bits
-
   let z_array_to_octets (arr : Z.t array) =
     let nb = Array.length arr in
-    let bits_arr = Array.map z_to_bits_be arr in
+    let bits_arr = Array.map Mirage_crypto_pk.Z_extra.to_octets_be arr in
     let len_arr = Array.map String.length bits_arr in
     let len = (2 * nb) + Array.fold_left ( + ) 0 len_arr in
     let b = Bytes.make len '\x00' in
@@ -42,32 +29,29 @@ module Binary_format_rsa = struct
       bits_arr;
     Bytes.unsafe_to_string b
 
-  let check = function
-    | false -> Error "rsa of_octets error, invalid data"
-    | true -> Ok ()
-
   let z_array_of_octets ~nb s =
     let s_len = String.length s in
-    let* () = check (s_len > 2 * nb) in
-    let pos = ref 0 in
-    let len_arr =
-      Array.init nb (fun _i ->
-          let len = String.get_uint16_be s !pos in
-          pos := !pos + 2;
-          len)
-    in
-    let* () =
+    if s_len <= 2 * nb then Error "rsa of_octets error"
+    else
+      let pos = ref 0 in
+      let len_arr =
+        Array.init nb (fun _i ->
+            let len = String.get_uint16_be s !pos in
+            pos := !pos + 2;
+            len)
+      in
       let len = (2 * nb) + Array.fold_left ( + ) 0 len_arr in
-      check (s_len = len)
-    in
-    let z_arr =
-      Array.init nb (fun i ->
-          let len = len_arr.(i) in
-          let z = z_of_bits_be s !pos len in
-          pos := !pos + len;
-          z)
-    in
-    Ok z_arr
+      if s_len <> len then Error "rsa of_octets error"
+      else
+        let z_arr =
+          Array.init nb (fun i ->
+              let len = len_arr.(i) in
+              let s = String.sub s !pos len in
+              let z = Mirage_crypto_pk.Z_extra.of_octets_be s in
+              pos := !pos + len;
+              z)
+        in
+        Ok z_arr
 
   let pub_to_octets ({ n; e } : Mirage_crypto_pk.Rsa.pub) =
     z_array_to_octets [| n; e |]