392 lines
15 KiB
Java
392 lines
15 KiB
Java
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.FileWriter;
|
|
import java.io.IOException;
|
|
import java.math.BigInteger;
|
|
import java.nio.charset.StandardCharsets;
|
|
import java.nio.file.Files;
|
|
import java.nio.file.Path;
|
|
import java.nio.file.StandardOpenOption;
|
|
import java.nio.file.attribute.PosixFilePermission;
|
|
import java.nio.file.attribute.PosixFilePermissions;
|
|
import java.security.*;
|
|
import java.security.spec.RSAKeyGenParameterSpec;
|
|
import java.util.Map;
|
|
import java.util.Set;
|
|
import java.util.concurrent.CompletableFuture;
|
|
import java.util.concurrent.locks.ReentrantLock;
|
|
|
|
/**
|
|
* Low-level client: key management, hybrid encryption, HTTP with retry, response decryption. Used by {@link SecureChatCompletion}.
|
|
*/
|
|
public class SecureCompletionClient {
|
|
|
|
/**
|
|
* NOMYO router base URL (trailing slash stripped).
|
|
*/
|
|
@Getter
|
|
private final String routerUrl;
|
|
|
|
/**
|
|
* Permit HTTP (non-HTTPS) URLs.
|
|
*/
|
|
@Getter
|
|
private final boolean allowHttp;
|
|
|
|
/**
|
|
* RSA key size in bits ({@link Constants#RSA_KEY_SIZE}).
|
|
*/
|
|
@Getter
|
|
private final int keySize;
|
|
|
|
/**
|
|
* Max retries on retryable errors.
|
|
*/
|
|
@Getter
|
|
private final int maxRetries;
|
|
|
|
/**
|
|
* Secure memory operations active.
|
|
*/
|
|
@Getter
|
|
private final boolean useSecureMemory;
|
|
|
|
/**
|
|
* Lock for double-checked key initialization.
|
|
*/
|
|
private final ReentrantLock keyInitLock = new ReentrantLock();
|
|
|
|
/**
|
|
* RSA private key ({@code null} until loaded/generated).
|
|
*/
|
|
@Getter
|
|
private PrivateKey privateKey;
|
|
|
|
/**
|
|
* PEM-encoded public key ({@code null} until loaded/generated).
|
|
*/
|
|
@Getter
|
|
private String publicPemKey;
|
|
|
|
/**
|
|
* Keys initialized.
|
|
*/
|
|
private volatile boolean keysInitialized = false;
|
|
|
|
/**
|
|
* Default settings: {@code https://api.nomyo.ai}, HTTPS-only, secure memory, 2 retries.
|
|
*/
|
|
public SecureCompletionClient() {
|
|
this(Constants.DEFAULT_BASE_URL, false, true, Constants.DEFAULT_MAX_RETRIES);
|
|
}
|
|
|
|
/**
|
|
* @param routerUrl NOMYO router base URL
|
|
* @param allowHttp permit HTTP URLs
|
|
* @param secureMemory enable memory locking/zeroing
|
|
* @param maxRetries 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;
|
|
}
|
|
|
|
private static String readFileContent(String filePath) throws IOException {
|
|
return Files.readString(Path.of(filePath));
|
|
}
|
|
|
|
/**
|
|
* Generates a 4096-bit RSA key pair (exponent 65537). Saves to disk if {@code saveToFile}.
|
|
*/
|
|
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) {
|
|
Path keyFolder = Path.of(keyDir);
|
|
if (!Files.exists(keyFolder)) {
|
|
try {
|
|
Files.createDirectories(keyFolder);
|
|
} catch (IOException e) {
|
|
throw new IOException("Failed to create key directory: " + keyDir, e);
|
|
}
|
|
}
|
|
|
|
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 (var writer = Files.newBufferedWriter(privateKeyPath)) {
|
|
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);
|
|
}
|
|
}
|
|
writer.write(privatePem);
|
|
}
|
|
}
|
|
|
|
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 (var writer = Files.newBufferedWriter(publicKeyPath)) {
|
|
writer.write(publicPem);
|
|
}
|
|
}
|
|
}
|
|
|
|
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 4096-bit RSA key pair and saves to the default directory.
|
|
*/
|
|
public void generateKeys(boolean saveToFile) {
|
|
generateKeys(saveToFile, Constants.DEFAULT_KEY_DIR, null);
|
|
}
|
|
|
|
/**
|
|
* Loads RSA private key from disk. If {@code publicPemKeyPath} is {@code null}, derives public key.
|
|
* Validates key size >= {@link Constants#MIN_RSA_KEY_SIZE} bits.
|
|
*
|
|
* @param privateKeyPath private key PEM path
|
|
* @param publicPemKeyPath optional public key PEM path
|
|
* @param password optional password for encrypted private key
|
|
*/
|
|
public void loadKeys(String privateKeyPath, String publicPemKeyPath, String password) {
|
|
Path keyPath = Path.of(privateKeyPath);
|
|
if (!Files.exists(keyPath)) {
|
|
throw new RuntimeException("Private key file not found: " + privateKeyPath);
|
|
}
|
|
|
|
String keyContent;
|
|
if (password != null && !password.isEmpty()) {
|
|
try {
|
|
keyContent = readFileContent(privateKeyPath);
|
|
} catch (IOException e) {
|
|
throw new RuntimeException("Failed to read private key file: " + e.getMessage(), e);
|
|
}
|
|
|
|
try {
|
|
keyContent = Pass2Key.decrypt("AES/GCM/NoPadding", keyContent, password);
|
|
} catch (NoSuchPaddingException | NoSuchAlgorithmException | BadPaddingException |
|
|
IllegalBlockSizeException | InvalidAlgorithmParameterException | InvalidKeyException e) {
|
|
System.out.println("Wrong password!");
|
|
return;
|
|
}
|
|
} else {
|
|
try {
|
|
keyContent = readFileContent(privateKeyPath);
|
|
} catch (IOException e) {
|
|
throw new RuntimeException("Failed to read private key file: " + e.getMessage(), e);
|
|
}
|
|
}
|
|
|
|
try {
|
|
this.privateKey = Pass2Key.convertStringToPrivateKey(keyContent);
|
|
} catch (Exception e) {
|
|
throw new RuntimeException("Failed to load private key: " + e.getMessage(), e);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Loads RSA private key from disk, deriving public key.
|
|
*/
|
|
public void loadKeys(String privateKeyPath, String password) {
|
|
loadKeys(privateKeyPath, null, password);
|
|
}
|
|
|
|
/**
|
|
* GET {@code {routerUrl}/pki/public_key}. Returns server PEM public key.
|
|
*/
|
|
public CompletableFuture<String> fetchServerPublicKey() {
|
|
throw new UnsupportedOperationException("Not yet implemented");
|
|
}
|
|
|
|
/**
|
|
* Hybrid encryption: AES-256-GCM for payload, RSA-OAEP-SHA256 for AES key wrapping.
|
|
*
|
|
* @param payload OpenAI-compatible chat parameters
|
|
* @return encrypted bytes (JSON package)
|
|
* @throws SecurityError if encryption fails or keys not loaded
|
|
*/
|
|
public CompletableFuture<byte[]> encryptPayload(Map<String, Object> payload) {
|
|
throw new UnsupportedOperationException("Not yet implemented");
|
|
}
|
|
|
|
/**
|
|
* Core hybrid encryption: AES-256-GCM encrypts {@code payloadBytes} with {@code aesKey}.
|
|
*/
|
|
public CompletableFuture<byte[]> doEncrypt(byte[] payloadBytes, byte[] aesKey) {
|
|
throw new UnsupportedOperationException("Not yet implemented");
|
|
}
|
|
|
|
/**
|
|
* Decrypts server response.
|
|
*/
|
|
public CompletableFuture<Map<String, Object>> decryptResponse(byte[] encryptedResponse, String payloadId) {
|
|
throw new UnsupportedOperationException("Not yet implemented");
|
|
}
|
|
|
|
/**
|
|
* encrypt → POST {routerUrl}/v1/chat/secure_completion → retry → decrypt → return.
|
|
* <p>Headers: Content-Type=octet-stream, X-Payload-ID, X-Public-Key, Authorization (Bearer), X-Security-Tier.
|
|
* Retryable: 429, 500, 502, 503, 504 + network errors. Backoff: 2^(attempt-1)s.
|
|
* <p>Status mapping: 200→return, 400→InvalidRequestError, 401→AuthenticationError, 403→ForbiddenError,
|
|
* 404→APIError, 429→RateLimitError, 500→ServerError, 503→ServiceUnavailableError,
|
|
* 502/504→APIError(retryable), network→APIConnectionError.
|
|
*
|
|
* @param payload OpenAI-compatible chat parameters
|
|
* @param payloadId unique payload identifier
|
|
* @param apiKey optional API key
|
|
* @param securityTier optional: "standard", "high", "maximum"
|
|
* @return decrypted response map
|
|
* @throws SecurityError encryption/decryption failure
|
|
* @throws APIConnectionError network error
|
|
* @throws InvalidRequestError HTTP 400
|
|
* @throws AuthenticationError HTTP 401
|
|
* @throws ForbiddenError HTTP 403
|
|
* @throws RateLimitError HTTP 429
|
|
* @throws ServerError HTTP 500
|
|
* @throws ServiceUnavailableError HTTP 503
|
|
* @throws APIError other errors
|
|
*/
|
|
public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId, String apiKey, String securityTier) {
|
|
throw new UnsupportedOperationException("Not yet implemented");
|
|
}
|
|
|
|
/**
|
|
* Without security tier.
|
|
*/
|
|
public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId, String apiKey) {
|
|
return sendSecureRequest(payload, payloadId, apiKey, null);
|
|
}
|
|
|
|
/**
|
|
* No API key or security tier.
|
|
*/
|
|
public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId) {
|
|
return sendSecureRequest(payload, payloadId, null, null);
|
|
}
|
|
|
|
/**
|
|
* Thread-safe key init via double-checked locking. Loads from disk if {@code keyDir} set, else generates.
|
|
*
|
|
* @param keyDir key directory or {@code null} for ephemeral
|
|
*/
|
|
public void ensureKeys(String keyDir) {
|
|
if (keysInitialized) return;
|
|
keyInitLock.lock();
|
|
try {
|
|
if (keysInitialized) return;
|
|
if (keyDir == null || keyDir.isEmpty()) {
|
|
generateKeys(false);
|
|
} else {
|
|
generateKeys(true);
|
|
}
|
|
keysInitialized = true;
|
|
} finally {
|
|
keyInitLock.unlock();
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Validates RSA key size >= {@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 = extractKeySize(key);
|
|
|
|
if (keySize < Constants.MIN_RSA_KEY_SIZE) {
|
|
throw new SecurityError("RSA key size " + keySize + " bits is below minimum " + Constants.MIN_RSA_KEY_SIZE + " bits");
|
|
}
|
|
}
|
|
|
|
private int extractKeySize(PrivateKey key) {
|
|
try {
|
|
var kf = KeyFactory.getInstance("RSA");
|
|
try {
|
|
var crtSpec = kf.getKeySpec(key, java.security.spec.RSAPrivateCrtKeySpec.class);
|
|
return crtSpec.getModulus().bitLength();
|
|
} catch (Exception ignored) {
|
|
var privSpec = kf.getKeySpec(key, java.security.spec.RSAPrivateKeySpec.class);
|
|
return privSpec.getModulus().bitLength();
|
|
}
|
|
} catch (Exception ignored) {
|
|
if (key.getEncoded() != null) {
|
|
return key.getEncoded().length * 8;
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
/**
|
|
* Maps HTTP status code to exception (200→null).
|
|
*/
|
|
public Exception mapHttpStatus(int statusCode, String responseBody) {
|
|
return switch (statusCode) {
|
|
case 200 -> null;
|
|
case 400 ->
|
|
new InvalidRequestError("Invalid request: " + (responseBody != null ? responseBody : "no body"));
|
|
case 401 ->
|
|
new AuthenticationError("Authentication failed: " + (responseBody != null ? responseBody : "no body"));
|
|
case 403 -> new ForbiddenError("Access forbidden: " + (responseBody != null ? responseBody : "no body"));
|
|
case 404 -> new APIError("Not found: " + (responseBody != null ? responseBody : "no body"));
|
|
case 429 -> new RateLimitError("Rate limit exceeded: " + (responseBody != null ? responseBody : "no body"));
|
|
case 500 -> new ServerError("Internal server error: " + (responseBody != null ? responseBody : "no body"));
|
|
case 503 ->
|
|
new ServiceUnavailableError("Service unavailable: " + (responseBody != null ? responseBody : "no body"));
|
|
case 502, 504 -> new APIError("Gateway error: " + (responseBody != null ? responseBody : "no body"));
|
|
default ->
|
|
new APIError("Unexpected status " + statusCode + ": " + (responseBody != null ? responseBody : "no body"));
|
|
};
|
|
}
|
|
|
|
/**
|
|
* URL-encodes PEM key for {@code X-Public-Key} header.
|
|
*/
|
|
public String urlEncodePublicKey(String pemKey) {
|
|
return java.net.URLEncoder.encode(pemKey, StandardCharsets.UTF_8);
|
|
}
|
|
|
|
/**
|
|
* Delegates to resource cleanup (stub).
|
|
*/
|
|
public void close() {
|
|
throw new UnsupportedOperationException("Not yet implemented");
|
|
}
|
|
}
|