From 0f815a2358d0f918a93ce2429e7b97b529ff4711 Mon Sep 17 00:00:00 2001
From: WilburZjh <wilburzjh@gmail.com>
Date: Tue, 20 Jun 2023 13:00:22 -0400
Subject: [PATCH] Support exporting plain SecretKey in FIPS mode

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 950f0ed247..63ca048481 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
@@ -342,6 +342,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 (attributes[0].getBoolean() || attributes[1].getBoolean() || (attributes[2].getBoolean() == false)) {
+                try {
+                    byte[] key = SunPKCS11.mysunpkcs11.exportKey(session.id(), attributes, keyID);
+                    SecretKey secretKey = new SecretKeySpec(key, algorithm);
+                    return new P11SecretKeyFIPS(session, keyID, algorithm, keyLength, attributes, 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,
                 attributes);
     }
@@ -494,6 +512,29 @@ byte[] getEncodedInternal() {
         }
     }
 
+    private static final class P11SecretKeyFIPS extends P11Key implements SecretKey {
+
+        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 {
         private static final long serialVersionUID = -7828241727014329084L;
         private volatile byte[] encoded;
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 2176ee12fe..de77baf91f 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
@@ -118,6 +118,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;
     }
@@ -503,10 +505,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);
         }