Repository: mirage/ocaml-tar
Pull request #105 (2022-12-15 17:36:34)

Implement allocate and set_partial


Patch:
diff --git a/mirage/tar_mirage.ml b/mirage/tar_mirage.ml
index a49dd7e..0ffa9a5 100644
--- a/mirage/tar_mirage.ml
+++ b/mirage/tar_mirage.ml
@@ -43,6 +43,10 @@ module Make_KV_RO (BLOCK : Mirage_block.S) = struct
     | #Mirage_kv.error as e -> Mirage_kv.pp_error ppf e
     | `Block b -> BLOCK.pp_error ppf b
 
+  let read t sector_start buffers =
+    Lwt_result.map_error (fun e -> `Block e)
+      (BLOCK.read t.b sector_start buffers)
+
   let get_node t key =
     let rec find e = function
       | [] -> Ok e
@@ -108,6 +112,20 @@ module Make_KV_RO (BLOCK : Mirage_block.S) = struct
   end
   module HR = Tar.HeaderReader(Lwt)(Reader)
 
+  (* [read_partial_sector t sector_start ~offset ~length dst]
+     reads a single sector and blits [length] bytes from [offset] into [dst]
+     with the same offset. *)
+  let read_partial_sector t sector_start ~offset ~length dst =
+    assert Int64.(add offset length <= of_int t.info.sector_size);
+    let length = Int64.to_int length and offset = Int64.to_int offset in
+    assert (Cstruct.length dst >= t.info.sector_size);
+    if length = 0 then Lwt_result.return () else
+    let ( >>>= ) = Lwt_result.bind in
+    let src = Cstruct.create t.info.sector_size in
+    read t sector_start [ src ] >>>= fun () ->
+    Cstruct.blit src offset dst offset length;
+    Lwt_result.return ()
+
   let get_partial t key ~offset ~length =
     match get_node t.map key with
     | Error e -> Lwt.return (Error e)
@@ -141,8 +159,8 @@ module Make_KV_RO (BLOCK : Mirage_block.S) = struct
           List.init (to_int n_sectors)
             (fun sec -> Cstruct.sub buf (sec * to_int sector_size) (to_int sector_size))
         in
-        BLOCK.read t.b start_sector tmps >|= function
-        | Error b -> Error (`Block b)
+        read t start_sector tmps >|= function
+        | Error _ as e -> e
         | Ok () ->
           let buf =
             Cstruct.sub buf (to_int start_padding) (to_int length_bytes)
@@ -213,6 +231,24 @@ module Make_KV_RO (BLOCK : Mirage_block.S) = struct
     in
     go map (Mirage_kv.Key.segments key)
 
+  let remove map key =
+    let rec go m = function
+      | [] -> assert false
+      | [hd] -> StringMap.remove hd m
+      | hd::tl ->
+        let hdr, m' = match StringMap.find_opt hd m with
+          | None -> Tar.Header.make hd 0L, StringMap.empty
+          | Some (Value _) -> assert false
+          | Some (Dict (hdr, m)) -> hdr, m
+        in
+        let m'' = go m' tl in
+        if StringMap.is_empty m'' then
+          StringMap.remove hd m
+        else
+          StringMap.add hd (Dict (hdr, m'')) m
+    in
+    go map (Mirage_kv.Key.segments key)
+
   let connect b =
     BLOCK.get_info b >>= fun info ->
     let ssize = info.Mirage_block.sector_size in
@@ -253,6 +289,21 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
 
   include Make_KV_RO(BLOCK)
 
+  type write_error = [ `Block of BLOCK.error | `Block_write of BLOCK.write_error | Mirage_kv.write_error | `Entry_already_exists | `Path_segment_is_a_value | `Append_only ]
+
+  let pp_write_error ppf = function
+   | `Block e -> Fmt.pf ppf "read error while writing: %a" BLOCK.pp_error e
+   | `Block_write e -> BLOCK.pp_write_error ppf e
+   | #Mirage_kv.write_error as e -> Mirage_kv.pp_write_error ppf e
+   | `Entry_already_exists -> Fmt.string ppf "entry already exists"
+   | `Path_segment_is_a_value -> Fmt.string ppf "path segment is a value"
+   | `Append_only -> Fmt.string ppf "append only"
+   | `Write_header msg -> Fmt.pf ppf "writing tar header failed: %s" msg
+
+  let write t sector_start buffers =
+    Lwt_result.map_error (fun e -> `Block_write e)
+      (BLOCK.write t.b sector_start buffers)
+
   let free t =
     Int64.(sub (mul (of_int t.info.sector_size) t.info.size_sectors)
              t.end_of_archive)
@@ -271,10 +322,13 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
     in
     find t.map (Mirage_kv.Key.segments key)
 
-  let header_of_key key len =
+  let header_of_key ?last_modified key len =
     let mod_time =
-      let ptime = Ptime.v (CLOCK.now_d_ps ()) in
-      Int64.of_float (Ptime.to_float_s ptime)
+      match last_modified with
+      | Some mod_time -> Int64.of_float (Ptime.to_float_s mod_time)
+      | None ->
+        let ptime = Ptime.v (CLOCK.now_d_ps ()) in
+        Int64.of_float (Ptime.to_float_s ptime)
     in
     Tar.Header.make ~mod_time (Mirage_kv.Key.to_string key) (Int64.of_int len)
 
@@ -296,6 +350,17 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
       let map = insert map key (Value (hdr, offset)) in
       Dict (root, map)
 
+  let update_remove map key =
+    match map with
+    | Value _ ->
+      (* if the root is a value we have done something very wrong. This should
+         be catched by [is_safe_to_set]. *)
+      assert false
+    | Dict (root, map) ->
+      (* [remove] may raise if [key] is [empty]. *)
+      let map = remove map key in
+      Dict (root, map)
+
   module Writer = struct
     type out_channel = {
       b: BLOCK.t;
@@ -324,16 +389,17 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
   end
   module HW = Tar.HeaderWriter(Lwt)(Writer)
 
-  type write_error = [ `Block of BLOCK.error | `Block_write of BLOCK.write_error | Mirage_kv.write_error | `Entry_already_exists | `Path_segment_is_a_value | `Append_only ]
-
-  let pp_write_error ppf = function
-   | `Block e -> Fmt.pf ppf "read error while writing: %a" BLOCK.pp_error e
-   | `Block_write e -> BLOCK.pp_write_error ppf e
-   | #Mirage_kv.write_error as e -> Mirage_kv.pp_write_error ppf e
-   | `Entry_already_exists -> Fmt.string ppf "entry already exists"
-   | `Path_segment_is_a_value -> Fmt.string ppf "path segment is a value"
-   | `Append_only -> Fmt.string ppf "append only"
-   | `Write_header msg -> Fmt.pf ppf "writing tar header failed: %s" msg
+  let write_header (t : t) header_start_bytes hdr =
+    let hw = Writer.{ b = t.b ; offset = header_start_bytes ; info = t.info } in
+    (* it is important we write at level [Ustar] at most as we assume the
+       header(s) taking up exactly 512 bytes. With [GNU] level extra blocks
+       may be used for long names. *)
+    Lwt.catch
+      (fun () -> HW.write ~level:Tar.Header.Ustar hdr hw >|= fun () -> Ok ())
+      (function
+        | Writer.Read e -> Lwt.return (Error (`Block e))
+        | Writer.Write e -> Lwt.return (Error (`Block_write e))
+        | exn -> raise exn)
 
   let set t key data =
     Lwt_mutex.with_lock t.write_lock (fun () ->
@@ -344,13 +410,17 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
           let* () = is_safe_to_set t key in
           let hdr = header_of_key key (Cstruct.length data) in
           let space_needed = space_needed hdr in
-          let* () = if free t >= space_needed then Ok () else Error `No_space in
+          let* () =
+            if free t >= Int64.(add space_needed (of_int Tar.Header.length)) then
+              Ok ()
+            else
+              Error `No_space
+          in
           Ok (hdr, space_needed)
         in
         Lwt.return r >>>= fun (hdr, space_needed) ->
         let open Int64 in
         let sector_size = of_int t.info.Mirage_block.sector_size in
-
         let data_start_bytes = sub t.end_of_archive (of_int Tar.Header.length) in
         let header_start_bytes = sub data_start_bytes (of_int Tar.Header.length) in
         let sentinel = mul 2L (of_int Tar.Header.length) in
@@ -360,33 +430,42 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
           div data_start_bytes sector_size,
           rem data_start_bytes sector_size
         in
-        let end_sector = div end_bytes sector_size in
-        let end_sector_offset = rem end_bytes sector_size in
+        let end_sector = div (add end_bytes sector_size) sector_size in
+        let last_sector_offset = rem end_bytes sector_size in
         let pad = Tar.Header.compute_zero_padding_length hdr in
 
-        (* Read slack for last sector *)
-        begin
-          if end_sector_offset = 0L then
-            Lwt.return_ok Cstruct.empty
-          else begin
-            let buf = Cstruct.create (to_int sector_size) in
-            BLOCK.read t.b end_sector [buf] >|= function
-            | Error e -> Error (`Block e)
-            | Ok () -> Ok (snd (Cstruct.split buf (to_int end_sector_offset)))
-          end
-        end >>>= fun slack ->
-        let data = Cstruct.concat [ data; Cstruct.create (pad + to_int sentinel); slack ] in
-        let first_sector, remaining_sectors =
-          let s =
-            Stdlib.min
-              (Cstruct.length data)
-              (to_int (sub sector_size data_start_sector_offset))
+        let data =
+          let slack =
+            if last_sector_offset = 0L then
+              0
+            else
+              to_int (sub sector_size last_sector_offset)
           in
-          Cstruct.split data s
+          Cstruct.concat [
+            Cstruct.create (to_int data_start_sector_offset);
+            data;
+            Cstruct.create (pad + to_int sentinel + slack);
+          ]
+        in
+        (* [data] is always at least one sector as the sentinel is always present *)
+        let first_sector, remaining_sectors = Cstruct.split data t.info.sector_size in
+        let last_sector =
+          (* sub on whole [data] as the first sector and last sector might be the same *)
+          Cstruct.sub data
+            (Cstruct.length data - t.info.sector_size)
+            t.info.sector_size
         in
+        (* blit in slack at the end if needed *)
+        begin if last_sector_offset = 0L then Lwt_result.return () else
+          read_partial_sector t (pred end_sector) last_sector
+            ~offset:last_sector_offset
+            ~length:(sub sector_size last_sector_offset)
+        end >>>= fun () ->
         (* to write robustly as we can:
-           - we write tar blocks 2..end,
-           - then the header AND the first tar block
+           - we write sectors 2..n,
+           - then the header,
+           - then we blit the first (data) sector as it may contain the header,
+           - finally we write the first (data) sector which contains the first tar data block.
         *)
         let remaining_sectors =
           (* XXX: this is to work around limitations in some block implementations *)
@@ -396,27 +475,14 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
                  (sector * to_int sector_size)
                  (to_int sector_size))
         in
-        Lwt_result.map_error (fun e -> `Block_write e)
-          (BLOCK.write t.b (succ data_start_sector) remaining_sectors) >>>= fun () ->
+        write t (succ data_start_sector) remaining_sectors >>>= fun () ->
         (* finally write header and first block *)
-        let hw = Writer.{ b = t.b ; offset = header_start_bytes ; info = t.info } in
-        (* it is important we write at level [Ustar] at most as we assume the
-           header(s) taking up exactly 512 bytes. With [GNU] level extra blocks
-           may be used for long names. *)
-        Lwt.catch
-          (fun () -> HW.write ~level:Tar.Header.Ustar hdr hw >|= fun () -> Ok ())
-          (function
-            | Writer.Read e -> Lwt.return (Error (`Block e))
-            | Writer.Write e -> Lwt.return (Error (`Block_write e))
-            | exn -> raise exn) >>>= fun () ->
-        let buf = Cstruct.create (to_int sector_size) in
-        Lwt_result.map_error (function e -> `Block e)
-          (BLOCK.read t.b data_start_sector [ buf ]) >>>= fun () ->
-        Cstruct.blit first_sector 0 buf (to_int data_start_sector_offset)
-          (Cstruct.length first_sector);
-        Lwt_result.map_error (fun e -> `Block_write e)
-          (BLOCK.write t.b data_start_sector [ buf ]) >>>= fun () ->
-        let tar_offset = Int64.div (sub t.end_of_archive 512L) 512L in
+        write_header t header_start_bytes hdr >>>= fun () ->
+        (* read in slack at beginning which could include the header *)
+        read_partial_sector t data_start_sector first_sector
+          ~offset:0L ~length:data_start_sector_offset >>>= fun () ->
+        write t data_start_sector [ first_sector ] >>>= fun () ->
+        let tar_offset = Int64.div data_start_bytes (of_int Tar.Header.length) in
         t.end_of_archive <- end_bytes;
         t.map <- update_insert t.map key hdr tar_offset;
         Lwt.return (Ok ()))
@@ -424,12 +490,144 @@ module Make_KV_RW (CLOCK : Mirage_clock.PCLOCK) (BLOCK : Mirage_block.S) = struc
   let remove _ _ =
     Lwt.return (Error `Append_only)
 
-  let rename _ ~source:_ ~dest:_ =
-    Lwt.return (Error `Append_only)
-
-  let set_partial _ _ ~offset:_ _ =
-    Lwt.return (Error `Append_only)
+  let rename t ~source ~dest =
+    let ( >>>= ) = Lwt_result.bind in
+    Lwt_mutex.with_lock t.write_lock (fun () ->
+        Lwt.return (is_safe_to_set t dest) >>>= fun () ->
+        Lwt.return begin match get_node t.map source with
+          | Ok Value (hdr, data_offset) -> Ok (hdr, data_offset)
+          | Ok Dict _ -> Error `Append_only
+          | Error _ as e -> e
+        end >>>= fun (hdr, data_offset) ->
+        let hdr = { hdr with Tar.Header.file_name = Mirage_kv.Key.to_string dest } in
+        write_header t Int64.(sub (mul data_offset (of_int Tar.Header.length)) (of_int Tar.Header.length)) hdr >>>= fun () ->
+        t.map <- update_insert t.map dest hdr data_offset;
+        t.map <- update_remove t.map source;
+        Lwt_result.return ())
+
+  let set_partial t key ~offset data =
+    let ( >>>= ) = Lwt_result.bind in
+    Lwt_mutex.with_lock t.write_lock (fun () ->
+        if Optint.Int63.(compare offset zero < 0) then
+          invalid_arg "Tar_mirage.set_partial: negative offset";
+        Lwt.return begin match get_node t.map key with
+          | Ok Value (hdr, data_offset) ->
+            Ok (hdr, Int64.(mul data_offset (of_int Tar.Header.length)))
+          | Ok Dict _ -> Error `Path_segment_is_a_value
+          | Error _ as e -> e
+        end >>>= fun (hdr, data_offset) ->
+        let offset = Optint.Int63.to_int64 offset in
+        let open Int64 in
+        let end_bytes = add offset (of_int (String.length data)) in
+        begin if end_bytes > hdr.file_size then
+            Lwt_result.fail `Append_only
+          else
+            Lwt_result.return ()
+        end >>>= fun () ->
+        (* compute the offsets into the archive *)
+        let end_bytes = add data_offset end_bytes in
+        let start_bytes = add data_offset offset in
+        let sector_size = of_int t.info.sector_size in
+        let start_sector_offset = rem start_bytes sector_size in
+        let end_sector = div (add end_bytes sector_size) sector_size in
+        let last_sector_offset = rem end_bytes sector_size in
+        (* allocate a buffer for what we need to write, and blit in data and slack
+           at first and last sector. *)
+        let data' =
+          let len =
+            add start_sector_offset
+              (add (of_int (String.length data))
+                 (if last_sector_offset = 0L then
+                    0L
+                  else
+                    sub sector_size last_sector_offset))
+          in
+          Cstruct.create (to_int len)
+        in
+        Cstruct.blit_from_string data 0 data'
+          (to_int start_sector_offset) (String.length data);
+        read_partial_sector t (div start_bytes sector_size) data'
+          ~offset:0L ~length:start_sector_offset >>>= fun () ->
+        let last_sector =
+          Cstruct.sub data' (Cstruct.length data' - t.info.sector_size)
+            t.info.sector_size
+        in
+        begin if last_sector_offset = 0L then Lwt_result.return () else
+          read_partial_sector t (pred end_sector) last_sector
+            ~offset:last_sector_offset
+            ~length:(sub sector_size last_sector_offset)
+        end >>>= fun () ->
+        (* XXX: this is to work around limitations in some block implementations *)
+        let data' =
+          List.init (Cstruct.length data' / t.info.sector_size)
+            (fun sector ->
+               Cstruct.sub data' (sector * t.info.sector_size) t.info.sector_size)
+        in
+        write t (div start_bytes sector_size) data')
 
-  let allocate _ _ ?last_modified:_ _ =
-    Lwt.return (Error `Append_only)
+  let allocate t key ?last_modified size =
+    Lwt_mutex.with_lock t.write_lock (fun () ->
+        let ( >>>= ) = Lwt_result.bind in
+        let r =
+          let ( let* ) = Result.bind in
+          (* XXX: map `Entry_already_exists to `Append_only ?! *)
+          let* () = is_safe_to_set t key in
+          let hdr = header_of_key ?last_modified key (Optint.Int63.to_int size) in
+          let space_needed = space_needed hdr in
+          let* () =
+            if free t >= Int64.(add space_needed (of_int Tar.Header.length)) then
+              Ok ()
+            else
+              Error `No_space
+          in
+          Ok (hdr, space_needed)
+        in
+        Lwt.return r >>>= fun (hdr, space_needed) ->
+        let open Int64 in
+        let sector_size = of_int t.info.Mirage_block.sector_size in
+        let header_start_bytes =
+          sub t.end_of_archive (of_int (2 * Tar.Header.length))
+        in
+        (* we don't have to zero the block before end_of_archive as it is
+           already zero due to the terminating double zero block (sentinel) *)
+        let to_zero_start_bytes = t.end_of_archive in
+        (* space_needed - 1 block + sentinel = space_needed + 1 block *)
+        let end_bytes = add t.end_of_archive
+            (add space_needed (of_int Tar.Header.length)) in
+        (* Compute the starting sector and ending sector *)
+        let to_zero_start_sector, to_zero_start_sector_offset =
+          div to_zero_start_bytes sector_size,
+          rem to_zero_start_bytes sector_size
+        in
+        let last_sector_offset = rem end_bytes sector_size in
+        let end_sector = div (add end_bytes sector_size) sector_size in
+        (* [num_to_zero_sectors] is at least 1 as we need to write at least one
+           zero block of the new sentinel. *)
+        let num_to_zero_sectors = to_int (sub end_sector to_zero_start_sector) in
+        let zero_sector = Cstruct.create t.info.Mirage_block.sector_size in
+        let data = Array.init num_to_zero_sectors (fun _ -> zero_sector) in
+        let nonzero_sector c =
+          (* we allocate a new buffer if [c] is [zero_sector], otherwise we can
+             reuse it in the case first and last sectors are the same. *)
+          if c != zero_sector then c else Cstruct.create (Cstruct.length c)
+        in
+        (* Read slack at start and end sector(s) *)
+        let () = data.(0) <- nonzero_sector data.(0) in
+        read_partial_sector t to_zero_start_sector data.(0)
+          ~offset:0L
+          ~length:to_zero_start_sector_offset
+        >>>= fun () ->
+        let last = nonzero_sector data.(num_to_zero_sectors - 1) in
+        let () = data.(num_to_zero_sectors - 1 ) <- last in
+        begin if last_sector_offset = 0L then Lwt_result.return () else
+          read_partial_sector t (pred end_sector) last
+            ~offset:last_sector_offset
+            ~length:(sub sector_size last_sector_offset)
+        end >>>= fun () ->
+        write t to_zero_start_sector (Array.to_list data) >>>= fun () ->
+        write_header t header_start_bytes hdr >>>= fun () ->
+        let tar_offset = div (sub t.end_of_archive (of_int Tar.Header.length)) (of_int Tar.Header.length) in
+        t.end_of_archive <- end_bytes;
+        t.map <- update_insert t.map key hdr tar_offset;
+        Lwt.return (Ok ()))
 end
diff --git a/tar-mirage.opam b/tar-mirage.opam
index aa69967..508c837 100644
--- a/tar-mirage.opam
+++ b/tar-mirage.opam
@@ -24,7 +24,7 @@ depends: [
   "optint"
   "ptime"
   "tar" {= version}
-  "mirage-block-unix" {with-test & >= "2.5.0"}
+  "mirage-block-unix" {with-test & >= "2.13.0"}
   "mirage-clock-unix" {with-test}
   "ounit2" {with-test}
   "ounit2-lwt" {with-test}


Test patch:
diff --git a/lib_test/allocate_set_partial_test.ml b/lib_test/allocate_set_partial_test.ml
new file mode 100644
index 0000000..4a707b6
--- /dev/null
+++ b/lib_test/allocate_set_partial_test.ml
@@ -0,0 +1,163 @@
+open OUnit2
+open Lwt.Infix
+
+module Int63 = Optint.Int63
+
+module B(Size : sig val sector_size : int end) = struct
+  include Block
+  include Size
+
+  let convert_path os path =
+    let ch = Unix.open_process_in (Printf.sprintf "cygpath -%c -- %s" (match os with `Mixed -> 'm' | `Unix -> 'u' | `Windows -> 'w') path) in
+    let line = input_line ch in
+    close_in ch;
+    line
+
+  let connect name =
+    let name = if Sys.win32 then convert_path `Windows name else name in
+    connect ~prefered_sector_size:(Some sector_size) name
+end
+
+module Block512 = B(struct let sector_size = 512 end)
+
+module Block4096 = B(struct let sector_size = 4096 end)
+
+module type BLOCK = sig
+  include module type of Block
+  val connect : string -> t Lwt.t
+  val sector_size : int
+end
+
+let int63_hdr_len = Int63.of_int Tar.Header.length
+
+module Test(B : BLOCK) = struct
+  module KV_RW = Tar_mirage.Make_KV_RW(Pclock)(B)
+
+  let kv_rw_error =
+    Lwt.wrap1 (Result.fold ~ok:Fun.id ~error:(Fmt.kstr failwith "%a" KV_RW.pp_error))
+
+  let kv_rw_write_error =
+    Lwt.wrap1 (Result.fold ~ok:Fun.id ~error:(Fmt.kstr failwith "%a" KV_RW.pp_write_error))
+
+  let connect_block test_ctxt =
+    let filename, ch = bracket_tmpfile ~prefix:"tar-allocate-set-partial-test" ~suffix:".tar" test_ctxt in
+    close_out ch;
+    B.connect filename
+
+  let resize b size =
+    B.resize b size >|=
+    Result.iter_error (Fmt.kstr failwith "%a" B.pp_write_error)
+
+  let allocate_empty_file test_ctxt =
+    connect_block test_ctxt >>= fun b ->
+    resize b 10240L >>= fun () ->
+    KV_RW.connect b >>= fun t ->
+    KV_RW.allocate t (Mirage_kv.Key.v "empty") Int63.zero >>=
+    kv_rw_write_error
+
+  let set_partial_no_file test_ctxt =
+    connect_block test_ctxt >>= fun b ->
+    resize b 10240L >>= fun () ->
+    KV_RW.connect b >>= fun t ->
+    KV_RW.set_partial t (Mirage_kv.Key.v "nonexistent")
+      ~offset:Int63.zero "" >>= function
+    | Ok () -> assert_failure "expected set_partial on nonexistent file to fail"
+    | Error _ -> Lwt.return_unit
+
+  let allocate_is_zeroed test_ctxt =
+    connect_block test_ctxt >>= fun b ->
+    resize b 10240L >>= fun () ->
+    KV_RW.connect b >>= fun t ->
+    let key = Mirage_kv.Key.v "zeroed" in
+    KV_RW.allocate t key int63_hdr_len >>=
+    kv_rw_write_error >>= fun () ->
+    KV_RW.get t key >>=
+    kv_rw_error >|=
+    assert_equal ~cmp:String.equal ~printer:String.escaped
+      (String.make Tar.Header.length '\000')
+
+  let allocate_two_one_byte_files_zeroed test_ctxt =
+    connect_block test_ctxt >>= fun b ->
+    resize b 10240L >>= fun () ->
+    KV_RW.connect b >>= fun t ->
+    let k1 = Mirage_kv.Key.v "zeroed" and k2 = Mirage_kv.Key.v "zeroed2" in
+    KV_RW.allocate t k1 Int63.one >>=
+    kv_rw_write_error >>= fun () ->
+    KV_RW.allocate t k2 Int63.one >>=
+    kv_rw_write_error >>= fun () ->
+    KV_RW.get t k1 >>=
+    kv_rw_error >>= fun s ->
+    assert_equal ~cmp:String.equal ~printer:String.escaped
+      "\000" s;
+    KV_RW.get t k2 >>=
+    kv_rw_error >|=
+    assert_equal ~cmp:String.equal ~printer:String.escaped
+      "\000"
+
+  let allocate_set_partial_first_byte test_ctxt =
+    connect_block test_ctxt >>= fun b ->
+    resize b 10240L >>= fun () ->
+    KV_RW.connect b >>= fun t ->
+    let k = Mirage_kv.Key.v "X" in
+    KV_RW.allocate t k Int63.(add one one) >>=
+    kv_rw_write_error >>= fun () ->
+    KV_RW.set_partial t k ~offset:Int63.zero "X" >>=
+    kv_rw_write_error >>= fun () ->
+    KV_RW.get t k >>=
+    kv_rw_error >|=
+    assert_equal ~cmp:String.equal ~printer:String.escaped
+      "X\000"
+
+  let rename_nonexistent_file test_ctxt =
+    connect_block test_ctxt >>= fun b ->
+    resize b 10240L >>= fun () ->
+    KV_RW.connect b >>= fun t ->
+    let source = Mirage_kv.Key.v "source"
+    and dest = Mirage_kv.Key.v "destination" in
+    KV_RW.rename t ~source ~dest >|= function
+    | Ok () -> assert_failure "Expected failure to rename nonexistent file"
+    | Error _ -> ()
+
+  let set_rename test_ctxt =
+    connect_block test_ctxt >>= fun b ->
+    resize b 10240L >>= fun () ->
+    KV_RW.connect b >>= fun t ->
+    let source = Mirage_kv.Key.v "source"
+    and dest = Mirage_kv.Key.v "destination" in
+    let s =
+      String.init (3 * Tar.Header.length)
+        (fun i -> "0123456789".[i mod 10])
+    in
+    KV_RW.set t source s >>=
+    kv_rw_write_error >>= fun () ->
+    KV_RW.rename t ~source ~dest >>=
+    kv_rw_write_error >>= fun () ->
+    KV_RW.get t dest >>=
+    kv_rw_error >|= fun s' ->
+    assert_equal ~cmp:String.equal ~printer:String.escaped
+      s s'
+end
+
+module Test512 = Test(Block512)
+module Test4096 = Test(Block4096)
+    
+
+let () =
+  let suite =
+    "tar-allocate-set-partial" >:::
+    [
+      "allocate empty file 512" >:: OUnitLwt.lwt_wrapper Test512.allocate_empty_file;
+      "allocate empty file 4096" >:: OUnitLwt.lwt_wrapper Test4096.allocate_empty_file;
+      "set_partial nonexistent file 512" >:: OUnitLwt.lwt_wrapper Test512.set_partial_no_file;
+      "set_partial nonexistent file 4096" >:: OUnitLwt.lwt_wrapper Test4096.set_partial_no_file;
+      "allocate is zeroed 512" >:: OUnitLwt.lwt_wrapper Test512.allocate_is_zeroed;
+      "allocate is zeroed 4096" >:: OUnitLwt.lwt_wrapper Test4096.allocate_is_zeroed;
+      "allocate two one-byte files 512" >:: OUnitLwt.lwt_wrapper Test512.allocate_two_one_byte_files_zeroed;
+      "allocate two one-byte files 4096" >:: OUnitLwt.lwt_wrapper Test4096.allocate_two_one_byte_files_zeroed;
+      "allocate and set first byte 512" >:: OUnitLwt.lwt_wrapper Test512.allocate_set_partial_first_byte;
+      "allocate and set first byte 4096" >:: OUnitLwt.lwt_wrapper Test4096.allocate_set_partial_first_byte;
+      "set and rename 512" >:: OUnitLwt.lwt_wrapper Test512.set_rename;
+      "set and rename 4096" >:: OUnitLwt.lwt_wrapper Test4096.set_rename;
+    ]
+  in
+  run_test_tt_main suite
diff --git a/lib_test/dune b/lib_test/dune
index 3f1fb0e..8f72855 100644
--- a/lib_test/dune
+++ b/lib_test/dune
@@ -1,5 +1,5 @@
 (tests
- (names parse_test write_test)
+ (names parse_test write_test allocate_set_partial_test)
  (package tar-mirage)
  (libraries
   mirage-block-unix
diff --git a/lib_test/parse_test.ml b/lib_test/parse_test.ml
index 0006422..17a388c 100644
--- a/lib_test/parse_test.ml
+++ b/lib_test/parse_test.ml
@@ -52,7 +52,7 @@ let header _test_ctxt =
 
 let set_difference a b = List.filter (fun a -> not(List.mem a b)) a
 
-let with_tar ?(level:Tar.Header.compatibility option) ?files ?(block_size = 512) test_ctxt f =
+let with_tar ?(level:Tar.Header.compatibility option) ?files ?(sector_size = 512) test_ctxt f =
   let format = match level with
     | None -> ""
     | Some format -> "--format=" ^ match format with
@@ -64,7 +64,7 @@ let with_tar ?(level:Tar.Header.compatibility option) ?files ?(block_size = 512)
   let tar_filename, ch = bracket_tmpfile ~prefix:"tar-test" ~suffix:".tar" test_ctxt in
   close_out ch;
   let tar_filename = if Sys.win32 then convert_path `Unix tar_filename else tar_filename in
-  let tar_block_size = block_size / 512 in
+  let tar_block_size = sector_size / 512 in
   let cmdline = Printf.sprintf "tar -cf %s -b %d %s %s" tar_filename tar_block_size format (String.concat " " files) in
   begin match Unix.system cmdline with
     | Unix.WEXITED 0 -> ()
@@ -155,7 +155,7 @@ let can_transform_tar test_ctxt =
 module Block4096 = struct
   include Block
 
-  let block_size = 4096
+  let sector_size = 4096
 
   let connect name =
     let name = if Sys.win32 then convert_path `Windows name else name in
@@ -165,13 +165,13 @@ end
 module type BLOCK = sig
   include Mirage_block.S
   val connect: string -> t Lwt.t
-  val block_size : int
+  val sector_size : int
 end
 
 module B = struct
   include Block
 
-  let block_size = 512
+  let sector_size = 512
 
   let connect name =
     let name = if Sys.win32 then convert_path `Windows name else name in
@@ -182,7 +182,8 @@ module Test(B: BLOCK) = struct
   let add_data_to_tar ?(level:Tar.Header.compatibility option) ?files test_ctxt f =
     let f tar_filename files =
       let size = Unix.(stat tar_filename).st_size in
-      Unix.truncate tar_filename (size + (min (2 * B.block_size) 4096));
+      let size = B.sector_size * ((pred size + 4096 + B.sector_size) / B.sector_size) in
+      Unix.truncate tar_filename size;
       B.connect tar_filename >>= fun b ->
       let module KV_RW = Tar_mirage.Make_KV_RW(Pclock)(B) in
       KV_RW.connect b >>= fun t ->
@@ -193,12 +194,14 @@ module Test(B: BLOCK) = struct
       let files = "barf" :: files in
       f tar_filename files
     in
-    with_tar ?level ?files ~block_size:B.block_size test_ctxt f
+    with_tar ?level ?files ~sector_size:B.sector_size test_ctxt f
 
   let add_more_data_to_tar ?(level:Tar.Header.compatibility option) ?files test_ctxt f =
     let f tar_filename files =
       let size = Unix.(stat tar_filename).st_size in
-      Unix.truncate tar_filename (size + (min (4 * B.block_size) 4096));
+      (* Add 4 KB rounding up to block size *)
+      let size = B.sector_size * ((pred size + 4096 + B.sector_size) / B.sector_size) in
+      Unix.truncate tar_filename size;
       B.connect tar_filename >>= fun b ->
       let module KV_RW = Tar_mirage.Make_KV_RW(Pclock)(B) in
       KV_RW.connect b >>= fun t ->
@@ -213,7 +216,7 @@ module Test(B: BLOCK) = struct
       let files = "barf" :: "barf2" :: files in
       f tar_filename files
     in
-    with_tar ?level ?files ~block_size:B.block_size test_ctxt f
+    with_tar ?level ?files ~sector_size:B.sector_size test_ctxt f
 
   let write_with_full_archive ?(level:Tar.Header.compatibility option) ?files test_ctxt =
     let f tar_filename files =
@@ -268,7 +271,7 @@ module Test(B: BLOCK) = struct
       ) files
 
   let can_read_through_BLOCK ?files test_ctxt =
-    with_tar ?files ~block_size:B.block_size test_ctxt check_tar
+    with_tar ?files ~sector_size:B.sector_size test_ctxt check_tar
 
   let write_test test_ctxt =
      add_data_to_tar test_ctxt check_tar