Java: add dataflow test for newly added KDF API

Author: IdrissRio

java/ql/test/library-tests/dataflow/kdf/KDFDataflowTest.java

@@ -0,0 +1,70 @@
+import javax.crypto.KDF;
+import javax.crypto.spec.HKDFParameterSpec;
+
+public class KDFDataflowTest {
+    public static void main(String[] args) throws Exception {
+        String userInput = args[0]; // source
+        byte[] taintedBytes = userInput.getBytes();
+
+        testBuilderPattern(taintedBytes);
+        testSeparateBuilder(taintedBytes);
+        testKDFWithSalt(taintedBytes);
+        testStaticParameterSpec(taintedBytes);
+        testCleanUsage();
+    }
+
+    public static void testBuilderPattern(byte[] taintedIKM) throws Exception {
+        HKDFParameterSpec.Builder builder = HKDFParameterSpec.ofExtract();
+        builder.addIKM(taintedIKM);
+        HKDFParameterSpec spec = builder.thenExpand("info".getBytes(), 32);
+
+        KDF kdf = KDF.getInstance("HKDF-SHA256");
+        byte[] result = kdf.deriveData(spec);
+        sink(result); // should flag
+    }
+
+    public static void testSeparateBuilder(byte[] taintedIKM) throws Exception {
+        HKDFParameterSpec.Builder builder1 = HKDFParameterSpec.ofExtract();
+        HKDFParameterSpec.Builder builder2 = builder1.addIKM(taintedIKM);
+        HKDFParameterSpec spec = builder2.thenExpand("info".getBytes(), 32);
+
+        KDF kdf = KDF.getInstance("HKDF-SHA256");
+        byte[] result = kdf.deriveData(spec);
+        sink(result); // should flag
+    }
+
+    public static void sink(Object o) {}
+
+    public static void testKDFWithSalt(byte[] taintedIKM) throws Exception {
+        HKDFParameterSpec.Builder builder = HKDFParameterSpec.ofExtract();
+        builder.addIKM(taintedIKM);
+        builder.addSalt("sensitive-salt".getBytes());
+        HKDFParameterSpec spec = builder.thenExpand("info".getBytes(), 32);
+
+        KDF kdf = KDF.getInstance("HKDF-SHA256");
+        byte[] result = kdf.deriveData(spec);
+        sink(result); // should flag
+    }
+
+    public static void testStaticParameterSpec(byte[] taintedIKM) throws Exception {
+        javax.crypto.spec.SecretKeySpec secretKey = new javax.crypto.spec.SecretKeySpec(taintedIKM, "AES");
+        HKDFParameterSpec spec = HKDFParameterSpec.expandOnly(
+            secretKey, "info".getBytes(), 32);
+
+        KDF kdf = KDF.getInstance("HKDF-SHA256");
+        byte[] result = kdf.deriveData(spec);
+        sink(result); // should flag
+    }
+
+    public static void testCleanUsage() throws Exception {
+        byte[] cleanKeyMaterial = "static-key-material".getBytes();
+
+        HKDFParameterSpec.Builder builder = HKDFParameterSpec.ofExtract();
+        builder.addIKM(cleanKeyMaterial); // clean input
+        HKDFParameterSpec spec = builder.thenExpand("info".getBytes(), 32);
+
+        KDF kdf = KDF.getInstance("HKDF-SHA256");
+        byte[] cleanResult = kdf.deriveData(spec);
+        sink(cleanResult); // should NOT flag - no taint source
+    }
+}

java/ql/test/library-tests/dataflow/kdf/options

@@ -0,0 +1 @@
+//semmle-extractor-options: --javac-args --enable-preview --release 25

java/ql/test/library-tests/dataflow/kdf/test.expected

@@ -0,0 +1,24 @@
+import java
+import semmle.code.java.dataflow.TaintTracking
+
+module Config implements DataFlow::ConfigSig {
+  predicate isSource(DataFlow::Node n) {
+    exists(ArrayAccess aa |
+      aa.getArray().(VarAccess).getVariable().hasName("args") and
+      n.asExpr() = aa
+    )
+  }
+
+  predicate isSink(DataFlow::Node n) {
+    exists(MethodCall ma |
+      ma.getMethod().hasName("sink") and
+      n.asExpr() = ma.getAnArgument()
+    )
+  }
+}
+
+module Flow = TaintTracking::Global<Config>;
+
+from DataFlow::Node src, DataFlow::Node sink
+where Flow::flow(src, sink)
+select src, sink