Do not destroy linked public keys (#335)

Co-authored-by: Greg Rubin <[email protected]>

src/com/amazon/corretto/crypto/provider/EvpEcPrivateKey.java

@@ -23,10 +23,12 @@ class EvpEcPrivateKey extends EvpEcKey implements ECPrivateKey, CanDerivePublicK
 
   @Override
   public EvpEcPublicKey getPublicKey() {
-    ephemeral =
-        false; // Once our internal key could be elsewhere, we can no longer safely release it when
-    // done
-    return new EvpEcPublicKey(internalKey);
+    // Once our internal key could be elsewhere, we can no longer safely release it when done
+    ephemeral = false;
+    sharedKey = true;
+    final EvpEcPublicKey result = new EvpEcPublicKey(internalKey);
+    result.sharedKey = true;
+    return result;
   }
 
   @Override

src/com/amazon/corretto/crypto/provider/EvpKey.java

@@ -29,8 +29,18 @@ abstract class EvpKey implements Key, Destroyable {
   protected final InternalKey internalKey;
   protected final EvpKeyType type;
   protected final boolean isPublicKey;
+  /**
+   * Indicates that the backing native key is used by another java object and thus must not be
+   * released by this one.
+   */
+  protected boolean sharedKey = false;
+  /**
+   * Indicates that this key is entirely managed within ACCP controlled code and thus we know when
+   * we're done with it and can release it.
+   */
   protected boolean ephemeral = false;
 
+  private volatile boolean isDestroyed = false;
   protected volatile byte[] encoded;
   protected volatile Integer cachedHashCode;
 
@@ -66,11 +76,13 @@ void releaseEphemeral() {
   // @CheckReturnValue // Restore once replacement for JSR-305 available
   <T, X extends Throwable> T use(final MiscInterfaces.ThrowingLongFunction<T, X> function)
       throws X {
+    assertNotDestroyed();
     return internalKey.use(function);
   }
 
   <X extends Throwable> void useVoid(final MiscInterfaces.ThrowingLongConsumer<X> function)
       throws X {
+    assertNotDestroyed();
     internalKey.useVoid(function);
   }
 
@@ -91,6 +103,7 @@ public byte[] getEncoded() {
   }
 
   protected byte[] internalGetEncoded() {
+    assertNotDestroyed();
     byte[] result = encoded;
     if (result == null) {
       synchronized (this) {
@@ -124,10 +137,14 @@ protected <T extends AlgorithmParameterSpec> T nativeParams(final Class<T> param
   }
 
   /**
-   * This method will be called by @{link #destroy()} after calling @{code internalKey.release()}.
+   * This method will be called by @{link #destroy()} after possibly calling @{code
+   * internalKey.release()}.
    */
   protected synchronized void destroyJavaState() {
-    // NOP
+    if (encoded != null) {
+      Arrays.fill(encoded, (byte) 0);
+    }
+    encoded = null;
   }
 
   @Override
@@ -198,17 +215,24 @@ public int hashCode() {
     return result;
   }
 
+  protected void assertNotDestroyed() {
+    if (isDestroyed) {
+      throw new IllegalStateException("Key has been destroyed");
+    }
+  }
+
   @Override
   public boolean isDestroyed() {
-    return internalKey.isReleased();
+    return isDestroyed;
   }
 
   @Override
   public synchronized void destroy() {
-    if (isDestroyed()) {
-      throw new IllegalStateException("Already destroyed");
+    assertNotDestroyed();
+    isDestroyed = true;
+    if (!sharedKey) {
+      internalKey.release();
     }
-    internalKey.release();
     destroyJavaState();
   }
 

src/com/amazon/corretto/crypto/provider/EvpRsaPrivateCrtKey.java

@@ -45,10 +45,12 @@ protected static EvpRsaPrivateKey buildProperKey(long ptr) {
 
   @Override
   public EvpRsaPublicKey getPublicKey() {
-    ephemeral =
-        false; // Once our internal key could be elsewhere, we can no longer safely release it when
-    // done
-    return new EvpRsaPublicKey(internalKey);
+    // Once our internal key could be elsewhere, we can no longer safely release it when done
+    ephemeral = false;
+    sharedKey = true;
+    final EvpRsaPublicKey result = new EvpRsaPublicKey(internalKey);
+    result.sharedKey = true;
+    return result;
   }
 
   @Override

src/com/amazon/corretto/crypto/provider/EvpRsaPrivateKey.java

@@ -47,6 +47,7 @@ protected synchronized void destroyJavaState() {
   @Override
   protected byte[] internalGetEncoded() {
     // RSA private keys in Java may lack CRT parameters and thus need custom serialization
+    assertNotDestroyed();
     byte[] result = encoded;
     if (result == null) {
       synchronized (this) {

tst/com/amazon/corretto/crypto/provider/test/EcGenTest.java

@@ -293,6 +293,12 @@ public void defaultParams() throws GeneralSecurityException {
     nativeGen.generateKeyPair();
   }
 
+  @Test
+  public void separateDestruction() throws Exception {
+    final KeyPair keyPair = nativeGen.generateKeyPair();
+    RsaGenTest.testSeparateDestruction(keyPair);
+  }
+
   @Test
   public void threadStorm() throws Throwable {
     final byte[] rngSeed = TestUtil.getRandomBytes(20);

tst/com/amazon/corretto/crypto/provider/test/EvpKeyFactoryTest.java

@@ -179,7 +179,7 @@ public static List<Arguments> ecPairsTranslation() {
     return result;
   }
 
-  private static long getRawPointer(Object evpKey) throws Exception {
+  static long getRawPointer(final Object evpKey) throws Exception {
     final Object internalKey = sneakyGetField(evpKey, "internalKey");
     final Object cell = sneakyGetField(internalKey, "cell");
     return (long) sneakyGetField(cell, "ptr");

tst/com/amazon/corretto/crypto/provider/test/RsaGenTest.java

@@ -8,6 +8,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.junit.jupiter.api.Assumptions.assumeFalse;
 
 import com.amazon.corretto.crypto.provider.AmazonCorrettoCryptoProvider;
@@ -248,6 +249,38 @@ public void threadStorm() throws Throwable {
     }
   }
 
+  @Test
+  public void separateDestruction() throws Exception {
+    final KeyPairGenerator generator = getGenerator();
+    generator.initialize(2048);
+    final KeyPair keyPair = generator.generateKeyPair();
+    testSeparateDestruction(keyPair);
+  }
+
+  static void testSeparateDestruction(final KeyPair kp) throws Exception {
+    // Make sure that the keys are backed by the same native object.
+    // Otherwise the test is invalid.
+    assertEquals(
+        EvpKeyFactoryTest.getRawPointer(kp.getPublic()),
+        EvpKeyFactoryTest.getRawPointer(kp.getPrivate()),
+        "Keys must be backed by same native object for test to be valid");
+    // Destroy the private key
+    kp.getPrivate().destroy();
+    // Getting encoded private key must fail and mention destruction
+    try {
+      kp.getPrivate().getEncoded();
+      fail("Expected exception");
+    } catch (final IllegalStateException ex) {
+      assertTrue(ex.getMessage().contains("destroy"), ex.getMessage());
+    }
+    // We must still be able to retrieve the public key
+    final byte[] encoded = kp.getPublic().getEncoded();
+    assertNotNull(encoded);
+    assertTrue(encoded.length > 0);
+    // Leading byte of an encoded key will never be zero
+    assertTrue(encoded[0] != 0);
+  }
+
   private static void assertConsistency(final RSAPublicKey pub, final RSAPrivateCrtKey priv)
       throws GeneralSecurityException {
     assertNotNull(pub);