Initial commit
This commit is contained in:
commit
8acf584d28
SSH key fingerprint:
SHA256:x4/RtnjUyuHkdvmwNDsWSfcfF1V5PNr3OpriZqOvCX8
24 changed files with 2408 additions and 0 deletions
513
src/main/java/ai/nomyo/SecureCompletionClient.java
Normal file
513
src/main/java/ai/nomyo/SecureCompletionClient.java
Normal file
|
|
@ -0,0 +1,513 @@
|
|||
package ai.nomyo;
|
||||
|
||||
import ai.nomyo.errors.*;
|
||||
import ai.nomyo.util.PEMConverter;
|
||||
import ai.nomyo.util.Pass2Key;
|
||||
import lombok.Getter;
|
||||
|
||||
import javax.crypto.*;
|
||||
import java.io.*;
|
||||
import java.math.BigInteger;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.attribute.PosixFilePermission;
|
||||
import java.nio.file.attribute.PosixFilePermissions;
|
||||
import java.security.*;
|
||||
import java.security.spec.PKCS8EncodedKeySpec;
|
||||
import java.security.spec.RSAKeyGenParameterSpec;
|
||||
import java.io.UnsupportedEncodingException;
|
||||
import java.util.Arrays;
|
||||
import java.util.Map;
|
||||
import java.util.Scanner;
|
||||
import java.util.Set;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import java.util.concurrent.locks.ReentrantLock;
|
||||
|
||||
/**
|
||||
* Low-level secure completion client for the NOMYO API.
|
||||
*
|
||||
* <p>This class handles key management, hybrid encryption, HTTP communication
|
||||
* with retry logic, and response decryption. It is the core of the NOMYO
|
||||
* Java client and is used internally by {@link SecureChatCompletion}.</p>
|
||||
*
|
||||
* <h3>Encryption Wire Format</h3>
|
||||
* <p>Encrypted payloads use hybrid encryption (AES-256-GCM + RSA-4096-OAEP-SHA256):</p>
|
||||
* <ul>
|
||||
* <li>A per-request 256-bit AES key encrypts the payload via AES-256-GCM</li>
|
||||
* <li>The AES key is encrypted via RSA-4096-OAEP-SHA256 using the server's public key</li>
|
||||
* <li>The result is a JSON package with base64-encoded ciphertext, nonce, tag, and encrypted AES key</li>
|
||||
* </ul>
|
||||
*
|
||||
* <h3>Key Lifecycle</h3>
|
||||
* <p>Client RSA keys are generated on first use (if not loaded from disk) and
|
||||
* reused across all subsequent calls until the client is discarded. Keys can
|
||||
* be persisted to disk via {@link #generateKeys(boolean, String, String)} or
|
||||
* loaded from disk via {@link #loadKeys(String, String, String)}.</p>
|
||||
*/
|
||||
public class SecureCompletionClient {
|
||||
|
||||
// ── Instance Attributes ─────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Base URL of the NOMYO router (trailing slash stripped).
|
||||
*/
|
||||
@Getter
|
||||
private final String routerUrl;
|
||||
|
||||
/**
|
||||
* Whether HTTP (non-HTTPS) URLs are permitted.
|
||||
*/
|
||||
@Getter
|
||||
private final boolean allowHttp;
|
||||
|
||||
/**
|
||||
* RSA key size in bits. Always {@link Constants#RSA_KEY_SIZE}.
|
||||
*/
|
||||
@Getter
|
||||
private final int keySize;
|
||||
|
||||
/**
|
||||
* Maximum number of retries for retryable errors.
|
||||
*/
|
||||
@Getter
|
||||
private final int maxRetries;
|
||||
|
||||
/**
|
||||
* Whether secure memory operations are active.
|
||||
*/
|
||||
@Getter
|
||||
private final boolean useSecureMemory;
|
||||
|
||||
/**
|
||||
* Lock for double-checked key initialization.
|
||||
*/
|
||||
private final ReentrantLock keyInitLock = new ReentrantLock();
|
||||
|
||||
/**
|
||||
* RSA private key, or {@code null} if not yet loaded/generated.
|
||||
*/
|
||||
@Getter
|
||||
private PrivateKey privateKey;
|
||||
|
||||
// ── Internal State ──────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* PEM-encoded public key string, or {@code null} if not yet loaded/generated.
|
||||
*/
|
||||
@Getter
|
||||
private String publicPemKey;
|
||||
|
||||
/**
|
||||
* Whether keys have been initialized.
|
||||
*/
|
||||
private volatile boolean keysInitialized = false;
|
||||
|
||||
/**
|
||||
* Constructs a {@code SecureCompletionClient} with default settings.
|
||||
*/
|
||||
public SecureCompletionClient() {
|
||||
this(Constants.DEFAULT_BASE_URL, false, true, Constants.DEFAULT_MAX_RETRIES);
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructs a {@code SecureCompletionClient} with the specified settings.
|
||||
*
|
||||
* @param routerUrl the NOMYO router base URL
|
||||
* @param allowHttp whether to permit HTTP (non-HTTPS) URLs
|
||||
* @param secureMemory whether to enable memory locking/zeroing
|
||||
* @param maxRetries number of retries on retryable errors
|
||||
*/
|
||||
public SecureCompletionClient(String routerUrl, boolean allowHttp, boolean secureMemory, int maxRetries) {
|
||||
this.routerUrl = routerUrl != null ? routerUrl.replaceAll("/+$", "") : Constants.DEFAULT_BASE_URL;
|
||||
this.allowHttp = allowHttp;
|
||||
this.useSecureMemory = secureMemory;
|
||||
this.keySize = Constants.RSA_KEY_SIZE;
|
||||
this.maxRetries = maxRetries;
|
||||
}
|
||||
|
||||
// ── Key Management ──────────────────────────────────────────────
|
||||
|
||||
private static String getEncryptedPrivateKeyFromFile(String privateKeyPath) {
|
||||
File myObj = new File(privateKeyPath);
|
||||
|
||||
StringBuilder builder = new StringBuilder();
|
||||
|
||||
try (Scanner myReader = new Scanner(myObj)) {
|
||||
while (myReader.hasNextLine()) {
|
||||
builder.append(myReader.nextLine());
|
||||
}
|
||||
} catch (FileNotFoundException e) {
|
||||
throw new RuntimeException("Tried to load private key from disk but no file found" + e.getMessage());
|
||||
}
|
||||
|
||||
return builder.toString();
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a new 4096-bit RSA key pair.
|
||||
*
|
||||
* <p>The public exponent is fixed at 65537. The generated key pair
|
||||
* is stored in memory. Use {@code saveToDir} to persist to disk.</p>
|
||||
*
|
||||
* @param saveToFile whether to save the keys to disk
|
||||
* @param keyDir directory to save keys (ignored if {@code saveToFile} is {@code false})
|
||||
* @param password optional password to encrypt the private key file
|
||||
*/
|
||||
public void generateKeys(boolean saveToFile, String keyDir, String password) {
|
||||
try {
|
||||
KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
|
||||
generator.initialize(new RSAKeyGenParameterSpec(Constants.RSA_KEY_SIZE, BigInteger.valueOf(Constants.RSA_PUBLIC_EXPONENT)));
|
||||
|
||||
KeyPair pair = generator.generateKeyPair();
|
||||
|
||||
String privatePem = PEMConverter.toPEM(pair.getPrivate().getEncoded(), true);
|
||||
String publicPem = PEMConverter.toPEM(pair.getPublic().getEncoded(), false);
|
||||
|
||||
if (saveToFile) {
|
||||
File keyFolder = new File(keyDir);
|
||||
if (!keyFolder.exists() && !keyFolder.mkdirs()) {
|
||||
throw new IOException("Failed to create key directory: " + keyDir);
|
||||
}
|
||||
|
||||
Path privateKeyPath = Path.of(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE);
|
||||
if (!Files.exists(privateKeyPath)) {
|
||||
Set<PosixFilePermission> filePermissions = PosixFilePermissions.fromString(Constants.PRIVATE_KEY_FILE_MODE);
|
||||
Files.createFile(privateKeyPath, PosixFilePermissions.asFileAttribute(filePermissions));
|
||||
|
||||
try (FileWriter fileWriter = new FileWriter(privateKeyPath.toFile())) {
|
||||
if (password == null || password.isEmpty()) {
|
||||
System.out.println("WARNING: Saving keys in plaintext!");
|
||||
} else {
|
||||
try {
|
||||
privatePem = Pass2Key.encrypt("AES/GCM/NoPadding", privatePem, password);
|
||||
} catch (NoSuchPaddingException | IllegalBlockSizeException | BadPaddingException |
|
||||
InvalidKeyException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
fileWriter.write(privatePem);
|
||||
fileWriter.flush();
|
||||
}
|
||||
}
|
||||
|
||||
Path publicKeyPath = Path.of(keyDir, Constants.DEFAULT_PUBLIC_KEY_FILE);
|
||||
if (!Files.exists(publicKeyPath)) {
|
||||
Set<PosixFilePermission> publicPermissions = PosixFilePermissions.fromString(Constants.PUBLIC_KEY_FILE_MODE);
|
||||
Files.createFile(publicKeyPath, PosixFilePermissions.asFileAttribute(publicPermissions));
|
||||
|
||||
try (FileWriter fileWriter = new FileWriter(publicKeyPath.toFile())) {
|
||||
fileWriter.write(publicPem);
|
||||
fileWriter.flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
this.privateKey = pair.getPrivate();
|
||||
this.publicPemKey = publicPem;
|
||||
|
||||
} catch (NoSuchAlgorithmException e) {
|
||||
throw new RuntimeException("RSA not available: " + e.getMessage(), e);
|
||||
} catch (InvalidAlgorithmParameterException e) {
|
||||
throw new RuntimeException(e);
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException("Failed to save keys: " + e.getMessage(), e);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a new 4096-bit RSA key pair and saves to the default directory.
|
||||
*
|
||||
* @param saveToFile whether to save the keys to disk
|
||||
*/
|
||||
public void generateKeys(boolean saveToFile) {
|
||||
generateKeys(saveToFile, Constants.DEFAULT_KEY_DIR, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads an RSA private key from disk.
|
||||
*
|
||||
* <p>If {@code publicPemKeyPath} is {@code null}, the public key is
|
||||
* derived from the loaded private key. Validates that the key size
|
||||
* is at least {@link Constants#MIN_RSA_KEY_SIZE} bits.</p>
|
||||
*
|
||||
* @param privateKeyPath path to the private key PEM file
|
||||
* @param publicPemKeyPath optional path to the public key PEM file
|
||||
* @param password optional password for the encrypted private key
|
||||
*/
|
||||
public void loadKeys(String privateKeyPath, String publicPemKeyPath, String password) {
|
||||
if (password != null && !password.isEmpty()) {
|
||||
String cipherText = getEncryptedPrivateKeyFromFile(privateKeyPath);
|
||||
|
||||
try {
|
||||
cipherText = Pass2Key.decrypt("AES/GCM/NoPadding", cipherText, password);
|
||||
} catch (NoSuchPaddingException | NoSuchAlgorithmException
|
||||
| BadPaddingException | IllegalBlockSizeException | InvalidAlgorithmParameterException |
|
||||
InvalidKeyException e) {
|
||||
System.out.println("Wrong password!");
|
||||
}
|
||||
|
||||
try {
|
||||
this.privateKey = Pass2Key.convertStringToPrivateKey(cipherText);
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads an RSA private key from disk, deriving the public key.
|
||||
*
|
||||
* @param privateKeyPath path to the private key PEM file
|
||||
* @param password optional password for the encrypted private key
|
||||
*/
|
||||
public void loadKeys(String privateKeyPath, String password) {
|
||||
loadKeys(privateKeyPath, null, password);
|
||||
}
|
||||
|
||||
// ── Server Key Fetching ─────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Fetches the server's RSA public key from the PKI endpoint.
|
||||
*
|
||||
* <p>Performs a GET request to {@code {routerUrl}/pki/public_key}
|
||||
* and returns the server PEM public key as a string. Validates that
|
||||
* the response parses as a valid PEM public key.</p>
|
||||
*
|
||||
* @return the server's PEM-encoded public key string
|
||||
* @throws SecurityError if the URL is not HTTPS and {@code allowHttp} is {@code false},
|
||||
* or if the response does not contain a valid PEM public key
|
||||
*/
|
||||
public CompletableFuture<String> fetchServerPublicKey() {
|
||||
throw new UnsupportedOperationException("Not yet implemented");
|
||||
}
|
||||
|
||||
// ── Encryption ──────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Encrypts a payload dict using hybrid encryption.
|
||||
*
|
||||
* <p>Serializes the payload to JSON, then encrypts it using:
|
||||
* <ol>
|
||||
* <li>A per-request 256-bit AES key (AES-256-GCM)</li>
|
||||
* <li>RSA-OAEP-SHA256 wrapping of the AES key with the server's public key</li>
|
||||
* </ol>
|
||||
* </p>
|
||||
*
|
||||
* @param payload the payload to encrypt (OpenAI-compatible chat parameters)
|
||||
* @return raw encrypted bytes (JSON package serialized to bytes)
|
||||
* @throws SecurityError if encryption fails or keys are not loaded
|
||||
*/
|
||||
public CompletableFuture<byte[]> encryptPayload(Map<String, Object> payload) {
|
||||
throw new UnsupportedOperationException("Not yet implemented");
|
||||
}
|
||||
|
||||
/**
|
||||
* Core hybrid encryption routine.
|
||||
*/
|
||||
public CompletableFuture<byte[]> doEncrypt(byte[] payloadBytes, byte[] aesKey) {
|
||||
throw new UnsupportedOperationException("Not yet implemented");
|
||||
}
|
||||
|
||||
// ── Decryption ──────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Decrypts a server response.
|
||||
*/
|
||||
public CompletableFuture<Map<String, Object>> decryptResponse(byte[] encryptedResponse, String payloadId) {
|
||||
throw new UnsupportedOperationException("Not yet implemented");
|
||||
}
|
||||
|
||||
// ── Secure Request Lifecycle ────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Full request lifecycle: encrypt → HTTP POST → retry → decrypt → return.
|
||||
*
|
||||
* <h3>Request Headers</h3>
|
||||
* <pre>
|
||||
* Content-Type: application/octet-stream
|
||||
* X-Payload-ID: {payloadId}
|
||||
* X-Public-Key: {urlEncodedPublicPemKey}
|
||||
* Authorization: Bearer {apiKey} (if apiKey is provided)
|
||||
* X-Security-Tier: {tier} (if securityTier is provided)
|
||||
* </pre>
|
||||
*
|
||||
* <h3>POST</h3>
|
||||
* {@code {routerUrl}/v1/chat/secure_completion} with encrypted payload as body.
|
||||
*
|
||||
* <h3>Retry Logic</h3>
|
||||
* <ul>
|
||||
* <li>Retryable status codes: {@code {429, 500, 502, 503, 504}}</li>
|
||||
* <li>Backoff: {@code 2^(attempt-1)} seconds (1s, 2s, 4s…)</li>
|
||||
* <li>Total attempts: {@code maxRetries + 1}</li>
|
||||
* <li>Network errors also retry</li>
|
||||
* <li>Non-retryable exceptions propagate immediately</li>
|
||||
* </ul>
|
||||
*
|
||||
* <h3>Status → Exception Mapping</h3>
|
||||
* <table>
|
||||
* <tr><th>Status</th><th>Result</th></tr>
|
||||
* <tr><td>200</td><td>Return decrypted response map</td></tr>
|
||||
* <tr><td>400</td><td>{@code InvalidRequestError}</td></tr>
|
||||
* <tr><td>401</td><td>{@code AuthenticationError}</td></tr>
|
||||
* <tr><td>403</td><td>{@code ForbiddenError}</td></tr>
|
||||
* <tr><td>404</td><td>{@code APIError}</td></tr>
|
||||
* <tr><td>429</td><td>{@code RateLimitError}</td></tr>
|
||||
* <tr><td>500</td><td>{@code ServerError}</td></tr>
|
||||
* <tr><td>503</td><td>{@code ServiceUnavailableError}</td></tr>
|
||||
* <tr><td>502/504</td><td>{@code APIError} (retryable)</td></tr>
|
||||
* <tr><td>other</td><td>{@code APIError} (non-retryable)</td></tr>
|
||||
* <tr><td>network error</td><td>{@code APIConnectionError}</td></tr>
|
||||
* </table>
|
||||
*
|
||||
* @param payload the payload to send (OpenAI-compatible chat parameters)
|
||||
* @param payloadId unique payload identifier
|
||||
* @param apiKey optional API key for authentication
|
||||
* @param securityTier optional security tier ({@code "standard"}, {@code "high"}, or {@code "maximum"})
|
||||
* @return the decrypted response map
|
||||
* @throws SecurityError if encryption/decryption fails
|
||||
* @throws APIConnectionError if a network error occurs
|
||||
* @throws InvalidRequestError if the API returns 400
|
||||
* @throws AuthenticationError if the API returns 401
|
||||
* @throws ForbiddenError if the API returns 403
|
||||
* @throws RateLimitError if the API returns 429
|
||||
* @throws ServerError if the API returns 500
|
||||
* @throws ServiceUnavailableError if the API returns 503
|
||||
* @throws APIError for other non-retryable errors
|
||||
*/
|
||||
public CompletableFuture<Map<String, Object>> sendSecureRequest(
|
||||
Map<String, Object> payload,
|
||||
String payloadId,
|
||||
String apiKey,
|
||||
String securityTier
|
||||
) {
|
||||
throw new UnsupportedOperationException("Not yet implemented");
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a secure request without a security tier.
|
||||
*
|
||||
* @param payload the payload to send
|
||||
* @param payloadId unique payload identifier
|
||||
* @param apiKey optional API key for authentication
|
||||
* @return the decrypted response map
|
||||
*/
|
||||
public CompletableFuture<Map<String, Object>> sendSecureRequest(
|
||||
Map<String, Object> payload,
|
||||
String payloadId,
|
||||
String apiKey
|
||||
) {
|
||||
return sendSecureRequest(payload, payloadId, apiKey, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a secure request with no API key or security tier.
|
||||
*
|
||||
* @param payload the payload to send
|
||||
* @param payloadId unique payload identifier
|
||||
* @return the decrypted response map
|
||||
*/
|
||||
public CompletableFuture<Map<String, Object>> sendSecureRequest(
|
||||
Map<String, Object> payload,
|
||||
String payloadId
|
||||
) {
|
||||
return sendSecureRequest(payload, payloadId, null, null);
|
||||
}
|
||||
|
||||
// ── Key Initialization ──────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Ensures RSA keys are loaded or generated.
|
||||
*
|
||||
* <p>Uses double-checked locking via {@link ReentrantLock} to ensure
|
||||
* thread-safe initialization. If {@code keyDir} is set, attempts to
|
||||
* load keys from disk first; if that fails, generates new keys.</p>
|
||||
*
|
||||
* @param keyDir directory to persist keys, or {@code null} for ephemeral
|
||||
*/
|
||||
public void ensureKeys(String keyDir) {
|
||||
if (keysInitialized) return;
|
||||
keyInitLock.lock();
|
||||
try {
|
||||
if (keysInitialized) return;
|
||||
// TODO: implement key loading/generation
|
||||
keysInitialized = true;
|
||||
} finally {
|
||||
keyInitLock.unlock();
|
||||
}
|
||||
}
|
||||
|
||||
// ── Key Validation ──────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Validates that an RSA key meets the minimum size requirement.
|
||||
*
|
||||
* @param key the RSA key to validate
|
||||
* @throws SecurityError if the key size is less than {@link Constants#MIN_RSA_KEY_SIZE} bits
|
||||
*/
|
||||
public void validateRsaKey(PrivateKey key) throws SecurityError {
|
||||
if (key == null) {
|
||||
throw new SecurityError("RSA key is null");
|
||||
}
|
||||
int keySize = key.getEncoded() != null ? key.getEncoded().length * 8 : 0;
|
||||
|
||||
System.out.println("Keysize: " + keySize);
|
||||
|
||||
if (keySize < Constants.MIN_RSA_KEY_SIZE) {
|
||||
throw new SecurityError(
|
||||
"RSA key size " + keySize + " bits is below minimum " + Constants.MIN_RSA_KEY_SIZE + " bits"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ── HTTP Status → Exception Mapping ─────────────────────────────
|
||||
|
||||
/**
|
||||
* Maps an HTTP status code to the appropriate exception.
|
||||
*/
|
||||
public Exception mapHttpStatus(int statusCode, String responseBody) {
|
||||
switch (statusCode) {
|
||||
case 200:
|
||||
return null;
|
||||
case 400:
|
||||
return new InvalidRequestError("Invalid request: " + (responseBody != null ? responseBody : "no body"));
|
||||
case 401:
|
||||
return new AuthenticationError("Authentication failed: " + (responseBody != null ? responseBody : "no body"));
|
||||
case 403:
|
||||
return new ForbiddenError("Access forbidden: " + (responseBody != null ? responseBody : "no body"));
|
||||
case 404:
|
||||
return new APIError("Not found: " + (responseBody != null ? responseBody : "no body"));
|
||||
case 429:
|
||||
return new RateLimitError("Rate limit exceeded: " + (responseBody != null ? responseBody : "no body"));
|
||||
case 500:
|
||||
return new ServerError("Internal server error: " + (responseBody != null ? responseBody : "no body"));
|
||||
case 503:
|
||||
return new ServiceUnavailableError("Service unavailable: " + (responseBody != null ? responseBody : "no body"));
|
||||
case 502:
|
||||
case 504:
|
||||
return new APIError("Gateway error: " + (responseBody != null ? responseBody : "no body"));
|
||||
default:
|
||||
return new APIError("Unexpected status " + statusCode + ": " + (responseBody != null ? responseBody : "no body"));
|
||||
}
|
||||
}
|
||||
|
||||
// ── URL Encoding ────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* URL-encodes a public key PEM string for use in the {@code X-Public-Key} header.
|
||||
*/
|
||||
public String urlEncodePublicKey(String pemKey) {
|
||||
return java.net.URLEncoder.encode(pemKey, StandardCharsets.UTF_8);
|
||||
}
|
||||
|
||||
// ── Getters ─────────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Closes the client and releases any resources.
|
||||
*/
|
||||
public void close() {
|
||||
throw new UnsupportedOperationException("Not yet implemented");
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue