codeql/java/ql/test/library-tests/dataflow/kdf/KDFDataflowTest.java at 728a4aff225e5af2be80f052bbec13e41c9e14e1 · github/codeql · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
import javax.crypto.KDF;
import javax.crypto.spec.HKDFParameterSpec;

public class KDFDataflowTest {
    public static String source(String label) {
        return "tainted";
    }

    public static void sink(Object o) {}

    public static void main(String[] args) throws Exception {
        String userInput = 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); // $ hasTaintFlow
    }

    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); // $ hasTaintFlow
    }

    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); // $ hasTaintFlow
    }

    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); // $ hasTaintFlow
    }

    public static void testCleanUsage() throws Exception {
        byte[] cleanKeyMaterial = "static-key-material".getBytes();

        HKDFParameterSpec.Builder builder = HKDFParameterSpec.ofExtract();
        builder.addIKM(cleanKeyMaterial);
        HKDFParameterSpec spec = builder.thenExpand("info".getBytes(), 32);

        KDF kdf = KDF.getInstance("HKDF-SHA256");
        byte[] cleanResult = kdf.deriveData(spec);
        sink(cleanResult); // Safe - no taint
    }

    public static void testThenExpand(byte[] cleanIKM) throws Exception {
        String userInput = source("");
        byte[] taintedInfo = userInput.getBytes();

        HKDFParameterSpec.Builder builder = HKDFParameterSpec.ofExtract();
        builder.addIKM(cleanIKM);
        HKDFParameterSpec spec = builder.thenExpand(taintedInfo, 32);

        KDF kdf = KDF.getInstance("HKDF-SHA256");
        byte[] result = kdf.deriveData(spec);
        sink(result); // $ hasTaintFlow
    }
}