Initial commit

src/main/java/ai/nomyo/SecureMemory.java

@@ -0,0 +1,237 @@
+package ai.nomyo;
+
+import lombok.Getter;
+import lombok.Setter;
+
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.util.Map;
+
+/**
+ * Cross-platform memory locking and secure zeroing utilities.
+ *
+ * <p>This module provides optional memory protection for sensitive
+ * cryptographic buffers. It fails gracefully if memory locking is
+ * unavailable on the current platform (e.g., Windows on some JVM
+ * configurations).</p>
+ *
+ * <h3>Protection Levels</h3>
+ * <ul>
+ *   <li><b>"full"</b> — Memory locking and secure zeroing both available</li>
+ *   <li><b>"zeroing_only"</b> — Only secure zeroing available</li>
+ *   <li><b>"none"</b> — No memory protection available</li>
+ * </ul>
+ *
+ * <h3>Usage</h3>
+ * <pre>{@code
+ * try (SecureBuffer buf = SecureMemory.secureBytearray(sensitiveData)) {
+ *     // buf is locked and ready to use
+ *     process(buf.getData());
+ * }
+ * // buf is automatically zeroed and unlocked on exit
+ * }</pre>
+ */
+public final class SecureMemory {
+
+    @Getter
+    @Setter
+    private static volatile boolean secureMemoryEnabled = true;
+
+    @Getter
+    private static final boolean HAS_MEMORY_LOCKING;
+    @Getter
+    private static final boolean HAS_SECURE_ZEROING;
+
+    static {
+        boolean locking = false;
+        boolean zeroing = false;
+        try {
+            locking = initMemoryLocking();
+            zeroing = true; // Secure zeroing is always available at the JVM level
+        } catch (Throwable t) {
+            // Degrade gracefully
+        }
+
+        HAS_MEMORY_LOCKING = locking;
+        HAS_SECURE_ZEROING = zeroing;
+    }
+
+    private static boolean initMemoryLocking() {
+        // FFM doesn't support memory locking at this point in time
+        // TODO: Bypass this with native libraries
+        return false;
+    }
+
+    /**
+     * Wraps a byte array with memory locking and guaranteed zeroing on exit.
+     *
+     * <p>Implements {@link AutoCloseable} for use with try-with-resources.
+     * The buffer is automatically zeroed and unlocked when closed, even
+     * if an exception occurs.</p>
+     */
+    public static class SecureBuffer implements AutoCloseable {
+
+        private final Arena arena;
+
+        @Getter
+        private final MemorySegment data;
+
+        @Getter
+        private final long size;
+
+        @Getter
+        private final long address;
+        private boolean locked;
+        private boolean closed;
+
+        /**
+         * Creates a new SecureBuffer wrapping the given data.
+         *
+         * @param data     the byte array to wrap
+         * @param lock     whether to attempt memory locking
+         */
+        public SecureBuffer(byte[] data, boolean lock) {
+            this.arena = Arena.ofConfined();
+            this.data = data != null ? this.arena.allocate(data.length) : MemorySegment.NULL;
+
+            if (data != null) {
+                this.data.asByteBuffer().put(data);
+            }
+
+            this.size = this.data.byteSize();
+            this.address = this.data.address();
+
+            this.locked = false;
+            this.closed = false;
+
+            if (lock && SecureMemory.secureMemoryEnabled) {
+                this.locked = lock();
+            }
+        }
+
+        /**
+         * Attempts to lock the buffer in memory, preventing swapping to disk.
+         *
+         * @return {@code true} if locking succeeded, {@code false} otherwise
+         */
+        public boolean lock() {
+            //data = data.asReadOnly();
+            return false;
+        }
+
+        /**
+         * Unlocks the buffer, allowing it to be swapped to disk.
+         *
+         * @return {@code true} if unlocking succeeded, {@code false} otherwise
+         */
+        public boolean unlock() {
+            if (!locked) return false;
+            locked = false;
+            return false;
+        }
+
+        /**
+         * Securely zeros the buffer contents.
+         */
+        public void zero() {
+            if (data != null) {
+                data.fill((byte) 0);
+            }
+        }
+
+        @Override
+        public void close() {
+            if (closed) return;
+
+            zero();
+            unlock();
+
+            arena.close();
+            closed = true;
+        }
+    }
+
+    /**
+     * Creates a SecureBuffer for the given data with memory locking.
+     *
+     * <p>This is the recommended way to handle sensitive data. The returned
+     * buffer should be used within a try-with-resources block to ensure
+     * secure zeroing on exit.</p>
+     *
+     * @param data   the sensitive data bytes
+     * @param lock   whether to attempt memory locking
+     * @return a new SecureBuffer
+     */
+    public static SecureBuffer secureByteArray(byte[] data, boolean lock) {
+        return new SecureBuffer(data, lock);
+    }
+
+    /**
+     * Creates a SecureBuffer for the given data with memory locking.
+     * Convenience variant that always attempts locking.
+     *
+     * @param data the sensitive data bytes
+     * @return a new SecureBuffer
+     */
+    public static SecureBuffer secureByteArray(byte[] data) {
+        return secureByteArray(data, true);
+    }
+
+    /**
+     * Creates a SecureBuffer for the given data with memory locking.
+     *
+     * <p><b>Deprecated:</b> Use {@link #secureByteArray(byte[])} instead.
+     * This method exists for compatibility with the Python reference.</p>
+     *
+     * @param data   the sensitive data bytes
+     * @param lock   whether to attempt memory locking
+     * @return a new SecureBuffer
+     * @deprecated Use {@link #secureByteArray(byte[])} instead
+     */
+    @Deprecated
+    public static SecureBuffer secureBytes(byte[] data, boolean lock) {
+        return new SecureBuffer(data, lock);
+    }
+
+    /**
+     * Creates a SecureBuffer for the given data with memory locking.
+     *
+     * <p><b>Deprecated:</b> Use {@link #secureByteArray(byte[])} instead.
+     * This method exists for compatibility with the Python reference.</p>
+     *
+     * @param data the sensitive data bytes
+     * @return a new SecureBuffer
+     * @deprecated Use {@link #secureByteArray(byte[])} instead
+     */
+    @Deprecated
+    public static SecureBuffer secureBytes(byte[] data) {
+        return secureBytes(data, true);
+    }
+
+    /**
+     * Returns information about the current memory protection capabilities.
+     *
+     * @return a map of protection capabilities
+     */
+    public static Map<String, Object> getMemoryProtectionInfo() {
+        String protectionLevel;
+        if (HAS_MEMORY_LOCKING && HAS_SECURE_ZEROING) {
+            protectionLevel = "full";
+        } else if (HAS_SECURE_ZEROING) {
+            protectionLevel = "zeroing_only";
+        } else {
+            protectionLevel = "none";
+        }
+
+        boolean supportsFull = HAS_MEMORY_LOCKING && HAS_SECURE_ZEROING && secureMemoryEnabled;
+
+        return Map.of(
+                "enabled", secureMemoryEnabled,
+                "protection_level", protectionLevel,
+                "has_memory_locking", HAS_MEMORY_LOCKING,
+                "has_secure_zeroing", HAS_SECURE_ZEROING,
+                "supports_full_protection", supportsFull,
+                "page_size", Constants.PAGE_SIZE
+        );
+    }
+}