(* 'df' command for virtual domains.

   (C) Copyright 2007 Richard W.M. Jones, Red Hat Inc.
   http://libvirt.org/

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.

   Support for Master Boot Record partition scheme.
*)

open Printf
open Unix
open ExtList

open Diskimage_impl

open Int63.Operators

let id = "mbr"

let sector_size = ~^512

(* Maximum number of extended partitions possible. *)
let max_extended_partitions = 100

(* The private data attached to a partitions structure. *)
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)

(* Device representing a single partition.  It just acts as an offset
 * into the underlying device.
 *
 * Notes:
 * (1) 'start'/'size' are measured in sectors.
 * (2) 'partno' is the partition number, starting at 1
 *     (cf. /dev/hda1 is the first partition).
 * (3) 'dev' is the underlying block device.
 * (4) natural blocksize to use is sector size.
 *)
class partition_device partno start size dev =
  let devname = dev#name in
  let name = sprintf "%s%d" devname partno in
  let start = start *^ sector_size in
  let size = size *^ sector_size in
object (self)
  inherit offset_device name start size sector_size dev
end

(** Probe the
    {{:http://en.wikipedia.org/wiki/Master_boot_record}master boot record}
    (if it is one) and read the partitions.

    @raise Not_found if it is not an MBR.
 *)
let rec probe dev =
  (* Read the first sector. *)
  let bits =
    try dev#read_bitstring ~^0 sector_size
    with exn -> raise Not_found in

  (* Does this match a likely-looking MBR? *)
  bitmatch bits with
  | { _ : 3568 : bitstring;		(* padding to byte offset 446 *)
      part0 : 128 : bitstring;		(* partitions *)
      part1 : 128 : bitstring;
      part2 : 128 : bitstring;
      part3 : 128 : bitstring;
      0x55 : 8; 0xAA : 8 } ->		(* MBR signature *)

      (* Parse the partition table entries. *)
      let primaries = List.map parse_mbr_entry [part0;part1;part2;part3] in

      (* Extended partitions are primary partitions with part_type 0x05,
       * containing extended partition data.
       *)
      let extendeds = List.map (
	function
	| ({ 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 : (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
      attach_private_data r privs;
      r

  | { _ } ->
      raise Not_found			(* not an MBR *)

(* Parse a single partition table entry.  See the table here:
 * http://en.wikipedia.org/wiki/Master_boot_record
 *)
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 }, null_priv

  | { ((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
      let first_lba = Int63.of_int32 first_lba in
      let part_size = Int63.of_int32 part_size 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 = 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;
	mbr_part_size = part_size;
      } in

      part, priv

  | { _ } ->
      { part_status = Malformed; part_type = 0;
	part_dev = null_device; part_content = `Unknown }, null_priv

(* Probe an extended partition. *)
and probe_extended_partition max dev start =
  if max > 0 then (
    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";
    []
  )

(*
(* Ugh, fake a UInt32 -> UInt64 conversion without sign extension, until
 * we get working UInt32/UInt64 modules in extlib.
 *)
and uint64_of_int32 u32 =
  let i64 = Int64.of_int32 u32 in
  if u32 >= 0l then i64
  else Int64.add i64 0x1_0000_0000_L
*)

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 ->
      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
		loop rest
	in
	loop rest
      )

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