TaintedFlow.h

#ifndef ROSE_BinaryAnalysis_TaintedFlow_H
#define ROSE_BinaryAnalysis_TaintedFlow_H
#include <featureTests.h>
#ifdef ROSE_ENABLE_BINARY_ANALYSIS

#include <Rose/BinaryAnalysis/DataFlow.h>
#include <Rose/Diagnostics.h>

#include <boost/shared_ptr.hpp>
#include <stdexcept>

namespace Rose {
namespace BinaryAnalysis {

class TaintedFlow {
public:
    enum Taintedness { BOTTOM, NOT_TAINTED, TAINTED, TOP };

    enum Approximation { UNDER_APPROXIMATE, OVER_APPROXIMATE };

    static Taintedness merge(Taintedness, Taintedness);

    typedef std::pair<DataFlow::Variable, Taintedness> VariableTaint;

    //                                  State
public:
    class State {
        typedef std::list<VariableTaint> VarTaintList;
        VarTaintList taints_;

    public:
        typedef boost::shared_ptr<State> Ptr;

    protected:
        // Initialize taintedness for all variables; this is protected because this is a reference-counted object
        State(const DataFlow::VariableList &variables, Taintedness taint) {
            for (const DataFlow::Variable &variable: variables)
                taints_.push_back(std::make_pair(variable, taint));
        }

    public:
        static State::Ptr instance(const DataFlow::VariableList &variables, Taintedness taint = BOTTOM) {
            return State::Ptr(new State(variables, taint));
        }

        virtual State::Ptr copy() const {
            return State::Ptr(new State(*this));
        }

        virtual ~State() {}

        Taintedness& lookup(const DataFlow::Variable&);

        bool setIfExists(const DataFlow::Variable&, Taintedness);

        bool merge(const State::Ptr&);

        const VarTaintList& variables() const { return taints_; }
        VarTaintList& variables() { return taints_; }
        void print(std::ostream&) const;
    };

    typedef State::Ptr StatePtr;

    //                                  Transfer function
protected:
    class TransferFunction {
        const DataFlow::VertexFlowGraphs &index_; // maps CFG vertex to data flow graph
        Approximation approximation_;
        SmtSolverPtr smtSolver_;
        Sawyer::Message::Facility &mlog;
    public:
        TransferFunction(const DataFlow::VertexFlowGraphs &index, Approximation approx, const SmtSolverPtr &solver,
                         Sawyer::Message::Facility &mlog)
            : index_(index), approximation_(approx), smtSolver_(solver), mlog(mlog) {}

        template<class CFG>
        StatePtr operator()(const CFG&, size_t cfgVertex, const StatePtr &in) {
            return (*this)(cfgVertex, in);
        }

        StatePtr operator()(size_t cfgVertex, const StatePtr &in);

        std::string toString(const StatePtr &in);
    };

    //                                  Merge function
protected:
    class MergeFunction {
    public:
        bool operator()(StatePtr &dst /*in,out*/, const StatePtr &src) const {
            ASSERT_not_null(src);
            if (!dst) {
                dst = src->copy();
                return true;                            // destination changed
            }
            return dst->merge(src);
        }
    };
    
    //                                  Data members
private:
    static Sawyer::Message::Facility mlog;
    Approximation approximation_;
    DataFlow dataFlow_;
    DataFlow::VertexFlowGraphs vertexFlowGraphs_;
    DataFlow::VariableList variableList_;
    bool vlistInitialized_;
    std::vector<StatePtr> results_;
    SmtSolverPtr smtSolver_;

public:
    explicit TaintedFlow(const InstructionSemantics::BaseSemantics::DispatcherPtr &userDispatcher)
        : approximation_(UNDER_APPROXIMATE), dataFlow_(userDispatcher), vlistInitialized_(false) {}

    static void initDiagnostics();

    Approximation approximation() const { return approximation_; }
    void approximation(Approximation a) { approximation_ = a; }
    SmtSolverPtr smtSolver() const { return smtSolver_; }
    void smtSolver(const SmtSolverPtr &solver) { smtSolver_ = solver; }
    template<class CFG>
    void computeFlowGraphs(const CFG &cfg, size_t cfgStartVertex) {
        using namespace Diagnostics;
        ASSERT_this();
        ASSERT_require(cfgStartVertex < cfg.nVertices());
        Stream mesg(mlog[WHERE] <<"computeFlowGraphs starting at CFG vertex " <<cfgStartVertex);
        vertexFlowGraphs_ = dataFlow_.buildGraphPerVertex(cfg, cfgStartVertex);
        variableList_ = dataFlow_.getUniqueVariables(vertexFlowGraphs_);
        results_.clear();
        vlistInitialized_ = true;
        mesg <<"; found " <<StringUtility::plural(variableList_.size(), "variables") <<"\n";
        if (mlog[DEBUG]) {
            for (const DataFlow::Variable &variable: variableList_)
                mlog[DEBUG] <<"  found variable: " <<variable <<"\n";
        }
    }

    const DataFlow::VertexFlowGraphs& vertexFlowGraphs() const {
        ASSERT_this();
        return vertexFlowGraphs_;
    }
    void vertexFlowGraphs(const DataFlow::VertexFlowGraphs &graphMap) {
        using namespace Diagnostics;
        ASSERT_this();
        vertexFlowGraphs_ = graphMap;
        variableList_ = dataFlow_.getUniqueVariables(vertexFlowGraphs_);
        vlistInitialized_ = true;
        results_.clear();
        mlog[WHERE] <<"vertexFlowGraphs set by user with " <<StringUtility::plural(variableList_.size(), "variables") <<"\n";
    }
    const DataFlow::VariableList& variables() const {
        ASSERT_this();
        ASSERT_require2(vlistInitialized_, "TaintedFlow::computeFlowGraphs must be called before TaintedFlow::variables");
        return variableList_;
    }

    StatePtr stateInstance(Taintedness taint) const {
        ASSERT_this();
        ASSERT_require2(vlistInitialized_, "TaintedFlow::computeFlowGraphs must be called before TaintedFlow::stateInstance");
        return State::instance(variableList_, taint);
    }

    template<class CFG>
    void runToFixedPoint(const CFG &cfg, size_t cfgStartVertex, const StatePtr &initialState) {
        using namespace Diagnostics;
        ASSERT_this();
        ASSERT_require(cfgStartVertex < cfg.nVertices());
        ASSERT_not_null(initialState);
        Stream mesg(mlog[WHERE] <<"runToFixedPoint starting at CFG vertex " <<cfgStartVertex);
        results_.clear();
        TransferFunction xfer(vertexFlowGraphs_, approximation_, smtSolver_, mlog);
        MergeFunction merge;
        DataFlow::Engine<CFG, StatePtr, TransferFunction, MergeFunction> dfEngine(cfg, xfer, merge);
        dfEngine.name("tainted-flow");
        dfEngine.runToFixedPoint(cfgStartVertex, initialState);
        results_ = dfEngine.getFinalStates();
        mesg <<"; results for " <<StringUtility::plural(results_.size(), "vertices", "vertex") <<"\n";
    }

    StatePtr getFinalState(size_t cfgVertexId) const {
        ASSERT_this();
        ASSERT_require(cfgVertexId < results_.size());
        return results_[cfgVertexId];
    }
};

std::ostream& operator<<(std::ostream &out, const TaintedFlow::State &state);

} // namespace
} // namespace

#endif
#endif