(* virt-dmesg
 * (C) Copyright 2008-2011 Red Hat Inc.
 *
 * 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.
 *)

open Printf

open Utils

(* This code tries 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
 * addresses of the ordinary ksyms, and it has some characteristics
 * which make it easy to detect in the memory.
 * 
 * kallsyms contains a complete list of symbols so is much more useful
 * than the basic list of exports.
 *)

let min_kallsyms_tabsize = 1_000L
let max_kallsyms_tabsize = 1_000_000L (* 433056 kallsyms in 64 bit F14 kernel *)

type kallsyms_compr =
  | Compressed of (string * int64) list * int64
  | Uncompressed of (string * int64) list

let search_kallsyms k map =
  let ksym_addrs = filter_map (
    fun ksym -> try Some (StringMap.find ksym map) with Not_found -> None
  ) Ksyms.common_ksyms in

  (* Search for those kernel addresses in the image.  We're looking
   * for the table kallsyms_addresses followed by kallsyms_num_syms
   * (number of symbols in the table).
   *)
  let ksym_addrs =
    List.concat (List.map (Kernel.find_all_pointers k) ksym_addrs) in

  (* Test each one to see if it's a candidate list of kernel
   * addresses followed by length of list.
   *)
  let kallsymtabs = filter_map (
    fun addr ->
      (* Search upwards from address until we find the length field.
       * If found, jump backwards by length and check all addresses.
       *)
      debug "testing candidate kallsyms at %Lx" addr;

      (* Addresses which can appear in kallsyms can be validly not
       * mapped.  This appears to only happen for 64 bit kernels.
       *)
      let is_kallsyms_valid_address =
        match k.Kernel.wordsize with
        | Kernel.Word32 -> Kernel.is_mapped k
        | Kernel.Word64 ->
            fun addr ->
              Kernel.is_mapped k addr
              || Kernel.addr_compare addr 0xffffffff_80000000_L >= 0
              || Kernel.addr_compare addr 0x00000000_00200000_L <= 0
      in

      let rec loop prev_addrp addr =
	let addrp = Kernel.follow_pointer k addr in

        (* kallsyms are sorted, so this gives us an additional heuristic *)
        let increasing =
          match prev_addrp with
          | None -> true
          | Some prev_addrp -> Kernel.addr_compare prev_addrp addrp <= 0 in

	if increasing && is_kallsyms_valid_address addrp then
	  (* continue up the table *)
	  loop (Some addrp) (Kernel.succ_word k addr)
	else
	  if addrp >= min_kallsyms_tabsize &&
	    addrp <= max_kallsyms_tabsize then (
	      (* addrp might be the symbol count.  Count backwards and
	       * check the full table.
	       *)
              debug "addrp (num_entries) = %Lx" addrp;

	      let num_entries = Int64.to_int addrp in
	      let entry_size = Kernel.bytes_of_wordsize k in
	      let start_addr =
		addr -^ Int64.of_int (entry_size * num_entries) in
	      let end_addr = addr in
	      let rec loop2 addr =
		if Kernel.addr_compare addr end_addr < 0 then (
		  let addrp = Kernel.follow_pointer k addr in
		  if is_kallsyms_valid_address addrp then
		    loop2 (Kernel.succ_word k addr)
		  else (
                    debug "rejected on second pass because addrp = %Lx" addrp;
		    None
                  )
		) else
		  (* ok! *)
		  let names_addr = Kernel.succ_word k end_addr in
		  debug "candidate kallsyms found at %Lx (names_addr at %Lx, num_entries %d)"
		    start_addr names_addr num_entries;
		  Some (start_addr, num_entries, names_addr)
	      in
	      loop2 start_addr
	    )
	  else (
            debug "rejected on first pass because addrp = %Lx" addrp;
	    None
          )
      in
      match loop None addr with
      | None -> None
      | Some (start_addr, num_entries, names_addr) ->
	  (* As an additional verification, check the list of
	   * kallsyms_names.
	   *)
	  try
	    (* If the first byte is '\000' and is followed by a
	     * C identifier, then this is old-school list of
	     * symbols with prefix compression as in 2.6.9.
	     * Otherwise Huffman-compressed kallsyms as in
	     * 2.6.25.
	     *)
	    if Kernel.get_byte k names_addr = 0 &&
	      Kernel.is_C_identifier k (names_addr+^1L) then (
		let names = ref [] in
		let prev = ref "" in
		let rec loop names_addr start_addr num =
		  if num > 0 then (
		    let prefix = Kernel.get_byte k names_addr in
		    let prefix = String.sub !prev 0 prefix in
		    let name =
		      Kernel.get_string k (names_addr+^1L) in
		    let len = String.length name in
		    let name = prefix ^ name in
		    prev := name;
		    let names_addr = names_addr +^ Int64.of_int len +^ 2L in
		    let sym_value = Kernel.follow_pointer k start_addr in
		    let start_addr = Kernel.succ_word k start_addr in
		    (*eprintf "%S -> %Lx\n" name sym_value;*)
		    names := (name, sym_value) :: !names;
		    loop names_addr start_addr (num-1)
		  )
		in
		loop names_addr start_addr num_entries;
		let names = List.rev !names in

		Some (start_addr, num_entries, names_addr,
		      Uncompressed names)
	      )
	    else ( (* new-style "compressed" names. *)
	      let compressed_names = ref [] in
	      let rec loop names_addr start_addr num =
		if num > 0 then (
		  let len = Kernel.get_byte k names_addr in
		  let name = Kernel.get_memory k (names_addr+^1L) len in
		  let names_addr = names_addr +^ Int64.of_int len +^ 1L in
		  let sym_value = Kernel.follow_pointer k start_addr in
		  let start_addr = Kernel.succ_word k start_addr in
		  compressed_names :=
		    (name, sym_value) :: !compressed_names;
		  loop names_addr start_addr (num-1)
		) else
		  names_addr
	      in
	      let markers_addr = loop names_addr start_addr num_entries in
	      let markers_addr = Kernel.succ_align k markers_addr in
	      let compressed_names = List.rev !compressed_names in

	      Some (start_addr, num_entries, names_addr,
		    Compressed (compressed_names, markers_addr))
	    )
	  with
	    Invalid_argument _ -> None (* bad names list *)
  ) ksym_addrs in

  debug "candidate kallsyms at:";
  List.iter (
    function
    | (start_addr, num_entries, names_addr, Uncompressed _) ->
	debug "\t%Lx %d entries names_addr=%Lx old-style"
	  start_addr num_entries names_addr
    | (start_addr, num_entries, names_addr,
       Compressed (_, markers_addr)) ->
	debug "\t%Lx %d entries names_addr=%Lx markers_addr=%Lx"
	  start_addr num_entries names_addr markers_addr
  ) kallsymtabs;

  (* Vote for the most popular symbol table candidate and
   * return the list of kallsyms.
   *)
  let freqs = frequency kallsymtabs in
  match freqs with
  | [] ->
      (* Can't find any kallsymtabs. *)
      raise Not_found

  | (_, (_, _, _, Uncompressed names)) :: _ ->
      names

  | (_, (start_addr, num_entries, names_addr,
	 Compressed (compressed_names, markers_addr))) :: _ ->
      (* Skip the markers and look for the token table. *)
      let num_markers = Int64.of_int ((num_entries + 255) / 256) in
      let marker_size = Int64.of_int (Kernel.bytes_of_wordsize k) in
      let tokens_addr = markers_addr +^ marker_size *^ num_markers in

      (* Now read out the compression tokens, which are just
       * 256 ASCIIZ strings that map bytes in the compression
       * names to substrings.
       *)
      let tokens = Array.make 256 "" in
      let rec loop i addr =
	if i < 256 then (
	  let str = Kernel.get_string k addr in
	  let len = String.length str in
	  let addr = addr +^ Int64.of_int (len+1) in
	  tokens.(i) <- str;
	  loop (i+1) addr
	)
      in
      loop 0 tokens_addr;

      (* Expand the compressed names using the tokens. *)
      filter_map (
	fun (name, sym_value) ->
	  let f c = tokens.(Char.code c) in
	  let name = replace_chars f name in
	  (* First character in uncompressed output is the symbol
	   * type, eg. 'T'/'t' for text etc.
	   *)
	  (* NOTE: Symbol names are NOT unique
	   * (eg. 'con_start' is both a function and data in
	   * some kernels).  XXX We need to handle this situation
	   * better.
	   *)
	  (*let typ = name.[0] in*)
	  let name = String.sub name 1 (String.length name - 1) in
	  (*eprintf "%S -> %Lx\n" name sym_value;*)
	  Some (name, sym_value)
      ) compressed_names