Fix sector alignment bug in formatting

THIS DESTROYS DATA

mirage/config.ml

@@ -1,35 +1,8 @@
+(* mirage >= 4.8.0 & < 4.9.0 *)
 open Mirage
 
-let setup = runtime_arg ~pos:__POS__ "Unikernel.K.setup"
-
-let ssh_key =
-  Runtime_arg.create ~pos:__POS__
-    {|let open Cmdliner in
-      let doc = Arg.info ~doc:"The private SSH key (rsa:<seed> or ed25519:<b64-key>)." ["ssh-key"] in
-      Arg.(value & opt (some string) None doc)|}
-
-let ssh_authenticator =
-  Runtime_arg.create ~pos:__POS__
-    {|let open Cmdliner in
-      let doc = Arg.info ~doc:"SSH authenticator." ["ssh-auth"] in
-      Arg.(value & opt (some string) None doc)|}
-
-let ssh_password =
-  Runtime_arg.create ~pos:__POS__
-    {|let open Cmdliner in
-      let doc = Arg.info ~doc:"The private SSH password." [ "ssh-password" ] in
-      Arg.(value & opt (some string) None doc)|}
-
-let tls_authenticator =
-  Runtime_arg.create ~pos:__POS__
-    {|let open Cmdliner in
-      let doc = "TLS host authenticator. See git_http in lib/mirage/mirage.mli for a description of the format." in
-      let doc = Arg.info ~doc ["tls-authenticator"] in
-      Arg.(value & opt (some string) None doc)|}
-
 let mirror =
   main "Unikernel.Make"
-    ~runtime_args:[ setup ]
     ~packages:[
       package ~min:"0.3.0" ~sublibs:[ "mirage" ] "paf" ;
       package "h2" ;
@@ -41,6 +14,8 @@ let mirror =
       package ~min:"3.0.0" ~sublibs:[ "gz" ] "tar" ;
       package ~min:"3.0.0" "tar-mirage" ;
       package ~max:"0.2.0" "mirage-block-partition" ;
+      package "gpt" ;
+      package "gptar" ~pin:"git+https://github.com/reynir/gptar.git" ;
       package "oneffs" ;
       package "digestif" ;
     ]
@@ -54,9 +29,9 @@ let block = block_of_file "tar"
 
 let git_client, alpn_client =
   let git = mimic_happy_eyeballs stack he dns in
-  merge_git_clients (git_ssh ~key:ssh_key ~authenticator:ssh_authenticator ~password:ssh_password tcp git)
+  merge_git_clients (git_ssh tcp git)
     (merge_git_clients (git_tcp tcp git)
-      (git_http ~authenticator:tls_authenticator tcp git)),
+      (git_http tcp git)),
   paf_client tcp (mimic_happy_eyeballs stack he dns)
 
 let () = register "mirror"

mirage/partitions.ml

@@ -0,0 +1,197 @@
+open Lwt.Syntax
+
+module Make(BLOCK : Mirage_block.S) = struct
+  module Part = Mirage_block_partition.Make(BLOCK)
+
+  include Part
+
+  type partitions = {
+    tar : Part.t ;
+    git_dump : Part.t ;
+    md5s : Part.t ;
+    sha512s : Part.t ;
+  }
+
+  (* I just made these ones up... *)
+  let tar_guid = Uuidm.of_string "53cd6812-46cc-474e-a141-30b3aed85f53" |> Option.get
+  let cache_guid = Uuidm.of_string "22ab9cf5-6e51-45c2-998a-862e23aab264" |> Option.get
+  let git_guid = Uuidm.of_string "30faa50a-4c9d-47ff-a1a5-ecfb3401c027" |> Option.get
+
+  (* GPT uses a 72 byte utf16be encoded string for partition names *)
+  let utf16be_of_ascii s =
+    String.init 72
+      (fun i ->
+         if i mod 2 = 0 && i / 2 < String.length s then
+           s.[i/2]
+         else
+           '\000')
+
+  let read_partition_table info block =
+    let b = Cstruct.create info.Mirage_block.sector_size in
+    (* We will ignore the protective MBR at lba [0L] *)
+    let* r = BLOCK.read block 1L [b] in
+    match r with
+    | Error e ->
+      Format.kasprintf failwith "Reading partition table: %a"
+        BLOCK.pp_error e
+    | Ok () ->
+      match Gpt.unmarshal b ~sector_size:info.Mirage_block.sector_size with
+      | Error e ->
+        Format.kasprintf failwith "Reading partition table: %s" e
+      | Ok (`Read_partition_table (lba, sectors), k) ->
+        let b = Cstruct.create (sectors * info.Mirage_block.sector_size) in
+        let* r = BLOCK.read block lba [b] in
+        match r with
+        | Error e ->
+          Format.kasprintf failwith "Reading partition table: %a"
+            BLOCK.pp_error e
+        | Ok () ->
+          match k b with
+          | Error e ->
+            Format.kasprintf failwith "Reading partition table: %s" e
+          | Ok gpt -> Lwt.return gpt
+
+  let connect block =
+    let* info = BLOCK.get_info block in
+    let* gpt = read_partition_table info block in
+    let tar, git_dump, md5s, sha512s =
+      match
+        List.fold_left
+          (fun (tar, git_dump, md5s, sha512s) p ->
+             if String.equal p.Gpt.Partition.name
+                 (utf16be_of_ascii "tar")
+             then
+               (Some p, git_dump, md5s, sha512s)
+             else if String.equal p.name
+                 (utf16be_of_ascii "git_dump")
+             then
+               (tar, Some p, md5s, sha512s)
+             else if String.equal p.name
+                 (utf16be_of_ascii "md5s")
+             then
+               (tar, git_dump, Some p, sha512s)
+             else if String.equal p.name
+                 (utf16be_of_ascii "sha512s")
+             then
+               (tar, git_dump, md5s, Some p)
+             else
+               Format.kasprintf failwith "Unknown partition %S" p.name)
+          (None, None, None, None)
+          gpt.partitions
+      with
+      | (Some tar, Some git_dump, Some md5s, Some sha512s) ->
+        (tar, git_dump, md5s, sha512s)
+      | _ ->
+        failwith "not all partitions found :("
+    in
+    let+ (_empty, p) = Part.connect 0L block in
+    let get_part part =
+      let len = Int64.(succ (sub part.Gpt.Partition.ending_lba part.starting_lba)) in
+      let (_before, after) = Part.subpartition part.starting_lba p in
+      let (part, _after) = Part.subpartition len after in
+      part
+    in
+    let tar = get_part tar and git_dump = get_part git_dump
+    and md5s = get_part md5s and sha512s = get_part sha512s in
+    { tar ; git_dump ; md5s ; sha512s }
+
+  let format block ~sectors_cache ~sectors_git =
+    let* { size_sectors; sector_size; _ } = BLOCK.get_info block in
+    let ( let*? ) = Lwt_result.bind in
+    (* ocaml-gpt uses a fixed size partition entries table. Create an empty GPT
+       to figure out the first usable LBA *)
+    let empty =
+      Gpt.make ~sector_size ~disk_sectors:size_sectors []
+      |> Result.get_ok
+    in
+    let*? () =
+      if size_sectors <
+         (* protective MBR + GPT header + GPT table *)
+         let ( + ) = Int64.add in
+         empty.first_usable_lba +
+         min 1L (Int64.of_int (2 * Tar.Header.length / sector_size)) + sectors_cache + sectors_cache + sectors_git
+         + 1L (* backup GPT header *) then
+        Lwt.return_error (`Msg "too small disk")
+      else Lwt_result.return ()
+    in
+    (* Current implementation of [Gpt.Partition.make] only returns [Ok _] or
+       raises [Invalid_argument _] :/ *)
+    let attributes = 1L in
+    let sha512s =
+      Gpt.Partition.make
+        ~name:(utf16be_of_ascii "sha512s")
+        ~type_guid:cache_guid
+        ~attributes
+        Int64.(succ (sub empty.last_usable_lba sectors_cache))
+        empty.last_usable_lba
+      |> Result.get_ok
+    in
+    let md5s =
+      Gpt.Partition.make
+        ~name:(utf16be_of_ascii "md5s")
+        ~type_guid:cache_guid
+        ~attributes
+        (Int64.sub sha512s.starting_lba sectors_cache)
+        (Int64.pred sha512s.starting_lba)
+      |> Result.get_ok
+    in
+    let git_dump =
+      Gpt.Partition.make
+        ~name:(utf16be_of_ascii "git_dump")
+        ~type_guid:git_guid
+        ~attributes
+        (Int64.sub md5s.starting_lba sectors_git)
+        (Int64.pred md5s.starting_lba)
+      |> Result.get_ok
+    in
+    let tar =
+      Gpt.Partition.make
+        ~name:(utf16be_of_ascii "tar")
+        ~type_guid:tar_guid
+        ~attributes
+        empty.first_usable_lba
+        (Int64.pred git_dump.starting_lba)
+      |> Result.get_ok
+    in
+    let gpt =
+      let partitions =
+        [ tar; git_dump; md5s; sha512s ]
+      in
+      Gpt.make ~sector_size ~disk_sectors:size_sectors partitions
+      |> Result.get_ok
+    in
+  let buf =
+    Cstruct.create (sector_size * (Int64.to_int gpt.first_usable_lba + 2 * Tar.Header.length))
+  in
+    Gptar.marshal_header ~sector_size buf gpt;
+    Gpt.marshal_partition_table ~sector_size
+      (Cstruct.shift buf (sector_size * Int64.to_int gpt.partition_entry_lba))
+    gpt;
+    let write block sector_start buffers =
+      BLOCK.write block sector_start buffers
+      |> Lwt_result.map_error (fun e -> `Block e)
+    in
+    let*? () =
+      write block 0L [ buf ]
+    in
+    (* Format the file systems by writing zeroes *)
+    (* For tar we need to zero (at least) the first 2*512 bytes so we round up
+       to the nearest sector alignment *)
+    let zeroes =
+      let sectors =
+        (2 * Tar.Header.length + sector_size - 1) / sector_size * sector_size
+      in
+      Cstruct.create sectors in
+    let*? () =
+      write block tar.starting_lba [ zeroes ]
+    in
+    (* For the OneFFS filesystems we just need to zero out the first sector *)
+    let zero_sector = Cstruct.create sector_size in
+    let*? () =
+      write block git_dump.starting_lba [ zero_sector ]
+    in
+    let*? () =
+      write block md5s.starting_lba [ zero_sector ]
+    in
+    write block sha512s.starting_lba [ zero_sector ]
+end

mirage/unikernel.ml

@@ -7,7 +7,7 @@ module K = struct
 
   let check =
     let doc = Arg.info ~doc:"Only check the cache" ["check"] in
-    Arg.(value & flag doc)
+    Mirage_runtime.register_arg Arg.(value & flag doc)
 
   let verify_sha256 =
     let doc = Arg.info
@@ -15,7 +15,7 @@ module K = struct
             re-build the other checksum caches."
       ["verify-sha256"]
     in
-    Arg.(value & flag doc)
+    Mirage_runtime.register_arg Arg.(value & flag doc)
 
   let remote =
     let doc = Arg.info
@@ -23,6 +23,7 @@ module K = struct
             https://github.com/ocaml/opam-repository.git"
       ["remote"]
     in
+    Mirage_runtime.register_arg
       Arg.(value & opt string "https://github.com/ocaml/opam-repository.git#master" doc)
 
   let parallel_downloads =
@@ -30,60 +31,37 @@ module K = struct
         ~doc:"Amount of parallel HTTP downloads"
         ["parallel-downloads"]
     in
-    Arg.(value & opt int 20 doc)
+    Mirage_runtime.register_arg Arg.(value & opt int 20 doc)
 
   let hook_url =
     let doc = Arg.info
         ~doc:"URL to conduct an update of the git repository" ["hook-url"]
     in
-    Arg.(value & opt string "update" doc)
+    Mirage_runtime.register_arg Arg.(value & opt string "update" doc)
 
   let port =
     let doc = Arg.info ~doc:"HTTP listen port." ["port"] in
-    Arg.(value & opt int 80 doc)
+    Mirage_runtime.register_arg Arg.(value & opt int 80 doc)
 
   let sectors_cache =
-    let doc = "Number of sectors reserved for each checksum cache (md5, sha512)." in
+    let doc = "Number of sectors reserved for each checksum cache (md5, sha512). Only used with --initialize-disk." in
     let doc = Arg.info ~doc ["sectors-cache"] in
-    Arg.(value & opt int64 Int64.(mul 4L 2048L) doc)
+    Mirage_runtime.register_arg Arg.(value & opt int64 Int64.(mul 4L 2048L) doc)
 
   let sectors_git =
-    let doc = "Number of sectors reserved for git dump." in
+    let doc = "Number of sectors reserved for git dump. Only used with --initialize-disk" in
     let doc = Arg.info ~doc ["sectors-git"] in
-    Arg.(value & opt int64 Int64.(mul 40L (mul 2L 1024L)) doc)
+    Mirage_runtime.register_arg Arg.(value & opt int64 Int64.(mul 40L (mul 2L 1024L)) doc)
+
+  let initialize_disk =
+    let doc = "Initialize the disk with a partition table. THIS IS DESTRUCTIVE!" in
+    let doc = Arg.info ~doc ["initialize-disk"] in
+    Mirage_runtime.register_arg Arg.(value & flag doc)
   
   let ignore_local_git =
     let doc = "Ignore restoring locally saved git repository." in
     let doc = Arg.info ~doc ["ignore-local-git"] in
-    Arg.(value & flag doc)
-
-  type t =
-    { check : bool
-    ; verify_sha256 : bool
-    ; remote : string
-    ; parallel_downloads : int
-    ; hook_url : string
-    ; port : int
-    ; sectors_cache : int64
-    ; sectors_git : int64
-    ; ignore_local_git : bool }
-
-  let v check verify_sha256 remote parallel_downloads hook_url port
-    sectors_cache sectors_git ignore_local_git =
-    { check; verify_sha256; remote; parallel_downloads; hook_url; port
-    ; sectors_cache; sectors_git; ignore_local_git }
-
-  let setup =
-    Term.(const v
-          $ check
-          $ verify_sha256
-          $ remote
-          $ parallel_downloads
-          $ hook_url
-          $ port
-          $ sectors_cache
-          $ sectors_git
-          $ ignore_local_git)
+    Mirage_runtime.register_arg Arg.(value & flag doc)
 end
 
 module Make
@@ -94,7 +72,7 @@ module Make
   (_ : sig end)
   (HTTP : Http_mirage_client.S) = struct
 
-  module Part = Mirage_block_partition.Make(BLOCK)
+  module Part = Partitions.Make(BLOCK)
   module KV = Tar_mirage.Make_KV_RW(Pclock)(Part)
   module Cache = OneFFS.Make(Part)
   module Store = Git_kv.Make(Pclock)
@@ -715,7 +693,7 @@ module Make
           and size = String.length data in
           let hdr = Tar.Header.make ~file_mode ~mod_time ~user_id ~group_id
             (Mirage_kv.Key.to_string path) (Int64.of_int size) in
-          Some (None, hdr, once data)
+          Some (Some Tar.Header.Ustar, hdr, once data)
         | Error _ -> None in
       let entries = Lwt_stream.filter_map_s to_entry entries in
       Lwt.return begin fun () -> Tar.High (High.inj (Lwt_stream.get entries >|= Result.ok)) end
@@ -724,7 +702,7 @@ module Make
       let now = Ptime.v (Pclock.now_d_ps ()) in
       let mtime = Option.value ~default:0 Ptime.(Span.to_int_s (to_span now)) in
       entries_of_git ~mtime store repo >>= fun entries ->
-      let t = Tar.out entries in
+      let t = Tar.out ~level:Ustar entries in
       let t = Tar_gz.out_gzipped ~level:4 ~mtime:(Int32.of_int mtime) Gz.Unix t in
       let buf = Buffer.create 1024 in
       to_buffer buf t >|= function
@@ -1011,29 +989,20 @@ stamp: %S
 
   module Paf = Paf_mirage.Make(Stack.TCP)
 
-  let start block _time _pclock stack git_ctx http_ctx
-    { K.check; verify_sha256; remote; parallel_downloads; hook_url
-    ; port; sectors_cache; sectors_git; ignore_local_git } =
-    BLOCK.get_info block >>= fun info ->
-    let git_start =
-      let cache_size = Int64.(mul 2L sectors_cache) in
-      Int64.(sub info.size_sectors (add cache_size sectors_git))
-    in
-    Part.connect git_start block >>= fun (kv, rest) ->
-    let git_dump, rest = Part.subpartition sectors_git rest in
-    let md5s, sha512s = Part.subpartition sectors_cache rest in
-    KV.connect kv >>= fun kv ->
+  let start_mirror { Part.tar; git_dump; md5s; sha512s } stack git_ctx http_ctx =
+    KV.connect tar >>= fun kv ->
+    Cache.connect git_dump >>= fun git_dump ->
     Cache.connect md5s >>= fun md5s ->
     Cache.connect sha512s >>= fun sha512s ->
-    Cache.connect git_dump >>= fun git_dump ->
     Logs.info (fun m -> m "Available bytes in tar storage: %Ld" (KV.free kv));
-    Disk.init ~verify_sha256 kv md5s sha512s >>= fun disk ->
-    if check then
+    Disk.init ~verify_sha256:(K.verify_sha256 ()) kv md5s sha512s >>= fun disk ->
+    let remote = K.remote () in
+    if K.check () then
       Lwt.return_unit
     else
       begin
         Logs.info (fun m -> m "Initializing git state. This may take a while...");
-        (if ignore_local_git then
+        (if K.ignore_local_git () then
            Lwt.return (Error ())
          else
            restore_git ~remote git_dump git_ctx) >>= function
@@ -1047,21 +1016,21 @@ stamp: %S
       Serve.commit_id git_kv >>= fun commit_id ->
       Logs.info (fun m -> m "git: %s" commit_id);
       Serve.create remote git_kv >>= fun serve ->
-      Paf.init ~port (Stack.tcp stack) >>= fun t ->
+      Paf.init ~port:(K.port ()) (Stack.tcp stack) >>= fun t ->
       let update () =
         Serve.update_git ~remote serve git_kv >>= function
         | None | Some [] -> Lwt.return_unit
         | Some _changes ->
           dump_git git_dump git_kv >>= fun () ->
-          download_archives parallel_downloads disk http_ctx git_kv
+          download_archives (K.parallel_downloads ()) disk http_ctx git_kv
       in
       let service =
         Paf.http_service
           ~error_handler:(fun _ ?request:_ _ _ -> ())
-          (Serve.dispatch serve disk hook_url update)
+          (Serve.dispatch serve disk (K.hook_url ()) update)
       in
       let `Initialized th = Paf.serve service t in
-      Logs.info (fun f -> f "listening on %d/HTTP" port);
+      Logs.info (fun f -> f "listening on %d/HTTP" (K.port ()));
       Lwt.async (fun () ->
           let rec go () =
             Time.sleep_ns (Duration.of_hour 1) >>= fun () ->
@@ -1069,6 +1038,23 @@ stamp: %S
             go ()
           in
           go ());
-      download_archives parallel_downloads disk http_ctx git_kv >>= fun () ->
+      download_archives (K.parallel_downloads ()) disk http_ctx git_kv >>= fun () ->
       (th >|= fun _v -> ())
+
+  let start block _time _pclock stack git_ctx http_ctx =
+    let initialize_disk = K.initialize_disk ()
+    and sectors_cache = K.sectors_cache ()
+    and sectors_git = K.sectors_git () in
+    if initialize_disk then
+      Part.format block ~sectors_cache ~sectors_git >>= function
+      | Ok () -> Lwt.return_unit
+      | Error `Msg e ->
+        Logs.err (fun m -> m "Error formatting disk: %s" e);
+        exit Mirage_runtime.argument_error
+      | Error `Block e ->
+        Logs.err (fun m -> m "Error formatting disk: %a" BLOCK.pp_write_error e);
+        exit 2
+    else
+      Part.connect block >>= fun parts ->
+      start_mirror parts stack git_ctx http_ctx
 end