Elementary -- and broken -- implementation of virt-mem capture.
[virt-mem.git] / lib / virt_mem.ml
index 3a42448..fcefa3c 100644 (file)
@@ -32,13 +32,25 @@ module MMap = Virt_mem_mmap
 let min_kallsyms_tabsize = 1_000L
 let max_kallsyms_tabsize = 250_000L
 
-let kernel_size = 0x100_0000
-let max_memory_peek = 0x1_000
+(* Make the kernel size around 16 MB, but just a bit smaller than
+ * maximum string length so we can still run this on a 32 bit platform.
+ *)
+let kernel_size =
+  if Sys.word_size = 32 then Sys.max_string_length
+  else 0x100_0000
+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
+
+type image_with_ksyms =
+    int option
+    * string
     * Virt_mem_utils.architecture
     * ([`Wordsize], [`Endian]) Virt_mem_mmap.t
     * (ksym -> MMap.addr)
@@ -47,9 +59,120 @@ type kallsyms_compr =
   | Compressed of (string * MMap.addr) list * MMap.addr
   | Uncompressed of (string * MMap.addr) list
 
-let start usage_msg =
-  (* Debug messages. *)
+(* When tools register themselves, they are added to this list.
+ * Later, we will alphabetize the list.
+ *)
+let tools = ref []
+
+(* Registration function used by the 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.
+ *)
+let main () =
+  (* Get the registered tools, alphabetically. *)
+  let tools = !tools in
+  let tools = List.sort ~cmp:(fun (a,_) (b,_) -> compare a b) tools in
+
+  (* Which tool did the user want to run?  Look at the executable
+   * name (eg. 'virt-dmesg' => tool == dmesg).  If we don't recognise
+   * the executable name then we must look for the first parameter
+   * which doesn't begin with a '-' character.
+   *
+   * Note that we must do all of this before using the OCaml Arg
+   * module to properly parse the command line (below), so that
+   * we can have a usage message ready.
+   *)
+  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" *)
+      try Filename.chop_extension prog with Invalid_argument _ -> prog in
+    let prog =                         (* eg. "dmesg" *)
+      if String.starts_with prog "virt-" then
+       String.sub prog 5 (String.length prog - 5)
+      else prog in
+    try Some (List.assoc prog tools), false
+    with Not_found ->
+      let arg1 =                       (* First non-option argument. *)
+       match Array.to_list Sys.argv with
+       | [] -> None
+       | _::args ->
+           let rec loop = function
+             | [] -> None
+             | a::args when String.length a > 0 && a.[0] = '-' -> loop args
+             | a::_ -> Some a
+           in
+           loop args in
+      match arg1 with
+      | None -> None, false
+      | Some prog ->                   (* Recognisable first argument? *)
+         let prog =
+           try Filename.chop_extension prog with Invalid_argument _ -> prog in
+         let prog =
+           if String.starts_with prog "virt-" then
+             String.sub prog 5 (String.length prog - 5)
+           else prog 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, _, 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 =
+         List.fold_left max 0 (List.map String.length (List.map fst tools)) in
+       let tools = List.map (fun (l,r) -> pad max_width l, r) tools in
+       let tools = List.map (fun (l,r) -> "  " ^ l ^ " - " ^ r) tools in
+       let tools = String.concat "\n" tools in
+
+       sprintf (f_"\
+
+virt-mem: Tools for providing information about virtual machines
+
+Currently available tools include:
+%s
+
+General usage is:
+  <tool> [-options] [domains...]
+
+To display extra help for a single tool, do:
+  virt-mem help <tool>
+
+Options:") tools
+
+                                        (* Tool-specific usage message. *)
+    | Some (name, summary, description, external_cmd, _, _, _, _) ->
+       let cmd =
+         if external_cmd then "virt-" ^ name else "virt-mem " ^ name in
+
+       sprintf (f_"\
+
+%s: %s
+
+Description:
+%s
+
+Options:") cmd summary description in
+
+  (* Now begin proper parsing of the command line arguments. *)
   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
@@ -65,9 +188,9 @@ let start usage_msg =
   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
 
@@ -91,44 +214,101 @@ let start usage_msg =
     | str -> def_text_addr := Int64.of_string str
   in
 
-  (* List of kernel images. *)
-  let images = ref [] in
-  let uri = ref "" in
-
+  (* Handle -t option. *)
   let memory_image filename =
     images :=
       (!def_wordsize, !def_endian, !def_architecture, !def_text_addr, filename)
     :: !images
   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";
-  ] in
-
-  let anon_fun str =
-    raise (Arg.Bad (sprintf (f_"%s: unknown parameter") str)) in
-  let usage_msg = usage_msg ^ s_"\n\nOPTIONS" in
-  Arg.parse argspec anon_fun usage_msg;
+  (* Handle --version option. *)
+  let version () =
+    printf "virt-mem %s\n" Virt_mem_version.version;
+
+    let major, minor, release =
+      let v, _ = Libvirt.get_version () in
+      v / 1_000_000, (v / 1_000) mod 1_000, v mod 1_000 in
+    printf "libvirt %d.%d.%d\n" major minor release;
+    exit 0
+  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, _, _, _, _, argcheck, beforeksyms, run =
+    match tool with
+    | Some t -> t
+    | None ->
+       prerr_endline (s_"\
+virt-mem: I could not work out which tool you are trying to run.
+Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
+       exit 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 (
@@ -143,14 +323,36 @@ let start usage_msg =
          );
          exit 1 in
 
-      (* List of active domains. *)
+      (* If we have a list of parameters, then it is the domain names / UUIDs /
+       * IDs ONLY that we wish to display.  Otherwise, display all active.
+       *)
       let doms =
-       let nr_active_doms = C.num_of_domains conn in
-       let active_doms =
-         Array.to_list (C.list_domains conn nr_active_doms) in
-       let active_doms =
-         List.map (D.lookup_by_id conn) active_doms in
-       active_doms in
+       if anon_args = [] then (
+         (* List of active domains. *)
+         let nr_active_doms = C.num_of_domains conn in
+         let active_doms =
+           Array.to_list (C.list_domains conn nr_active_doms) in
+         List.map (D.lookup_by_id conn) active_doms
+       ) else (
+         List.map (
+           fun arg ->
+             let dom =
+               try D.lookup_by_uuid_string conn arg
+               with _ ->
+                 try D.lookup_by_name conn arg
+                 with _ ->
+                   try D.lookup_by_id conn (int_of_string arg)
+                   with _ ->
+                     failwith (sprintf (f_"%s: unknown domain (not a UUID, name or ID of any active domain)") arg) in
+
+             (* XXX Primitive test to see if the domain is active. *)
+             let is_active = try D.get_id dom >= 0 with _ -> false in
+             if not is_active then
+               failwith (sprintf (f_"%s: domain is not running") arg);
+
+             dom
+         ) anon_args
+       ) in
 
       (* Get their XML. *)
       let xmls = List.map (fun dom -> dom, D.get_xml_desc dom) doms in
@@ -170,6 +372,7 @@ let start usage_msg =
 
       List.map (
        fun (dom, _) ->
+         let id = D.get_id dom in
          let name = D.get_name dom in
 
          let wordsize =
@@ -196,8 +399,9 @@ let start usage_msg =
 
          if !def_text_addr = 0L then
            failwith
-             (sprintf (f_"%s: use -T to define kernel load address for this image")
-                name);
+             (sprintf (f_"%s: use -T to define kernel load address for this image") name);
+
+         let start_t = gettimeofday () in
 
          (* Read the kernel memory.
           * Maximum 64K can be read over remote connections.
@@ -214,6 +418,12 @@ let start usage_msg =
          in
          loop 0;
 
+         if debug then (
+           let end_t = gettimeofday () in
+           eprintf "timing: downloading kernel took %f seconds\n%!"
+             (end_t -. start_t)
+         );
+
          (* Map the virtual memory. *)
          let mem = MMap.of_string str !def_text_addr in
 
@@ -221,9 +431,13 @@ let start usage_msg =
          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
+    ) else (
+      (* One or more -t options passed. *)
+      if anon_args <> [] then
+       failwith (s_"virt-mem: if -t given on command line, then no domain arguments should be listed");
+
       List.map (
        fun (wordsize, endian, arch, text_addr, filename) ->
          (* Quite a lot of limitations on the kernel images we can
@@ -265,12 +479,27 @@ let start usage_msg =
          let mem = MMap.set_wordsize mem wordsize in
          let mem = MMap.set_endian mem endian in
 
-         (filename, arch, mem)
-      ) images in
+         ((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
@@ -292,12 +521,22 @@ let start usage_msg =
        (* Searching for <NUL>string<NUL> *)
        let common_ksyms_nul = List.map (sprintf "\000%s\000") common_ksyms in
 
+       let start_t = gettimeofday () in
+
        (* Search for these strings in the memory image. *)
        let ksym_strings = List.map (MMap.find_all mem) common_ksyms_nul in
        let ksym_strings = List.concat ksym_strings in
        (* Adjust found addresses to start of the string (skip <NUL>). *)
        let ksym_strings = List.map Int64.succ ksym_strings in
 
+       if debug then (
+         let end_t = gettimeofday () in
+         eprintf "timing: searching for common_ksyms took %f seconds\n%!"
+           (end_t -. start_t)
+       );
+
+       let start_t = gettimeofday () in
+
        (* For any we found, try to look up the symbol table
         * base addr and size.
         *)
@@ -359,11 +598,11 @@ let start usage_msg =
        let ksymtabs = List.filter (fun (_, size) -> size > 64L) ksymtabs in
 
        if debug then (
-         printf "%s: candidate symbol tables at:\n" name;
+         eprintf "%s: candidate symbol tables at:\n" name;
          List.iter (
            fun (addr, size) ->
-             printf "\t%Lx\t%Lx\t%!" addr size;
-             printf "first symbol: %s\n%!"
+             eprintf "\t%Lx\t%Lx\t%!" addr size;
+             eprintf "first symbol: %s\n%!"
                (MMap.get_string mem
                   (MMap.follow_pointer mem
                      (MMap.succ_long mem addr)))
@@ -382,13 +621,13 @@ let start usage_msg =
 
          | (_, (ksymtab_addr, ksymtab_size)) :: _ ->
              if debug then
-               printf
+               eprintf
                  "%s: Kernel symbol table found at %Lx, size %Lx bytes\n%!"
                  name ksymtab_addr ksymtab_size;
 
              (* 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
 
@@ -414,6 +653,14 @@ let start usage_msg =
              lookup_ksym
        in
 
+       if debug then (
+         let end_t = gettimeofday () in
+         eprintf "timing: searching for ordinary ksyms took %f seconds\n%!"
+           (end_t -. start_t)
+       );
+
+       let start_t = gettimeofday () in
+
        (* Now try to find the /proc/kallsyms table.  This is in an odd
         * compressed format (but not a very successful compression
         * format).  However if it exists we know that it will contain
@@ -444,7 +691,7 @@ let start usage_msg =
             * If found, jump backwards by length and check all addresses.
             *)
            if debug then
-             printf "%s: testing candidate kallsyms at %Lx\n" name addr;
+             eprintf "%s: testing candidate kallsyms at %Lx\n" name addr;
            let rec loop addr =
              let addrp = MMap.follow_pointer mem addr in
              if MMap.is_mapped mem addrp then
@@ -471,7 +718,7 @@ let start usage_msg =
                      (* ok! *)
                      let names_addr = MMap.succ_long mem end_addr in
                      if debug then
-                       printf "%s: candidate kallsyms found at %Lx (names_addr at %Lx, num_entries %d)\n"
+                       eprintf "%s: candidate kallsyms found at %Lx (names_addr at %Lx, num_entries %d)\n"
                          name start_addr names_addr num_entries;
                      Some (start_addr, num_entries, names_addr)
                  in
@@ -508,7 +755,7 @@ let start usage_msg =
                        let names_addr = names_addr +^ Int64.of_int len +^ 2L in
                        let sym_value = MMap.follow_pointer mem start_addr in
                        let start_addr = MMap.succ_long mem start_addr in
-                       (*printf "%S -> %Lx\n" name sym_value;*)
+                       (*eprintf "%S -> %Lx\n" name sym_value;*)
                        names := (name, sym_value) :: !names;
                        loop names_addr start_addr (num-1)
                      )
@@ -546,15 +793,15 @@ let start usage_msg =
        ) ksym_addrs in
 
        if debug then (
-         printf "%s: candidate kallsyms at:\n" name;
+         eprintf "%s: candidate kallsyms at:\n" name;
          List.iter (
            function
            | (start_addr, num_entries, names_addr, Uncompressed _) ->
-               printf "\t%Lx %d entries names_addr=%Lx old-style\n%!"
+               eprintf "\t%Lx %d entries names_addr=%Lx old-style\n%!"
                  start_addr num_entries names_addr
            | (start_addr, num_entries, names_addr,
               Compressed (_, markers_addr)) ->
-               printf "\t%Lx %d entries names_addr=%Lx markers_addr=%Lx\n%!"
+               eprintf "\t%Lx %d entries names_addr=%Lx markers_addr=%Lx\n%!"
                  start_addr num_entries names_addr markers_addr
          ) kallsymtabs
        );
@@ -619,7 +866,7 @@ let start usage_msg =
                   *)
                  (*let typ = name.[0] in*)
                  let name = String.sub name 1 (String.length name - 1) in
-                 (*printf "%S -> %Lx\n" name sym_value;*)
+                 (*eprintf "%S -> %Lx\n" name sym_value;*)
                  Some (name, sym_value)
              ) compressed_names in
 
@@ -632,6 +879,12 @@ let start usage_msg =
 
              lookup_ksym in
 
+       if debug then (
+         let end_t = gettimeofday () in
+         eprintf "timing: searching for kallsyms took %f seconds\n%!"
+           (end_t -. start_t)
+       );
+
        (* Just wrap the lookup_ksym call in something which prints
         * the query when debug is set.
         *)
@@ -640,10 +893,10 @@ let start usage_msg =
            let lookup_ksym sym =
              try
                let value = lookup_ksym sym in
-               printf "lookup_ksym %S = %Lx\n%!" sym value;
+               eprintf "lookup_ksym %S = %Lx\n%!" sym value;
                value
              with Not_found ->
-               printf "lookup_ksym %S failed\n%!" sym;
+               eprintf "lookup_ksym %S failed\n%!" sym;
                raise Not_found
            in
            lookup_ksym
@@ -651,7 +904,12 @@ let start usage_msg =
            lookup_ksym
        in
 
-       ((name, arch, mem, lookup_ksym) : image)
+       ((domid, name, arch, mem, lookup_ksym) : image_with_ksyms)
     ) images in
 
-  debug, images
+  (* Run the tool's main function. *)
+  (match run with
+   | None -> ()
+   | Some run ->
+       run debug images
+  )