diskzip/diskzip.ml

   1 (* 'diskzip' command for intelligently compressing disk images.
   2    (C) Copyright 2007 Richard W.M. Jones, Red Hat Inc.
   3    http://libvirt.org/
   4
   5    This program is free software; you can redistribute it and/or modify
   6    it under the terms of the GNU General Public License as published by
   7    the Free Software Foundation; either version 2 of the License, or
   8    (at your option) any later version.
   9
  10    This program is distributed in the hope that it will be useful,
  11    but WITHOUT ANY WARRANTY; without even the implied warranty of
  12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13    GNU General Public License for more details.
  14
  15    You should have received a copy of the GNU General Public License
  16    along with this program; if not, write to the Free Software
  17    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
  18  *)
  19
  20 open Unix
  21 open Printf
  22
  23 open Int63.Operators
  24 open Diskzip_gettext.Gettext
  25
  26 type output = File of string | Dir of string
  27 type extcompress = BZip2 | GZip | External of string
  28
  29 let rec main () =
  30   (* Program name changes behaviour. *)
  31   let compressing =
  32     let name = Sys.argv.(0) in
  33     let name = Filename.basename name in (* just the executable name *)
  34     let name =                           (* remove .opt or .exe *)
  35       try Filename.chop_extension name
  36       with Invalid_argument("Filename.chop_extension") -> name in
  37     let name = String.lowercase name in
  38     match name with
  39     | "diskzcat" -> false
  40     | "diskzip" -> true
  41     | _ ->
  42         eprintf
  43           (f_"diskzip: unknown executable name '%s', assuming 'diskzip'\n")
  44           name;
  45         true in
  46   let compressing = ref compressing in
  47
  48   (* Command line argument parsing. *)
  49   let version () =
  50     printf "diskzip\n"; (* XXX version XXX *)
  51     exit 0
  52   in
  53
  54   let output = ref None in
  55   let set_output path =
  56     if !output <> None then (
  57       prerr_endline (s_"diskzip: '-o' option cannot appear more than once");
  58       exit 2
  59     );
  60     try
  61       let statbuf = stat path in
  62       if statbuf.st_kind = S_DIR then
  63         output := Some (Dir path)
  64       else
  65         output := Some (File path)
  66     with
  67     (* No such file or directory, assume it's a file output. *)
  68     | Unix_error (ENOENT, _, _) -> output := Some (File path)
  69   in
  70
  71   (* By default we don't use any external compression program. *)
  72   let extcompress = ref None in
  73   let set_extcompress t () =
  74     if !extcompress <> None then (
  75       prerr_endline (s_"diskzip: '-z' or '-j' cannot appear more than once");
  76       exit 2
  77     );
  78     extcompress := Some t
  79   in
  80
  81   let force = ref false in
  82
  83   let argspec = Arg.align [
  84     "-d", Arg.Clear compressing,
  85       " " ^ s_ "Uncompress (default: depends on executable name)";
  86     "--debug", Arg.Set Diskimage.debug,
  87       " " ^ s_ "Debug mode (default: false)";
  88     "-f", Arg.Set force,
  89       " " ^ s_"Force compress even if stdout looks like a tty";
  90     "-j", Arg.Unit (set_extcompress BZip2),
  91       " " ^ s_"Pipe the output/input through bzip2";
  92     "-o", Arg.String set_output,
  93       "path " ^ s_"Set the output filename or directory name";
  94     "-p", Arg.String (fun prog -> set_extcompress (External prog) ()),
  95       "prog " ^ s_"Pipe the output/input through external program";
  96     "--version", Arg.Unit version,
  97       " " ^ s_"Display version and exit";
  98     "-z", Arg.Unit (set_extcompress GZip),
  99       " " ^ s_"Pipe the output/input through gzip";
 100   ] in
 101
 102   let args = ref [] in
 103   let anon_fun str = args := str :: !args in
 104   let usage_msg = s_"diskzip: Intelligently compress disk images
 105
 106 SUMMARY
 107   diskzip [-options] disk.img [disk.img ...] > output.dz
 108   diskzcat [-options] output.dz > disk.img
 109
 110 OPTIONS" in
 111
 112   Arg.parse argspec anon_fun usage_msg;
 113
 114   (* Turn refs back into normal values. *)
 115   let compressing = !compressing in
 116   let extcompress = !extcompress in
 117   let output = !output in
 118   let force = !force in
 119   let args = !args in
 120
 121   (* Check the arguments make sense. *)
 122   if compressing && output <> None then (
 123     prerr_endline (s_"diskzip: '-o' option cannot be used when compressing");
 124     exit 2
 125   );
 126   if compressing && args = [] then (
 127     prerr_endline (s_"diskzip: no input");
 128     exit 2
 129   );
 130   if compressing && not force && isatty stdout then (
 131     prerr_endline (s_"diskzip: compressed data not written to a terminal, use '-f' to force");
 132     exit 2
 133   );
 134
 135   (* Run the compression or decompression functions. *)
 136   if compressing then
 137     go_compress extcompress args
 138   else
 139     go_decompress ?output extcompress args
 140
 141 (* Do compression. *)
 142 and go_compress extcompress images =
 143   (* Create a Diskimage machine description from the requested images.  This
 144    * also checks that everything we need is readable.
 145    *)
 146   let machine =
 147     Diskimage.open_machine "diskzip" (List.map (fun n -> (n,n)) images) in
 148
 149   (* Scan the images for filesystems. *)
 150   let machine = Diskimage.scan_machine machine in
 151
 152   (* Create ownership tables. *)
 153   let ownership = Diskimage.create_ownership machine in
 154
 155   (* Redirect output through external pipe if asked. *)
 156   (match extcompress with
 157    | None -> ()
 158    | Some prog ->
 159        let prog, progargs =
 160          match prog with
 161          | BZip2 -> "bzip2", [|"bzip2"; "-c"|]
 162          | GZip -> "gzip", [|"gzip"; "-c"|]
 163          | External prog -> "sh", [|"sh"; "-c"; prog |] in
 164        let rfd, wfd = pipe () in
 165        let pid = fork () in
 166        if pid = 0 then (                (* child *)
 167          close wfd;
 168          dup2 rfd stdin;
 169          close rfd;
 170          execvp prog progargs
 171        ) else (                         (* parent *)
 172          close rfd;
 173          dup2 wfd stdout;
 174          close wfd
 175        )
 176   );
 177
 178   (* Iterate over the disks. *)
 179   List.iter (
 180     fun { Diskimage.d_name = name; d_dev = disk } ->
 181       let blocksize = disk#blocksize in
 182       let size = disk#size in           (* Size in bytes. *)
 183       let nr_blocks = size /^ blocksize in (* Number of disk sectors. *)
 184
 185       if !Diskimage.debug then
 186         eprintf "Writing disk %s (%s sectors) ...\n%!"
 187           disk#name (Int63.to_string nr_blocks);
 188
 189       (* Get the lookup function for this disk. *)
 190       let lookup_offset =
 191         Diskimage.get_owners_lookup machine ownership disk in
 192
 193       (* Convenience function to look up a block and test freeness. *)
 194       let block_is_free blk =
 195         let offset = blk *^ blocksize in
 196         Diskimage.offset_is_free (lookup_offset offset)
 197       in
 198
 199       (* Look up owners for each sector in turn. *)
 200       let rec loop blk =
 201         if blk < nr_blocks then (
 202           (* The current sector (blk) is either free or not free.  Look
 203            * for a stretch of sectors which are the same.
 204            *)
 205           let current_free = block_is_free blk in
 206           let rec find_end blk =
 207             if blk < nr_blocks then (
 208               if block_is_free blk = current_free then
 209                 find_end (Int63.succ blk)
 210               else
 211                 blk
 212             ) else
 213               nr_blocks                 (* End of the disk image. *)
 214           in
 215           let end_blk = find_end (Int63.succ blk) in
 216
 217           (* Current stretch is from blk .. end_blk-1. *)
 218           if !Diskimage.debug then
 219             eprintf "  %s stretch %s to %s-1 (%s bytes)\n"
 220               (if current_free then "free" else "used")
 221               (Int63.to_string blk) (Int63.to_string end_blk)
 222               (Int63.to_string ((end_blk-^blk) *^ blocksize));
 223
 224
 225
 226
 227           loop end_blk
 228         )
 229       in
 230       loop ~^0
 231   ) machine.Diskimage.m_disks
 232
 233
 234
 235
 236
 237
 238
 239
 240
 241
 242
 243
 244 and go_decompress ?output extcompress args =
 245   (* Read the input, which may be a single named file, or a series of
 246    * files (we just concatenate them).  We may have to feed the input
 247    * through an external program.
 248    *)
 249   let () =
 250     match args with
 251     | [] -> ()                          (* Reading from stdin. *)
 252     | [file] ->                         (* Read the named file. *)
 253         let fd = openfile file [O_RDONLY] 0 in
 254         dup2 fd stdin;
 255         close fd
 256     | files ->                          (* Concatenate files. *)
 257         let rfd, wfd = pipe () in
 258         let pid = fork () in
 259         if pid = 0 then (               (* child *)
 260           close rfd;
 261           dup2 wfd stdout;
 262           close wfd;
 263           execvp "cat" (Array.of_list ("cat" :: "--" :: files))
 264         ) else (                        (* parent *)
 265           close wfd;
 266           dup2 rfd stdin;
 267           close rfd
 268         )
 269   in
 270   (match extcompress with
 271    | None -> ()
 272    | Some prog ->
 273        let prog, progargs =
 274          match prog with
 275          | BZip2 -> "bzip2", [|"bzip2"; "-cd"|]
 276          | GZip -> "gzip", [|"gzip"; "-cd"|]
 277          | External prog -> "sh", [|"sh"; "-c"; prog |] in
 278        let rfd, wfd = pipe () in
 279        let pid = fork () in
 280        if pid = 0 then (                (* child *)
 281          close rfd;
 282          dup2 wfd stdout;
 283          close wfd;
 284          execvp prog progargs
 285        ) else (                         (* parent *)
 286          close wfd;
 287          dup2 rfd stdin;
 288          close rfd
 289        )
 290   )
 291
 292 (*
 293   let header = read_header () in
 294   XXX
 295
 296 *)
 297
 298
 299
 300
 301
 302
 303
 304
 305 (*
 306 (* Since we have the wonderful pa_bitmatch, might as well use it to
 307  * define a robust binary format for the compressed files.
 308  *)
 309 and write_header ... =
 310   let bs = BITSTRING {
 311     0xD152 : 16; 0x01 : 8; 0x00 : 8;    (* file magic, version 1.0 *)
 312     nr_disks : 8;                       (* number of disks being packed *)
 313
 314
 315
 316
 317   } in
 318
 319 and read_header () =
 320   (* Diskzip headers are limited to overall max size of 1024 bytes. *)
 321   let bs = Bitmatch.bitstring_of_file_descr_max stdin 1024 in
 322
 323   bitmatch bs with
 324   | { 0xD152 : 16;                      (* file magic *)
 325       0x01 : 8; (_ as minor) : 8;       (* major, minor versions *)
 326     } ->
 327
 328   (* Is this a later version (major != 1)? *)
 329   | { 0xD152 : 16;                      (* file magic *)
 330       (_ as major) : 8; (_ as minor) : 8 } when major <> 1 ->
 331       eprintf (f_"diskzip: archive version %d.%d, this program only understands version 1.x")
 332         major minor;
 333       exit 1
 334
 335   (* If it looks like gzip or bzip2, exit with an informative error. *)
 336   | { 0o37 : 8; 0o213 : 8 } ->          (* gzip *)
 337       prerr_endline (s_"diskzip: This looks like a gzip archive. Did you mean to pass the '-z' option?");
 338       exit 1
 339   | { "BZh" : 24 : string } ->          (* bzip2 *)
 340       prerr_endline (s_"diskzip: This looks like a bzip2 archive. Did you mean to pass the '-j' option?");
 341       exit 1
 342
 343   (* If it looks like a disk image (MBR), give an error. *)
 344   | { _ : 4080 : bitstring; 0x55 : 8; 0xAA : 8 } ->
 345       prerr_endline (s_"diskzip: This looks like a disk image. Did you mean to compress it?");
 346       exit 1
 347
 348   | { _ } ->
 349       prerr_endline (s_"diskzip: Not a diskzip archive.");
 350       exit 1
 351 *)
 352
 353 let () = main ()