Block free/used for NTFS working

diff --git a/lib/diskimage_ntfs.ml b/lib/diskimage_ntfs.ml
index d7eb7d6..e4e9036 100644 (file)
--- a/lib/diskimage_ntfs.ml
+++ b/lib/diskimage_ntfs.ml
@@ -27,14 +27,11 @@ open Int63.Operators
 
 let id = "ntfs"
 
-(* Private data functions. *)
-let attach_private_data, get_private_data =
-  private_data_functions (fun {fs_cb = {fs_cb_uq = u}} -> u)
-
 (* Type of the private data, basically all the metadata that we
  * read from the NTFS volume.
  *)
 type ntfs_fs = {
+  ntfs_dev : device;                   (* Device. *)
   ntfs_blocksize : int63;              (* Blocksize (cluster size) *)
   ntfs_mft_lcn : int63;                        (* MFT location (bytes) *)
   ntfs_mft_size : int63;               (* MFT size (bytes) *)
@@ -62,6 +59,11 @@ and ntfs_runentry =
     (* VCN start,size => LCN / None if sparse hole *)
     (int63 * int63)   *  int63 option
 
+(* Private data functions. *)
+let attach_private_data, get_private_data =
+  private_data_functions (fun {fs_cb = {fs_cb_uq = u}} -> u)
+
+(* Probe for an NTFS filesystem on this device. *)
 let rec probe dev =
   let fs = probe_superblock dev in
   fs
@@ -109,6 +111,7 @@ and probe_superblock dev =
          dev#name blocksize volume_serial_number;
 
       let blocksize = Int63.of_int blocksize in
+      let number_of_sectors = Int63.of_int64 number_of_sectors in
 
       (* The blocksize of the filesystem is likely to be quite different
        * from that of the underlying device, so create an overlay device
@@ -122,14 +125,38 @@ and probe_superblock dev =
 
       let mft = parse_mft dev mft_lcn mft_size in
 
-      let priv = {
+      let ntfs = {
+       ntfs_dev = fs_dev;
        ntfs_blocksize = blocksize;
        ntfs_mft_lcn = mft_lcn;
        ntfs_mft_size = mft_size;
        ntfs_mft_records = mft
       } in
 
-      raise Not_found                  (* XXX *)
+      (* Query free space.  I cannot find any metadata in the NTFS
+       * structures which records free space directly, so instead we
+       * need to read the $Bitmap::$Data (bitmap of allocated LCNs).
+       *)
+      let blocks_used, blocks_avail = parse_bitmap_freespace ntfs in
+
+      (* Create a filesystem structure. *)
+      let fs = {
+       fs_cb = callbacks ();
+       fs_dev = fs_dev;
+       fs_blocksize = blocksize;
+       fs_blocks_total = number_of_sectors *^ bytes_per_sector /^ blocksize;
+       fs_is_swap = false;
+       fs_blocks_reserved = ~^0;       (* XXX MFT, bitmap are "reserved" *)
+       fs_blocks_avail = blocks_avail;
+       fs_blocks_used = blocks_used;
+       fs_inodes_total = ~^0;          (* XXX MFT records are like inodes *)
+       fs_inodes_reserved = ~^0;
+       fs_inodes_avail = ~^0;
+       fs_inodes_used = ~^0;
+      } in
+
+      attach_private_data fs ntfs;
+      fs
 
   | { _ } -> raise Not_found           (* Not an NTFS boot sector. *)
 
@@ -245,10 +272,9 @@ and parse_attrs attrs mft_record =
       pad : (Int32.to_int attr_size - 8) * 8 : bitstring;
       rest : -1 : bitstring } ->
 
-      if !debug then (
-       eprintf "cannot parse MFT attribute entry\n%!";
-       Bitmatch.hexdump_bitstring Pervasives.stderr attrs
-      );
+      if !debug then
+       eprintf "cannot parse MFT attribute entry, attr_type = %lx\n%!"
+         attr_type;
 
       parse_attrs rest mft_record
 
@@ -430,6 +456,90 @@ and ucs2_to_utf8 name len =
   done;
   outstr
 
+(* Parse $Bitmap::$Data to get free/used.  Returns (used, free) blocks. *)
+and parse_bitmap_freespace ntfs =
+  (* Can throw Not_found - allow that to escape because we don't
+   * expect an NTFS filesystem without this magic file.
+   *)
+  let file = find_system_file ntfs "$Bitmap" in
+
+  (* Count used/free bits. *)
+  let used = ref ~^0 and free = ref ~^0 in
+  iter_blocks ntfs file (
+    fun lcn vcn data ->
+      for i = 0 to String.length data - 1 do
+       let c = Char.code data.[i] in
+       if c = 0 then                   (* common cases *)
+         free := !free +^ ~^8
+       else if c = 0xff then
+         used := !used +^ ~^8
+       else (                          (* uncommon case: count the bits *)
+         let m = ref 0x80 in
+         while !m > 0 do
+           if c land !m <> 0 then
+             used := !used +^ ~^1
+           else
+             free := !free +^ ~^1;
+           m := !m lsr 1
+         done
+       )
+      done
+  );
+  (!used, !free)
+
+and find_system_file { ntfs_mft_records = mft_records } fname =
+  let rec loop =
+    function 
+    | [] -> raise Not_found
+    | ({ ntfs_filename = Some { ntfs_name = name } } as file) :: _
+       when name = fname ->
+       file
+    | _ :: rest -> loop rest
+  in
+  loop mft_records
+
+and iter_blocks { ntfs_blocksize = blocksize; ntfs_dev = dev }
+    { ntfs_data = data } f =
+  match data with
+  | None -> ()                         (* No $Data attribute. *)
+  | Some { ntfs_data_size = data_size; ntfs_runlist = runlist } ->
+      let rec loop data_size = function
+       | [] -> ()
+
+       (* Run of vcnsize clusters. *)
+       | ((vcnstart, vcnsize), Some lcn) :: rest ->
+           let data_size = ref data_size in
+           let lcn = ref lcn in
+           let vcn = ref vcnstart in
+           let vcnsize = ref vcnsize in
+           while !vcnsize > ~^0 && !data_size > ~^0 do
+             let size = min blocksize !data_size in
+             let data = dev#read (!lcn *^ blocksize) size in
+             f (Some !lcn) !vcn data;
+             lcn := !lcn +^ ~^1;
+             vcn := !vcn +^ ~^1;
+             vcnsize := !vcnsize -^ ~^1;
+             data_size := !data_size -^ size
+           done;
+           loop !data_size rest
+
+       (* Sparse hole. *)
+       | ((vcnstart, vcnsize), None) :: rest ->
+           let data_size = ref data_size in
+           let vcn = ref vcnstart in
+           let vcnsize = ref vcnsize in
+           while !vcnsize > ~^0 && !data_size > ~^0 do
+             let size = min blocksize !data_size in
+             let data = String.make size '\000' in
+             f None !vcn data;
+             vcn := !vcn +^ ~^1;
+             vcnsize := !vcnsize -^ ~^1;
+             data_size := !data_size -^ size
+           done;
+           loop !data_size rest
+      in
+      loop data_size runlist
+
 and offset_is_free _ _ = false
 
 and callbacks =