Support for writing diskzip files.
[virt-df.git] / diskzip / diskzip.ml
index b4a0a5a..55094e7 100644 (file)
  *)
 
 open Unix
+open ExtList
+open ExtString
 open Printf
 
+open Int63.Operators
 open Diskzip_gettext.Gettext
 
 type output = File of string | Dir of string
 type extcompress = BZip2 | GZip | External of string
 
+let max_disks = 32
+let max_image_name = 256
+
 let rec main () =
   (* Program name changes behaviour. *)
   let compressing =
     let name = Sys.argv.(0) in
-    let name = Filename.basename name in       (* just the executable name *)
-    let name = Filename.chop_extension name in (* remove .opt or .exe *)
+    let name = Filename.basename name in (* just the executable name *)
+    let name =                          (* remove .opt or .exe *)
+      try Filename.chop_extension name
+      with Invalid_argument("Filename.chop_extension") -> name in
     let name = String.lowercase name in
     match name with
     | "diskzcat" -> false
     | "diskzip" -> true
-    | name ->
+    | _ ->
        eprintf
          (f_"diskzip: unknown executable name '%s', assuming 'diskzip'\n")
-         name in
+         name;
+       true in
   let compressing = ref compressing in
 
   (* Command line argument parsing. *)
@@ -145,6 +154,9 @@ and go_compress extcompress images =
   (* Scan the images for filesystems. *)
   let machine = Diskimage.scan_machine machine in
 
+  (* Create ownership tables. *)
+  let ownership = Diskimage.create_ownership machine in
+
   (* Redirect output through external pipe if asked. *)
   (match extcompress with
    | None -> ()
@@ -166,16 +178,77 @@ and go_compress extcompress images =
         dup2 wfd stdout;
         close wfd
        )
-  )
-
-
-
-
-
-
-
+  );
 
+  (* Write the main header. *)
+  write_header machine.Diskimage.m_disks;
+
+  (* Iterate over the disks. *)
+  List.iteri (
+    fun disknum { Diskimage.d_name = name; d_dev = disk } ->
+      let blocksize = disk#blocksize in
+      let size = disk#size in          (* Size in bytes. *)
+      let nr_blocks = size /^ blocksize in (* Number of disk sectors. *)
+
+      if !Diskimage.debug then
+       eprintf "Writing disk %s (%s sectors) ...\n%!"
+         disk#name (Int63.to_string nr_blocks);
+
+      (* Get the lookup function for this disk. *)
+      let lookup_offset =
+       Diskimage.get_owners_lookup machine ownership disk in
+
+      (* Convenience function to look up a block and test freeness. *)
+      let block_is_free blk =
+       let offset = blk *^ blocksize in
+       Diskimage.offset_is_free (lookup_offset offset)
+      in
+
+      (* Look up owners for each sector in turn. *)
+      let rec loop start_blk =
+       if start_blk < nr_blocks then (
+         (* The current sector (start_blk) is either free or not free.
+          * Look for a stretch of sectors which are the same.
+          *)
+         let current_free = block_is_free start_blk in
+         let rec find_end blk =
+           if blk < nr_blocks then (
+             if block_is_free blk = current_free then
+               find_end (Int63.succ blk)
+             else
+               blk
+           ) else
+             nr_blocks                 (* End of the disk image. *)
+         in
+         let end_blk = find_end (Int63.succ start_blk) in
+
+         let len_blks = end_blk -^ start_blk in
+
+         let start_offset = start_blk *^ blocksize in
+         let len_bytes = len_blks *^ blocksize in
+
+         (* Current stretch is from start_blk .. end_blk-1. *)
+         if !Diskimage.debug then
+           eprintf "  %s stretch %s to %s-1 (%s bytes)\n%!"
+             (if current_free then "free" else "used")
+             (Int63.to_string start_blk) (Int63.to_string end_blk)
+             (Int63.to_string (len_blks *^ blocksize));
+
+         (* Write the stretch to stdout. *)
+         write_stretch_header disknum current_free start_offset len_bytes;
+
+         (* Write the data (note: we only need to write it if
+          * it's not marked as free!).
+          *)
+         if not current_free then write_stretch disk start_offset len_bytes;
+
+         loop end_blk
+       )
+      in
+      loop ~^0
+  ) machine.Diskimage.m_disks;
 
+  write_trailer ()
 
 and go_decompress ?output extcompress args =
   (* Read the input, which may be a single named file, or a series of
@@ -238,20 +311,75 @@ and go_decompress ?output extcompress args =
 
 
 
-(*
 (* Since we have the wonderful pa_bitmatch, might as well use it to
  * define a robust binary format for the compressed files.
+ *
+ * These functions are in matched pairs "write_foo" / "read_foo" so
+ * you can check that the write and read protocols agree.
  *)
-and write_header ... =
-  let bs = BITSTRING {
-    0xD152 : 16; 0x01 : 8; 0x00 : 8;   (* file magic, version 1.0 *)
-    nr_disks : 8;                      (* number of disks being packed *)
-    
-
+and write_header disks =
+  let nr_disks = List.length disks in
+  assert (nr_disks > 0);
 
-
-  } in
-  
+  (* Don't allow large numbers of disks. *)
+  if nr_disks > max_disks then (
+    eprintf (f_"diskzip: maximum number of disk images is limited by the current image format to %d") nr_disks;
+    exit 2
+  );
+  let names = List.map (
+    fun { Diskimage.d_name = name } ->
+      let name =
+       try
+         let i = 1 + String.rindex name '/' in
+         String.sub name i (String.length name - i)
+       with
+         Invalid_string | Not_found -> name in
+
+      let contains_dotdot =
+       try ignore (String.find name ".."); true
+       with Invalid_string | Not_found -> false in
+      if contains_dotdot then (
+       prerr_endline (s_"diskzip: disk image names cannot contain \"..\"");
+       exit 2
+      );
+
+      (* Don't allow very long names. *)
+      if String.length name > max_image_name then (
+       eprintf (f_"diskzip: maximum length of disk image name is limited by the current image format to %d bytes") max_image_name;
+       exit 2
+      );
+
+      name
+  ) disks in
+
+  (* Header followed by names. *)
+  let header =
+    let header =
+      let bs = BITSTRING {
+       0xD152 : 16; 0x01 : 8; 0x00 : 8; (* file magic, version 1.0 *)
+       nr_disks : 8                    (* number of disks being packed *)
+      } in
+      let len = Bitmatch.bitstring_length bs in
+      assert (len land 7 = 0);
+      Bitmatch.string_of_bitstring bs in
+    let names = List.map (
+      fun name ->
+       let bs =
+         BITSTRING {
+           String.length name : 16;
+           name : -1 : string
+         } in
+       let len = Bitmatch.bitstring_length bs in
+       assert (len land 7 = 0);
+       Bitmatch.string_of_bitstring bs
+    ) names in
+
+    (* Construct the final header. *)
+    header ^ String.concat "" names in
+
+  ignore (write stdout header 0 (String.length header))
+
+(*  
 and read_header () =
   (* Diskzip headers are limited to overall max size of 1024 bytes. *)
   let bs = Bitmatch.bitstring_of_file_descr_max stdin 1024 in
@@ -286,4 +414,42 @@ and read_header () =
       exit 1
 *)
 
+and write_stretch_header disknum free start_offset len_bytes =
+  let start_offset = Int63.to_int64 start_offset in
+  let len_bytes = Int63.to_int64 len_bytes in
+
+  let bs = BITSTRING {
+    (* Stretch header magic.  Allows us to find synchronization errors. *)
+    0xD1525555_l : 32;
+    free : 1; 0 : 7;                   (* Flags. *)
+    disknum : 8;                       (* Disk number. *)
+    start_offset : 64;                 (* Start offset in disk image (bytes)*)
+    len_bytes : 64                     (* Length in bytes. *)
+  } in
+
+  let str = Bitmatch.string_of_bitstring bs in
+  ignore (write stdout str 0 (String.length str))
+
+and write_stretch dev start_offset len_bytes =
+  (* Limit size of blocks that we write, since some implementations
+   * of device class cannot handle large blocks.  In any case I tested
+   * this and writing large blocks isn't any faster.
+   *)
+  let blocksize = ~^65536 in
+  let rec loop offset len =
+    if len > ~^0 then (
+      let n = min blocksize len in
+      ignore (write stdout (dev#read offset n) 0 (Int63.to_int n));
+      loop (offset+^n) (len-^n)
+    )
+  in
+  loop start_offset len_bytes
+
+and write_trailer () =
+  let bs = BITSTRING {
+    0xD152FFFF_l : 32
+  } in
+  let str = Bitmatch.string_of_bitstring bs in
+  ignore (write stdout str 0 (String.length str))
+
 let () = main ()