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

@@ -261,35 +261,39 @@
         if (algorithm.endsWith("inP1363Format")) {
             this.p1363Format = true;
         }
     }
 
-    private void ensureInitialized() {
-        token.ensureValid();
-        if (initialized == false) {
-            initialize();
-        }
-    }
-
-    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;
-        }
-        try {
+        } else {
             // "cancel" operation by finishing it
             // XXX make sure all this always works correctly
-            if (mode == M_SIGN) {
                 try {
+                if (mode == M_SIGN) {
                     if (type == T_UPDATE) {
                         token.p11.C_SignFinal(session.id(), 0);
                     } else {
                         byte[] digest;
                         if (type == T_DIGEST) {

@@ -297,15 +301,11 @@
                         } else { // T_RAW
                             digest = buffer;
                         }
                         token.p11.C_Sign(session.id(), digest);
                     }
-                } catch (PKCS11Exception e) {
-                    throw new ProviderException("cancel failed", e);
-                }
             } else { // M_VERIFY
-                try {
                     byte[] signature;
                     if (keyAlgorithm.equals("DSA")) {
                         signature = new byte[40];
                     } else {
                         signature = new byte[(p11Key.length() + 7) >> 3];

@@ -319,22 +319,29 @@
                         } else { // T_RAW
                             digest = buffer;
                         }
                         token.p11.C_Verify(session.id(), digest, signature);
                     }
-                } catch (PKCS11Exception e) {
-                    // will fail since the signature is incorrect
-                    // XXX check error code
                 }
+            } catch (PKCS11Exception e) {
+                throw new ProviderException("cancel failed", e);
             }
-        } finally {
-            session = token.releaseSession(session);
         }
     }
 
     // assumes current state is initialized == false
-    private void initialize() {
+    private void ensureInitialized() {
+        if (initialized) {
+            return;
+        }
+        if (p11Key == null) {
+            throw new ProviderException(
+                    "Operation cannot be performed without calling engineInit first");
+        }
+        try {
+            token.ensureValid();
+            p11Key.incNativeKeyRef();
         try {
             if (session == null) {
                 session = token.getOpSession();
             }
             if (mode == M_SIGN) {

@@ -342,21 +349,22 @@
                         new CK_MECHANISM(mechanism), p11Key.keyID);
             } else {
                 token.p11.C_VerifyInit(session.id(),
                         new CK_MECHANISM(mechanism), p11Key.keyID);
             }
-            initialized = true;
-        } catch (PKCS11Exception e) {
-            // release session when initialization failed
+            } catch (Throwable t) {
+                p11Key.decNativeKeyRef();
             session = token.releaseSession(session);
-            throw new ProviderException("Initialization failed", e);
+                throw t;
         }
-        if (bytesProcessed != 0) {
-            bytesProcessed = 0;
-            if (md != null) {
+            initialized = true;
+            if (bytesProcessed != 0 && md != null) {
                 md.reset();
             }
+            bytesProcessed = 0;
+        } catch (PKCS11Exception e) {
+            throw new ProviderException("Initialization failed", e);
         }
     }
 
     private void checkKeySize(String keyAlgo, Key key)
         throws InvalidKeyException {

@@ -449,14 +457,14 @@
         }
         // Need to check key length whenever a new key is set
         if (publicKey != p11Key) {
             checkKeySize(keyAlgorithm, publicKey);
         }
-        cancelOperation();
+        reset(true);
         mode = M_VERIFY;
         p11Key = P11KeyFactory.convertKey(token, publicKey, keyAlgorithm);
-        initialize();
+        ensureInitialized();
     }
 
     // see JCA spec
     @Override
     protected void engineInitSign(PrivateKey privateKey)

@@ -466,14 +474,14 @@
         }
         // Need to check RSA key length whenever a new key is set
         if (privateKey != p11Key) {
             checkKeySize(keyAlgorithm, privateKey);
         }
-        cancelOperation();
+        reset(true);
         mode = M_SIGN;
         p11Key = P11KeyFactory.convertKey(token, privateKey, keyAlgorithm);
-        initialize();
+        ensureInitialized();
     }
 
     // see JCA spec
     @Override
     protected void engineUpdate(byte b) throws SignatureException {

@@ -515,12 +523,11 @@
                 } else {
                     token.p11.C_VerifyUpdate(session.id(), 0, b, ofs, len);
                 }
                 bytesProcessed += len;
             } catch (PKCS11Exception e) {
-                initialized = false;
-                session = token.releaseSession(session);
+                reset(true);
                 throw new ProviderException(e);
             }
             break;
         case T_DIGEST:
             md.update(b, ofs, len);

@@ -639,15 +646,14 @@
                 }
             }
         } catch (PKCS11Exception pe) {
             throw new ProviderException(pe);
         } catch (SignatureException | ProviderException e) {
-            cancelOperation();
+            reset(true);
             throw e;
         } finally {
-            initialized = false;
-            session = token.releaseSession(session);
+            reset(false);
         }
     }
 
     // see JCA spec
     @Override

@@ -708,15 +714,14 @@
             if (errorCode == CKR_DATA_LEN_RANGE) {
                 return false;
             }
             throw new ProviderException(pe);
         }  catch (SignatureException | ProviderException e) {
-            cancelOperation();
+            reset(true);
             throw e;
         } finally {
-            initialized = false;
-            session = token.releaseSession(session);
+            reset(false);
         }
     }
 
     private byte[] pkcs1Pad(byte[] data) {
         try {