Add input and output nodes and fix cross product

java/ql/lib/experimental/Quantum/JCA.qll

@@ -59,14 +59,22 @@ module JCAModel {
     Expr getProviderArg() { result = this.getArgument(1) }
   }
 
-  private class JCACipherOperationCall extends Call {
-    JCACipherOperationCall() {
+  private class CipherOperationCall extends MethodCall {
+    CipherOperationCall() {
       exists(string s | s in ["doFinal", "wrap", "unwrap"] |
-        this.getCallee().hasQualifiedName("javax.crypto", "Cipher", s)
+        this.getMethod().hasQualifiedName("javax.crypto", "Cipher", s)
       )
     }
 
-    DataFlow::Node getMessageArg() { result.asExpr() = this.getArgument(0) }
+    Expr getInput() { result = this.getArgument(0) }
+
+    Expr getOutput() {
+      result = this.getArgument(3)
+      or
+      this.getMethod().getReturnType().hasName("byte[]") and result = this
+    }
+
+    DataFlow::Node getMessageArg() { result.asExpr() = this.getInput() }
   }
 
   /**
@@ -304,7 +312,7 @@ module JCAModel {
     predicate isSink(DataFlow::Node sink, FlowState state) { none() }
 
     predicate isSink(DataFlow::Node sink) {
-      exists(JCACipherOperationCall c | c.getQualifier() = sink.asExpr())
+      exists(CipherOperationCall c | c.getQualifier() = sink.asExpr())
     }
 
     predicate isAdditionalFlowStep(
@@ -330,7 +338,7 @@ module JCAModel {
   class CipherOperationInstance extends Crypto::CipherOperationInstance instanceof Call {
     Crypto::CipherOperationSubtype mode;
     CipherGetInstanceToCipherOperationFlow::PathNode sink;
-    JCACipherOperationCall doFinalize;
+    CipherOperationCall doFinalize;
     CipherGetInstanceAlgorithmArg consumer;
 
     CipherOperationInstance() {
@@ -350,11 +358,15 @@ module JCAModel {
       result = sink.getState().(InitializedCipherModeFlowState).getInitCall().getNonceArg()
     }
 
-    override Crypto::ArtifactConsumer getMessageConsumer() {
+    override Crypto::CipherInputConsumer getInputConsumer() {
       result = doFinalize.getMessageArg().asExpr()
     }
 
     override Crypto::AlgorithmConsumer getAlgorithmConsumer() { result = consumer }
+
+    override Crypto::CipherOutputArtifactInstance getOutputArtifact() {
+      result = doFinalize.getOutput()
+    }
   }
 
   /**
@@ -481,15 +493,27 @@ module JCAModel {
     }
   }
 
-  class CipherInitCallNonceArgConsumer extends Crypto::NonceArtifactConsumer instanceof Expr {
+  class CipherInitCallNonceArgConsumer extends NonceArtifactConsumer instanceof Expr {
     CipherInitCallNonceArgConsumer() { this = any(CipherInitCall call).getNonceArg() }
 
     override DataFlow::Node getInputNode() { result.asExpr() = this }
   }
 
-  class CipherInitCallKeyConsumer extends Crypto::ArtifactConsumer instanceof Expr {
+  class CipherInitCallKeyConsumer extends Crypto::ArtifactConsumer {
     CipherInitCallKeyConsumer() { this = any(CipherInitCall call).getKeyArg() }
 
     override DataFlow::Node getInputNode() { result.asExpr() = this }
   }
+
+  class CipherMessageInputConsumer extends Crypto::CipherInputConsumer {
+    CipherMessageInputConsumer() { this = any(CipherOperationCall call).getMessageArg().asExpr() }
+
+    override DataFlow::Node getInputNode() { result.asExpr() = this }
+  }
+
+  class CipherOperationCallOutput extends CipherOutputArtifact {
+    CipherOperationCallOutput() { this = any(CipherOperationCall call).getOutput() }
+
+    override DataFlow::Node getOutputNode() { result.asExpr() = this }
+  }
 }

java/ql/lib/experimental/Quantum/Language.qll

@@ -69,8 +69,6 @@ class ConstantDataSource extends Crypto::GenericConstantOrAllocationSource insta
   override DataFlow::Node getOutputNode() { result.asExpr() = this }
 
   override predicate flowsTo(Crypto::FlowAwareElement other) {
-    other instanceof NonceArtifactInstance and
-    // limit to only nonces for now
     // TODO: separate config to avoid blowing up data-flow analysis
     GenericDataSourceUniversalFlow::flow(this.getOutputNode(), other.getInputNode())
   }
@@ -113,29 +111,10 @@ abstract class AdditionalFlowInputStep extends DataFlow::Node {
   final DataFlow::Node getInput() { result = this }
 }
 
-module ArtifactUniversalFlowConfig implements DataFlow::ConfigSig {
-  predicate isSource(DataFlow::Node source) {
-    source = any(Crypto::ArtifactElement artifact).getOutputNode()
-  }
-
-  predicate isSink(DataFlow::Node sink) {
-    sink = any(Crypto::FlowAwareElement other).getInputNode()
-  }
-
-  predicate isBarrierIn(DataFlow::Node node) {
-    node = any(Crypto::FlowAwareElement element).getOutputNode()
-  }
-
-  predicate isAdditionalFlowStep(DataFlow::Node node1, DataFlow::Node node2) {
-    node1.(AdditionalFlowInputStep).getOutput() = node2
-  }
-}
-
 module ArtifactUniversalFlow = DataFlow::Global<ArtifactUniversalFlowConfig>;
 
-class NonceArtifactInstance extends Crypto::NonceArtifactInstance {
-  NonceArtifactInstance() { this instanceof Crypto::NonceArtifactConsumer }
-
+class NonceArtifactConsumer extends Crypto::NonceArtifactInstance instanceof Crypto::NonceArtifactConsumer
+{
   override predicate flowsTo(Crypto::FlowAwareElement other) {
     ArtifactUniversalFlow::flow(this.getOutputNode(), other.getInputNode())
   }
@@ -149,6 +128,27 @@ class NonceArtifactInstance extends Crypto::NonceArtifactInstance {
   }
 }
 
+class CipherInputConsumer extends Crypto::CipherInputArtifactInstance instanceof Crypto::CipherInputConsumer
+{
+  override predicate flowsTo(Crypto::FlowAwareElement other) {
+    ArtifactUniversalFlow::flow(this.getOutputNode(), other.getInputNode())
+  }
+
+  override DataFlow::Node getOutputNode() { none() }
+
+  override DataFlow::Node getInputNode() {
+    result = this.(Crypto::CipherInputArtifactInstance).getInputNode()
+  }
+}
+
+abstract class CipherOutputArtifact extends Crypto::CipherOutputArtifactInstance {
+  override predicate flowsTo(Crypto::FlowAwareElement other) {
+    ArtifactUniversalFlow::flow(this.getOutputNode(), other.getInputNode())
+  }
+
+  override DataFlow::Node getInputNode() { none() }
+}
+
 /**
  * Generic data source to node input configuration
  */
@@ -161,6 +161,32 @@ module GenericDataSourceUniversalFlowConfig implements DataFlow::ConfigSig {
     sink = any(Crypto::FlowAwareElement other).getInputNode()
   }
 
+  predicate isBarrierOut(DataFlow::Node node) {
+    node = any(Crypto::FlowAwareElement element).getInputNode()
+  }
+
+  predicate isBarrierIn(DataFlow::Node node) {
+    node = any(Crypto::FlowAwareElement element).getOutputNode()
+  }
+
+  predicate isAdditionalFlowStep(DataFlow::Node node1, DataFlow::Node node2) {
+    node1.(AdditionalFlowInputStep).getOutput() = node2
+  }
+}
+
+module ArtifactUniversalFlowConfig implements DataFlow::ConfigSig {
+  predicate isSource(DataFlow::Node source) {
+    source = any(Crypto::ArtifactElement artifact).getOutputNode()
+  }
+
+  predicate isSink(DataFlow::Node sink) {
+    sink = any(Crypto::FlowAwareElement other).getInputNode()
+  }
+
+  predicate isBarrierOut(DataFlow::Node node) {
+    node = any(Crypto::FlowAwareElement element).getInputNode()
+  }
+
   predicate isBarrierIn(DataFlow::Node node) {
     node = any(Crypto::FlowAwareElement element).getOutputNode()
   }