From 4a1e3c915d9ad34ea618f789745737d1ae0ce9dc Mon Sep 17 00:00:00 2001 From: Oracle Date: Thu, 30 Apr 2026 14:17:44 +0200 Subject: [PATCH] Fix SecureBuffer data access --- .../java/ai/nomyo/SecureCompletionClient.java | 26 ++++++++++--------- src/main/java/ai/nomyo/SecureMemory.java | 14 ++++++---- src/main/java/ai/nomyo/util/Pass2Key.java | 14 +++++----- src/test/java/ai/nomyo/SecureMemoryTest.java | 7 ++--- 4 files changed, 34 insertions(+), 27 deletions(-) diff --git a/src/main/java/ai/nomyo/SecureCompletionClient.java b/src/main/java/ai/nomyo/SecureCompletionClient.java index 85b3778..1d702de 100644 --- a/src/main/java/ai/nomyo/SecureCompletionClient.java +++ b/src/main/java/ai/nomyo/SecureCompletionClient.java @@ -17,6 +17,7 @@ import java.net.*; import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; @@ -345,7 +346,7 @@ public class SecureCompletionClient { Cipher cipher; try { cipher = Cipher.getInstance("AES/GCM/NoPadding"); - cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(aesKey.getEncoded(), "AES"), new GCMParameterSpec(Constants.GCM_TAG_SIZE * Byte.SIZE, secureNonce.getData().asByteBuffer().array())); + cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(aesKey.getEncoded(), "AES"), new GCMParameterSpec(Constants.GCM_TAG_SIZE * Byte.SIZE, secureNonce.getAsByteArray())); } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidAlgorithmParameterException | InvalidKeyException e) { throw new RuntimeException(new SecurityError("AES-GCM cipher initialization failed: " + e.getMessage(), e)); @@ -354,7 +355,7 @@ public class SecureCompletionClient { byte[] ciphertext; try { - ciphertext = cipher.doFinal(securePayload.getData().asByteBuffer().array()); + ciphertext = cipher.doFinal(securePayload.getAsByteArray()); } catch (IllegalBlockSizeException | BadPaddingException e) { throw new RuntimeException(new SecurityError("AES-GCM encryption failed: " + e.getMessage(), e)); } @@ -399,7 +400,7 @@ public class SecureCompletionClient { byte[] encryptedAESKey; try (SecureBuffer secureAesKeyEncoded = SecureMemory.secureByteArray(aesKey.getEncoded())) { - encryptedAESKey = rsa.doFinal(secureAesKeyEncoded.getData().asByteBuffer().array()); + encryptedAESKey = rsa.doFinal(secureAesKeyEncoded.getAsByteArray()); } catch (IllegalBlockSizeException | BadPaddingException e) { throw new RuntimeException(new SecurityError("RSA-OAEP key wrapping failed: " + e.getMessage(), e)); } @@ -777,10 +778,10 @@ public class SecureCompletionClient { Cipher rsaCipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding"); rsaCipher.init(Cipher.DECRYPT_MODE, this.privateKey, oaepParams); - byte[] aesKeyBytes = rsaCipher.doFinal(secureEncryptedAESKey.getData().asByteBuffer().array()); + byte[] aesKeyBytes = rsaCipher.doFinal(secureEncryptedAESKey.getAsByteArray()); try (SecureBuffer secureAesKeyBytes = SecureMemory.secureByteArray(aesKeyBytes)) { - SecretKeySpec aesKey = new SecretKeySpec(secureAesKeyBytes.getData().asByteBuffer().array(), "AES"); + SecretKeySpec aesKey = new SecretKeySpec(secureAesKeyBytes.getAsByteArray(), "AES"); // Decrypt payload byte[] ciphertext = Base64.getDecoder().decode(encryptedPayload.get("ciphertext").getAsString()); @@ -790,19 +791,20 @@ public class SecureCompletionClient { try (SecureBuffer secureCiphertext = SecureMemory.secureByteArray(ciphertext); SecureBuffer secureNonce = SecureMemory.secureByteArray(nonce); SecureBuffer secureTag = SecureMemory.secureByteArray(tag)) { Cipher aesCipher = Cipher.getInstance("AES/GCM/NoPadding"); - aesCipher.init(Cipher.DECRYPT_MODE, aesKey, new GCMParameterSpec(Constants.GCM_TAG_SIZE * 8, secureNonce.getData().asByteBuffer().array())); + aesCipher.init(Cipher.DECRYPT_MODE, aesKey, new GCMParameterSpec(Constants.GCM_TAG_SIZE * 8, secureNonce.getAsByteArray())); - // Combine ciphertext (without tag) and tag for decryption using SecureBuffer - try (SecureBuffer secureCiphertextWithTag = SecureMemory.secureByteArray(new byte[ciphertext.length + tag.length])) { - secureCiphertextWithTag.getData().asByteBuffer().put(secureCiphertext.getData().asByteBuffer().array()); - secureCiphertextWithTag.getData().asByteBuffer().put(secureTag.getData().asByteBuffer().array()); + // Combine ciphertext (without tag) and tag for decryption using SecureBuffer + try (SecureBuffer secureCiphertextWithTag = SecureMemory.secureByteArray(new byte[ciphertext.length + tag.length])) { + ByteBuffer combinedBuf = secureCiphertextWithTag.getData().asByteBuffer(); + combinedBuf.put(secureCiphertext.getAsByteArray()); + combinedBuf.put(secureTag.getAsByteArray()); - byte[] plaintextBytes = aesCipher.doFinal(secureCiphertextWithTag.getData().asByteBuffer().array()); + byte[] plaintextBytes = aesCipher.doFinal(secureCiphertextWithTag.getAsByteArray()); // Parse JSON response Map response; try (SecureBuffer securePlaintext = SecureMemory.secureByteArray(plaintextBytes)) { - Object parsed = gson.fromJson(new String(securePlaintext.getData().asByteBuffer().array(), StandardCharsets.UTF_8), Object.class); + Object parsed = gson.fromJson(new String(securePlaintext.getAsByteArray(), StandardCharsets.UTF_8), Object.class); @SuppressWarnings("unchecked") Map resultMap = (Map) parsed; response = resultMap != null ? resultMap : new HashMap<>(); } diff --git a/src/main/java/ai/nomyo/SecureMemory.java b/src/main/java/ai/nomyo/SecureMemory.java index b8ba7c2..d06c486 100644 --- a/src/main/java/ai/nomyo/SecureMemory.java +++ b/src/main/java/ai/nomyo/SecureMemory.java @@ -5,7 +5,6 @@ import lombok.Setter; import java.lang.foreign.*; import java.lang.invoke.MethodHandle; -import java.nio.ByteBuffer; import java.util.Arrays; import java.util.Map; @@ -145,7 +144,7 @@ public final class SecureMemory { * @param lock whether to attempt memory locking */ public SecureBuffer(byte[] data, boolean lock) { - this.arena = Arena.ofConfined(); + this.arena = Arena.ofShared(); this.data = data != null ? this.arena.allocate(data.length) : MemorySegment.NULL; if (data != null) { @@ -164,6 +163,10 @@ public final class SecureMemory { } } + public byte[] getAsByteArray() { + return this.data.toArray(ValueLayout.JAVA_BYTE); + } + /** * Locks buffer in memory (prevents disk swapping). Returns false if unavailable. */ @@ -174,7 +177,7 @@ public final class SecureMemory { try { long result = (long) MLOCK_HANDLE.invokeExact( - MemorySegment.ofAddress(this.address), + this.address, this.size); this.locked = result == 0; } catch (Throwable t) { @@ -194,7 +197,7 @@ public final class SecureMemory { try { long result = (long) MUNLOCK_HANDLE.invokeExact( - MemorySegment.ofAddress(this.address), + this.address, this.size); locked = result != 0; return result == 0; @@ -208,7 +211,7 @@ public final class SecureMemory { * Securely zeros buffer contents. */ public void zero() { - if (data != null) { + if (data != MemorySegment.NULL) { data.fill((byte) 0); } } @@ -220,6 +223,7 @@ public final class SecureMemory { zero(); unlock(); + arena.close(); closed = true; } diff --git a/src/main/java/ai/nomyo/util/Pass2Key.java b/src/main/java/ai/nomyo/util/Pass2Key.java index 7100d24..e82420e 100644 --- a/src/main/java/ai/nomyo/util/Pass2Key.java +++ b/src/main/java/ai/nomyo/util/Pass2Key.java @@ -54,17 +54,17 @@ public final class Pass2Key { byte[] iv = new byte[GCM_IV_LENGTH]; RANDOM.nextBytes(iv); try (SecureBuffer secureSalt = SecureMemory.secureByteArray(salt); SecureBuffer secureIv = SecureMemory.secureByteArray(iv)) { - GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, secureIv.getData().asByteBuffer().array()); + GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, secureIv.getAsByteArray()); byte[] ciphertext = encryptWithCipher(algorithm, key, spec, input); try (SecureBuffer secureCiphertext = SecureMemory.secureByteArray(ciphertext)) { - payload = assemblePayloadGcm(secureSalt.getData().asByteBuffer().array(), secureIv.getData().asByteBuffer().array(), secureCiphertext.getData().asByteBuffer().array()); + payload = assemblePayloadGcm(secureSalt.getAsByteArray(), secureIv.getAsByteArray(), secureCiphertext.getAsByteArray()); } } } else { try (SecureBuffer secureSalt = SecureMemory.secureByteArray(salt)) { byte[] ciphertext = encryptWithCipher(algorithm, key, input); try (SecureBuffer secureCiphertext = SecureMemory.secureByteArray(ciphertext)) { - payload = assemblePayloadSalt(secureSalt.getData().asByteBuffer().array(), secureCiphertext.getData().asByteBuffer().array()); + payload = assemblePayloadSalt(secureSalt.getAsByteArray(), secureCiphertext.getAsByteArray()); } } } @@ -92,13 +92,13 @@ public final class Pass2Key { byte[] iv = java.util.Arrays.copyOfRange(decoded, SALT_LENGTH, SALT_LENGTH + GCM_IV_LENGTH); byte[] ciphertext = java.util.Arrays.copyOfRange(decoded, SALT_LENGTH + GCM_IV_LENGTH, decoded.length); try (SecureBuffer secureSalt = SecureMemory.secureByteArray(salt); SecureBuffer secureIv = SecureMemory.secureByteArray(iv); SecureBuffer secureCiphertext = SecureMemory.secureByteArray(ciphertext)) { - GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, secureIv.getData().asByteBuffer().array()); - result = decryptWithCipher(algorithm, key, spec, secureCiphertext.getData().asByteBuffer().array()); + GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, secureIv.getAsByteArray()); + result = decryptWithCipher(algorithm, key, spec, secureCiphertext.getAsByteArray()); } } else { byte[] ciphertext = java.util.Arrays.copyOfRange(decoded, SALT_LENGTH, decoded.length); try (SecureBuffer secureSalt = SecureMemory.secureByteArray(salt); SecureBuffer secureCiphertext = SecureMemory.secureByteArray(ciphertext)) { - result = decryptWithCipher(algorithm, key, secureCiphertext.getData().asByteBuffer().array()); + result = decryptWithCipher(algorithm, key, secureCiphertext.getAsByteArray()); } } @@ -172,7 +172,7 @@ public final class Pass2Key { try (SecureBuffer secureKey = SecureMemory.secureByteArray(decodedKey)) { // Create a PKCS8EncodedKeySpec object - PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(secureKey.getData().asByteBuffer().array()); + PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(secureKey.getAsByteArray()); // Get an instance of the KeyFactory for the desired algorithm (e.g., RSA) KeyFactory keyFactory = KeyFactory.getInstance("RSA"); diff --git a/src/test/java/ai/nomyo/SecureMemoryTest.java b/src/test/java/ai/nomyo/SecureMemoryTest.java index 790cdcc..ac14bbf 100644 --- a/src/test/java/ai/nomyo/SecureMemoryTest.java +++ b/src/test/java/ai/nomyo/SecureMemoryTest.java @@ -257,11 +257,12 @@ class SecureMemoryTest { @Test @DisplayName("SecureBuffer should preserve data contents") void secureBuffer_shouldPreserveData() { - byte[] original = new byte[]{(byte) 0x00, (byte) 0xFF, (byte) 0xAA, (byte) 0x55, 42}; + byte[] expected = new byte[]{(byte) 0x00, (byte) 0xFF, (byte) 0xAA, (byte) 0x55, 42}; + byte[] original = expected.clone(); try (SecureMemory.SecureBuffer buffer = SecureMemory.secureByteArray(original)) { - byte[] retrieved = new byte[original.length]; + byte[] retrieved = new byte[expected.length]; buffer.getData().asByteBuffer().get(retrieved); - assertArrayEquals(original, retrieved, "Data should be preserved in buffer"); + assertArrayEquals(expected, retrieved, "Data should be preserved in buffer"); } }