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

@@ -194,11 +194,11 @@
         }
         implInit(opmode, key);
     }
 
     private void implInit(int opmode, Key key) throws InvalidKeyException {
-        cancelOperation();
+        reset(true);
         p11Key = P11KeyFactory.convertKey(token, key, algorithm);
         boolean encrypt;
         if (opmode == Cipher.ENCRYPT_MODE) {
             encrypt = true;
         } else if (opmode == Cipher.DECRYPT_MODE) {

@@ -233,29 +233,41 @@
         outputSize = n;
         buffer = new byte[n];
         maxInputSize = ((padType == PAD_PKCS1 && encrypt) ?
                             (n - PKCS1_MIN_PADDING_LENGTH) : n);
         try {
-            initialize();
+            ensureInitialized();
         } catch (PKCS11Exception e) {
             throw new InvalidKeyException("init() failed", e);
         }
     }
 
-    private void cancelOperation() {
-        token.ensureValid();
-        if (initialized == false) {
+    // reset the states to the pre-initialized values
+    private void reset(boolean doCancel) {
+        if (!initialized) {
             return;
         }
         initialized = false;
-        if ((session == null) || (token.explicitCancel == false)) {
+        try {
+            if (session == null) {
             return;
         }
+            if (doCancel && token.explicitCancel) {
+                cancelOperation();
+            }
+        } finally {
+            p11Key.decNativeKeyRef();
+            session = token.releaseSession(session);
+        }
+    }
+
+    private void cancelOperation() {
+        token.ensureValid();
         if (session.hasObjects() == false) {
             session = token.killSession(session);
             return;
-        }
+        } else {
         try {
             PKCS11 p11 = token.p11;
             int inLen = maxInputSize;
             int outLen = buffer.length;
             switch (mode) {

@@ -281,19 +293,23 @@
             }
         } catch (PKCS11Exception e) {
             // XXX ensure this always works, ignore error
         }
     }
+    }
 
     private void ensureInitialized() throws PKCS11Exception {
-        token.ensureValid();
-        if (initialized == false) {
-            initialize();
+        if (initialized) {
+            return;
         }
+        if (p11Key == null) {
+            throw new ProviderException(
+                    "Operation cannot be performed without calling engineInit first");
     }
-
-    private void initialize() throws PKCS11Exception {
+        token.ensureValid();
+        p11Key.incNativeKeyRef();
+        try {
         if (session == null) {
             session = token.getOpSession();
         }
         PKCS11 p11 = token.p11;
         CK_MECHANISM ckMechanism = new CK_MECHANISM(mechanism);

@@ -311,12 +327,17 @@
             p11.C_VerifyRecoverInit(session.id(), ckMechanism, p11Key.keyID);
             break;
         default:
             throw new AssertionError("internal error");
         }
-        bufOfs = 0;
+        } catch (Throwable t) {
+            p11Key.decNativeKeyRef();
+            session = token.releaseSession(session);
+            throw t;
+        }
         initialized = true;
+        bufOfs = 0;
     }
 
     private void implUpdate(byte[] in, int inOfs, int inLen) {
         try {
             ensureInitialized();

@@ -375,12 +396,11 @@
             return n;
         } catch (PKCS11Exception e) {
             throw (BadPaddingException)new BadPaddingException
                 ("doFinal() failed").initCause(e);
         } finally {
-            initialized = false;
-            session = token.releaseSession(session);
+            reset(false);
         }
     }
 
     // see JCE spec
     protected byte[] engineUpdate(byte[] in, int inOfs, int inLen) {

@@ -452,15 +472,19 @@
             }
         }
         Session s = null;
         try {
             s = token.getOpSession();
+            p11Key.incNativeKeyRef();
+            sKey.incNativeKeyRef();
             return token.p11.C_WrapKey(s.id(), new CK_MECHANISM(mechanism),
                 p11Key.keyID, sKey.keyID);
         } catch (PKCS11Exception e) {
             throw new InvalidKeyException("wrap() failed", e);
         } finally {
+            p11Key.decNativeKeyRef();
+            sKey.decNativeKeyRef();
             token.releaseSession(s);
         }
     }
 
     // see JCE spec

@@ -526,15 +550,21 @@
                             new CK_ATTRIBUTE(CKA_CLASS, CKO_SECRET_KEY),
                             new CK_ATTRIBUTE(CKA_KEY_TYPE, keyType),
                         };
                     attributes = token.getAttributes(
                             O_IMPORT, CKO_SECRET_KEY, keyType, attributes);
-                    long keyID = token.p11.C_UnwrapKey(s.id(),
+                    p11Key.incNativeKeyRef();
+                    long keyID;
+                    try {
+                        keyID = token.p11.C_UnwrapKey(s.id(),
                             new CK_MECHANISM(mechanism), p11Key.keyID,
                             wrappedKey, attributes);
+                    } finally {
+                        p11Key.decNativeKeyRef();
+                    }
                     secretKey = P11Key.secretKey(s, keyID,
-                            algorithm, 48 << 3, attributes);
+                            algorithm, 48 << 3, attributes, true);
                 } catch (PKCS11Exception e) {
                     if (isTlsRsaPremasterSecret) {
                         failover = e;
                     } else {
                         throw new InvalidKeyException("unwrap() failed", e);

@@ -579,11 +609,11 @@
                     CKK_GENERIC_SECRET, new CK_ATTRIBUTE[0]);
             long keyID = token.p11.C_GenerateKey(session.id(),
                     new CK_MECHANISM(CKM_SSL3_PRE_MASTER_KEY_GEN, version),
                     attributes);
             newKey = P11Key.secretKey(session,
-                    keyID, "TlsRsaPremasterSecret", 48 << 3, attributes);
+                    keyID, "TlsRsaPremasterSecret", 48 << 3, attributes, true);
         } catch (PKCS11Exception e) {
             throw new ProviderException(
                     "Could not generate premaster secret", e);
         }