dwarf-synthesis/DwarfSynth/Simplest.ml

321 lines
10 KiB
OCaml

open Std
module CFG = BStd.Graphs.Ir
type memory_offset = int64
type cfa_pos =
RspOffset of memory_offset
| RbpOffset of memory_offset
| CfaLostTrack
type cfa_change = CfaChange of BStd.word * cfa_pos
type cfa_changes_fde = cfa_change list
module StrMap = Map.Make(String)
type cfa_changes = cfa_changes_fde StrMap.t
module TIdMap = Map.Make(BStd.Tid)
exception InvalidSub
let pp_cfa_pos ppx = function
| RspOffset off -> Format.fprintf ppx "RSP + (%s)@." (Int64.to_string off)
| RbpOffset off -> Format.fprintf ppx "RBP + (%s)@." (Int64.to_string off)
| CfaLostTrack -> Format.fprintf ppx "??@."
let pp_cfa_change ppx = function CfaChange(addr, cfa_pos) ->
Format.fprintf ppx "%a: %a" BStd.Word.pp_hex addr pp_cfa_pos cfa_pos
let pp_cfa_changes_fde ppx = List.iter (pp_cfa_change ppx)
let pp_cfa_changes ppx =
StrMap.iter (fun fde_name entry ->
Format.fprintf ppx "%s@\n====@\n@\n%a@." fde_name
pp_cfa_changes_fde entry)
let pp_option_of sub_pp ppx = function
| None -> Format.fprintf ppx "None"
| Some x -> Format.fprintf ppx "Some %a" sub_pp x
let addr_of term =
(** Get the address of a term *)
match BStd.Term.get_attr term BStd.address with
| None -> assert false
| Some addr -> addr
let interpret_var_expr c_var offset expr = BStd.Bil.(
let closed_form = BStd.Exp.substitute
(var c_var)
(int (BStd.Word.of_int64 offset))
expr
in
let res = BStd.Exp.eval closed_form in
match res with
| Imm value ->
Some (BStd.Word.to_int64_exn @@ BStd.Word.signed value)
| _ -> None
)
let process_def def (cur_offset: memory_offset)
: ((cfa_change * memory_offset) option) =
let lose_track addr =
Some (CfaChange(addr, CfaLostTrack), Int64.zero)
in
(match Regs.X86_64.of_var (BStd.Def.lhs def) with
| Some reg when reg = Regs.X86_64.rsp ->
let exp = BStd.Def.rhs def in
let free_vars = BStd.Exp.free_vars exp in
let free_x86_regs = Regs.X86_64.map_varset free_vars in
(match Regs.DwRegOptSet.cardinal free_x86_regs with
| 1 ->
let free_var = Regs.DwRegOptSet.choose free_x86_regs in
(match free_var with
| Some dw_var when dw_var = Regs.X86_64.rsp ->
let bil_var = (match BStd.Var.Set.choose free_vars with
| None -> assert false
| Some x -> x) in
let interpreted = interpret_var_expr bil_var cur_offset exp in
(match interpreted with
| None -> lose_track (addr_of def)
| Some interp_val ->
let gap = Int64.sub interp_val cur_offset in
let new_offset = Int64.sub cur_offset gap in
Some (CfaChange(addr_of def, RspOffset(new_offset)), new_offset)
)
| _ -> lose_track (addr_of def)
)
| _ -> lose_track @@ addr_of def
)
| _ -> None)
let process_jmp jmp (cur_offset: memory_offset)
: ((cfa_change * memory_offset) option) =
let gen_change off =
let new_offset = Int64.add cur_offset (Int64.of_int off) in
Some (CfaChange(addr_of jmp, RspOffset(new_offset)), new_offset)
in
match (BStd.Jmp.kind jmp) with
| BStd.Call call -> gen_change (-8)
| BStd.Ret ret -> gen_change (8)
| _ -> None
let sym_of_blk blk : cfa_changes_fde =
(** Extracts the symbolic CFA changes of a block. These changes assume that
at the beginning of the block, CFA = RspOffset(0) and will be offset
after *)
let apply_offset (accu, cur_offset) = function
| None -> (accu, cur_offset)
| Some (change, n_offset) -> (change :: accu, n_offset)
in
let fold_elt (accu, cur_offset) elt = match elt with
| `Def(def) ->
apply_offset (accu, cur_offset) @@ process_def def cur_offset
| `Jmp(jmp) ->
apply_offset (accu, cur_offset) @@ process_jmp jmp cur_offset
| _ -> (accu, cur_offset)
in
let elts_seq = BStd.Blk.elts blk in
let out, end_offset = BStd.Seq.fold elts_seq
~init:([], Int64.zero)
~f:fold_elt in
out
let end_offset (changelist: cfa_changes_fde): memory_offset option =
match changelist with
| CfaChange(_, RspOffset(x)) :: _ -> Some x
| _ -> None
exception Inconsistent of BStd.tid
let rec dfs_propagate changemap propagated parent_val node graph =
let c_tid = BStd.Term.tid @@ CFG.Node.label node in
match TIdMap.find_opt c_tid propagated with
| Some x ->
if x = parent_val then
(* Already propagated and consistent, all fine *)
propagated
else
(* Already propagated with a different value: inconsistency *)
raise (Inconsistent c_tid)
| None ->
let n_propagated = TIdMap.add c_tid parent_val propagated in
let outwards = CFG.Node.outputs node graph in
let self_entry = TIdMap.find c_tid changemap in
let offset = (match end_offset self_entry with
Some x -> x
| None -> Int64.zero) in
let cur_val = Int64.add parent_val offset in
BStd.Seq.fold outwards
~init:n_propagated
~f:(fun accu edge ->
dfs_propagate changemap accu cur_val (CFG.Edge.dst edge) graph)
let get_entry_blk graph =
let entry =
BStd.Seq.find (CFG.nodes graph)
(fun node -> BStd.Seq.is_empty @@ CFG.Node.inputs node graph)
in match entry with
| None -> assert false
| Some x -> x
let term_addr term =
BStd.Term.get_attr term BStd.address
let same_keys map1 map2 =
let exists_in_other other key _ =
TIdMap.mem key other in
TIdMap.for_all (exists_in_other map2) map1
&& TIdMap.for_all (exists_in_other map1) map2
let of_sub sub : cfa_changes_fde =
(** Extracts the `cfa_changes_fde` of a subroutine *)
Format.eprintf "Sub %s...@." @@ BStd.Sub.name sub ;
let initial_cfa_rsp_offset = Int64.of_int 8 in
let store_sym accu blk =
let blk = CFG.Node.label blk in
let res = sym_of_blk blk in
TIdMap.add (BStd.Term.tid blk) res accu
in
let node_addr nd = term_addr @@ CFG.Node.label nd in
let merge_corrected blk_tid changes offset = match (changes, offset) with
| Some changes, Some offset ->
Some (
List.map (fun (CfaChange(addr, pos)) -> match pos with
RspOffset(off) -> CfaChange(addr,
RspOffset(Int64.add off offset))
| RbpOffset(off) -> CfaChange(addr, RbpOffset(off))
| CfaLostTrack -> CfaChange(addr, CfaLostTrack)
)
changes
)
| _ -> None
in
let cfg = BStd.Sub.to_cfg sub in
let tid_match = BStd.Seq.fold (CFG.nodes cfg)
~init:TIdMap.empty
~f:(fun accu node ->
let tid = BStd.Term.tid @@ CFG.Node.label node in
TIdMap.add tid node accu)
in
let blk_sym = BStd.Seq.fold
~init:TIdMap.empty
~f:store_sym
@@ CFG.nodes cfg
in
let entry_blk = get_entry_blk cfg in
let offset_map = dfs_propagate
blk_sym (TIdMap.empty) initial_cfa_rsp_offset entry_blk cfg in
let corrected = TIdMap.merge merge_corrected blk_sym offset_map in
let is_connex = same_keys tid_match corrected in
if not is_connex then
raise InvalidSub ;
let tid_list = TIdMap.bindings tid_match in
let sorted_blk = List.sort (fun (tid1, bl1) (tid2, bl2) ->
let res = match (node_addr bl1, node_addr bl2) with
| Some addr1, Some addr2 -> compare addr1 addr2
| Some _, None -> 1
| None, Some _ -> -1
| None, None -> compare tid1 tid2
in
-res)
tid_list
in
let out = List.fold_left
(fun accu blk ->
let changes = TIdMap.find blk corrected in
List.fold_left (fun accu chg -> chg::accu)
accu
changes
)
[]
(List.map (fun (x, y) -> x) sorted_blk) in
let sub_addr = (match term_addr sub with
| Some x -> x
| None -> assert false) in
let init = [
CfaChange(sub_addr, RspOffset(initial_cfa_rsp_offset))
] in
init @ out
let cleanup_fde (fde_changes: cfa_changes_fde) : cfa_changes_fde =
(** Cleanup the result of `of_sub`.
Merges entries at the same address, propagates track lost *)
let fold_one (accu, (last_commit:cfa_pos option), in_flight, lost_track) = function
| CfaChange(addr, cfa_change) as cur_change -> (
match lost_track, in_flight, cfa_change with
| true, _, _ ->
(* Already lost track: give up *)
(accu, last_commit, None, lost_track)
| false, _, CfaLostTrack ->
(* Just lost track: give up the operation on flight as well *)
(cur_change :: accu, None, None, true)
| _, Some CfaChange(flight_addr, flight_chg), _
when flight_addr = addr ->
(* On flight address matches current address: continue flying *)
accu, last_commit, Some cur_change, false
| _, Some CfaChange(_, in_flight_inner_pos), _
when last_commit = Some in_flight_inner_pos ->
(* Doesn't match anymore, but there was some operation in flight,
which has the same result as what was last committed. Discard. *)
(accu, last_commit, Some cur_change, false)
| _, Some (CfaChange(_, in_flight_inner_pos) as in_flight_inner), _ ->
(* Doesn't match anymore, but there was some operation in flight:
commit it, put the new one in flight *)
(in_flight_inner :: accu, Some in_flight_inner_pos,
Some cur_change, false)
| _, None, _ ->
(* No operation in flight: put the new one in flight *)
(accu, last_commit, Some cur_change, false)
)
in
let extract_end_value (accu, _, in_flight, lost_track) =
List.rev @@ match lost_track, in_flight with
| true, _ -> accu
| false, None -> accu
| false, Some x -> x :: accu
in
extract_end_value
@@ List.fold_left fold_one ([], None, None, false) fde_changes
let of_prog prog : cfa_changes =
(** Extracts the `cfa_changes` of a program *)
let fold_step accu sub =
(try
let res = (*cleanup_fde @@ *)of_sub sub in
StrMap.add (BStd.Sub.name sub) res accu
with InvalidSub -> accu)
in
let subroutines = BStd.Term.enum BStd.sub_t prog in
BStd.Seq.fold subroutines
~init:StrMap.empty
~f:fold_step
let of_proj proj : cfa_changes =
(** Extracts the `cfa_changes` of a project *)
let prog = BStd.Project.program proj in
of_prog prog