From d02eebf02d4697b2d07cab78517b218d4fe90e2f Mon Sep 17 00:00:00 2001
From: "Richard W.M. Jones" <rjones@redhat.com>
Date: Thu, 1 Jan 1970 00:00:00 +0000
Subject: [PATCH] Use private data to report correct offset_is_free.

---
 lib/diskimage_mbr.ml | 104 +++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 75 insertions(+), 29 deletions(-)

diff --git a/lib/diskimage_mbr.ml b/lib/diskimage_mbr.ml
index 0de443b..8e1c5c0 100644
--- a/lib/diskimage_mbr.ml
+++ b/lib/diskimage_mbr.ml
@@ -35,6 +35,16 @@ let sector_size = ~^512
 (* Maximum number of extended partitions possible. *)
 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 *)
+}
+
+let attach_private_data, get_private_data =
+  private_data_functions (fun {parts_cb = {parts_cb_uq = u}} -> u)
+
 (* Device representing a single partition.  It just acts as an offset
  * into the underlying device.
  *
@@ -89,9 +99,15 @@ let rec probe dev =
 	| part -> []
       ) primaries in
       let extendeds = List.concat extendeds in
-      primaries @ extendeds
 *)
-      { parts_cb = callbacks; parts_dev = dev; parts = primaries }
+
+      let parts = primaries (* @ extendeds *) in
+      let privs = List.concat (List.map snd parts) in
+      let parts = List.map fst parts in
+
+      let r = { parts_cb = callbacks (); parts_dev = dev; parts = parts } in
+      attach_private_data r privs;
+      r
 
   | { _ } ->
       raise Not_found			(* not an MBR *)
@@ -103,34 +119,40 @@ and parse_mbr_entry dev i 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 }, []
 
   | { ((0|0x80) as bootable) : 8; first_chs : 24;
       part_type : 8; last_chs : 24;
       first_lba : 32 : unsigned, littleendian;
       part_size : 32 : unsigned, littleendian } ->
+
       let bootable = if bootable = 0 then Nonbootable else Bootable in
-      make_mbr_entry bootable dev (i+1) part_type first_lba part_size
+      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_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;
+      } in
+
+      part, [priv]
 
   | { _ } ->
       { part_status = Malformed; part_type = 0;
-	part_dev = null_device; part_content = `Unknown }
-
-and make_mbr_entry part_status dev partno part_type first_lba part_size =
-  let first_lba = Int63.of_int32 first_lba in
-  let part_size = Int63.of_int32 part_size in
-  (*
-    XXX Used to be:
-  let first_lba = uint63_of_int32 first_lba in
-  let part_size = uint63_of_int32 part_size in
-  *)
-  if !debug then
-    eprintf "make_mbr_entry: first_lba = %s part_size = %s\n%!"
-      (Int63.to_string first_lba) (Int63.to_string part_size);
-  { part_status = part_status;
-    part_type = part_type;
-    part_dev = new partition_device partno first_lba part_size dev;
-    part_content = `Unknown }
+	part_dev = null_device; part_content = `Unknown }, []
 
 (*
 This code worked previously, but now needs some love ...
@@ -185,15 +207,39 @@ and uint64_of_int32 u32 =
   else Int64.add i64 0x1_0000_0000_L
 *)
 
-(* XXX We don't currently keep enough data in the parts structure
- * to allow us to reconstruct missing partition table entries.
- *)
-and offset_is_free _ _ = false
+and offset_is_free parts offset =
+  let privs = get_private_data parts in
+
+  (* The first partition is somehow privileged in that we assume
+   * everything before this is not free.  Usually this is the first
+   * 63 sectors containing the MBR itself and sectors which should
+   * be blank but in reality contain all sorts of stupid hacks like
+   * alternate partitioning schemes.
+   *)
+  match privs with
+  | [] -> false
+  | { mbr_part_start = start; mbr_part_size = size } :: rest ->
+      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 ->
+	      if start <= offset && offset < start +^ size then
+		false
+	      else
+		loop rest
+	in
+	loop rest
+      )
 
-and callbacks = {
-  parts_cb_name = id;
-  parts_cb_offset_is_free = offset_is_free;
-}
+and callbacks =
+  let i = ref 0 in
+  fun () -> {
+    parts_cb_uq = (incr i; !i);
+    parts_cb_name = id;
+    parts_cb_offset_is_free = offset_is_free;
+  }
 
 (* Register the plugin. *)
 let () = register_plugin ~partitioner:probe id
-- 
1.8.3.1