Properly inline the int63 impl on 64 bit platforms.
[virt-df.git] / 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 module Bitmap = Diskzip_bitmap
26
27 type output = File of string | Dir of string
28 type extcompress = BZip2 | GZip | External of string
29
30 let rec main () =
31   (* Program name changes behaviour. *)
32   let compressing =
33     let name = Sys.argv.(0) in
34     let name = Filename.basename name in (* just the executable name *)
35     let name =                           (* remove .opt or .exe *)
36       try Filename.chop_extension name
37       with Invalid_argument("Filename.chop_extension") -> name in
38     let name = String.lowercase name in
39     match name with
40     | "diskzcat" -> false
41     | "diskzip" -> true
42     | name ->
43         eprintf
44           (f_"diskzip: unknown executable name '%s', assuming 'diskzip'\n")
45           name 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   (* Create ownership bitmap for each disk. *)
156   List.iter (
157     fun { Diskimage.d_name = name; d_dev = disk } ->
158       let blocksize = disk#blocksize in
159       let size = disk#size in           (* Size in bytes. *)
160       let nr_blocks = size /^ blocksize in (* Number of disk sectors. *)
161
162       if !Diskimage.debug then
163         eprintf "Creating bitmap for %s (%s sectors) ...\n%!"
164           disk#name (Int63.to_string nr_blocks);
165
166       (* Create an empty bitmap, one bit per sector. *)
167       let bitmap = Bitmap.create nr_blocks in
168
169       (* Get the lookup function for this disk. *)
170       let lookup = Diskimage.get_owners_lookup machine ownership disk in
171
172       (* Lookup each sector. *)
173       Bitmap.iter_set (
174         fun blk _ ->
175           let owners = lookup blk in
176           false
177       ) bitmap
178   ) machine.Diskimage.m_disks;
179
180   (* Redirect output through external pipe if asked. *)
181   (match extcompress with
182    | None -> ()
183    | Some prog ->
184        let prog, progargs =
185          match prog with
186          | BZip2 -> "bzip2", [|"bzip2"; "-c"|]
187          | GZip -> "gzip", [|"gzip"; "-c"|]
188          | External prog -> "sh", [|"sh"; "-c"; prog |] in
189        let rfd, wfd = pipe () in
190        let pid = fork () in
191        if pid = 0 then (                (* child *)
192          close wfd;
193          dup2 rfd stdin;
194          close rfd;
195          execvp prog progargs
196        ) else (                         (* parent *)
197          close rfd;
198          dup2 wfd stdout;
199          close wfd
200        )
201   )
202
203
204
205
206
207
208
209
210
211
212 and go_decompress ?output extcompress args =
213   (* Read the input, which may be a single named file, or a series of
214    * files (we just concatenate them).  We may have to feed the input
215    * through an external program.
216    *)
217   let () =
218     match args with
219     | [] -> ()                          (* Reading from stdin. *)
220     | [file] ->                         (* Read the named file. *)
221         let fd = openfile file [O_RDONLY] 0 in
222         dup2 fd stdin;
223         close fd
224     | files ->                          (* Concatenate files. *)
225         let rfd, wfd = pipe () in
226         let pid = fork () in
227         if pid = 0 then (               (* child *)
228           close rfd;
229           dup2 wfd stdout;
230           close wfd;
231           execvp "cat" (Array.of_list ("cat" :: "--" :: files))
232         ) else (                        (* parent *)
233           close wfd;
234           dup2 rfd stdin;
235           close rfd
236         )
237   in
238   (match extcompress with
239    | None -> ()
240    | Some prog ->
241        let prog, progargs =
242          match prog with
243          | BZip2 -> "bzip2", [|"bzip2"; "-cd"|]
244          | GZip -> "gzip", [|"gzip"; "-cd"|]
245          | External prog -> "sh", [|"sh"; "-c"; prog |] in
246        let rfd, wfd = pipe () in
247        let pid = fork () in
248        if pid = 0 then (                (* child *)
249          close rfd;
250          dup2 wfd stdout;
251          close wfd;
252          execvp prog progargs
253        ) else (                         (* parent *)
254          close wfd;
255          dup2 rfd stdin;
256          close rfd
257        )
258   )
259
260 (*
261   let header = read_header () in
262   XXX
263
264 *)
265
266
267
268
269
270
271
272
273 (*
274 (* Since we have the wonderful pa_bitmatch, might as well use it to
275  * define a robust binary format for the compressed files.
276  *)
277 and write_header ... =
278   let bs = BITSTRING {
279     0xD152 : 16; 0x01 : 8; 0x00 : 8;    (* file magic, version 1.0 *)
280     nr_disks : 8;                       (* number of disks being packed *)
281     
282
283
284
285   } in
286   
287 and read_header () =
288   (* Diskzip headers are limited to overall max size of 1024 bytes. *)
289   let bs = Bitmatch.bitstring_of_file_descr_max stdin 1024 in
290
291   bitmatch bs with
292   | { 0xD152 : 16;                      (* file magic *)
293       0x01 : 8; (_ as minor) : 8;       (* major, minor versions *)
294     } ->
295
296   (* Is this a later version (major != 1)? *)
297   | { 0xD152 : 16;                      (* file magic *)
298       (_ as major) : 8; (_ as minor) : 8 } when major <> 1 ->
299       eprintf (f_"diskzip: archive version %d.%d, this program only understands version 1.x")
300         major minor;
301       exit 1
302
303   (* If it looks like gzip or bzip2, exit with an informative error. *)
304   | { 0o37 : 8; 0o213 : 8 } ->          (* gzip *)
305       prerr_endline (s_"diskzip: This looks like a gzip archive. Did you mean to pass the '-z' option?");
306       exit 1
307   | { "BZh" : 24 : string } ->          (* bzip2 *)
308       prerr_endline (s_"diskzip: This looks like a bzip2 archive. Did you mean to pass the '-j' option?");
309       exit 1
310
311   (* If it looks like a disk image (MBR), give an error. *)
312   | { _ : 4080 : bitstring; 0x55 : 8; 0xAA : 8 } ->
313       prerr_endline (s_"diskzip: This looks like a disk image. Did you mean to compress it?");
314       exit 1
315
316   | { _ } ->
317       prerr_endline (s_"diskzip: Not a diskzip archive.");
318       exit 1
319 *)
320
321 let () = main ()