DalvikBytecodeAnalysis/src/main/scala/analysis/Environment.scala

package analysis

import com.microsoft.z3._
import dex.DexType.{ObjectType, PrimitiveType, VoidType}
import dex.{DexApp, DexClass, DexField, DexMethod}
import domain.{AbstractDomain, AbstractDomainsCollection}

/**
  * Environment for the analysis.
  */
class Environment(
    val ctx: Context,
    val domains: AbstractDomainsCollection,
    val app: DexApp,
    val preAnalysisResult: PreAnalysisResult
) {

  val z3Obj = new Z3Object(ctx, domains)

  /**
    * Abstract registers.
    */
  object R {

    /**
      * Index of the ret register.
      */
    val retIdx: Int = -1

    /**
      * Approximation of a register.
      */
    class RApprox(
        val numApprox: BoolExpr,
        val referenceApprox: BoolExpr,
        val taintApprox: BoolExpr
    ) {}

    /**
      * Register at a program point.
      */
    class RAtPP(
        private val regCount: Int,
        val numbers: Int => domains.primitiveDomain.Element,
        val references: Int => domains.referenceDomain.Element,
        val taints: Int => domains.taintDomain.Element,
    ) {

      /**
        * Approximate the register to top (taint excluded).
        * @param i index of the register.
        * @return a register approximation.
        */
      def top(i: Int): RApprox =
        new RApprox(
          domains.primitiveDomain.approx(
            domains.primitiveDomain.top,
            this.numbers(i),
          ),
          domains.referenceDomain.approx(
            domains.referenceDomain.top,
            this.references(i),
          ),
          ctx.mkTrue()
        )

      /**
        * Approximate a register to another.
        * @param i index of the register that must be approximated.
        * @param j index of the register that must approximates.
        * @param r register at the desired program point.
        * @return a register approximation.
        */
      def <=(i: Int, j: Int)(r: RAtPP): RApprox =
        new RApprox(
          domains.primitiveDomain.approx(
            this.numbers(i),
            r.numbers(j)
          ),
          domains.referenceDomain.approx(
            this.references(i),
            r.references(j)
          ),
          domains.taintDomain.approx(
            this.taints(i),
            r.taints(j)
          )
        )

      /**
        * Approximate all the registers between two program points.
        * @param r registers at the desider program point.
        * @return registers approximations.
        */
      def <=(r: RAtPP): List[RApprox] =
        (-1 until this.regCount)
          .map(i => this.<=(i, i)(r))
          .toList

      /**
        * Approximate all the registers between two program points except for the register at the given index.
        * @param n index of the register to exclude.
        * @param r registers at the desider program point.
        * @return registers approximations.
        */
      def <=(n: Int)(r: RAtPP): List[RApprox] =
        (-1 until this.regCount)
          .filter(i => i != n)
          .map(i => this.<=(i, i)(r))
          .toList

    }

    /**
      * Returns the registers at the given program points.
      * @param pp the program point.
      * @return the registers.
      */
    def apply(pp: PP): RAtPP =
      new RAtPP(
        pp.m.registersCount,
        rTypeAtPP(pp, domains.primitiveDomain, "R-prim"),
        rTypeAtPP(pp, domains.referenceDomain, "R-ref"),
        rTypeAtPP(pp, domains.taintDomain, "R-taint")
      )

    /**
      * Returns a function that maps a register index to the given domain Element.
      * @param pp program point of the registers.
      * @param domain the chosen domain.
      * @param prefix prefix to distinguish between Elements of different domains.
      * @tparam E the domain Element type.
      * @return function that maps the index to the Element.
      */
    private def rTypeAtPP[E](
        pp: PP,
        domain: AbstractDomain,
        prefix: String
    ): Int => E =
      (i) =>
        ctx
          .mkApp(
            ctx.mkFuncDecl(
              s"$prefix-${pp.c.name}-${pp.m.name}-${pp.pc}",
              Array[Sort](ctx.mkIntSort()),
              domain.sort
            ),
            ctx.mkInt(i)
          )
          .asInstanceOf[E]

    /**
      * Convert a list of registers approximations into a list of Z3 constraints.
      */
    implicit def toBoolExpr(approxMap: List[RApprox]): List[BoolExpr] =
      approxMap
        .flatMap(v => toBoolExpr(v))

    /**
      * Convert a register approximation into a list of Z3 constraints.
      */
    implicit def toBoolExpr(approx: RApprox): List[BoolExpr] =
      List(approx.numApprox, approx.referenceApprox, approx.taintApprox)

  }

  /**
    * Abstract heap.
    */
  object H {

    /**
      * Returns the heap that maps references to objects.
      */
    def obj(): FuncDecl = {
      val z3Obj = new Z3Object(ctx, domains)
      ctx.mkFuncDecl(
        "H",
        Array[Sort](ctx.mkStringSort()),
        z3Obj.sort
      )
    }
  }

  /**
    * Abstract static heap.
    */
  object SH {

    /**
      * Returns the class field value.
      * @param dexClass the class.
      * @param dexField the field.
      * @return a tuple where the first element is the value and the second the taint.
      */
    def apply(dexClass: DexClass, dexField: DexField): (Expr, Expr) = {
      (
        dexField.dexType match {
          case PrimitiveType() =>
            const(dexClass, dexField, domains.primitiveDomain.sort, "SH-prim")
          case ObjectType(_) =>
            const(dexClass, dexField, domains.referenceDomain.sort, "SH-ref")
        },
        const(dexClass, dexField, domains.taintDomain.sort, "SH-taint")
      )
    }

    /**
      * Returns the constant that represent the class field.
      * @param dexClass the class.
      * @param dexField the field.
      * @param sort the type of the constant domain.
      * @param prefix prefix to distinguish between the same constant in different domains.
      * @return the constant.
      */
    private def const(
        dexClass: DexClass,
        dexField: DexField,
        sort: Sort,
        prefix: String
    ): Expr = {
      ctx.mkConst(s"$prefix-${dexClass.name}-${dexField.name}", sort)
    }
  }

  /**
    * Method return value.
    */
  object M {

    /**
      * Value returned by the given class method.
      * @param dexClass the class.
      * @param dexMethod the method.
      * @return the value.
      */
    def apply(dexClass: DexClass, dexMethod: DexMethod): Option[Expr] = {
      dexMethod.returnType match {
        case PrimitiveType() =>
          Option(
            typeReturn(
              dexClass,
              dexMethod,
              domains.primitiveDomain.sort,
              "M-prim"
            )
          )
        case ObjectType(_) =>
          Option(
            typeReturn(
              dexClass,
              dexMethod,
              domains.referenceDomain.sort,
              "M-ref"
            )
          )
        case VoidType() => Option.empty
      }
    }

    /**
      * Method result in the given domain.
      */
    private def typeReturn(
        dexClass: DexClass,
        dexMethod: DexMethod,
        sort: Sort,
        prefix: String
    ): Expr = {
      ctx.mkConst(
        s"$prefix-${dexClass.name}-${dexMethod.name}",
        sort
      )
    }

  }

  /**
    * Method invocation helper.
    */
  object Dispatch {

    /**
      * Returns a map where the key is the class that has the given
      * method, either because it implements it or because it is inherited,
      * and the value is the corresponding class with the method implementation.
      */
    def apply(dexMethod: String,
              definingClass: String): Map[DexClass, DexClass] = {
      dispatchAux(dexMethod, definingClass).toMap
    }

    /**
      * Tail recursive search of method implementations.
      */
    private def dispatchAux(
        dexMethod: String,
        definingClass: String): Seq[(DexClass, DexClass)] = {
      val defClass = app.classes.find(c => c.name == definingClass)
      val impl: Option[DexClass] =
        defClass match {
          case Some(c) => firstImpl(dexMethod, c)
          case None    => Option.empty
        }
      val impls = app.classes
        .filter(c => c.superClass == definingClass)
        .flatMap(
          c => dispatchAux(dexMethod, c.name)
        )
        .toList
      impl match {
        case Some(c) => (defClass.get -> c) :: impls
        case None    => impls
      }
    }

    /**
      * Returns the first implementation (if present) of the given method when invoked to the given class.
      */
    private def firstImpl(dexMethod: String,
                          dexClass: DexClass): Option[DexClass] = {
      dexClass.methods.find(m => m.name == dexMethod) match {
        case Some(_) => Option(dexClass)
        case None =>
          app.classes.find(c => c.name == dexClass.superClass) match {
            case Some(c) => firstImpl(dexMethod, c)
            case None    => Option.empty
          }
      }
    }

  }

}