propagates stub resolver results to program term attributes

plugins/stub_resolver/stub_resolver.ml

@@ -11,28 +11,31 @@ open KB.Syntax
 type groups = int Tid.Map.t
 type names = String.Set.t Int.Map.t
 
-type t = {
+type state = {
   groups : groups;
   names  : names;
   next   : int;
   stubs  : Tid.Set.t;
 }
 
-let tids = Knowledge.Domain.mapping (module Tid) "tids"
-    ~equal:Tid.equal
-    ~inspect:sexp_of_tid
+module Class = struct
+  type t
 
-let resolver = Knowledge.Class.declare
-    ~package "stub-resolver" ()
+  let t : (t,unit) KB.cls = Knowledge.Class.declare ~package "stubs" ()
 
-let result = Knowledge.Class.property
-    resolver "result" tids
-    ~persistent:(Knowledge.Persistent.of_binable (module struct
-                   type t = tid Tid.Map.t
-                   [@@deriving bin_io]
-                 end))
-    ~package
-    ~desc:"The mapping from stubs to real symbols"
+  let links = Knowledge.Class.property t "stub-refs"
+      ~package
+      ~desc:"Describes unambiguous connections between stubs \
+             and their implementations" @@
+    Knowledge.Domain.mapping (module Tid) "links"
+      ~equal:Tid.equal
+      ~inspect:sexp_of_tid
+
+  let stubs = KB.Class.property t "stub-tids"
+      ~package
+      ~desc:"The set of identified stubs" @@
+    KB.Domain.powerset (module Tid) "tids"
+end
 
 let empty = {
   groups = Map.empty (module Tid);
@@ -42,9 +45,14 @@ let empty = {
 }
 
 let is_stub sub =
-  KB.collect (Value.Tag.slot Sub.stub) (Term.tid sub) >>= function
-  | None -> KB.return false
-  | Some () -> KB.return true
+  if Term.has_attr sub Sub.stub then KB.return true
+  else match Term.get_attr sub address with
+    | None -> KB.return true
+    | Some addr ->
+      Theory.Label.for_addr (Word.to_bitvec addr) >>= fun sub ->
+      KB.collect (Value.Tag.slot Sub.stub) sub >>= function
+      | None -> KB.return false
+      | Some () -> KB.return true
 
 let aliases_of_sub s = KB.collect Theory.Label.aliases (Term.tid s)
 
@@ -126,19 +134,26 @@ let find_pairs t =
   unambiguous_pairs t.stubs
 
 let resolve prog =
-  Knowledge.Seq.fold ~init:empty
-    (Term.to_sequence sub_t prog) ~f:add >>|
-  find_pairs
+  Term.to_sequence sub_t prog |>
+  Knowledge.Seq.fold ~init:empty ~f:add >>| fun state ->
+  state, find_pairs state
 
 let provide prog =
-  Knowledge.Object.create resolver >>= fun obj ->
-  resolve prog >>= fun links ->
-  KB.provide result obj links >>= fun () ->
+  Knowledge.Object.create Class.t >>= fun obj ->
+  resolve prog >>= fun ({stubs},links) ->
+  KB.sequence [
+    KB.provide Class.links obj links;
+    KB.provide Class.stubs obj stubs;
+  ] >>= fun () ->
   KB.return obj
 
 let run prog =
-  match Knowledge.run resolver (provide prog) (Toplevel.current ()) with
-  | Ok (v,_) -> Knowledge.Value.get result v
+  match Knowledge.run Class.t (provide prog) (Toplevel.current ()) with
+  | Ok (v,_) -> v
   | Error cnf ->
     error "%a\n" Knowledge.Conflict.pp cnf;
-    Map.empty (module Tid)
+    KB.Value.empty Class.t
+
+type t = (Class.t,unit) KB.cls KB.Value.t
+let links = KB.Value.get Class.links
+let stubs = KB.Value.get Class.stubs

plugins/stub_resolver/stub_resolver.mli

@@ -3,6 +3,12 @@ open Bap.Std
 open Bap_core_theory
 open Bap_knowledge
 
+type t
+
 (** [run prog] - returns the mapping from
     stubs to implementations *)
-val run : program term -> tid Tid.Map.t
+val run : program term -> t
+
+val stubs : t -> Set.M(Tid).t
+
+val links : t -> tid Map.M(Tid).t

plugins/stub_resolver/stub_resolver_main.ml

@@ -316,8 +316,14 @@ let detect_stubs_by_signatures () : unit =
         Option.some_if (matches sigs mem) ())
 
 let update prog =
-  let links = Stub_resolver.run prog in
+  let resolver = Stub_resolver.run prog in
+  let stubs = Stub_resolver.stubs resolver
+  and links = Stub_resolver.links resolver in
   (object inherit Term.mapper
+    method! map_sub sub =
+      if Set.mem stubs (Term.tid sub)
+      then Term.set_attr sub Sub.stub ()
+      else sub
     method! map_jmp jmp =
       match Jmp.alt jmp with
       | None -> jmp
@@ -329,10 +335,10 @@ let update prog =
           | _ -> jmp
   end)#run prog
 
-let main = Project.map_program ~f:update
+let abi_pass = Project.map_program ~f:update
 
 let () = Extension.declare ~doc @@ fun ctxt ->
-  Bap_abi.register_pass main;
+  Bap_abi.register_pass abi_pass;
   mark_plt_as_stub ();
   detect_stubs_by_signatures ();
   Stubs.prepare ctxt;

plugins/stub_resolver/stub_resolver_tests.ml

@@ -121,7 +121,7 @@ let run name symbols expected should_fail _ctxt =
           Map.add_exn tids
             (tid_for_name_exn prog stub)
             (tid_for_name_exn prog impl)) in
-  let pairs = Stub_resolver.run prog in
+  let pairs = Stub_resolver.(links@@run prog) in
   let equal = Map.equal Tid.equal expected pairs in
   assert_bool name (equal || should_fail)