make autocloseable and add cleaner

jpos · Aug 21, 2023 · e5f2961 · e5f2961
1 parent b30c77d
commit e5f2961
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 1 deletion.
diff --git a/jpos/src/main/java/org/jpos/security/SensitiveString.java b/jpos/src/main/java/org/jpos/security/SensitiveString.java
@@ -22,24 +22,30 @@
 
 import javax.crypto.*;
 import javax.crypto.spec.IvParameterSpec;
+import java.lang.ref.Cleaner;
 import java.nio.ByteBuffer;
 import java.security.*;
+import java.util.Arrays;
 import java.util.Random;
 import java.util.function.Supplier;
 
-public class SensitiveString implements Supplier<String> {
+public class SensitiveString implements Supplier<String>, AutoCloseable {
     private SecretKey key;
     private byte[] encoded;
     private static Random rnd;
     private static final String AES = "AES/CBC/PKCS5Padding";
 
+    private static final Cleaner cleaner = Cleaner.create();
+    private Cleaner.Cleanable cleanable;
+
     static {
         rnd = new SecureRandom();
     }
 
     public SensitiveString(String s) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException, NoSuchProviderException, InvalidAlgorithmParameterException {
         key = generateKey();
         encoded = encrypt(s.getBytes());
+        cleanable = cleaner.register(this, this::clean);
     }
 
     @Override
@@ -86,12 +92,25 @@ private byte[] randomIV() {
         return b;
     }
 
+    public void clean () {
+        byte[] b = encoded;
+        encoded = null;
+        Arrays.fill (b, (byte) 0);
+    }
+
     @Override
     public String get() {
+        if (encoded == null)
+            throw new IllegalStateException ("SensitiveString not available");
         try {
             return new String(decrypt(encoded));
         } catch (NoSuchPaddingException | NoSuchAlgorithmException | BadPaddingException | InvalidKeyException | NoSuchProviderException | IllegalBlockSizeException | InvalidAlgorithmParameterException e) {
             throw new AssertionError(e.getMessage());
         }
     }
+
+    @Override
+    public void close() throws Exception {
+        clean();
+    }
 }
diff --git a/jpos/src/test/java/org/jpos/security/SensitiveStringTest.java b/jpos/src/test/java/org/jpos/security/SensitiveStringTest.java
@@ -18,9 +18,20 @@
 
 package org.jpos.security;
 
+import org.jpos.space.SpaceProxy;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import javax.crypto.BadPaddingException;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.security.NoSuchProviderException;
+
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class SensitiveStringTest {
     @Test
@@ -40,4 +51,13 @@ public void testSSEquals() throws Exception {
         SensitiveString ss1 = new SensitiveString(s);
         assertEquals (ss0, ss1, "Equals should be true");
     }
+
+    @Test
+    public void testClean() throws Exception {
+        SensitiveString ss = new SensitiveString("abc");
+        try (ss) {
+            assertEquals(ss.get(), "abc", "Equals should be true");
+        };
+        assertThrows(IllegalStateException.class, ss::get);
+    }
 }