X-Git-Url: http://git.annexia.org/?a=blobdiff_plain;f=lib%2Fvirt_mem.ml;h=d1c4d2bbd30fdf4a334d414b97f1a1aa6a654a21;hb=c4e61cb8704aef23207d6186ab1e00f3d77ac5d7;hp=1f21a66ea1a2628eca4d54d761e89cb722588492;hpb=b70c967911e197b74d6d7ad98e3df9240d82572f;p=virt-mem.git

diff --git a/lib/virt_mem.ml b/lib/virt_mem.ml
index 1f21a66..d1c4d2b 100644
--- a/lib/virt_mem.ml
+++ b/lib/virt_mem.ml
@@ -43,9 +43,16 @@ let max_memory_peek = 65536 (* XXX Use D.max_peek function *)
 type ksym = string
 
 type image =
-    string
+    int option
+    * string
     * Virt_mem_utils.architecture
-    * ([`Wordsize], [`Endian]) Virt_mem_mmap.t
+    * ([`Wordsize], [`Endian], [`HasMapping]) Virt_mem_mmap.t
+
+type image_with_ksyms =
+    int option
+    * string
+    * Virt_mem_utils.architecture
+    * ([`Wordsize], [`Endian], [`HasMapping]) Virt_mem_mmap.t
     * (ksym -> MMap.addr)
 
 type kallsyms_compr =
@@ -55,23 +62,16 @@ type kallsyms_compr =
 (* When tools register themselves, they are added to this list.
  * Later, we will alphabetize the list.
  *)
-let tools = ref [
-  "capture", (
-    "capture",
-    s_"capture memory image for post-mortem analysis",
-    s_"Capture a memory image to a file for later post-mortem
-analysis.  Use the '-o memoryimage' option to specify the
-output file.
-
-Other tools can load the memory image using the '-t' option.",
-    false,
-    (fun _ _ -> ())
-  );
-]
+let tools = ref []
 
 (* Registration function used by the tools. *)
-let register name summary description is_cmd run_fn =
-  tools := (name, (name, summary, description, is_cmd, run_fn)) :: !tools
+let register ?(external_cmd = true) ?(extra_args = [])
+    ?argcheck ?beforeksyms ?run
+    name summary description =
+  tools :=
+    (name, (name, summary, description, external_cmd, extra_args,
+	    argcheck, beforeksyms, run))
+  :: !tools
 
 (* Main program, called from mem/virt_mem_main.ml when all the
  * tools have had a chance to register themselves.
@@ -90,7 +90,7 @@ let main () =
    * module to properly parse the command line (below), so that
    * we can have a usage message ready.
    *)
-  let tool =
+  let tool, ignore_first_anon_arg =
     let prog = Sys.executable_name in	(* eg. "/usr/bin/virt-dmesg.opt" *)
     let prog = Filename.basename prog in(* eg. "virt-dmesg.opt" *)
     let prog =				(* eg. "virt-dmesg" *)
@@ -99,7 +99,7 @@ let main () =
       if String.starts_with prog "virt-" then
 	String.sub prog 5 (String.length prog - 5)
       else prog in
-    try Some (List.assoc prog tools)
+    try Some (List.assoc prog tools), false
     with Not_found ->
       let arg1 =			(* First non-option argument. *)
 	match Array.to_list Sys.argv with
@@ -112,7 +112,7 @@ let main () =
 	    in
 	    loop args in
       match arg1 with
-      | None -> None
+      | None -> None, false
       | Some prog ->			(* Recognisable first argument? *)
 	  let prog =
 	    try Filename.chop_extension prog with Invalid_argument _ -> prog in
@@ -120,16 +120,17 @@ let main () =
 	    if String.starts_with prog "virt-" then
 	      String.sub prog 5 (String.length prog - 5)
 	    else prog in
-	  (try Some (List.assoc prog tools) with Not_found -> None) in
+	  (try Some (List.assoc prog tools), true
+	   with Not_found -> None, false) in
 
   (* Make a usage message. *)
   let usage_msg =
     match tool with
     | None ->				(* Generic usage message. *)
 	let tools = List.map (
-	  fun (name, (_, summary, _, is_cmd, _)) ->
-	    if is_cmd then "virt-"^name, summary
-	    else           "virt-mem "^name, summary
+	  fun (name, (_, summary, _, external_cmd, _, _, _, _)) ->
+	    if external_cmd then "virt-"^name, summary
+	    else                 "virt-mem "^name, summary
 	) tools in
 	(* Maximum width of field in the left hand column. *)
 	let max_width =
@@ -154,8 +155,9 @@ To display extra help for a single tool, do:
 Options:") tools
 
                                         (* Tool-specific usage message. *)
-    | Some (name, summary, description, is_cmd, _) ->
-	let cmd = if is_cmd then "virt-" ^ name else "virt-mem " ^ name in
+    | Some (name, summary, description, external_cmd, _, _, _, _) ->
+	let cmd =
+	  if external_cmd then "virt-" ^ name else "virt-mem " ^ name in
 
 	sprintf (f_"\
 
@@ -167,9 +169,10 @@ Description:
 Options:") cmd summary description in
 
   (* Now begin proper parsing of the command line arguments. *)
-
-  (* Debug messages. *)
   let debug = ref false in
+  let images = ref [] in
+  let uri = ref "" in
+  let anon_args = ref [] in
 
   (* Default wordsize. *)
   let def_wordsize = ref None in
@@ -185,9 +188,9 @@ Options:") cmd summary description in
   let set_endian = function
     | "auto" -> def_endian := None
     | "le" | "little" | "littleendian" | "intel" ->
-	def_endian := Some Bitmatch.LittleEndian
+	def_endian := Some Bitstring.LittleEndian
     | "be" | "big" | "bigendian" | "motorola" ->
-	def_endian := Some Bitmatch.BigEndian
+	def_endian := Some Bitstring.BigEndian
     | str -> failwith (sprintf (f_"set_endian: %s: unknown endianness") str)
   in
 
@@ -211,17 +214,14 @@ Options:") cmd summary description in
     | str -> def_text_addr := Int64.of_string str
   in
 
-  (* List of kernel images. *)
-  let images = ref [] in
-  let uri = ref "" in
-  let anon_args = ref [] in
-
+  (* Handle -t option. *)
   let memory_image filename =
     images :=
       (!def_wordsize, !def_endian, !def_architecture, !def_text_addr, filename)
     :: !images
   in
 
+  (* Handle --version option. *)
   let version () =
     printf "virt-mem %s\n" Virt_mem_version.version;
 
@@ -232,41 +232,67 @@ Options:") cmd summary description in
     exit 0
   in
 
-  let argspec = Arg.align [
-    "-A", Arg.String set_architecture,
-      "arch " ^ s_"Set kernel architecture, endianness and word size";
-    "-E", Arg.String set_endian,
-      "endian " ^ s_"Set kernel endianness";
-    "-T", Arg.String set_text_addr,
-      "addr " ^ s_"Set kernel text address";
-    "-W", Arg.String set_wordsize,
-      "addr " ^ s_"Set kernel word size";
-    "-c", Arg.Set_string uri,
-      "uri " ^ s_ "Connect to URI";
-    "--connect", Arg.Set_string uri,
-      "uri " ^ s_ "Connect to URI";
-    "--debug", Arg.Set debug,
-      " " ^ s_"Debug mode (default: false)";
-    "-t", Arg.String memory_image,
-      "image " ^ s_"Use saved kernel memory image";
-    "--version", Arg.Unit version,
-      " " ^ s_"Display version and exit";
-  ] in
-
+  (* Function to collect up any anonymous args (domain names/IDs). *)
   let anon_arg str = anon_args := str :: !anon_args in
+
+  (* Construct the argspec.
+   * May include extra arguments specified by the tool.
+   *)
+  let argspec =
+    let extra_args = match tool with
+      | None -> []
+      | Some (_, _, _, _, extra_args, _, _, _) -> extra_args in
+    let argspec = [
+      "-A", Arg.String set_architecture,
+        "arch " ^ s_"Set kernel architecture, endianness and word size";
+      "-E", Arg.String set_endian,
+        "endian " ^ s_"Set kernel endianness";
+      "-T", Arg.String set_text_addr,
+        "addr " ^ s_"Set kernel text address";
+      "-W", Arg.String set_wordsize,
+        "addr " ^ s_"Set kernel word size";
+      "-c", Arg.Set_string uri,
+        "uri " ^ s_ "Connect to URI";
+      "--connect", Arg.Set_string uri,
+        "uri " ^ s_ "Connect to URI";
+      "--debug", Arg.Set debug,
+        " " ^ s_"Debug mode (default: false)";
+      "-t", Arg.String memory_image,
+        "image " ^ s_"Use saved kernel memory image";
+      "--version", Arg.Unit version,
+        " " ^ s_"Display version and exit";
+    ] @ extra_args in
+
+    (* Sort options alphabetically on first alpha character. *)
+    let cmp (a,_,_) (b,_,_) =
+      let chars = "-" in
+      let a = String.strip ~chars a and b = String.strip ~chars b in
+      compare a b
+    in
+    let argspec = List.sort ~cmp argspec in
+    (* Make the options line up nicely. *)
+    Arg.align argspec in
+
+  (* Parse the command line.  This will exit if --version or --help found. *)
   Arg.parse argspec anon_arg usage_msg;
 
   let images = !images in
   let debug = !debug in
   let uri = if !uri = "" then None else Some !uri in
+
+  (* Discard the first anonymous argument if, above, we previously
+   * found it contained the tool name.
+   *)
   let anon_args = List.rev !anon_args in
+  let anon_args =
+    if ignore_first_anon_arg then List.tl anon_args else anon_args in
 
   (* At this point, either --help was specified on the command line
    * (and so the program has exited) or we must have determined tool,
    * or the user didn't give us a valid tool (eg. "virt-mem foobar").
    * Detect that final case now and give an error.
    *)
-  let name, _, _, _, run_fn =
+  let name, _, _, _, _, argcheck, beforeksyms, run =
     match tool with
     | Some t -> t
     | None ->
@@ -277,6 +303,12 @@ Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
   in
   if debug then eprintf "tool = %s\n%!" name;
 
+  (* Optional argument checking in the tool. *)
+  (match argcheck with
+   | None -> ()
+   | Some argcheck -> argcheck debug
+  );
+
   (* Get the kernel images. *)
   let images =
     if images = [] then (
@@ -340,6 +372,7 @@ Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
 
       List.map (
 	fun (dom, _) ->
+	  let id = D.get_id dom in
 	  let name = D.get_name dom in
 
 	  let wordsize =
@@ -398,7 +431,7 @@ Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
 	  let mem = MMap.set_wordsize mem wordsize in
 	  let mem = MMap.set_endian mem endian in
 
-	  (name, arch, mem)
+	  ((Some id, name, arch, mem) : image)
       ) xmls
     ) else (
       (* One or more -t options passed. *)
@@ -446,13 +479,27 @@ Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
 	  let mem = MMap.set_wordsize mem wordsize in
 	  let mem = MMap.set_endian mem endian in
 
-	  (filename, arch, mem)
+	  ((None, filename, arch, mem) : image)
       ) images
     ) in
 
+  (* Optional callback into the tool before we start looking for
+   * kernel symbols.
+   *)
+  (match beforeksyms with
+   | None -> ()
+   | Some beforeksyms -> beforeksyms debug images
+  );
+
+  (* If there is no run function, then there is no point continuing
+   * with the rest of the program (kernel symbol analysis) ...
+   *)
+  if run = None then exit 0;
+
+  (* Now kernel symbol analysis starts ... *)
   let images =
     List.map (
-      fun (name, arch, mem) ->
+      fun (domid, name, arch, mem) ->
 	(* Look for some common entries in the exported symbol table and
 	 * from that find the symbol table itself.  These are just
 	 * supposed to be symbols which are very likely to be present
@@ -580,7 +627,7 @@ Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
 
 	      (* Load the whole symbol table as a bitstring. *)
 	      let ksymtab =
-		Bitmatch.bitstring_of_string
+		Bitstring.bitstring_of_string
 		  (MMap.get_bytes mem ksymtab_addr
 		     (Int64.to_int ksymtab_size)) in
 
@@ -857,8 +904,12 @@ Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
 	    lookup_ksym
 	in
 
-	((name, arch, mem, lookup_ksym) : image)
+	((domid, name, arch, mem, lookup_ksym) : image_with_ksyms)
     ) images in
 
-  (* Run the actual tool. *)
-  run_fn debug images
+  (* Run the tool's main function. *)
+  (match run with
+   | None -> ()
+   | Some run ->
+       run debug images
+  )