path_tainter.ml

open Core_kernel
open Bap.Std
open Bap_primus.Std
open Bap_taint.Std
open Format
include Self()

type state = {
  path_taints : Taint.Object.Set.t
}


let state = Primus.Machine.State.declare
    ~uuid:"10d552e4-9c4d-40ce-8bfe-50faed839dcf"
    ~name:"path-tainter"
    (fun _ -> {path_taints = Taint.Object.Set.empty})

module Tracker(Machine : Primus.Machine.S) = struct
  open Machine.Syntax
  module Tracker = Taint.Tracker.Make(Machine)
  module Eval = Primus.Interpreter.Make(Machine)

  let on_jmp (cnd,_) =
    Tracker.lookup cnd Taint.Rel.direct >>= fun taints ->
    Machine.Local.update state ~f:(fun {path_taints} -> {
          path_taints = Set.union path_taints taints;
        })

  let on_load addr =
    Tracker.lookup addr Taint.Rel.indirect >>= fun taints ->
    Machine.Local.get state >>= fun {path_taints} ->
    Eval.pc >>= fun pc ->

    (* simple policy *)
    if not (Set.is_empty taints) && not (Set.is_empty path_taints)
    then eprintf "violation: %a: load(%a)@\n%!"
        Addr.pp pc Primus.Value.pp addr;
    Machine.return ()

  let init () = Machine.sequence Primus.Interpreter.[
      jumping >>> on_jmp;
      loading >>> on_load;
    ]
end

let () = Config.when_ready (fun _ ->
    Primus.Machine.add_component (module Tracker))