Skip to content

Instantly share code, notes, and snippets.

@ivg
Created August 25, 2020 22:51
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save ivg/bbd1ba6d9362e778e723a467a3f8c8b0 to your computer and use it in GitHub Desktop.
Save ivg/bbd1ba6d9362e778e723a467a3f8c8b0 to your computer and use it in GitHub Desktop.
A simple example that shows how to explore the knowledge base.
open Core_kernel
open Bap_main
open Bap_knowledge
open Bap_core_theory
open Bap.Std
open KB.Syntax
let zero_collector = object
inherit [Tid.Set.t] Term.visitor
method! enter_def t zeros = match Def.rhs t with
| Bil.Int x when Word.is_zero x -> Set.add zeros (Term.tid t)
| _ -> zeros
end
let collect_zeros =
KB.objects Theory.Program.cls >>=
KB.Seq.fold ~init:Tid.Set.empty ~f:(fun tids obj ->
KB.collect Theory.Program.Semantics.slot obj >>| fun sema ->
KB.Value.get Term.slot sema |>
List.fold ~init:tids ~f:(fun tids blk ->
zero_collector#visit_blk blk tids))
type collector
let collector = KB.Class.declare "zero-collector" ()
let tids = KB.Domain.powerset (module Tid) "tids"
let result = KB.Class.property collector "result" tids
let create : collector KB.Object.t knowledge =
KB.Object.create collector >>= fun obj ->
collect_zeros >>= fun zeros ->
KB.provide result obj zeros >>| fun () ->
obj
let print_zeros state =
match Knowledge.run collector create state with
| Error _ -> ()
| Ok (collector,_state) ->
let zeros = KB.Value.get result collector in
Set.iter zeros ~f:(Format.printf "%a@\n" Tid.pp)
let main _ =
print_zeros @@ Toplevel.current ()
let () =
Extension.declare @@ fun _ctxt ->
Project.register_pass' main;
Ok ()
@ivg
Copy link
Author

ivg commented Aug 25, 2020

Building

bapbuild knowledge_explorer.plugin && bapbundle install knowledge_explorer.plugin

Using

bap ./exe -pknowledge-explorer

Description

In this example, we will build a program that will scan through the knowledge base and collect all terms that are assigned to zero. For no other reason than the didactic one, we will use BIR to find such objects.

First of all, since we would like to analyze the BIL IR we will Term.visitor to collect all terms that are assigned to zero,

let zero_collector = object
  inherit [Tid.Set.t] Term.visitor
  method! enter_def t zeros = match Def.rhs t with
    | Bil.Int x when Word.is_zero x -> Set.add zeros (Term.tid t)
    | _ -> zeros
end

Now, we will use the Knowledge.objects function to collect all zeros assigned terms,

let collect_zeros =
  KB.objects Theory.Program.cls >>=
  KB.Seq.fold ~init:Tid.Set.empty ~f:(fun tids obj ->
      KB.collect Theory.Program.Semantics.slot obj >>| fun sema ->
      KB.Value.get Term.slot sema |>
      List.fold ~init:tids ~f:(fun tids blk ->
          zero_collector#visit_blk blk tids))

I.e., for each object in the knowledge base, we compute its semantics, then extract the IR property, which is represented as a list of IR blks, and fold over it using our IR visitor.

Sometimes that is enough, you can write this computation and run it as a part of your overall analysis. Moreover, we will soon provide an extension point which will let you register your knowledge computation and run it, the same as passes are run on the old-style project. But for didactic purposes, let's learn how to run the knowledge monad manually.

The run function has a rather strange type,

  val run : ('k,'s) cls -> 'k obj t -> state -> (('k,'s) cls value * state, conflict) result

instead of a more natural (for state monads) 'a monad -> state -> ('a * state,error) Result.t we have to provide a class and a computation that returns an object just to get the value of that object as the result (along with the modified state).
The reason for that is that knowledge monad computes objects, not values and in general a knowledge computation should be treated more like a statement or command, rather than expression. And the whole computation is about the side-effects, not computing the value. Therefore, in order to extract some value from the knowledge computation, we must create an object that bears that value as a property. Most likely we will need to create a separate class for that. Pretty much like in Java we have to create a class just to run the main function :)

So, let's define a class and its properties,

type collector
let collector = KB.Class.declare "zero-collector" ()
let tids = KB.Domain.powerset (module Tid) "tids"
let result = KB.Class.property collector "result" tids

let create : collector KB.Object.t knowledge =
  KB.Object.create collector >>= fun obj ->
  collect_zeros >>= fun zeros ->
  KB.provide result obj zeros >>| fun () ->
  obj

We declared the collector class that has only one property result that has type tids, which has the structure of the powerset domain. We also implemented the constructor for that class, which collects all zero assigned terms. Now we can run it,

let print_zeros state =
  match Knowledge.run collector create state with
  | Error _ -> ()
  | Ok (collector,_state) ->
    let zeros = KB.Value.get result collector in
    Set.iter zeros ~f:(Format.printf "%a@\n" Tid.pp)

But to run we need the state, i.e., the knowledge base itself. When we will finish with #1169, we wouldn't need this. In fact, even now we can load the knowledge and store the knowledge. But we don't want to do this, and rely on the disassemble command that will disassemble the binary and provide the knowledge. So consider the following as a sort of workaround. We register a pass, then extract the knowledge from the toplevel and run our print_zeros function.

let main _ =
  print_zeros @@ Toplevel.current ()

let () =
  Extension.declare @@ fun _ctxt ->
  Project.register_pass' main;
  Ok ()

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment