Adding a sketch for a CipherOperation concept to model encryption/decryption operations.

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

@@ -83,9 +83,9 @@ module JCAModel {
   class CipherGetInstanceAlgorithmArg extends Crypto::EncryptionAlgorithmInstance,
     Crypto::ModeOfOperationAlgorithmInstance, Crypto::PaddingAlgorithmInstance instanceof Expr
   {
-    CipherGetInstanceAlgorithmArg() {
-      exists(CipherGetInstanceCall call | this = call.getArgument(0))
-    }
+    CipherGetInstanceCall call;
+
+    CipherGetInstanceAlgorithmArg() { this = call.getAlgorithmArg() }
 
     /**
      * Returns the `StringLiteral` from which this argument is derived, if known.
@@ -94,26 +94,26 @@ module JCAModel {
       AlgorithmStringToFetchFlow::flow(DataFlow::exprNode(result),
         DataFlow::exprNode(this.(Expr).getAChildExpr*()))
     }
+
+    CipherGetInstanceCall getCall() { result = call }
   }
 
   // TODO: what if encrypt/decrypt mode isn't known
   private module CipherGetInstanceToFinalizeConfig implements DataFlow::StateConfigSig {
-    class FlowState = string;
+    class FlowState = Crypto::TCipherOperationMode;
 
     predicate isSource(DataFlow::Node src, FlowState state) {
-      state = "UNKNOWN" and
+      state = Crypto::UnknownCipherOperationMode() and
       src.asExpr() instanceof CipherGetInstanceCall
     }
 
     predicate isSink(DataFlow::Node sink, FlowState state) {
-      state in ["ENCRYPT", "DECRYPT", "UNKNOWN"] and
       exists(CipherDoFinalCall c | c.getQualifier() = sink.asExpr())
     }
 
     predicate isAdditionalFlowStep(
       DataFlow::Node node1, FlowState state1, DataFlow::Node node2, FlowState state2
     ) {
-      state1 in ["UNKNOWN", "ENCRYPT", "DECRYPT"] and
       exists(CipherInitCall c |
         c.getQualifier() = node1.asExpr() and
         // TODO: not taking into consideration if the mode traces to this arg
@@ -121,10 +121,16 @@ module JCAModel {
           c.getModeArg() = fa and
           (
             fa.getField().getName() = "ENCRYPT_MODE" and
-            state2 = "ENCRYPT"
+            state2 = Crypto::EncryptionMode()
             or
             fa.getField().getName() = "DECRYPT_MODE" and
-            state2 = "DECRYPT"
+            state2 = Crypto::DecryptionMode()
+            or
+            fa.getField().getName() = "WRAP_MODE" and
+            state2 = Crypto::EncryptionMode()
+            or
+            fa.getField().getName() = "UNWRAP_MODE" and
+            state2 = Crypto::DecryptionMode()
           )
         )
       ) and
@@ -135,16 +141,29 @@ module JCAModel {
   module CipherGetInstanceToFinalizeFlow =
     DataFlow::GlobalWithState<CipherGetInstanceToFinalizeConfig>;
 
-  // TODO: what if the mode is UNKNOWN?
-  class CipherEncryptionOperation extends Crypto::EncryptionOperationInstance instanceof Call {
+  class CipherEncryptionOperation extends Crypto::CipherOperationInstance instanceof Call {
+    Crypto::TCipherOperationMode mode;
+    Crypto::EncryptionAlgorithmInstance algorithm;
+
     CipherEncryptionOperation() {
-      exists(CipherGetInstanceToFinalizeFlow::PathNode sink, CipherDoFinalCall c |
-        CipherGetInstanceToFinalizeFlow::flowPath(_, sink) and
-        sink.getNode().asExpr() = c.getQualifier() and
-        sink.getState() = "ENCRYPT" and
-        this = c
+      exists(
+        CipherGetInstanceToFinalizeFlow::PathNode sink,
+        CipherGetInstanceToFinalizeFlow::PathNode src, CipherGetInstanceCall getCipher,
+        CipherDoFinalCall doFinalize, CipherGetInstanceAlgorithmArg arg
+      |
+        CipherGetInstanceToFinalizeFlow::flowPath(src, sink) and
+        src.getNode().asExpr() = getCipher and
+        sink.getNode().asExpr() = doFinalize.getQualifier() and
+        sink.getState() = mode and
+        this = doFinalize and
+        arg.getCall() = getCipher and 
+        algorithm = arg
       )
     }
+
+    override Crypto::EncryptionAlgorithmInstance getAlgorithm() { result = algorithm }
+
+    override Crypto::TCipherOperationMode getCipherOperationMode() { result = mode }
   }
 
   /**

shared/cryptography/codeql/cryptography/Model.qll

@@ -79,7 +79,11 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
 
   abstract class KeyDerivationAlgorithmInstance extends LocatableElement { }
 
-  abstract class EncryptionOperationInstance extends LocatableElement { }
+  abstract class CipherOperationInstance extends LocatableElement {
+    abstract EncryptionAlgorithmInstance getAlgorithm();
+
+    abstract TCipherOperationMode getCipherOperationMode();
+  }
 
   abstract class EncryptionAlgorithmInstance extends LocatableElement { }
 
@@ -115,7 +119,7 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
     // Operations (e.g., hashing, encryption)
     THashOperation(HashOperationInstance e) or
     TKeyDerivationOperation(KeyDerivationOperationInstance e) or
-    TEncryptionOperation(EncryptionOperationInstance e) or
+    TCipherOperation(CipherOperationInstance e) or
     TKeyEncapsulationOperation(KeyEncapsulationOperationInstance e) or
     // Algorithms (e.g., SHA-256, AES)
     TEncryptionAlgorithm(EncryptionAlgorithmInstance e) or
@@ -238,13 +242,14 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
      */
     abstract Algorithm getAlgorithm();
 
-    /**
-     * Gets the name of this operation, e.g., "hash" or "encrypt".
-     */
-    abstract string getOperationType();
-
-    final override string getInternalType() { result = this.getOperationType() }
-
+    // TODO: I only removed this because I want the operation type to be non-string
+    // since for CipherOperations the user will have to pick the right type,
+    // and I want to force them to use a type that is restricted. In this case to a TCipherOperationType
+    // /**
+    //  * Gets the name of this operation, e.g., "hash" or "encrypt".
+    //  */
+    // abstract string getOperationType();
+    // final override string getInternalType() { result = this.getOperationType() }
     override NodeBase getChild(string edgeName) {
       result = super.getChild(edgeName)
       or
@@ -290,8 +295,7 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
    */
   abstract class HashOperation extends Operation, THashOperation {
     abstract override HashAlgorithm getAlgorithm();
-
-    override string getOperationType() { result = "HashOperation" }
+    //override string getOperationType() { result = "HashOperation" }
   }
 
   newtype THashType =
@@ -401,8 +405,7 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
     final override Location getLocation() {
       exists(LocatableElement le | this = TKeyDerivationOperation(le) and result = le.getLocation())
     }
-
-    override string getOperationType() { result = "KeyDerivationOperation" }
+    //override string getOperationType() { result = "KeyDerivationOperation" }
   }
 
   /**
@@ -681,15 +684,31 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
     abstract override string getRawAlgorithmName();
   }
 
+  newtype TCipherOperationMode =
+    EncryptionMode() or
+    DecryptionMode() or
+    UnknownCipherOperationMode()
+
   /**
    * An encryption operation that processes plaintext to generate a ciphertext.
    * This operation takes an input message (plaintext) of arbitrary content and length
    * and produces a ciphertext as the output using a specified encryption algorithm (with a mode and padding).
    */
-  abstract class EncryptionOperation extends Operation, TEncryptionOperation {
-    override string getOperationType() { result = "EncryptionOperation" }
+  // NOTE FOR NICK: making this concrete here as I don't think users need to worry about making/extending these operations, just instances
+  class CipherOperation extends Operation, TCipherOperation {
+    CipherOperationInstance instance;
 
-    abstract override EncryptionAlgorithm getAlgorithm();
+    CipherOperation() { this = TCipherOperation(instance) }
+
+    override Location getLocation() { result = instance.getLocation() }
+
+    final TCipherOperationMode getCipherOperationMode() {
+      result = instance.getCipherOperationMode()
+    }
+
+    final override EncryptionAlgorithm getAlgorithm() { result = instance.getAlgorithm() }
+
+    override string getInternalType() { result = "CipherOperation" }
     // /**
     //  * Gets the initialization vector associated with this encryption operation.
     //  *