This is an automated email from the ASF dual-hosted git repository.
pvillard pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new e6266cb123 NIFI-15032 Fixed Key encoding for Snowflake Key Pair
Authentication
e6266cb123 is described below
commit e6266cb123aa4246202bf89b6100b106be9108d2
Author: exceptionfactory <[email protected]>
AuthorDate: Thu Oct 2 15:04:36 2025 -0500
NIFI-15032 Fixed Key encoding for Snowflake Key Pair Authentication
- Added PEM formatting and Base64 encoding for Private Key
Signed-off-by: Pierre Villard <[email protected]>
This closes #10375.
---
.../nifi-snowflake-services/pom.xml | 5 +
.../service/SnowflakeComputingConnectionPool.java | 15 ++-
.../SnowflakeComputingConnectionPoolTest.java | 113 +++++++++++++++++++++
3 files changed, 131 insertions(+), 2 deletions(-)
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
index acff34c3db..a0ef5738c6 100644
---
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
@@ -79,5 +79,10 @@
<artifactId>nifi-kerberos-user-service-api</artifactId>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-mock</artifactId>
+ <scope>test</scope>
+ </dependency>
</dependencies>
</project>
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowflakeComputingConnectionPool.java
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowflakeComputingConnectionPool.java
index e03838b64c..5547c871cd 100644
---
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowflakeComputingConnectionPool.java
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowflakeComputingConnectionPool.java
@@ -49,6 +49,7 @@ import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.snowflake.service.util.ConnectionUrlFormat;
import org.apache.nifi.snowflake.service.util.ConnectionUrlFormatParameters;
+import java.nio.charset.StandardCharsets;
import java.security.PrivateKey;
import java.sql.Connection;
import java.sql.Driver;
@@ -168,6 +169,8 @@ public class SnowflakeComputingConnectionPool extends
AbstractDBCPConnectionPool
private static final String AUTHENTICATOR_SNOWFLAKE_JWT = "SNOWFLAKE_JWT";
+ private static final String PEM_CONTENT_FORMAT = "-----BEGIN PRIVATE
KEY-----%n%s%n-----END PRIVATE KEY-----%n";
+
private volatile OAuth2AccessTokenProvider accessTokenProvider;
private static final List<PropertyDescriptor> PROPERTY_DESCRIPTORS =
List.of(
@@ -310,8 +313,7 @@ public class SnowflakeComputingConnectionPool extends
AbstractDBCPConnectionPool
if (privateKeyServiceProperty.isSet()) {
final PrivateKeyService privateKeyService =
privateKeyServiceProperty.asControllerService(PrivateKeyService.class);
final PrivateKey privateKey = privateKeyService.getPrivateKey();
- final byte[] privateKeyEncoded = privateKey.getEncoded();
- final String privateKeyBase64 =
Base64.getEncoder().encodeToString(privateKeyEncoded);
+ final String privateKeyBase64 = getPrivateKeyBase64(privateKey);
connectionProperties.put(PRIVATE_KEY_BASE64.getPropertyKey(),
privateKeyBase64);
connectionProperties.put(AUTHENTICATOR.getPropertyKey(),
AUTHENTICATOR_SNOWFLAKE_JWT);
}
@@ -362,4 +364,13 @@ public class SnowflakeComputingConnectionPool extends
AbstractDBCPConnectionPool
context.getProperty(SNOWFLAKE_CLOUD_TYPE).evaluateAttributeExpressions().getValue()
);
}
+
+ private String getPrivateKeyBase64(final PrivateKey privateKey) {
+ final byte[] privateKeyEncoded = privateKey.getEncoded();
+ final Base64.Encoder encoder = Base64.getEncoder();
+ final String privateKeyEncodedBase64 =
encoder.encodeToString(privateKeyEncoded);
+ final String pemPrivateKey =
PEM_CONTENT_FORMAT.formatted(privateKeyEncodedBase64);
+ final byte[] pemPrivateKeyBinary =
pemPrivateKey.getBytes(StandardCharsets.UTF_8);
+ return encoder.encodeToString(pemPrivateKeyBinary);
+ }
}
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/SnowflakeComputingConnectionPoolTest.java
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/SnowflakeComputingConnectionPoolTest.java
new file mode 100644
index 0000000000..671dc46fcf
--- /dev/null
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/SnowflakeComputingConnectionPoolTest.java
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.snowflake.service;
+
+import net.snowflake.client.core.SFSessionProperty;
+import org.apache.nifi.components.PropertyDescriptor;
+import org.apache.nifi.key.service.api.PrivateKeyService;
+import org.apache.nifi.util.MockConfigurationContext;
+import org.apache.nifi.util.MockProcessContext;
+import org.apache.nifi.util.NoOpProcessor;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.nio.charset.StandardCharsets;
+import java.security.GeneralSecurityException;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.PrivateKey;
+import java.util.Base64;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+class SnowflakeComputingConnectionPoolTest {
+
+ private static final String PRIVATE_KEY_SERVICE_ID =
PrivateKeyService.class.getSimpleName();
+
+ private static final String RSA_ALGORITHM = "RSA";
+
+ private static final String SNOWFLAKE_JWT = "SNOWFLAKE_JWT";
+
+ private static final String PEM_HEADER = "-----BEGIN PRIVATE KEY-----";
+
+ private static final String PEM_HEADER_BASE64 =
Base64.getEncoder().encodeToString(PEM_HEADER.getBytes(StandardCharsets.UTF_8));
+
+ private static final String PEM_CONTENT_FORMAT = "%s%n%s%n-----END PRIVATE
KEY-----%n";
+
+ private static PrivateKey privateKey;
+
+ private static String pemPrivateKeyBase64;
+
+ @Mock
+ private PrivateKeyService privateKeyService;
+
+ private SnowflakeComputingConnectionPool pool;
+
+ @BeforeAll
+ static void setPrivateKey() throws GeneralSecurityException {
+ final KeyPairGenerator keyPairGenerator =
KeyPairGenerator.getInstance(RSA_ALGORITHM);
+ final KeyPair keyPair = keyPairGenerator.generateKeyPair();
+ privateKey = keyPair.getPrivate();
+ final byte[] privateKeyEncoded = privateKey.getEncoded();
+ final String privateKeyEncodedBase64 =
Base64.getEncoder().encodeToString(privateKeyEncoded);
+ final String pemPrivateKey = PEM_CONTENT_FORMAT.formatted(PEM_HEADER,
privateKeyEncodedBase64);
+ pemPrivateKeyBase64 =
Base64.getEncoder().encodeToString(pemPrivateKey.getBytes(StandardCharsets.UTF_8));
+ }
+
+ @BeforeEach
+ void setPool() {
+ pool = new SnowflakeComputingConnectionPool();
+ }
+
+ @Test
+ void testGetConnectionPropertiesPrivateKeyService() {
+ final Map<PropertyDescriptor, String> properties = new HashMap<>();
+ properties.put(SnowflakeComputingConnectionPool.PRIVATE_KEY_SERVICE,
PRIVATE_KEY_SERVICE_ID);
+
+ final MockProcessContext controllerServiceLookup = new
MockProcessContext(new NoOpProcessor());
+
when(privateKeyService.getIdentifier()).thenReturn(PRIVATE_KEY_SERVICE_ID);
+ controllerServiceLookup.addControllerService(privateKeyService,
Map.of(), null);
+
+ final Map<String, String> environmentVariables = Map.of();
+ final MockConfigurationContext context = new
MockConfigurationContext(properties, controllerServiceLookup,
environmentVariables);
+
+ when(privateKeyService.getPrivateKey()).thenReturn(privateKey);
+ final Map<String, String> connectionProperties =
pool.getConnectionProperties(context);
+ assertNotNull(connectionProperties);
+ assertFalse(connectionProperties.isEmpty());
+
+ final String authenticator =
connectionProperties.get(SFSessionProperty.AUTHENTICATOR.getPropertyKey());
+ assertEquals(SNOWFLAKE_JWT, authenticator);
+
+ final String privateKeyBase64 =
connectionProperties.get(SFSessionProperty.PRIVATE_KEY_BASE64.getPropertyKey());
+ assertNotNull(privateKeyBase64);
+
+ assertTrue(privateKeyBase64.startsWith(PEM_HEADER_BASE64), "PEM Header
encoded with Bas64 not found");
+ assertEquals(pemPrivateKeyBase64, privateKeyBase64);
+ }
+}
