Add code for formatting a disk

mirage/partitions.ml

@@ -12,6 +12,20 @@ module Make(BLOCK : Mirage_block.S) = struct
     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] *)
@@ -38,14 +52,6 @@ module Make(BLOCK : Mirage_block.S) = struct
           | Ok gpt -> Lwt.return gpt
 
   let connect block =
-    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')
-    in
     let* info = BLOCK.get_info block in
     let* gpt = read_partition_table info block in
     let tar, git_dump, md5s, sha512s =
@@ -88,4 +94,79 @@ module Make(BLOCK : Mirage_block.S) = struct
     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
+    let gpt =
+      let partitions =
+        (* 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
+        [ 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;
+    BLOCK.write block 0L [ buf ]
+    |> Lwt_result.map_error (fun e -> `Block e)
 end

mirage/unikernel.ml

@@ -44,15 +44,20 @@ module K = struct
     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
     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
     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
@@ -984,8 +989,7 @@ stamp: %S
 
   module Paf = Paf_mirage.Make(Stack.TCP)
 
-  let start block _time _pclock stack git_ctx http_ctx =
-    Part.connect block >>= fun { Part.tar ; git_dump; md5s ; sha512s } ->
+  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 ->
@@ -1036,4 +1040,21 @@ stamp: %S
           go ());
       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