Verifies a safety condition on all paths

path_check.ml

open Core_kernel.Std
open Bap.Std
open Graphlib.Std
open Format

include Self()

module CG = Graphs.Callgraph
module CFG = Graphs.Tid
module DAG = Graphlib.Make(Tid)(Unit)

type pcfg = {
  scc : tid partition;
  cfg : CFG.t;
}


type callsite = {
  caller : tid;
  callee : tid;
}

type t = {
  srcs : Tid.Set.t;
  dsts : Tid.Set.t;
  cg : CG.t;
}

type no_src_or_dst = [`No_src | `No_dst]
type must_not = [`Not_present | no_src_or_dst]
type may_not = [`Maybe_not]
type must = [`Calls]
type may = [`Equiv of tid group | `Pdom of (tid * tid)]
type decision = [must_not | may_not | may | must]


let remove_loops nodes cfg =
  let groups = Partition.groups nodes in
  Seq.cartesian_product groups groups |>
  Seq.filter  ~f:(fun (g1,g2) ->
      Seq.exists (Group.enum g1) ~f:(fun n1 ->
          Seq.exists (Group.enum g2) ~f:(fun n2 ->
              CFG.Node.has_edge n1 n2 cfg))) |>
  Seq.fold ~init:DAG.empty ~f:(fun g (g1,g2) ->
      let n1 = Group.top g1 and n2 = Group.top g2 in
      let edge = DAG.Edge.create n1 n2 () in
      DAG.Edge.insert edge g)

let insert_exit_node cfg =
  let exit = Tid.create () in
  exit, DAG.nodes cfg |> Seq.fold ~init:cfg ~f:(fun cfg n ->
      if DAG.Node.degree ~dir:`Out n cfg = 0
      then
        let edge = DAG.Edge.create n exit () in
        DAG.Edge.insert edge cfg
      else cfg)

let compute_pdom {scc; cfg} =
  let cfg = remove_loops scc cfg in
  let exit,cfg = insert_exit_node cfg in
  Graphlib.dominators (module DAG) ~rev:true cfg exit

let find_by_name prog name =
  Term.enum sub_t prog |>
  Seq.find_map ~f:(fun s ->
      Option.some_if (Sub.name s = name) (Term.tid s))

let calls_of_blk blk =
  Term.enum jmp_t blk |> Seq.filter_map ~f:(fun j ->
      match Jmp.kind j with
      | Goto _ | Ret _ | Int (_,_) -> None
      | Call dst -> match Call.target dst with
        | Direct tid -> Some {caller = Term.tid blk; callee=tid}
        | Indirect _ -> None)

let calls_of_program prog =
  Term.enum sub_t prog |>
  Seq.concat_map ~f:(fun sub ->
      Term.enum blk_t sub |>
      Seq.concat_map ~f:calls_of_blk)

let reaches cg target {callee} =
  Graphlib.is_reachable (module CG) cg callee target

let reaching cg callsites target =
  Seq.filter callsites ~f:(reaches cg target) |>
  Seq.fold ~init:Tid.Set.empty ~f:(fun cs {caller} ->
      Set.add cs caller)

let cfg_of_sub sub =
  let cfg = Sub.to_graph sub in
  { cfg; scc = Graphlib.strong_components (module CFG) cfg}

let partition prog =
  Term.enum sub_t prog |> Seq.map ~f:cfg_of_sub

let cross_product {srcs; dsts} =
  Seq.cartesian_product
    (Set.to_sequence srcs)
    (Set.to_sequence dsts)

let has_equiv p cfgs =
  cross_product p |> Seq.find_map ~f:(fun (src,dst) ->
      if Tid.equal src dst then None
      else Seq.find_map cfgs ~f:(fun {scc} ->
          if Partition.equiv scc src dst
          then Partition.group scc src
          else None))

let has_pdom p pdoms =
  cross_product p |> Seq.find ~f:(fun (src,dst) ->
      Seq.exists pdoms ~f:(fun pdom ->
          Tree.is_ancestor_of pdom ~child:src dst))

let verify_graph prog p =
  let cfgs = Seq.memoize (partition prog) in
  match has_equiv p cfgs with
  | Some x -> `Equiv x
  | None ->
    let pdoms = Seq.map cfgs ~f:compute_pdom |> Seq.memoize in
    match has_pdom p pdoms with
    | None -> `Maybe_not
    | Some x -> `Pdom x

let verify prog src dst : decision =
  match find_by_name prog src, find_by_name prog dst with
  | None,None    -> `Not_present
  | None, Some _ -> `No_src
  | Some _, None -> `No_dst
  | Some src, Some dst ->
    let cg = Program.to_graph prog in
    if Graphlib.is_reachable (module CG) cg src dst
    then `Calls
    else
      let calls = calls_of_program prog in
      verify_graph prog {
        srcs = reaching cg calls src;
        dsts = reaching cg calls dst;
        cg;
      }

let pp_decision src dst (d : decision) =
  let sol = match d with
    | #must_not -> "can't be called"
    | #may -> "may be called"
    | #must -> "is called"
    | #may_not -> "is unlikely to be called" in
  let desc = match d with
    | `Maybe_not -> "there is no such path, unless there is an indirect call"
    | `Calls -> sprintf "%s directly calls %s" src dst
    | `Not_present -> "both functions are not present in the binary"
    | `Equiv g ->
      asprintf "there are callsites to these functions that belong \
                to the same loop, e.g., %a" (Group.pp Tid.pp) g
    | `Pdom (cs,cd) ->
      asprintf
        "there is a callsite to %s that is postdominated by %s \
        e.g., %a and %a" src dst Tid.pp cs Tid.pp cd
    | #no_src_or_dst as t ->
      sprintf "%s is not present in the binary" @@ match t with
      | `No_src -> src
      | `No_dst -> dst in
  printf "%s %s before %s because %s@\n" src sol dst desc


let main src dst proj =
  let prog = Project.program proj in
  pp_decision src dst (verify prog src dst)


module Cmdline = struct
  open Config
  let src = param string "src" ~doc:"Name of the source function"
  let dst = param string "dst" ~doc:"Name of the sink function"
      ~synonyms:["sink"]

  ;;

  when_ready (fun {get=(!!)} ->
      Project.register_pass' (main !!src !!dst))

  ;;

  manpage [
    `S "DESCRIPTION";
    `P
      "Checks whether there is an execution path that contains a call
    to the $(v,src) function before a call to the $(v,dst)
    parameter. Outputs a possible (shortest) chain of calls that can
    lead to the policy violation."
  ]


end