import cpp
import semmle.code.cpp.dataflow.new.TaintTracking


    newtype ConnStateType = 
        STATE0() or
        STATE1() or
        STATE2()

    predicate allStates(ConnStateType s) {
        s = 
        STATE0() or
        s = STATE1()
        or s = STATE2()
    }

    class ConnState extends ConnStateType {
        final string toString() {
            this = STATE0() and result = "STATE0"
            or
            this = STATE1() and result = "STATE1"
            or this = STATE2() and result = "STATE2"
        }
        Location getLocation() { result instanceof UnknownLocation }
    }
    predicate isSetState(Function m, int arg,  ConnState sqt) {
        m.getName() = "set"
        and (
        arg = 0  and sqt = STATE0()
        or
        arg = 1 and sqt = STATE1()
        or
        arg = 2 and sqt = STATE2()
        )
    }
    predicate isCallSetState(FunctionCall fc, ConnState sqt) {
        isSetState(fc.getTarget(), fc.getArgument(1).getValue().toInt(), sqt)
    }

class ConnFlowStateTransformerNode extends DataFlow::Node {
    ConnState next_state;
    ConnFlowStateTransformerNode() {
        exists(FunctionCall fc|
            isCallSetState(fc, next_state )
            and this.asIndirectArgument() = fc.getArgument(0)
        )
        or
        exists(FunctionCall fc|
            isCallSetState(fc, next_state )
            and this.asExpr() = fc.getArgument(0)
        )   
        or
        exists(FunctionCall fc|
            isCallSetState(fc, next_state )
            and this.asDefiningArgument() = fc.getArgument(0)
        )                   
    }    
    ConnState transform(ConnState flowstate){
        flowstate != next_state and result = next_state
    }
}

predicate isAdditionalStepImpl(DataFlow::Node n1, DataFlow::Node n2){
    // from event to event.conn
    exists(FieldAccess fa |
      n1.asIndirectExpr() = fa.getQualifier() and
      n2.asExpr() = fa
      and fa.getTarget().getName() = "conn"
    )
    or 
    exists(FieldAccess fa |
      n1.asIndirectExpr() = fa.getQualifier() and
      n2.asIndirectExpr() = fa
      and fa.getTarget().getName() = "conn"
    )   
}
module TestConfig implements DataFlow::StateConfigSig {
  class FlowState = ConnState;
  predicate isSource(DataFlow::Node node, FlowState sig) {
    exists(FunctionCall fc |
            fc.getTarget().getName()= "get_data"
            and fc = node.asExpr()
            and allStates(sig)
    )
  }

  predicate isSink(DataFlow::Node sink, FlowState state) {
    exists(FunctionCall fc |
            fc.getTarget().getName()= "sink"
            and fc.getArgument(0) = sink.asExpr()
            and allStates(state)
    ) 
  }

  predicate isAdditionalFlowStep(DataFlow::Node n1, DataFlow::Node n2) {
        isAdditionalStepImpl(n1,n2)
  }
  predicate isAdditionalFlowStep(
    DataFlow::Node node1, FlowState prevState, DataFlow::Node node2, FlowState succState
  ) {
    node2.(ConnFlowStateTransformerNode).transform(prevState) = succState
    and DataFlow::simpleLocalFlowStep(node1, node2, _)
  }

  predicate isBarrier(DataFlow::Node node, FlowState state) { 
    node.(ConnFlowStateTransformerNode).transform(state) != state      
  }  
}

module FlowTest = TaintTracking::GlobalWithState<TestConfig>;
int explorationLimit() { result = 4 }
module MyPartialFlowFwd = FlowTest::FlowExplorationFwd<explorationLimit/0>;