(* '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 LVM2 PVs.
*)

open Printf
open ExtList

open Diskimage_impl
open Diskimage_lvm2_metadata

open Int63.Operators

let id = "LVM2"

let sector_size_int = 512
let sector_size = ~^sector_size_int

(*let attach_private_data, get_private_data =
  private_data_functions (fun {lvm_cb = {lvm_cb_uq = u}} -> u)*)

(*----------------------------------------------------------------------*)
(* Block device which can do linear maps, same as the kernel dm-linear.c *)
class linear_map_device name extent_size segments =
  (* The segments are passed containing (start_extent, extent_count, ...)
   * but it's easier to deal with (start_extent, end_extent, ...) so
   * rewrite them.
   *)
  let segments = List.map
    (fun (start_extent, extent_count, dev, pvoffset) ->
       (start_extent, start_extent +^ extent_count, dev, pvoffset)
    ) segments in

  (* Calculate the size of the device (in bytes).  Note that because
   * of the random nature of the mapping this doesn't imply that we can
   * satisfy any read request up to the full size.
   *)
  let size_in_extents =
    List.fold_left max ~^0
      (List.map (fun (_, end_extent, _, _) -> end_extent) segments) in
  let size = size_in_extents *^ extent_size in
object (self)
  inherit device
  method name = name
  method size = size

  (* The natural blocksize for LVM devices is the extent size.
   * NB. Throws a runtime exception if the extent size is bigger
   * than an int (only likely to matter on 32 bit).
   *)
  method blocksize = extent_size

  method private map i =
    if i < ~^0 || i >= size_in_extents then
      invalid_arg "linear_map_device: read outside device";

    let rec loop = function
      | [] ->
	  None
      | (start_extent, end_extent, dev, pvoffset) :: rest ->
	  if start_extent <= i && i < end_extent then (
	    let dev_offset = (pvoffset +^ i) *^ extent_size in
	    Some (start_extent, end_extent, dev, dev_offset, pvoffset)
	  ) else
	    loop rest
    in
    loop segments

  (* Map block (extent) i to the underlying device. *)
  method map_block i =
    match self#map i with
    | Some (_, _, dev, dev_offset, _) -> [dev, dev_offset]
    | None -> []

  (* Continguous span. *)
  method contiguous offset =
    let offset_in_extents = offset /^ extent_size in

    (* Get the segment that this offset lies in. *)
    match self#map offset_in_extents with
    | Some (_, end_extent, dev, dev_offset, _) ->
	(* Contiguous bytes up to the end of this extent. *)
	end_extent *^ extent_size -^ offset
    | None -> ~^0

  (* NB. Use the superclass #read method. *)
end

(*----------------------------------------------------------------------*)
(* Probe to see if it's an LVM2 PV. *)
let rec probe dev =
  try
    let uuid, _, _ = read_pv_label dev in
    if !debug then
      eprintf "LVM2 detected PV UUID %s\n%!" uuid;
    { pv_cb = callbacks (); pv_uuid = uuid; pv_dev = dev }
  with exn ->
    if !debug then prerr_endline (Printexc.to_string exn);
    raise Not_found

and read_pv_label dev =
  (* Load the first 8 sectors.  I found by experimentation that
   * the second sector contains the header ("LABELONE" etc) and
   * the nineth sector contains some additional information about
   * the location of the current metadata.
   *)
  let bits = dev#read_bitstring ~^0 (~^9 *^ sector_size) in

  bitmatch bits with
  | {
      (* sector 0 *)
      _ : sector_size_int*8 : bitstring;

      (* sector 1 *)
      "LABELONE" : 64 : string;		(* "LABELONE" *)
      _ : 128 : bitstring;		(* Seems to contain something. *)
      "LVM2 001" : 64 : string;		(* "LVM2 001" *)
      uuid : 256 : string;		(* PV UUID *)
      _ : (sector_size_int-64)*8 : bitstring;(* to end of second sector *)

      (* sectors 2-7 *)
      _ : sector_size_int*8 * 6 : bitstring;

      (* sector 8 *)
      _ : 320 : bitstring;		(* start of sector 8 *)
      metadata_offset : 32 : littleendian; (* metadata offset *)
      _ : 32 : bitstring;
      metadata_length : 32 : littleendian (* length of metadata (bytes) *)
    } ->

      (* Metadata offset is relative to end of PV label. *)
      let metadata_offset = Int63.of_int32 metadata_offset +^ ~^0x1000 in
      (* Metadata length appears to include the trailing \000 which
       * we don't want, so subtract 1 to get the true length.
       *)
      let metadata_length = Int63.of_int32 metadata_length -^ ~^1 in

      (* Check the metadata offset and length are sensible. *)
      if metadata_offset <= ~^0x1200 || metadata_offset >= dev#size
	|| metadata_length <= ~^0
	|| metadata_offset +^ metadata_length >= dev#size then
	  invalid_arg "LVM2: read_metadata: bad metadata offset or length";

      uuid, metadata_offset, metadata_length

  | { _ } ->
      invalid_arg
	(sprintf "LVM2: read_pv_label: %s: not an LVM2 physical volume"
	   dev#name)

(*----------------------------------------------------------------------*)
(* We are passed a list of devices which we previously identified
 * as PVs belonging to us.  From these produce a list of all LVs
 * (as devices) and return them.  Note that we don't try to detect
 * what is on these LVs - that will be done in the main code.
 *)
and list_lvs pvs =
  (* Read the PV label (again) for each PV, and this time also
   * read out the metadata, which is a big block of text.
   *)
  let pvsmap = List.map (
    fun { pv_dev = dev } ->
      let uuid, metadata_offset, metadata_length = read_pv_label dev in
      let metadata = dev#read metadata_offset metadata_length in

      if !debug then
	eprintf "list_lvs: metadata for PV %s (offset %s len %s):\n%s\n%!"
	  dev#name
	  (Int63.to_string metadata_offset) (Int63.to_string metadata_length)
	  metadata;

      (uuid, (metadata, dev))
  ) pvs in

  (* Parse the metadata using the external lexer/parser. *)
  let pvsmap = List.map (
    fun (uuid, (metadata, dev)) ->
      uuid, (Diskimage_lvm2_lexer.parse_lvm2_metadata_from_string metadata,
	     dev)
  ) pvsmap in

  (* Print the parsed metadata. *)
  if !debug then
    List.iter (
      fun (uuid, (metadata, dev)) ->
	eprintf "metadata for PV UUID %s on %s:\n" uuid dev#name;
	output_metadata stderr metadata
    ) pvsmap;

  (* Scan for volume groups.  The first entry in the metadata
   * appears to be the volume group name.  This gives us a
   * list of VGs and the metadata for each underlying PV.
   *)
  let vgnames =
    List.filter_map (
      function
      | pvuuid, (((vgname, Metadata vgmeta) :: _), dev) ->
	  Some (vgname, (pvuuid, vgmeta))
      | _ -> None
    ) pvsmap in

  let cmp ((a:string),_) ((b:string),_) = compare a b in
  let vgnames = List.sort ~cmp vgnames in
  let vgs = group_by vgnames in

  (* Note that the metadata is supposed to be duplicated
   * identically across all PVs (for redundancy purposes).
   * In theory we should check this and use the 'seqno'
   * field to find the latest metadata if it doesn't match,
   * but in fact we don't check this.
   *)
  let vgs = List.map (
    fun (vgname, metas) ->
      let pvuuids = List.map fst metas in
      let _, vgmeta = List.hd metas in (* just pick any metadata *)
      vgname, (pvuuids, vgmeta)) vgs in

  (* Print the VGs. *)
  if !debug then
    List.iter (
      fun (vgname, (pvuuids, vgmeta)) ->
	eprintf "VG %s is on PVs: %s\n%!" vgname (String.concat "," pvuuids)
    ) vgs;

  (* Some useful getter functions.  If these can't get a value
   * from the metadata or if the type is wrong they raise Not_found.
   *)
  let rec get_int63 field meta =
    match List.assoc field meta with
    | Int i -> i
    | _ -> raise Not_found
  and get_int_bounded field meta max =
    match List.assoc field meta with
    | Int i when i >= ~^0 && i <= Int63.of_int max -> Int63.to_int i
    | _ -> raise Not_found
  and get_string field meta =
    match List.assoc field meta with
    | String s -> s
    | _ -> raise Not_found
  and get_meta field meta =
    match List.assoc field meta with
    | Metadata md -> md
    | _ -> raise Not_found
  and get_stripes field meta =		(* List of (string,int) pairs. *)
    match List.assoc field meta with
    | List xs ->
	let rec loop = function
	  | [] -> []
	  | String pvname :: Int offset :: xs ->
	      (pvname, offset) :: loop xs
	  | _ -> raise Not_found
	in
	loop xs
    | _ -> raise Not_found
  in

  (* The volume groups refer to the physical volumes using their
   * own naming system ("pv0", "pv1", etc.) instead of PV UUIDs.
   *
   * Each PV also has a start (in sectors) & count (in extents)
   * of the writable area (the bit after the superblock and metadata)
   * which normally starts at sector 384.
   *
   * Create a PV device (simple offset + size) and a map from PV
   * names to these devices.
   *)
  let vgs = List.map (
    fun (vgname, (pvuuids, vgmeta)) ->
      let pvdevs, extent_size =
	try
	  (* NB: extent_size is in sectors here - we convert to bytes. *)
	  let extent_size =
	    get_int_bounded "extent_size" vgmeta (1024*1024) in
	  let extent_size = Int63.of_int extent_size in
	  let extent_size = extent_size *^ sector_size in

	  (* Get the physical_volumes section of the metadata. *)
	  let pvdevs = get_meta "physical_volumes" vgmeta in

	  List.filter_map (
	    function
	    | (pvname, Metadata meta) ->
		(* Get the UUID. *)
		let pvuuid = get_string "id" meta in
		let pvuuid = canonical_uuid pvuuid in

		(* Get the underlying physical device. *)
		let _, dev = List.assoc pvuuid pvsmap in

		(* Construct a PV device. *)
		let pe_start = get_int63 "pe_start" meta in
		let pe_start = pe_start *^ sector_size in
		let pe_count = get_int63 "pe_count" meta in
		let pe_count = pe_count *^ extent_size in
		let pvdev =
		  new offset_device
		    pvuuid (* name *)
		    pe_start pe_count (* start, size in bytes *)
		    (* don't really have a natural block size ... *)
		    extent_size
		    dev (* underlying device *) in

		Some (pvname, pvdev)
	    | _ ->
		None
	  ) pvdevs, extent_size
	with
	  (* Something went wrong - just return an empty map. *)
	  Not_found -> [], ~^0 in
      (vgname, (pvuuids, vgmeta, pvdevs, extent_size))
  ) vgs in

  (* Scan for logical volumes.  Each VG contains several LVs.
   * This gives us a list of LVs within each VG (hence extends
   * the vgs variable).
   *)
  let vgs = List.map (
    fun (vgname, (pvuuids, vgmeta, pvdevs, extent_size)) ->
      let lvs =
	try
	  let lvs = get_meta "logical_volumes" vgmeta in
	  let lvs = List.filter_map (
	    function
	    | lvname, Metadata lvmeta ->
		(try
		   let segment_count =
		     get_int_bounded "segment_count" lvmeta 1024 in

		   (* Get the segments for this LV. *)
		   let segments = range 1 (segment_count+1) in
		   let segments =
		     List.map
		       (fun i -> get_meta ("segment" ^ string_of_int i) lvmeta)
		       segments in

		   let segments =
		     List.map (
		       fun segmeta ->
			 let start_extent =
			   get_int63 "start_extent" segmeta in
			 let extent_count =
			   get_int63 "extent_count" segmeta in
			 let segtype = get_string "type" segmeta in

			 (* Can only handle striped segments at the
			  * moment. XXX
			  *)
			 if segtype <> "striped" then raise Not_found;

			 let stripe_count =
			   get_int_bounded "stripe_count" segmeta 1024 in
			 let stripes = get_stripes "stripes" segmeta in

			 if List.length stripes <> stripe_count then
			   raise Not_found;

			 (* Can only handle linear striped segments at
			  * the moment. XXX
			  *)
			 if stripe_count <> 1 then raise Not_found;
			 let pvname, pvoffset = List.hd stripes in

			 (start_extent, extent_count, pvname, pvoffset)
		     ) segments in

		   Some (lvname, segments)
		 with
		   (* Something went wrong with segments - omit this LV. *)
		   Not_found -> None)
	    | _ -> None
	  ) lvs in

	  lvs
	with
	  Not_found ->
	    (* Something went wrong - assume no LVs found. *)
	    [] in
      (vgname, (pvuuids, vgmeta, pvdevs, extent_size, lvs))
  ) vgs in

  (* Print the LVs. *)
  if !debug then (
    List.iter (
      fun (vgname, (pvuuids, vgmeta, pvdevs, extent_size, lvs)) ->
	eprintf "VG %s: (extent_size = %s bytes)\n" vgname
	  (Int63.to_string extent_size);
	List.iter (
	  fun (lvname, segments) ->
	    eprintf "  %s/%s:\n" vgname lvname;
	    List.iter (
	      fun (start_extent, extent_count, pvname, pvoffset) ->
		eprintf "    start %s count %s at %s:%s\n"
		  (Int63.to_string start_extent)
		  (Int63.to_string extent_count)
		  pvname (Int63.to_string pvoffset)
	    ) segments
	) lvs
    ) vgs;
    flush stderr
  );

(*  List.iter (fun pv -> attach_private_data pv vgs) pvs; *)

  (* Finally we can set up devices for the LVs. *)
  let lvs =
    List.map (
      fun (vgname, (pvuuids, vgmeta, pvdevs, extent_size, lvs)) ->
	try
	  List.map (
	    fun (lvname, segments) ->
	      let name = vgname ^ "/" ^ lvname in
	      let segments = List.map (
		fun (start_extent, extent_count, pvname, pvoffset) ->
		  (* Get the PV device. *)
		  let pvdev = List.assoc pvname pvdevs in

		  (* Extents                 mapped to:             *)
		  (start_extent, extent_count,          pvdev, pvoffset)
	      ) segments in

	      (* Create a linear mapping device. *)
	      let lv_dev = new linear_map_device name extent_size segments in

	      { lv_dev = lv_dev }
	  ) lvs
	with
	  Not_found -> []
    ) vgs in
  let lvs = List.concat lvs in

  (* Return the list of LV devices. *)
  lvs

(* XXX We need to reparse the metadata in a different way in
 * order to calculate this.  Need to generalize metadata handling.
 *)
and offset_is_free _ _ = false

and callbacks =
  let i = ref 0 in
  fun () -> {
    lvm_cb_uq = (incr i; !i);
    lvm_cb_name = id;
    lvm_cb_list_lvs = list_lvs;
    lvm_cb_offset_is_free = offset_is_free;
  }

(* Register the plugin. *)
let () = register_plugin ~lvm:probe id