Refactor getPrivateKey method

ballerina-platform · Mar 22, 2024 · e2ca5f9 · e2ca5f9
1 parent 53d18df
commit e2ca5f9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 9 deletions.
diff --git a/ballerina/tests/sign_verify_test.bal b/ballerina/tests/sign_verify_test.bal
@@ -419,6 +419,20 @@ isolated function testDecodeMlDsa65PrivateKeyError() returns Error? {
     }
 }
 
+@test:Config {}
+isolated function testDecodeMlKem768PrivateKeyError() returns Error? {
+    KeyStore keyStore = {
+        path: KEYSTORE_PATH,
+        password: "ballerina"
+    };
+    PrivateKey|Error privateKey = decodeMlKem768PrivateKeyFromKeyStore(keyStore, "ballerina", "ballerina");
+    if privateKey is Error {
+        test:assertEquals(privateKey.message(), "Not a valid ML-KEM-768 key");
+    } else {
+        test:assertFail("Expected error not found");
+    }
+}
+
 @test:Config {}
 isolated function testDecodeEcPublicKeyError() returns Error? {
     KeyStore keyStore = {

diff --git a/native/src/main/java/io/ballerina/stdlib/crypto/nativeimpl/Decode.java b/native/src/main/java/io/ballerina/stdlib/crypto/nativeimpl/Decode.java
@@ -172,7 +172,7 @@ public static Object decodeMlKem768PrivateKeyFromKeyFile(BString keyFilePath, Ob
         if (Security.getProvider(BouncyCastlePQCProvider.PROVIDER_NAME) == null) {
             Security.addProvider(new BouncyCastlePQCProvider());
         }
-        Object decodedPrivateKey = getPrivateKey(keyFilePath, keyPassword, BouncyCastlePQCProvider.PROVIDER_NAME);
+        Object decodedPrivateKey = getPrivateKey(keyFilePath, keyPassword);
         if (decodedPrivateKey instanceof PrivateKey privateKey) {
             return buildMlKem768PrivateKeyRecord(privateKey);
         }
@@ -183,14 +183,10 @@ private static Object getPrivateKey(BString keyFilePath, Object keyPassword) {
         if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
             Security.addProvider(new BouncyCastleProvider());
         }
-        return getPrivateKey(keyFilePath, keyPassword, BouncyCastleProvider.PROVIDER_NAME);
-    }
-
-    private static Object getPrivateKey(BString keyFilePath, Object keyPassword, String provider) {
         File privateKeyFile = new File(keyFilePath.getValue());
         try (PEMParser pemParser = new PEMParser(new FileReader(privateKeyFile, StandardCharsets.UTF_8))) {
             Object obj = pemParser.readObject();
-            JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(provider);
+            JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(BouncyCastleProvider.PROVIDER_NAME);
             PrivateKeyInfo privateKeyInfo;
             if (obj instanceof PEMEncryptedKeyPair) {
                 if (keyPassword == null) {
@@ -221,9 +217,6 @@ private static Object getPrivateKey(BString keyFilePath, Object keyPassword, Str
         } catch (FileNotFoundException e) {
             return CryptoUtils.createError("Key file not found at: " + privateKeyFile.getAbsoluteFile());
         } catch (PKCSException | IOException e) {
-            if (!BouncyCastleProvider.PROVIDER_NAME.equalsIgnoreCase(provider)) {
-                return getPrivateKey(keyFilePath, keyPassword, BouncyCastleProvider.PROVIDER_NAME);
-            }
             return CryptoUtils.createError("Unable to do private key operations: " + e.getMessage());
         }
     }