Add JCA key (generation) modelling

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

@@ -84,7 +84,7 @@ module JCAModel {
     predicate isSource(DataFlow::Node src) { src.asExpr() instanceof CipherStringLiteral }
 
     predicate isSink(DataFlow::Node sink) {
-      exists(CipherGetInstanceCall call | sink.asExpr() = call.getAlgorithmArg())
+      exists(Crypto::AlgorithmValueConsumer consumer | sink = consumer.getInputNode())
     }
   }
 
@@ -102,13 +102,13 @@ module JCAModel {
   class CipherStringLiteralAlgorithmInstance extends Crypto::CipherAlgorithmInstance,
     Crypto::ModeOfOperationAlgorithmInstance, Crypto::PaddingAlgorithmInstance instanceof CipherStringLiteral
   {
-    CipherGetInstanceAlgorithmArg consumer;
+    Crypto::AlgorithmValueConsumer consumer;
 
     CipherStringLiteralAlgorithmInstance() {
-      AlgorithmStringToFetchFlow::flow(DataFlow::exprNode(this), DataFlow::exprNode(consumer))
+      AlgorithmStringToFetchFlow::flow(DataFlow::exprNode(this), consumer.getInputNode())
     }
 
-    CipherGetInstanceAlgorithmArg getConsumer() { result = consumer }
+    Crypto::AlgorithmValueConsumer getConsumer() { result = consumer }
 
     override Crypto::ModeOfOperationAlgorithmInstance getModeOfOperationAlgorithm() {
       result = this and exists(this.getRawModeAlgorithmName()) // TODO: provider defaults
@@ -411,6 +411,19 @@ module JCAModel {
     }
   }
 
+  // e.g., getPublic or getPrivate
+  class KeyPairGetKeyCall extends MethodCall {
+    KeyPairGetKeyCall() {
+      this.getCallee().hasQualifiedName("java.security", "KeyPair", "getPublic")
+      or
+      this.getCallee().hasQualifiedName("java.security", "KeyPair", "getPrivate")
+    }
+
+    DataFlow::Node getInputNode() { result.asExpr() = this.getQualifier() }
+
+    DataFlow::Node getOutputNode() { result.asExpr() = this }
+  }
+
   predicate additionalFlowSteps(DataFlow::Node node1, DataFlow::Node node2) {
     exists(IvParameterSpecGetIvCall m |
       node1.asExpr() = m.getQualifier() and
@@ -421,12 +434,17 @@ module JCAModel {
       node1 = n.getInputNode() and
       node2 = n.getOutputNode()
     )
+    or
+    exists(KeyPairGetKeyCall call |
+      node1 = call.getInputNode() and
+      node2 = call.getOutputNode()
+    )
   }
 
-  class NonceAdditionalFlowInputStep extends AdditionalFlowInputStep {
+  class ArtifactAdditionalFlowStep extends AdditionalFlowInputStep {
     DataFlow::Node output;
 
-    NonceAdditionalFlowInputStep() { additionalFlowSteps(this, output) }
+    ArtifactAdditionalFlowStep() { additionalFlowSteps(this, output) }
 
     override DataFlow::Node getOutput() { result = output }
   }
@@ -605,4 +623,72 @@ module JCAModel {
       )
     }
   }
+
+  class KeyGeneratorCallAlgorithmValueConsumer extends Crypto::AlgorithmValueConsumer {
+    KeyGeneratorGetInstanceCall call;
+
+    KeyGeneratorCallAlgorithmValueConsumer() { this = call.getAlgorithmArg() }
+
+    override DataFlow::Node getInputNode() { result.asExpr() = this }
+
+    override Crypto::AlgorithmInstance getAKnownAlgorithmSource() {
+      result.(CipherStringLiteralAlgorithmInstance).getConsumer() = this
+    }
+  }
+
+  // flow from instance created by getInstance to generateKey
+  module KeyGeneratorGetInstanceToGenerateConfig implements DataFlow::ConfigSig {
+    predicate isSource(DataFlow::Node src) {
+      exists(KeyGeneratorGetInstanceCall call | src.asExpr() = call)
+    }
+
+    predicate isSink(DataFlow::Node sink) {
+      exists(KeyGeneratorGenerateCall call | sink.asExpr() = call.(MethodCall).getQualifier())
+    }
+  }
+
+  module KeyGeneratorGetInstanceToGenerateFlow =
+    DataFlow::Global<KeyGeneratorGetInstanceToGenerateConfig>;
+
+  class KeyGeneratorGetInstanceCall extends MethodCall {
+    KeyGeneratorGetInstanceCall() {
+      this.getCallee().hasQualifiedName("javax.crypto", "KeyGenerator", "getInstance")
+      or
+      this.getCallee().hasQualifiedName("java.security", "KeyPairGenerator", "getInstance")
+    }
+
+    Expr getAlgorithmArg() { result = super.getArgument(0) }
+
+    predicate flowsTo(KeyGeneratorGenerateCall sink) {
+      KeyGeneratorGetInstanceToGenerateFlow::flow(DataFlow::exprNode(this),
+        DataFlow::exprNode(sink.(MethodCall).getQualifier()))
+    }
+  }
+
+  class KeyGeneratorGenerateCall extends Crypto::KeyGenerationOperationInstance instanceof MethodCall
+  {
+    Crypto::KeyArtifactType type;
+
+    KeyGeneratorGenerateCall() {
+      this.getCallee().hasQualifiedName("javax.crypto", "KeyGenerator", "generateKey") and
+      type instanceof Crypto::TSymmetricKeyType
+      or
+      this.getCallee().hasQualifiedName("java.security", "KeyPairGenerator", "generateKeyPair") and
+      type instanceof Crypto::TAsymmetricKeyType
+    }
+
+    override DataFlow::Node getOutputKeyArtifact() { result.asExpr() = this }
+
+    override Crypto::KeyArtifactType getOutputKeyType() { result = type }
+
+    override Crypto::AlgorithmValueConsumer getAnAlgorithmValueConsumer() {
+      exists(KeyGeneratorGetInstanceCall getInstance |
+        getInstance.flowsTo(this) and result = getInstance.getAlgorithmArg()
+      )
+    }
+
+    Crypto::AlgorithmInstance getAKnownAlgorithm() {
+      result = this.getAnAlgorithmValueConsumer().getAKnownAlgorithmSource()
+    }
+  }
 }

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

@@ -30,6 +30,12 @@ module CryptoInput implements InputSig<Language::Location> {
     result = node.asExpr() or
     result = node.asParameter()
   }
+
+  predicate artifactOutputFlowsToGenericInput(
+    DataFlow::Node artifactOutput, DataFlow::Node otherInput
+  ) {
+    ArtifactUniversalFlow::flow(artifactOutput, otherInput)
+  }
 }
 
 /**
@@ -70,16 +76,20 @@ class GenericRemoteDataSource extends Crypto::GenericRemoteDataSource {
   override string getAdditionalDescription() { result = this.toString() }
 }
 
-class ConstantDataSource extends Crypto::GenericConstantOrAllocationSource instanceof Literal {
-  override DataFlow::Node getOutputNode() { result.asExpr() = this }
-
-  override predicate flowsTo(Crypto::FlowAwareElement other) {
-    // TODO: separate config to avoid blowing up data-flow analysis
-    GenericDataSourceUniversalFlow::flow(this.getOutputNode(), other.getInputNode())
-  }
-
-  override string getAdditionalDescription() { result = this.toString() }
-}
+/*
+ * class ConstantDataSource extends Crypto::GenericConstantOrAllocationSource instanceof Literal {
+ *  ConstantDataSource() { not this instanceof Crypto::KnownElement }
+ *
+ *  override DataFlow::Node getOutputNode() { result.asExpr() = this }
+ *
+ *  override predicate flowsTo(Crypto::FlowAwareElement other) {
+ *    // TODO: separate config to avoid blowing up data-flow analysis
+ *    GenericDataSourceUniversalFlow::flow(this.getOutputNode(), other.getInputNode())
+ *  }
+ *
+ *  override string getAdditionalDescription() { result = this.toString() }
+ * }
+ */
 
 /**
  * Random number generation, where each instance is modelled as the expression

java/ql/src/experimental/Quantum/TestCipher.ql

@@ -7,11 +7,12 @@ import experimental.Quantum.Language
 from
   Crypto::CipherOperationNode op, Crypto::CipherAlgorithmNode a,
   Crypto::ModeOfOperationAlgorithmNode m, Crypto::PaddingAlgorithmNode p,
-  Crypto::NonceArtifactNode nonce
+  Crypto::NonceArtifactNode nonce, Crypto::KeyArtifactNode k
 where
   a = op.getAKnownCipherAlgorithm() and
   m = a.getModeOfOperation() and
   p = a.getPaddingAlgorithm() and
-  nonce = op.getANonce()
+  nonce = op.getANonce() and
+  k = op.getAKey()
 select op, op.getCipherOperationSubtype(), a, a.getRawAlgorithmName(), m, m.getRawAlgorithmName(),
-  p, p.getRawAlgorithmName(), nonce
+  p, p.getRawAlgorithmName(), nonce, k, k.getSourceElement()

java/ql/src/experimental/Quantum/TestCipherKey.ql

@@ -0,0 +1,17 @@
+/**
+ * @name "PQC Test"
+ */
+
+import experimental.Quantum.Language
+
+from Crypto::CipherOperationNode op, Crypto::CipherAlgorithmNode a, Crypto::KeyArtifactNode k
+where
+  a = op.getAKnownCipherAlgorithm() and
+  k = op.getAKey()
+select op, op.getCipherOperationSubtype(), a, a.getRawAlgorithmName(), k, k.getSourceNode()
+/*
+ * from Crypto::CipherOperationNode op
+ * where op.getLocation().getFile().getBaseName() = "AsymmetricEncryptionMacHybridCryptosystem.java"
+ * select op, op.getAKey().getSourceNode()
+ */
+

shared/cryptography/codeql/cryptography/Model.qll

@@ -22,6 +22,10 @@ signature module InputSig<LocationSig Location> {
   class UnknownLocation instanceof Location;
 
   LocatableElement dfn_to_element(DataFlowNode node);
+
+  predicate artifactOutputFlowsToGenericInput(
+    DataFlowNode artifactOutput, DataFlowNode otherFlowAwareInput
+  );
 }
 
 module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
@@ -51,26 +55,22 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
       )
   }
 
-  NodeBase getPassthroughNodeChild(NodeBase node) {
-    result = node.(CipherInputNode).getChild(_) or
-    result = node.(NonceArtifactNode).getChild(_)
-  }
+  NodeBase getPassthroughNodeChild(NodeBase node) { result = node.getChild(_) }
 
   predicate isPassthroughNode(NodeBase node) {
-    node instanceof CipherInputNode or
-    node instanceof NonceArtifactNode
+    node.asElement() instanceof ArtifactConsumerAndInstance
   }
 
   predicate nodes_graph_impl(NodeBase node, string key, string value) {
     not node.isExcludedFromGraph() and
-    not isPassthroughNode(node) and
+    not isPassthroughNode(node) and // TODO: punt to fix known unknowns for passthrough nodes
     (
       key = "semmle.label" and
       value = node.toString()
       or
       // CodeQL's DGML output does not include a location
       key = "Location" and
-      value = node.getLocation().toString()
+      value = "demo" // node.getLocation().toString()
       or
       // Known unknown edges should be reported as properties rather than edges
       node = node.getChild(key) and
@@ -305,22 +305,17 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
   // in order for a flowsTo to be defined. At the per-modeling-instance level, extend that language-level class!
   abstract class OutputArtifactInstance extends ArtifactInstance {
     override predicate isConsumerArtifact() { none() }
+
+    override DataFlowNode getInputNode() { none() }
   }
 
-  abstract class DigestArtifactInstance extends OutputArtifactInstance {
-    final override DataFlowNode getInputNode() { none() }
-  }
+  abstract class DigestArtifactInstance extends OutputArtifactInstance { }
 
   abstract class RandomNumberGenerationInstance extends OutputArtifactInstance {
     // TODO: input seed?
-    final override DataFlowNode getInputNode() { none() }
   }
 
-  abstract class CipherOutputArtifactInstance extends ArtifactInstance {
-    final override DataFlowNode getInputNode() { none() }
-
-    override predicate isConsumerArtifact() { none() }
-  }
+  abstract class CipherOutputArtifactInstance extends OutputArtifactInstance { }
 
   // Artifacts that may be outputs or inputs
   newtype TKeyArtifactType =
@@ -338,13 +333,30 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
     }
   }
 
-  abstract class KeyArtifactInstance extends ArtifactInstance {
+  abstract private class KeyArtifactInstance extends ArtifactInstance {
     abstract KeyArtifactType getKeyType();
   }
 
-  final class KeyArtifactConsumer extends ArtifactConsumerAndInstance, KeyArtifactInstance {
+  final class KeyArtifactOutputInstance extends KeyArtifactInstance, OutputArtifactInstance {
+    KeyCreationOperationInstance creator;
+
+    KeyArtifactOutputInstance() { Input::dfn_to_element(creator.getOutputKeyArtifact()) = this }
+
+    final KeyCreationOperationInstance getCreator() { result = creator }
+
+    override KeyArtifactType getKeyType() { result = creator.getOutputKeyType() }
+
+    override DataFlowNode getOutputNode() { result = creator.getOutputKeyArtifact() }
+
+    override predicate flowsTo(FlowAwareElement other) {
+      Input::artifactOutputFlowsToGenericInput(this.getOutputNode(), other.getInputNode())
+    }
+  }
+
+  final class KeyArtifactConsumer extends KeyArtifactInstance, ArtifactConsumerAndInstance {
     DataFlowNode inputNode;
 
+    // TODO: key type hint? e.g. hint: private || public
     KeyArtifactConsumer() {
       exists(CipherOperationInstance op | inputNode = op.getKeyConsumer()) and
       this = Input::dfn_to_element(inputNode)
@@ -483,6 +495,33 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
 
   abstract class KeyDerivationAlgorithmInstance extends AlgorithmInstance { }
 
+  abstract private class KeyCreationOperationInstance extends OperationInstance {
+    abstract string getKeyCreationTypeDescription();
+
+    /**
+     * Gets the key artifact produced by this operation.
+     */
+    abstract DataFlowNode getOutputKeyArtifact();
+
+    /**
+     * Gets the key artifact type produced.
+     */
+    abstract KeyArtifactType getOutputKeyType();
+
+    /**
+     * Gets the key size of the key produced by this operation.
+     */
+    string getKeySize() { none() } // TODO: punt, might need a generic value consumer?
+  }
+
+  abstract class KeyGenerationOperationInstance extends KeyCreationOperationInstance {
+    final override string getKeyCreationTypeDescription() { result = "KeyGeneration" }
+  }
+
+  abstract class KeyLoadOperationInstance extends KeyCreationOperationInstance {
+    final override string getKeyCreationTypeDescription() { result = "KeyLoad" }
+  }
+
   private signature class AlgorithmInstanceType instanceof AlgorithmInstance;
 
   module AlgorithmInstanceOrValueConsumer<AlgorithmInstanceType Alg> {
@@ -519,6 +558,8 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
     TKeyDerivationOperation(KeyDerivationOperationInstance e) or
     TCipherOperation(CipherOperationInstance e) or
     TKeyEncapsulationOperation(KeyEncapsulationOperationInstance e) or
+    // Key Creation Operations
+    TKeyCreationOperation(KeyCreationOperationInstance e) or
     // Algorithms (e.g., SHA-256, AES)
     TCipherAlgorithm(CipherAlgorithmInstanceOrValueConsumer e) or
     TEllipticCurveAlgorithm(EllipticCurveAlgorithmInstance e) or
@@ -662,7 +703,7 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
      *
      * If a child class defines this predicate as `none()`, no relationship will be reported.
      */
-    string getSourceNodeRelationship() { result = "Source" }
+    string getSourceNodeRelationship() { result = "Source" } // TODO: revisit why this exists
 
     override NodeBase getChild(string edgeName) {
       result = super.getChild(edgeName)
@@ -737,6 +778,33 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
 
     override LocatableElement asElement() { result = instance }
 
+    /**
+     * Gets the algorithm or unknown source nodes consumed as an algorithm associated with this operation.
+     */
+    NodeBase getAnAlgorithmOrUnknown() {
+      result = this.getAKnownAlgorithm() or
+      result =
+        this.asElement().(OperationInstance).getAnAlgorithmValueConsumer().getAnUnknownSourceNode()
+    }
+
+    /**
+     * Gets a known algorithm associated with this operation
+     */
+    CipherAlgorithmNode getAKnownAlgorithm() {
+      result =
+        this.asElement().(OperationInstance).getAnAlgorithmValueConsumer().getAKnownSourceNode()
+    }
+
+    override NodeBase getChild(string edgeName) {
+      result = super.getChild(edgeName)
+      or
+      // [KNOWN_OR_UNKNOWN]
+      edgeName = "Algorithm" and
+      if exists(this.getAnAlgorithmOrUnknown())
+      then result = this.getAnAlgorithmOrUnknown()
+      else result = this
+    }
+
     override predicate properties(string key, string value, Location location) {
       super.properties(key, value, location)
       or
@@ -787,6 +855,16 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
     }
   }
 
+  final class KeyCreationOperationNode extends OperationNode, TKeyCreationOperation {
+    KeyCreationOperationInstance instance;
+
+    KeyCreationOperationNode() { this = TKeyCreationOperation(instance) }
+
+    override LocatableElement asElement() { result = instance }
+
+    override string getInternalType() { result = instance.getKeyCreationTypeDescription() }
+  }
+
   newtype TCipherOperationSubtype =
     TEncryptionMode() or
     TDecryptionMode() or
@@ -867,6 +945,11 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
       result.asElement() = this.asElement().(CipherOperationInstance).getOutputArtifact()
     }
 
+    KeyArtifactNode getAKey() {
+      result.asElement() =
+        Input::dfn_to_element(this.asElement().(CipherOperationInstance).getKeyConsumer())
+    }
+
     override NodeBase getChild(string key) {
       result = super.getChild(key)
       or
@@ -891,6 +974,10 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
       if exists(this.getAnOutputArtifact())
       then result = this.getAnOutputArtifact()
       else result = this
+      or
+      // [KNOWN_OR_UNKNOWN]
+      key = "Key" and
+      if exists(this.getAKey()) then result = this.getAKey() else result = this
     }
 
     override predicate properties(string key, string value, Location location) {