From 05b4c9852a0c7fd9d81b7f40ed05c95d0a8461dd Mon Sep 17 00:00:00 2001
From: Chris Conlon <chris@wolfssl.com>
Date: Tue, 19 Mar 2024 14:55:20 -0600
Subject: [PATCH] JCE: prepend zero byte to DH shared secret if less than prime
 length

---
 .../provider/jce/WolfCryptKeyAgreement.java   | 39 +++++++++++++++++++
 .../jce/test/WolfCryptKeyAgreementTest.java   | 11 ++++++
 2 files changed, 50 insertions(+)

diff --git a/src/main/java/com/wolfssl/provider/jce/WolfCryptKeyAgreement.java b/src/main/java/com/wolfssl/provider/jce/WolfCryptKeyAgreement.java
index 1451e4b..63a5f7d 100644
--- a/src/main/java/com/wolfssl/provider/jce/WolfCryptKeyAgreement.java
+++ b/src/main/java/com/wolfssl/provider/jce/WolfCryptKeyAgreement.java
@@ -21,6 +21,7 @@
 
 package com.wolfssl.provider.jce;
 
+import java.util.Arrays;
 import javax.crypto.KeyAgreementSpi;
 import javax.crypto.SecretKey;
 import javax.crypto.ShortBufferException;
@@ -194,6 +195,33 @@ public class WolfCryptKeyAgreement extends KeyAgreementSpi {
             secret = new byte[len];
             System.arraycopy(tmp, 0, secret, 0, len);
 
+            /* DH shared secrets can vary in length depending on if they are
+             * padded or not at the beginning with zero bytes to make a total
+             * output size matching the prime length.
+             *
+             * Native wolfCrypt does not prepend zero bytes to DH shared
+             * secrets, following RFC 5246 (8.1.2) which instructs to strip
+             * leading zero bytes.
+             *
+             * Sun KeyAgreement DH implementations as of after Java 8
+             * prepend zero bytes if total length is not equal to prime length.
+             * This was changed with OpenJDK bug fix JDK-7146728.
+             *
+             * BouncyCastle also behaves the same way, prepending zero bytes
+             * if total secret size is not prime length. This follows
+             * RFC 2631 (2.1.2).
+             *
+             * To match Sun and BC behavior, we will follow the same here if
+             * running on a Java version later than Java 8.
+             */
+            if (this.type == KeyAgreeType.WC_DH) {
+                tmp = new byte[this.primeLen];
+                Arrays.fill(tmp, (byte)0);
+                System.arraycopy(secret, 0, tmp,
+                    tmp.length - secret.length, secret.length);
+                secret = tmp.clone();
+            }
+
         } catch (ShortBufferException e) {
             zeroArray(tmp);
             zeroArray(secret);
@@ -360,6 +388,12 @@ public class WolfCryptKeyAgreement extends KeyAgreementSpi {
                 this.dh.setParams(paramP, paramG);
 
                 primeLen = paramP.length;
+
+                /* prime may have leading zero */
+                if (paramP[0] == 0x00) {
+                    primeLen--;
+                }
+
                 return;
 
             } else {
@@ -383,6 +417,11 @@ public class WolfCryptKeyAgreement extends KeyAgreementSpi {
 
         primeLen = paramP.length;
 
+        /* prime may have leading zero */
+        if (paramP[0] == 0x00) {
+            primeLen--;
+        }
+
         /* import private key */
         dhPriv = dhKey.getX().toByteArray();
         if (dhPriv == null) {
diff --git a/src/test/java/com/wolfssl/provider/jce/test/WolfCryptKeyAgreementTest.java b/src/test/java/com/wolfssl/provider/jce/test/WolfCryptKeyAgreementTest.java
index 41271c6..4d316f8 100644
--- a/src/test/java/com/wolfssl/provider/jce/test/WolfCryptKeyAgreementTest.java
+++ b/src/test/java/com/wolfssl/provider/jce/test/WolfCryptKeyAgreementTest.java
@@ -314,6 +314,17 @@ public class WolfCryptKeyAgreementTest {
             byte secretA[] = aKeyAgree.generateSecret();
             byte secretB[] = bKeyAgree.generateSecret();
 
+            /* Older versions of Java did not prepend a zero byte to shared
+             * secrets that were smaller than the prime length. This was
+             * changed in SunJCE as of JDK-7146728, but we may be running this
+             * test on an older version that does not prepend the zero byte.
+             * Since wolfJCE does prepend the zero byte, for the sake of this
+             * interop test, we strip the zero byte from wolfJCE's secret
+             * if lengths are different and try to compare that. */
+            if (secretB.length == (secretA.length - 1)) {
+                secretA = Arrays.copyOfRange(secretA, 1, secretA.length);
+            }
+
             if (secretA.length != secretB.length) {
                 int i = 0;
                 System.out.println("secretA.length != secretB.length");