unikernel/duniverse/ocaml-dns/dnssec/base32.ml

@@ -0,0 +1,236 @@
+(* RFC 4648, Section 7 - Base 32 with extended hex alphabet *)
+
+let make_alphabet alphabet =
+  if String.length alphabet <> 32
+  then invalid_arg "Length of alphabet must be 32" ;
+  if String.contains alphabet '='
+  then invalid_arg "Alphabet can not contain padding character" ;
+  let emap =
+    Array.init (String.length alphabet) (fun i -> Char.code alphabet.[i]) in
+  let dmap = Array.make 256 (-1) in
+  String.iteri (fun idx chr -> dmap.(Char.code chr) <- idx) alphabet ;
+  (emap, dmap)
+
+let alphabet = make_alphabet "0123456789ABCDEFGHIJKLMNOPQRSTUV"
+
+let pad_char = '='
+let pad_int = int_of_char pad_char
+
+let encode ?(pad = true) str =
+  let len = String.length str in
+  (* since String.get_uint8 is OCaml >= 4.13 only *)
+  let str = Bytes.unsafe_of_string str in
+  let out_len = (len + 4) / 5 * 8 in
+  let out = Bytes.make out_len pad_char in
+  let o1 b1 = b1 lsr 3
+  and o2 b1 b2 = (b1 land 0x07) lsl 2 + b2 lsr 6
+  and o3 b2 = (b2 land 0x3E) lsr 1
+  and o4 b2 b3 = ((b2 land 0x01) lsl 4) + b3 lsr 4
+  and o5 b3 b4 = (b3 land 0x0F) lsl 1 + b4 lsr 7
+  and o6 b4 = (b4 land 0x7c) lsr 2
+  and o7 b4 b5 = (b4 land 0x03) lsl 3 + b5 lsr 5
+  and o8 b5 = b5 land 0x1F
+  in
+  let emit b1 b2 b3 b4 b5 off =
+    List.iteri (fun idx v -> Bytes.set_uint8 out (off + idx) ((fst alphabet).(v)))
+      [ o1 b1; o2 b1 b2; o3 b2; o4 b2 b3; o5 b3 b4; o6 b4; o7 b4 b5; o8 b5 ]
+  in
+  let rec enc s_off d_off =
+    if s_off = len then
+      (* case 1 *) 0
+    else if s_off = len - 1 then
+      (* case 2 - 6 padding = *)
+      let b1 = Bytes.get_uint8 str s_off in
+      let p1 = o1 b1 and p2 = o2 b1 0 in
+      Bytes.set_uint8 out d_off ((fst alphabet).(p1));
+      Bytes.set_uint8 out (d_off + 1) ((fst alphabet).(p2));
+      6
+    else if s_off = len - 2 then
+      (* case 3 - 4 padding = *)
+      let b1 = Bytes.get_uint8 str s_off
+      and b2 = Bytes.get_uint8 str (s_off + 1)
+      in
+      let p1 = o1 b1 and p2 = o2 b1 b2 and p3 = o3 b2 and p4 = o4 b2 0 in
+      Bytes.set_uint8 out d_off ((fst alphabet).(p1));
+      Bytes.set_uint8 out (d_off + 1) ((fst alphabet).(p2));
+      Bytes.set_uint8 out (d_off + 2) ((fst alphabet).(p3));
+      Bytes.set_uint8 out (d_off + 3) ((fst alphabet).(p4));
+      4
+    else if s_off = len - 3 then
+      (* case 4 - 3 padding = *)
+      let b1 = Bytes.get_uint8 str s_off
+      and b2 = Bytes.get_uint8 str (s_off + 1)
+      and b3 = Bytes.get_uint8 str (s_off + 2)
+      in
+      let p1 = o1 b1 and p2 = o2 b1 b2 and p3 = o3 b2 and p4 = o4 b2 b3 and p5 = o5 b3 0 in
+      Bytes.set_uint8 out d_off ((fst alphabet).(p1));
+      Bytes.set_uint8 out (d_off + 1) ((fst alphabet).(p2));
+      Bytes.set_uint8 out (d_off + 2) ((fst alphabet).(p3));
+      Bytes.set_uint8 out (d_off + 3) ((fst alphabet).(p4));
+      Bytes.set_uint8 out (d_off + 4) ((fst alphabet).(p5));
+      3
+    else if s_off = len - 4 then
+      (* case 5 - 1 padding = *)
+      let b1 = Bytes.get_uint8 str s_off
+      and b2 = Bytes.get_uint8 str (s_off + 1)
+      and b3 = Bytes.get_uint8 str (s_off + 2)
+      and b4 = Bytes.get_uint8 str (s_off + 3)
+      in
+      let p1 = o1 b1 and p2 = o2 b1 b2 and p3 = o3 b2 and p4 = o4 b2 b3 and p5 = o5 b3 b4 and p6 = o6 b4 and p7 = o7 b4 0 in
+      Bytes.set_uint8 out d_off ((fst alphabet).(p1));
+      Bytes.set_uint8 out (d_off + 1) ((fst alphabet).(p2));
+      Bytes.set_uint8 out (d_off + 2) ((fst alphabet).(p3));
+      Bytes.set_uint8 out (d_off + 3) ((fst alphabet).(p4));
+      Bytes.set_uint8 out (d_off + 4) ((fst alphabet).(p5));
+      Bytes.set_uint8 out (d_off + 5) ((fst alphabet).(p6));
+      Bytes.set_uint8 out (d_off + 6) ((fst alphabet).(p7));
+      1
+    else
+      let b1 = Bytes.get_uint8 str s_off in
+      let b2 = Bytes.get_uint8 str (s_off + 1) in
+      let b3 = Bytes.get_uint8 str (s_off + 2) in
+      let b4 = Bytes.get_uint8 str (s_off + 3) in
+      let b5 = Bytes.get_uint8 str (s_off + 4) in
+      emit b1 b2 b3 b4 b5 d_off;
+      enc (s_off + 5) (d_off + 8)
+  in
+  let padding_bytes = enc 0 0 in
+  let out_s = Bytes.unsafe_to_string out in
+  if pad then out_s else String.sub out_s 0 (out_len - padding_bytes)
+
+(* RFC 4648 section 10
+   BASE32-HEX("") = ""
+   BASE32-HEX("f") = "CO======"
+   BASE32-HEX("fo") = "CPNG===="
+   BASE32-HEX("foo") = "CPNMU==="
+   BASE32-HEX("foob") = "CPNMUOG="
+   BASE32-HEX("fooba") = "CPNMUOJ1"
+   BASE32-HEX("foobar") = "CPNMUOJ1E8======"
+*)
+
+let decode ?(unpadded = false) str =
+  let ( let* ) = Result.bind in
+  let* str =
+    let lmod8 = String.length str mod 8 in
+    if lmod8 > 0 then
+      if unpadded then
+        Ok (str ^ String.make (8 - lmod8) pad_char)
+      else
+        Error (`Msg "invalid input length (not divisible by 8)")
+    else
+      Ok str
+  in
+  let len = String.length str in
+  let str = Bytes.unsafe_of_string str in
+  let out_len = len / 8 * 5 in (* max length *)
+  let out = Bytes.create out_len in
+  let o1 b1 b2 = b1 lsl 3 + b2 lsr 2
+  and o2 b2 b3 b4 = (b2 land 0x03) lsl 6 + b3 lsl 1 + b4 lsr 4
+  and o3 b4 b5 = (b4 land 0x0F) lsl 4 + b5 lsr 1
+  and o4 b5 b6 b7 = (b5 land 0x01) lsl 7 + b6 lsl 2 + b7 lsr 3
+  and o5 b7 b8 = (b7 land 0x07) lsl 5 + b8
+  in
+  let c ~off idx =
+    let r = (snd alphabet).(idx) in
+    if r = -1 then
+      Error (`Msg ("bad encoding at " ^ string_of_int off))
+    else
+      Ok r
+  in
+  let emit s_off v1 v2 v3 v4 v5 v6 v7 v8 off =
+    let* b1 = c ~off:s_off v1 in
+    let* b2 = c ~off:(s_off + 1) v2 in
+    let* b3 = c ~off:(s_off + 2) v3 in
+    let* b4 = c ~off:(s_off + 3) v4 in
+    let* b5 = c ~off:(s_off + 4) v5 in
+    let* b6 = c ~off:(s_off + 5) v6 in
+    let* b7 = c ~off:(s_off + 6) v7 in
+    let* b8 = c ~off:(s_off + 7) v8 in
+    Bytes.set_uint8 out off (o1 b1 b2);
+    Bytes.set_uint8 out (off + 1) (o2 b2 b3 b4);
+    Bytes.set_uint8 out (off + 2) (o3 b4 b5);
+    Bytes.set_uint8 out (off + 3) (o4 b5 b6 b7);
+    Bytes.set_uint8 out (off + 4) (o5 b7 b8);
+    Ok ()
+  in
+  let rec dec s_off d_off =
+    if s_off = len then
+      Ok (0, 0)
+    else
+      let v1 = Bytes.get_uint8 str s_off
+      and v2 = Bytes.get_uint8 str (s_off + 1)
+      and v3 = Bytes.get_uint8 str (s_off + 2)
+      and v4 = Bytes.get_uint8 str (s_off + 3)
+      and v5 = Bytes.get_uint8 str (s_off + 4)
+      and v6 = Bytes.get_uint8 str (s_off + 5)
+      and v7 = Bytes.get_uint8 str (s_off + 6)
+      and v8 = Bytes.get_uint8 str (s_off + 7)
+      in
+      if v3 = pad_int then
+        let* b1 = c ~off:s_off v1 in
+        let* b2 = c ~off:(s_off + 1) v2 in
+        let p1 = o1 b1 b2 in
+        Bytes.set_uint8 out d_off p1;
+        Ok (6, 4)
+      else if v5 = pad_int then
+        let* b1 = c ~off:s_off v1 in
+        let* b2 = c ~off:(s_off + 1) v2 in
+        let* b3 = c ~off:(s_off + 2) v3 in
+        let* b4 = c ~off:(s_off + 3) v4 in
+        let p1 = o1 b1 b2
+        and p2 = o2 b2 b3 b4
+        in
+        Bytes.set_uint8 out d_off p1;
+        Bytes.set_uint8 out (d_off + 1) p2;
+        Ok (4, 3)
+      else if v6 = pad_int then
+        let* b1 = c ~off:s_off v1 in
+        let* b2 = c ~off:(s_off + 1) v2 in
+        let* b3 = c ~off:(s_off + 2) v3 in
+        let* b4 = c ~off:(s_off + 3) v4 in
+        let* b5 = c ~off:(s_off + 4) v5 in
+        let p1 = o1 b1 b2
+        and p2 = o2 b2 b3 b4
+        and p3 = o3 b4 b5
+        in
+        Bytes.set_uint8 out d_off p1;
+        Bytes.set_uint8 out (d_off + 1) p2;
+        Bytes.set_uint8 out (d_off + 2) p3;
+        Ok (3, 2)
+      else if v8 = pad_int then
+        let* b1 = c ~off:s_off v1 in
+        let* b2 = c ~off:(s_off + 1) v2 in
+        let* b3 = c ~off:(s_off + 2) v3 in
+        let* b4 = c ~off:(s_off + 3) v4 in
+        let* b5 = c ~off:(s_off + 4) v5 in
+        let* b6 = c ~off:(s_off + 5) v6 in
+        let* b7 = c ~off:(s_off + 6) v7 in
+        let p1 = o1 b1 b2
+        and p2 = o2 b2 b3 b4
+        and p3 = o3 b4 b5
+        and p4 = o4 b5 b6 b7
+        in
+        Bytes.set_uint8 out d_off p1;
+        Bytes.set_uint8 out (d_off + 1) p2;
+        Bytes.set_uint8 out (d_off + 2) p3;
+        Bytes.set_uint8 out (d_off + 3) p4;
+        Ok (1, 1)
+      else
+        let* () = emit s_off v1 v2 v3 v4 v5 v6 v7 v8 d_off in
+        dec (s_off + 8) (d_off + 5)
+  in
+  let* (pad_bytes, to_remove) = dec 0 0 in
+  let rec check_pad = function
+    | 0 -> Ok ()
+    | n ->
+      if Bytes.get_uint8 str (len - n) = pad_int then
+        check_pad (n - 1)
+      else
+        Error (`Msg ("expected pad character at " ^ (string_of_int (len - n))))
+  in
+  let* () = check_pad pad_bytes in
+  let out_str = Bytes.unsafe_to_string out in
+  if to_remove > 0 then
+    Ok (String.sub out_str 0 (out_len - to_remove))
+  else
+    Ok out_str