Add zeroDek function to protect against JIT

Author: coltfred

src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKey.java

@@ -301,8 +301,7 @@ public CompletableFuture<BatchResult<PlaintextDocument>> decryptBatch(
   @Override
   public void close() {
     if (closed.compareAndSet(false, true)) {
-      // Zero out the DEK bytes for security
-      Arrays.fill(dek, (byte) 0);
+      zeroDek(dek);
       // Report operations to TSP
       int encrypts = encryptCount.get();
       int decrypts = decryptCount.get();
@@ -311,4 +310,16 @@ public void close() {
       }
     }
   }
+
+  /**
+   * Zero a DEK byte array with a subsequent access to prevent the JIT from eliminating the fill as
+   * a dead store. The volatile write after the fill ensures the zeroing is not optimized away.
+   */
+  @SuppressWarnings("unused")
+  private static volatile byte ZERO_FENCE;
+
+  static void zeroDek(byte[] dek) {
+    Arrays.fill(dek, (byte) 0);
+    ZERO_FENCE = dek[0];
+  }
 }

src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java

@@ -12,7 +12,6 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
-import java.util.Arrays;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.function.Function;
@@ -493,7 +492,7 @@ private CompletableFuture<CachedKey> newCachedKeyFromUnwrap(String edek,
     return this.encryptionService.unwrapKey(edek, metadata).thenApply(dekBytes -> {
       CachedKey cachedKey = new CachedKey(dekBytes, edek, this.encryptionExecutor,
           this.secureRandom, this.encryptionService, metadata);
-      Arrays.fill(dekBytes, (byte) 0);
+      CachedKey.zeroDek(dekBytes);
       return cachedKey;
     });
   }
@@ -503,7 +502,7 @@ private CompletableFuture<CachedKey> newCachedKeyFromWrap(DocumentMetadata metad
       byte[] dekBytes = wrappedKey.getDekBytes();
       CachedKey cachedKey = new CachedKey(dekBytes, wrappedKey.getEdek(), this.encryptionExecutor,
           this.secureRandom, this.encryptionService, metadata);
-      Arrays.fill(dekBytes, (byte) 0);
+      CachedKey.zeroDek(dekBytes);
       return cachedKey;
     });
   }