From 3b866585a269f18881692d69f6f94e0c1f699b65 Mon Sep 17 00:00:00 2001
From: WilburZjh <wilburzjh@gmail.com>
Date: Mon, 26 Jun 2023 11:07:54 -0400
Subject: [PATCH] Support exporting plain SecretKey in FIPS mode

We support exporting and importing plain
RSA/EC private keys and importing secret
keys as AES in FIPS mode for now. We also
need to support exporting secret keys in
FIPS mode. Currently, while user generating
an AES secretkey in FIPS mode through
KeyGenerator, this AES key's encode is NULL.
Therefore, we modify the P11Key.java and
SunPKCS11.java files to enable the support
for exporting secret keys inclusing AES
and TripleDES.

Signed-off-by: Jinhang Zhang <Jinhang.Zhang@ibm.com>
---
 .../classes/sun/security/pkcs11/P11Key.java   | 41 +++++++++++++++++++
 .../sun/security/pkcs11/SunPKCS11.java        |  4 ++
 2 files changed, 45 insertions(+)

diff --git a/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/P11Key.java b/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/P11Key.java
index 3b72ffadb33..7d4907b803c 100644
--- a/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/P11Key.java
+++ b/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/P11Key.java
@@ -358,6 +358,24 @@ static SecretKey secretKey(Session session, long keyID, String algorithm,
                     new CK_ATTRIBUTE(CKA_SENSITIVE),
                     new CK_ATTRIBUTE(CKA_EXTRACTABLE),
         });
+
+        if ((SunPKCS11.mysunpkcs11 != null) && !SunPKCS11.isExportWrapKey.get()
+            && ("AES".equals(algorithm) || "TripleDES".equals(algorithm))
+        ) {
+            if (attrs[0].getBoolean() || attrs[1].getBoolean() || (attrs[2].getBoolean() == false)) {
+                try {
+                    byte[] key = SunPKCS11.mysunpkcs11.exportKey(session.id(), attrs, keyID);
+                    SecretKey secretKey = new SecretKeySpec(key, algorithm);
+                    return new P11SecretKeyFIPS(session, keyID, algorithm, keyLength, attrs, secretKey);
+                } catch (PKCS11Exception e) {
+                    // Attempt failed, create a P11SecretKey object.
+                    if (debug != null) {
+                        debug.println("Attempt failed, creating a SecretKey object for " + algorithm);
+                    }
+                }
+            }
+        }
+
         return new P11SecretKey(session, keyID, algorithm, keyLength, attrs);
     }
 
@@ -480,6 +498,29 @@ byte[] getEncodedInternal() {
         }
     }
 
+    private static final class P11SecretKeyFIPS extends P11Key implements SecretKey {
+        @Serial
+        private static final long serialVersionUID = -9186806495402041696L;
+        private final SecretKey key;
+
+        P11SecretKeyFIPS(Session session, long keyID, String algorithm,
+                int keyLength, CK_ATTRIBUTE[] attributes, SecretKey key) {
+            super(SECRET, session, keyID, algorithm, keyLength, attributes);
+            this.key = key;
+        }
+
+        @Override
+        public String getFormat() {
+            return "RAW";
+        }
+
+        @Override
+        byte[] getEncodedInternal() {
+            return key.getEncoded();
+        }
+
+    }
+
     private static class P11SecretKey extends P11Key implements SecretKey {
         @Serial
         private static final long serialVersionUID = -7828241727014329084L;
diff --git a/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/SunPKCS11.java b/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/SunPKCS11.java
index 83e2a1dccbf..8be24e2e444 100644
--- a/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/SunPKCS11.java
+++ b/src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/SunPKCS11.java
@@ -128,6 +128,8 @@ public final class SunPKCS11 extends AuthProvider {
     // FIPS mode.
     static SunPKCS11 mysunpkcs11;
 
+    static final ThreadLocal<Boolean> isExportWrapKey = ThreadLocal.withInitial(() -> Boolean.FALSE);
+
     Token getToken() {
         return token;
     }
@@ -506,10 +508,12 @@ byte[] exportKey(long hSession, CK_ATTRIBUTE[] attributes, long keyId) throws PK
 
         try {
             long genKeyId = token.p11.C_GenerateKey(wrapKeyGenSession.id(), new CK_MECHANISM(CKM_AES_KEY_GEN), wrapKeyAttributes);
+            isExportWrapKey.set(Boolean.TRUE);
             wrapKey = (P11Key)P11Key.secretKey(wrapKeyGenSession, genKeyId, "AES", 256 >> 3, null);
         } catch (PKCS11Exception e) {
             throw e;
         } finally {
+            isExportWrapKey.set(Boolean.FALSE);
             token.releaseSession(wrapKeyGenSession);
         }