Make dispose idempotent and reject use after dispose

src/main/c/KeyEncapsulation.c

@@ -29,7 +29,10 @@ JNIEXPORT void JNICALL Java_org_openquantumsafe_KeyEncapsulation_free_1KEM
   (JNIEnv *env, jobject obj)
 {
     OQS_KEM *kem = (OQS_KEM *) getHandle(env, obj, "native_kem_handle_");
+    if (kem == NULL) return;
     OQS_KEM_free(kem);
+    // Clear the handle so that the freed object cannot be used again
+    setHandle(env, obj, NULL, "native_kem_handle_");
 }
 
 /*

src/main/c/Signature.c

@@ -29,7 +29,10 @@ JNIEXPORT void JNICALL Java_org_openquantumsafe_Signature_free_1sig
   (JNIEnv *env, jobject obj)
 {
     OQS_SIG *sig = (OQS_SIG *) getHandle(env, obj, "native_sig_handle_");
+    if (sig == NULL) return;
     OQS_SIG_free(sig);
+    // Clear the handle so that the freed object cannot be used again
+    setHandle(env, obj, NULL, "native_sig_handle_");
 }
 
 /*

src/main/java/org/openquantumsafe/KeyEncapsulation.java

@@ -163,12 +163,23 @@ public void dispose_KEM() {
         free_KEM();
     }
 
+    /**
+     * \brief Throw IllegalStateException if dispose_KEM() has freed the native
+     * object.
+     */
+    private void check_not_disposed() throws IllegalStateException {
+        if (native_kem_handle_ == 0) {
+            throw new IllegalStateException("KeyEncapsulation has been disposed");
+        }
+    }
+
 
     /**
      * \brief Invoke native generate_keypair method using the PK and SK lengths
      * from alg_details_. Check return value and if != 0 throw Exception.
      */
     public byte[] generate_keypair() throws RuntimeException {
+        check_not_disposed();
         int rv_ = generate_keypair(this.public_key_, this.secret_key_);
         if (rv_ != 0) throw new RuntimeException("Cannot generate keypair");
         return this.public_key_;
@@ -179,6 +190,7 @@ public byte[] generate_keypair() throws RuntimeException {
      * from alg_details_. Check return value and if != 0 throw Exception.
      */
     public byte[] generate_keypair(byte[] seed) throws RuntimeException {
+        check_not_disposed();
         if (seed.length != alg_details_.length_keypair_seed) {
             throw new RuntimeException("Incorrect seed length");
         }
@@ -192,20 +204,23 @@ public byte[] generate_keypair(byte[] seed) throws RuntimeException {
      * \brief Return seed length
      */
     public long get_keypair_seed_length() {
+        check_not_disposed();
         return alg_details_.length_keypair_seed;
     }
 
     /**
      * \brief Return public key
      */
     public byte[] export_public_key() {
+        check_not_disposed();
         return this.public_key_;
     }
 
     /**
     * \brief Return secret key
      */
     public byte[] export_secret_key() {
+        check_not_disposed();
         return this.secret_key_;
     }
 
@@ -216,6 +231,7 @@ public byte[] export_secret_key() {
      */
     public Pair<byte[], byte[]> encap_secret(byte[] public_key)
                                                     throws RuntimeException {
+        check_not_disposed();
         if (public_key.length != alg_details_.length_public_key) {
             throw new RuntimeException("Incorrect public key length");
         }
@@ -236,6 +252,7 @@ public Pair<byte[], byte[]> encap_secret(byte[] public_key)
      * \return Shared secret
      */
     public byte[] decap_secret(byte[] ciphertext) throws RuntimeException {
+        check_not_disposed();
         if (ciphertext.length != alg_details_.length_ciphertext) {
             throw new RuntimeException("Incorrect ciphertext length");
         }
@@ -255,6 +272,7 @@ public byte[] decap_secret(byte[] ciphertext) throws RuntimeException {
      * initialized, initialize it and fill it using native C code.
      */
     public void print_KeyEncapsulation() {
+        check_not_disposed();
         if (alg_details_ == null) {
             alg_details_ = get_KEM_details();
         }
@@ -267,6 +285,7 @@ public void print_KeyEncapsulation() {
      * is not initialized, initialize it and fill it using native C code.
      */
     public void print_details() {
+        check_not_disposed();
         if (alg_details_ == null) {
             alg_details_ = get_KEM_details();
         }

src/main/java/org/openquantumsafe/Signature.java

@@ -200,11 +200,22 @@ public void dispose_sig() {
         free_sig();
     }
 
+    /**
+     * \brief Throw IllegalStateException if dispose_sig() has freed the native
+     * object.
+     */
+    private void check_not_disposed() throws IllegalStateException {
+        if (native_sig_handle_ == 0) {
+            throw new IllegalStateException("Signature has been disposed");
+        }
+    }
+
     /**
      * \brief Invoke native generate_keypair method using the PK and SK lengths
      * from alg_details_. Check return value and if != 0 throw RuntimeException.
      */
     public byte[] generate_keypair() throws RuntimeException {
+        check_not_disposed();
         int rv_ = generate_keypair(this.public_key_, this.secret_key_);
         if (rv_ != 0) throw new RuntimeException("Cannot generate keypair");
         return this.public_key_;
@@ -214,13 +225,15 @@ public byte[] generate_keypair() throws RuntimeException {
      * \brief Return public key
      */
     public byte[] export_public_key() {
+        check_not_disposed();
         return this.public_key_;
     }
 
     /**
     * \brief Return secret key
      */
     public byte[] export_secret_key() {
+        check_not_disposed();
         return this.secret_key_;
     }
 
@@ -230,6 +243,7 @@ public byte[] export_secret_key() {
      * \return signature
      */
     public byte[] sign(byte[] message) throws RuntimeException {
+        check_not_disposed();
         if (this.secret_key_.length != alg_details_.length_secret_key) {
             throw new RuntimeException("Incorrect secret key length, " +
                                     "make sure you specify one in the " +
@@ -256,6 +270,7 @@ public byte[] sign(byte[] message) throws RuntimeException {
      */
     public boolean verify(byte[] message, byte[] signature, byte[] public_key)
                                                     throws RuntimeException {
+        check_not_disposed();
         if (public_key.length != alg_details_.length_public_key) {
             throw new RuntimeException("Incorrect public key length");
         }
@@ -273,6 +288,7 @@ public boolean verify(byte[] message, byte[] signature, byte[] public_key)
      * \return signature
      */
     public byte[] sign(byte[] message, byte[] ctx) throws RuntimeException {
+        check_not_disposed();
         if (this.secret_key_.length != alg_details_.length_secret_key) {
             throw new RuntimeException("Incorrect secret key length, " +
                                     "make sure you specify one in the " +
@@ -303,6 +319,7 @@ public byte[] sign(byte[] message, byte[] ctx) throws RuntimeException {
      */
     public boolean verify(byte[] message, byte[] signature, byte[] ctx, byte[] public_key)
                                                     throws RuntimeException {
+        check_not_disposed();
         if (public_key.length != alg_details_.length_public_key) {
             throw new RuntimeException("Incorrect public key length");
         }
@@ -320,6 +337,7 @@ public boolean verify(byte[] message, byte[] signature, byte[] ctx, byte[] publi
      * initialized, initialize it and fill it using native C code.
      */
     public void print_signature() {
+        check_not_disposed();
         if (alg_details_ == null) {
             alg_details_ = get_sig_details();
         }
@@ -331,6 +349,7 @@ public void print_signature() {
      * not initialized, initialize it and fill it using native C code.
      */
     public void print_details() {
+        check_not_disposed();
         if (alg_details_ == null) {
             alg_details_ = get_sig_details();
         }

src/test/java/org/openquantumsafe/KEMTest.java

@@ -3,6 +3,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.function.Executable;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
@@ -57,6 +58,46 @@ public void testAllKEMs(String kem_name) {
         System.out.println(sb.toString());
     }
 
+    /**
+     * Test that disposing of a KEM twice is harmless
+     */
+    @Test
+    public void testDisposeKEMTwice() {
+        KeyEncapsulation client = new KeyEncapsulation("ML-KEM-512");
+        client.generate_keypair();
+
+        client.dispose_KEM();
+        client.dispose_KEM();
+    }
+
+    /**
+     * Test that a KEM cannot be used after it has been disposed of
+     */
+    @Test
+    public void testUseAfterDisposeKEMExpectedException() {
+        KeyEncapsulation client = new KeyEncapsulation("ML-KEM-512");
+        byte[] client_public_key = client.generate_keypair();
+        byte[] ciphertext = client.encap_secret(client_public_key).getLeft();
+        byte[] seed = Rand.randombytes(client.get_keypair_seed_length());
+
+        client.dispose_KEM();
+
+        assertDisposed(() -> client.generate_keypair());
+        assertDisposed(() -> client.generate_keypair(seed));
+        assertDisposed(() -> client.encap_secret(client_public_key));
+        assertDisposed(() -> client.decap_secret(ciphertext));
+        assertDisposed(() -> client.get_keypair_seed_length());
+        assertDisposed(() -> client.export_public_key());
+        assertDisposed(() -> client.export_secret_key());
+        assertDisposed(() -> client.print_KeyEncapsulation());
+        assertDisposed(() -> client.print_details());
+    }
+
+    private static void assertDisposed(Executable executable) {
+        IllegalStateException e = Assertions.assertThrows(IllegalStateException.class, executable);
+        Assertions.assertEquals("KeyEncapsulation has been disposed", e.getMessage());
+    }
+
     /**
      * Test KEMs with derandomized keypair generation.
      */

src/test/java/org/openquantumsafe/SigTest.java

@@ -3,6 +3,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.function.Executable;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 
@@ -64,6 +65,46 @@ public void testAllSigs(String sig_name) {
         System.out.println(sb.toString());
     }
 
+    /**
+     * Test that disposing of a Sig twice is harmless
+     */
+    @Test
+    public void testDisposeSigTwice() {
+        Signature signer = new Signature("ML-DSA-44");
+        signer.generate_keypair();
+
+        signer.dispose_sig();
+        signer.dispose_sig();
+    }
+
+    /**
+     * Test that a Sig cannot be used after it has been disposed of
+     */
+    @Test
+    public void testUseAfterDisposeSigExpectedException() {
+        byte[] context = "01234567890".getBytes();
+        Signature signer = new Signature("ML-DSA-44");
+        byte[] signer_public_key = signer.generate_keypair();
+        byte[] signature = signer.sign(message);
+
+        signer.dispose_sig();
+
+        assertDisposed(() -> signer.generate_keypair());
+        assertDisposed(() -> signer.sign(message));
+        assertDisposed(() -> signer.verify(message, signature, signer_public_key));
+        assertDisposed(() -> signer.sign(message, context));
+        assertDisposed(() -> signer.verify(message, signature, context, signer_public_key));
+        assertDisposed(() -> signer.export_public_key());
+        assertDisposed(() -> signer.export_secret_key());
+        assertDisposed(() -> signer.print_signature());
+        assertDisposed(() -> signer.print_details());
+    }
+
+    private static void assertDisposed(Executable executable) {
+        IllegalStateException e = Assertions.assertThrows(IllegalStateException.class, executable);
+        Assertions.assertEquals("Signature has been disposed", e.getMessage());
+    }
+
     /**
      * Test Sigs with context.
      */