src/jdk.crypto.cryptoki/share/classes/sun/security/pkcs11/P11KeyAgreement.java

@@ -207,13 +207,19 @@
                 new CK_ATTRIBUTE(CKA_CLASS, CKO_SECRET_KEY),
                 new CK_ATTRIBUTE(CKA_KEY_TYPE, CKK_GENERIC_SECRET),
             };
             attributes = token.getAttributes
                 (O_GENERATE, CKO_SECRET_KEY, CKK_GENERIC_SECRET, attributes);
-            long keyID = token.p11.C_DeriveKey(session.id(),
+            privateKey.incNativeKeyRef();
+            long keyID;
+            try {
+                keyID = token.p11.C_DeriveKey(session.id(),
                 new CK_MECHANISM(mechanism, publicValue), privateKey.keyID,
                 attributes);
+            } finally {
+                privateKey.decNativeKeyRef();
+            }
             attributes = new CK_ATTRIBUTE[] {
                 new CK_ATTRIBUTE(CKA_VALUE)
             };
             token.p11.C_GetAttributeValue(session.id(), keyID, attributes);
             byte[] secret = attributes[0].getByteArray();

@@ -331,20 +337,26 @@
                 new CK_ATTRIBUTE(CKA_CLASS, CKO_SECRET_KEY),
                 new CK_ATTRIBUTE(CKA_KEY_TYPE, keyType),
             };
             attributes = token.getAttributes
                 (O_GENERATE, CKO_SECRET_KEY, keyType, attributes);
-            long keyID = token.p11.C_DeriveKey(session.id(),
+            privateKey.incNativeKeyRef();
+            long keyID;
+            try {
+                keyID = token.p11.C_DeriveKey(session.id(),
                 new CK_MECHANISM(mechanism, publicValue), privateKey.keyID,
                 attributes);
+            } finally {
+                privateKey.decNativeKeyRef();
+            }
             CK_ATTRIBUTE[] lenAttributes = new CK_ATTRIBUTE[] {
                 new CK_ATTRIBUTE(CKA_VALUE_LEN),
             };
             token.p11.C_GetAttributeValue(session.id(), keyID, lenAttributes);
             int keyLen = (int)lenAttributes[0].getLong();
             SecretKey key = P11Key.secretKey
-                        (session, keyID, algorithm, keyLen << 3, attributes);
+                        (session, keyID, algorithm, keyLen << 3, attributes, true);
             if ("RAW".equals(key.getFormat())) {
                 // Workaround for Solaris bug 6318543.
                 // Strip leading zeroes ourselves if possible (key not sensitive).
                 // This should be removed once the Solaris fix is available
                 // as here we always retrieve the CKA_VALUE even for tokens