From 704a06e1fa5d706fa1a06af33a1e67b8d8b76935 Mon Sep 17 00:00:00 2001
From: "REDMOND\\brodes" <brodes@microsoft.com>
Date: Thu, 2 Oct 2025 11:45:13 -0400
Subject: [PATCH] Crypto: Update JCA PBKDF2 modeling: 1) add further
 inheritance structures to make the inheritance decomposition and caveats
 clearer, and 2) use getConsumer to establish the hash and hmac consumer.
 Update the Model to expect hash node types specifically for HMAC
 getHashALgorithmOrUnknown.

---
 java/ql/lib/experimental/quantum/JCA.qll      | 38 +++++++++++++------
 .../codeql/quantum/experimental/Model.qll     |  2 +-
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/java/ql/lib/experimental/quantum/JCA.qll b/java/ql/lib/experimental/quantum/JCA.qll
index de2afdc49aa..a79210c6d23 100644
--- a/java/ql/lib/experimental/quantum/JCA.qll
+++ b/java/ql/lib/experimental/quantum/JCA.qll
@@ -1265,23 +1265,39 @@ module JCAModel {
     override int getFixedDigestLength() { exists(hash_name_to_type_known(hashName, result)) }
   }
 
-  //TODO: handle PBE
-  class Pbkdf2AlgorithmStringLiteral extends KdfAlgorithmStringLiteral,
-    Crypto::Pbkdf2AlgorithmInstance, Crypto::HmacAlgorithmInstance
+  //TODO: handle PBE "with" cases
+  class Pbkdf2WithHmac_Pbkdf2AlgorithmInstance extends Crypto::Pbkdf2AlgorithmInstance,
+    KdfAlgorithmStringLiteral, // this is a parent already, but extending to have immediate access to 'getConsumer()'
+    Pbkdf2WithHmac_KeyOperationAlgorithmStringLiteral
   {
-    Pbkdf2AlgorithmStringLiteral() { super.getKdfType() instanceof Crypto::PBKDF2 }
+    override Crypto::AlgorithmValueConsumer getHmacAlgorithmValueConsumer() {
+      result = this.getConsumer()
+    }
+  }
 
-    override Crypto::AlgorithmValueConsumer getHmacAlgorithmValueConsumer() { result = this }
+  // NOTE: must use instanceof to avoid non-monotonic recursion
+  class Pbkdf2WithHmac_HmacAlgorithmInstance extends Crypto::HmacAlgorithmInstance instanceof Pbkdf2WithHmac_KeyOperationAlgorithmStringLiteral
+  {
+    override Crypto::AlgorithmValueConsumer getHashAlgorithmValueConsumer() {
+      result = this.(KdfAlgorithmStringLiteral).getConsumer()
+    }
 
-    override Crypto::AlgorithmValueConsumer getHashAlgorithmValueConsumer() { result = this }
+    override int getKeySizeFixed() {
+      // already defined by parent key operation algorithm, but extending an instance
+      // still requires we override this method
+      result = super.getKeySizeFixed()
+    }
 
-    override int getKeySizeFixed() { none() }
-
-    override Crypto::ConsumerInputDataFlowNode getKeySizeConsumer() { none() }
+    override Crypto::ConsumerInputDataFlowNode getKeySizeConsumer() {
+      // already defined by parent key operation algorithm, but extending an instance
+      // still requires we override this method
+      result = super.getKeySizeConsumer()
+    }
 
     override string getRawAlgorithmName() {
-      // Note: hard coding "hmac" since that should be the only option
-      result = "Hmac"
+      // already defined by parent key operation algorithm, but extending an instance
+      // still requires we override this method
+      result = super.getRawAlgorithmName()
     }
 
     override Crypto::KeyOpAlg::AlgorithmType getAlgorithmType() {
diff --git a/shared/quantum/codeql/quantum/experimental/Model.qll b/shared/quantum/codeql/quantum/experimental/Model.qll
index 218b1e02001..a09b96cd550 100644
--- a/shared/quantum/codeql/quantum/experimental/Model.qll
+++ b/shared/quantum/codeql/quantum/experimental/Model.qll
@@ -1633,7 +1633,7 @@ module CryptographyBase<LocationSig Location, InputSig<Location> Input> {
 
     override string getInternalType() { result = "HMACAlgorithm" }
 
-    NodeBase getHashAlgorithmOrUnknown() {
+    HashAlgorithmNode getHashAlgorithmOrUnknown() {
       result.asElement() = hmacInstance.getHashAlgorithmValueConsumer().getASource()
     }