Rewrote MBR parsing so it supports extended partitions (again).

diff --git a/lib/diskimage_mbr.ml b/lib/diskimage_mbr.ml
index 8e1c5c0..6839380 100644 (file)
--- a/lib/diskimage_mbr.ml
+++ b/lib/diskimage_mbr.ml
@@ -36,11 +36,11 @@ let sector_size = ~^512
 let max_extended_partitions = 100
 
 (* The private data attached to a partitions structure. *)
-type private_t = mbr_part list         (* list of partitions that we found *)
-and mbr_part = {
-  mbr_part_start : int63;              (* start of partition in bytes *)
-  mbr_part_size : int63;               (* size of partition in bytes *)
+type mbr_priv = {
+  mbr_part_start : int63;              (* start of partition in SECTORS *)
+  mbr_part_size : int63;               (* size of partition in SECTORS *)
 }
+let null_priv = { mbr_part_start = ~^0; mbr_part_size = ~^0 }
 
 let attach_private_data, get_private_data =
   private_data_functions (fun {parts_cb = {parts_cb_uq = u}} -> u)
@@ -86,23 +86,35 @@ let rec probe dev =
       0x55 : 8; 0xAA : 8 } ->          (* MBR signature *)
 
       (* Parse the partition table entries. *)
-      let primaries =
-       List.mapi (parse_mbr_entry dev) [part0;part1;part2;part3] in
+      let primaries = List.map parse_mbr_entry [part0;part1;part2;part3] in
 
-(*
-      (* Read extended partition data. *)
+      (* Extended partitions are primary partitions with part_type 0x05,
+       * containing extended partition data.
+       *)
       let extendeds = List.map (
        function
-       | { part_type = 0x05 } as part ->
-           probe_extended_partition
-             max_extended_partitions fd part part.part_lba_start
+       | ({ part_type = 0x05 }, { mbr_part_start = start }) ->
+           probe_extended_partition max_extended_partitions dev start
        | part -> []
       ) primaries in
       let extendeds = List.concat extendeds in
-*)
 
-      let parts = primaries (* @ extendeds *) in
-      let privs = List.concat (List.map snd parts) in
+      let parts : (partition * mbr_priv) list = primaries @ extendeds in
+
+      (* Create partition devices for all partitions that aren't null. *)
+      let parts = List.mapi (
+       fun partno (part, priv) ->
+         let partno = partno+1 in
+         match part with
+         | { part_status = (Bootable|Nonbootable) } as part ->
+             let start = priv.mbr_part_start and size = priv.mbr_part_size in
+             let dev = new partition_device partno start size dev in
+             { part with part_dev = dev }, priv
+         | _ -> (part, priv)
+      ) parts in
+
+      (* Separate out the private data and attach it. *)
+      let privs = List.map snd parts in
       let parts = List.map fst parts in
 
       let r = { parts_cb = callbacks (); parts_dev = dev; parts = parts } in
@@ -115,11 +127,11 @@ let rec probe dev =
 (* Parse a single partition table entry.  See the table here:
  * http://en.wikipedia.org/wiki/Master_boot_record
  *)
-and parse_mbr_entry dev i bits =
+and parse_mbr_entry bits =
   bitmatch bits with
   | { 0l : 32; 0l : 32; 0l : 32; 0l : 32 } ->
       { part_status = NullEntry; part_type = 0;
-       part_dev = null_device; part_content = `Unknown }, []
+       part_dev = null_device; part_content = `Unknown }, null_priv
 
   | { ((0|0x80) as bootable) : 8; first_chs : 24;
       part_type : 8; last_chs : 24;
@@ -129,73 +141,72 @@ and parse_mbr_entry dev i bits =
       let bootable = if bootable = 0 then Nonbootable else Bootable in
       let first_lba = Int63.of_int32 first_lba in
       let part_size = Int63.of_int32 part_size in
-      let partno = i+1 in
 
       if !debug then
        eprintf "parse_mbr_entry: first_lba = %s part_size = %s\n%!"
          (Int63.to_string first_lba) (Int63.to_string part_size);
 
       let part = {
-       part_status = bootable;
-       part_type = part_type;
-       part_dev = new partition_device partno first_lba part_size dev;
+       part_status = bootable; part_type = part_type;
+       part_dev = null_device;         (* This gets overwritten by probe. *)
        part_content = `Unknown;
       } in
 
       (* Extra private data which we'll use to calculate free offsets. *)
       let priv = {
-       mbr_part_start = first_lba *^ sector_size;
-       mbr_part_size = part_size *^ sector_size;
+       mbr_part_start = first_lba;
+       mbr_part_size = part_size;
       } in
 
-      part, [priv]
+      part, priv
 
   | { _ } ->
       { part_status = Malformed; part_type = 0;
-       part_dev = null_device; part_content = `Unknown }, []
-
-(*
-This code worked previously, but now needs some love ...
-XXX
+       part_dev = null_device; part_content = `Unknown }, null_priv
 
 (* Probe an extended partition. *)
-and probe_extended_partition max fd epart sect =
+and probe_extended_partition max dev start =
   if max > 0 then (
-    (* Offset of the first EBR. *)
-    let ebr_offs = sect *^ sector_size in
-    (* EBR Signature? *)
-    LargeFile.lseek fd (ebr_offs +^ 510L) SEEK_SET;
-    let str = String.create 2 in
-    if read fd str 0 2 <> 2 || str.[0] != '\x55' || str.[1] != '\xAA' then
-      [] (* Not EBR *)
-    else (
-      (* Read the extended partition table entries (just 2 of them). *)
-      LargeFile.lseek fd (ebr_offs +^ 446L) SEEK_SET;
-      let str = String.create 32 in
-      if read fd str 0 32 <> 32 then
-       failwith (s_ "error reading extended partition")
-      else (
-       (* Extract partitions from the data. *)
-       let part1, part2 =
-         match List.map (get_partition str) [ 0; 16 ] with
-         | [p1;p2] -> p1,p2
-         | _ -> failwith (s_ "probe_extended_partition: internal error") in
-       (* First partition entry has offset to the start of this partition. *)
-       let part1 = { part1 with
-                       part_lba_start = sect +^ part1.part_lba_start } in
-       (* Second partition entry is zeroes if end of list, otherwise points
-        * to the next partition.
-        *)
-       if part2.part_status = NullEntry then
-         [part1]
-       else
-         part1 :: probe_extended_partition
-                    (max-1) fd epart (sect +^ part2.part_lba_start)
+    try
+      (* Get the partition table (like a boot sector). *)
+      let bits = dev#read_bitstring (start *^ sector_size) sector_size in
+
+      (bitmatch bits with
+       | { _ : 3568 : bitstring;      (* padding to byte offset 446 *)
+          part : 128 : bitstring;    (* this partition *)
+          next : 128 : bitstring;    (* pointer to next extended partition *)
+          _ : 128 : bitstring;       (* ignored - should be zero *)
+          _ : 128 : bitstring;
+          0x55 : 8; 0xAA : 8 } ->    (* MBR signature *)
+
+          let (part, ppriv) = parse_mbr_entry part in
+
+          (* The first partition's LBA is actually offset relative
+           * to the current sector.
+           *)
+          let ppriv =
+            { ppriv with mbr_part_start = ppriv.mbr_part_start + start } in
+
+          let (next, npriv) = parse_mbr_entry next in
+
+          if next.part_status = NullEntry then
+            [ part, ppriv ]            (* End of list. *)
+          else (
+            let start_of_next = start + npriv.mbr_part_start in
+            (part, ppriv) ::
+              probe_extended_partition (max-1) dev start_of_next
+          )
+
+       | { _ } ->
+          invalid_arg "mbr: invalid extended partition table"
       )
-    )
+    with exn ->
+      prerr_endline (Printexc.to_string exn);
+      []
+  ) else (
+    prerr_endline "mbr: too many extended partitions";
+    []
   )
-  else []
-*)
 
 (*
 (* Ugh, fake a UInt32 -> UInt64 conversion without sign extension, until
@@ -219,12 +230,16 @@ and offset_is_free parts offset =
   match privs with
   | [] -> false
   | { mbr_part_start = start; mbr_part_size = size } :: rest ->
+      let start = start *^ sector_size in
+      let size = size *^ sector_size in
       if offset < start +^ size then
        false
       else (
        let rec loop = function
          | [] -> true                  (* not in a partition, must be free *)
          | { mbr_part_start = start; mbr_part_size = size } :: rest ->
+             let start = start *^ sector_size in
+             let size = size *^ sector_size in
              if start <= offset && offset < start +^ size then
                false
              else