Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-618478: support passing private key as base64 encoded string (private_key_base64) as alternative to (private_key_file) #1810

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 21 additions & 5 deletions src/main/java/net/snowflake/client/core/SFLoginInput.java
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ public class SFLoginInput {
private OCSPMode ocspMode;
private HttpClientSettingsKey httpClientKey;
private String privateKeyFile;
private String privateKeyFilePwd;
private String privateKeyBase64;
private String privateKeyPwd;
private String inFlightCtx; // Opaque string sent for Snowsight account activation

private boolean disableConsoleLogin = true;
Expand Down Expand Up @@ -325,22 +326,37 @@ SFLoginInput setPrivateKey(PrivateKey privateKey) {
return this;
}

SFLoginInput setPrivateKeyBase64(String privateKeyBase64) {
this.privateKeyBase64 = privateKeyBase64;
return this;
}

SFLoginInput setPrivateKeyFile(String privateKeyFile) {
this.privateKeyFile = privateKeyFile;
return this;
}

SFLoginInput setPrivateKeyFilePwd(String privateKeyFilePwd) {
this.privateKeyFilePwd = privateKeyFilePwd;
SFLoginInput setPrivateKeyPwd(String privateKeyPwd) {
this.privateKeyPwd = privateKeyPwd;
return this;
}

String getPrivateKeyFile() {
return privateKeyFile;
}

String getPrivateKeyFilePwd() {
return privateKeyFilePwd;
String getPrivateKeyBase64() {
return privateKeyBase64;
}

String getPrivateKeyPwd() {
return privateKeyPwd;
}

boolean isPrivateKeyProvided() {
return (getPrivateKey() != null
|| getPrivateKeyFile() != null
|| getPrivateKeyBase64() != null);
}

public String getApplication() {
Expand Down
20 changes: 16 additions & 4 deletions src/main/java/net/snowflake/client/core/SFSession.java
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ public class SFSession extends SFBaseSession {
private String idToken;
private String mfaToken;
private String privateKeyFileLocation;
private String privateKeyBase64;
private String privateKeyPassword;
private PrivateKey privateKey;

Expand Down Expand Up @@ -452,7 +453,14 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
}
break;

case PRIVATE_KEY_BASE64:
if (propertyValue != null) {
privateKeyBase64 = (String) propertyValue;
}
break;

case PRIVATE_KEY_FILE_PWD:
case PRIVATE_KEY_PWD:
if (propertyValue != null) {
privateKeyPassword = (String) propertyValue;
}
Expand Down Expand Up @@ -583,7 +591,7 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
connectionPropertiesMap.get(SFSessionProperty.TRACING),
connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE),
SFLoggerUtil.isVariableProvided(
(String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD)),
(String) connectionPropertiesMap.getOrDefault(SFSessionProperty.PRIVATE_KEY_PWD, connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD) )),
connectionPropertiesMap.get(SFSessionProperty.ENABLE_DIAGNOSTICS),
connectionPropertiesMap.get(SFSessionProperty.DIAGNOSTICS_ALLOWLIST_FILE),
sessionParametersMap.get(CLIENT_STORE_TEMPORARY_CREDENTIAL),
Expand Down Expand Up @@ -631,8 +639,9 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
.setSessionParameters(sessionParametersMap)
.setPrivateKey((PrivateKey) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY))
.setPrivateKeyFile((String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE))
.setPrivateKeyFilePwd(
(String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD))
.setPrivateKeyBase64((String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_BASE64))
.setPrivateKeyPwd(
(String) connectionPropertiesMap.getOrDefault(SFSessionProperty.PRIVATE_KEY_PWD, connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD) ))
.setApplication((String) connectionPropertiesMap.get(SFSessionProperty.APPLICATION))
.setServiceName(getServiceName())
.setOCSPMode(getOCSPMode())
Expand Down Expand Up @@ -750,7 +759,10 @@ private boolean isSnowflakeAuthenticator() {
Map<SFSessionProperty, Object> connectionPropertiesMap = getConnectionPropertiesMap();
String authenticator = (String) connectionPropertiesMap.get(SFSessionProperty.AUTHENTICATOR);
PrivateKey privateKey = (PrivateKey) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY);
return (authenticator == null && privateKey == null && privateKeyFileLocation == null)
return (authenticator == null
&& privateKey == null
&& privateKeyFileLocation == null
&& privateKeyBase64 == null)
|| ClientAuthnDTO.AuthenticatorType.SNOWFLAKE.name().equalsIgnoreCase(authenticator);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,14 @@ public enum SFSessionProperty {
VALIDATE_DEFAULT_PARAMETERS("validateDefaultParameters", false, Boolean.class),
INJECT_WAIT_IN_PUT("inject_wait_in_put", false, Integer.class),
PRIVATE_KEY_FILE("private_key_file", false, String.class),
/**
* @deprecated Use {@link #PRIVATE_KEY_PWD} for clarity. The given password will be used to decrypt
* the private key value independent of whether that value is supplied as a file or base64 string
*/
@Deprecated()
PRIVATE_KEY_FILE_PWD("private_key_file_pwd", false, String.class),
PRIVATE_KEY_BASE64("private_key_base64", false, String.class),
PRIVATE_KEY_PWD("private_key_pwd", false, String.class),
CLIENT_INFO("snowflakeClientInfo", false, String.class),
ALLOW_UNDERSCORES_IN_HOST("allowUnderscoresInHost", false, Boolean.class),

Expand Down
37 changes: 33 additions & 4 deletions src/main/java/net/snowflake/client/core/SessionUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,7 @@ private static ClientAuthnDTO.AuthenticatorType getAuthenticator(SFLoginInput lo
// authenticator is null, then jdbc will decide authenticator depends on
// if privateKey is specified or not. If yes, authenticator type will be
// SNOWFLAKE_JWT, otherwise it will use SNOWFLAKE.
return (loginInput.getPrivateKey() != null || loginInput.getPrivateKeyFile() != null)
return loginInput.isPrivateKeyProvided()
? ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT
: ClientAuthnDTO.AuthenticatorType.SNOWFLAKE;
}
Expand Down Expand Up @@ -421,7 +421,8 @@ private static SFLoginOutput newSession(
new SessionUtilKeyPair(
loginInput.getPrivateKey(),
loginInput.getPrivateKeyFile(),
loginInput.getPrivateKeyFilePwd(),
loginInput.getPrivateKeyBase64(),
loginInput.getPrivateKeyPwd(),
loginInput.getAccountName(),
loginInput.getUserName());

Expand Down Expand Up @@ -676,7 +677,8 @@ private static SFLoginOutput newSession(
new SessionUtilKeyPair(
loginInput.getPrivateKey(),
loginInput.getPrivateKeyFile(),
loginInput.getPrivateKeyFilePwd(),
loginInput.getPrivateKeyBase64(),
loginInput.getPrivateKeyPwd(),
loginInput.getAccountName(),
loginInput.getUserName());

Expand Down Expand Up @@ -1723,6 +1725,7 @@ public static void resetOCSPUrlIfNecessary(String serverUrl) throws IOException
*
* @param privateKey private key
* @param privateKeyFile path to private key file
* @param privateKeyBase64 base64 encoded content of the private key file
* @param privateKeyFilePwd password for private key file
* @param accountName account name
* @param userName user name
Expand All @@ -1732,13 +1735,39 @@ public static void resetOCSPUrlIfNecessary(String serverUrl) throws IOException
public static String generateJWTToken(
ets marked this conversation as resolved.
Show resolved Hide resolved
PrivateKey privateKey,
String privateKeyFile,
String privateKeyBase64,
String privateKeyFilePwd,
String accountName,
String userName)
throws SFException {
SessionUtilKeyPair s =
new SessionUtilKeyPair(
privateKey, privateKeyFile, privateKeyFilePwd, accountName, userName);
privateKey, privateKeyFile, privateKeyBase64, privateKeyFilePwd, accountName, userName);
return s.issueJwtToken();
}

/**
* Helper function to generate a JWT token
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please add link in javadoc to function above

*
* @param privateKey private key
* @param privateKeyFile path to private key file
* @param privateKeyFilePwd password for private key file
* @param accountName account name
* @param userName user name
* @return JWT token
* @throws SFException if Snowflake error occurs
*/
@Deprecated()
public static String generateJWTToken(
PrivateKey privateKey,
String privateKeyFile,
String privateKeyFilePwd,
String accountName,
String userName)
throws SFException {
SessionUtilKeyPair s =
new SessionUtilKeyPair(
privateKey, privateKeyFile, null, privateKeyFilePwd, accountName, userName);
return s.issueJwtToken();
}

Expand Down
78 changes: 54 additions & 24 deletions src/main/java/net/snowflake/client/core/SessionUtilKeyPair.java
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import java.io.FileReader;
import java.io.IOException;
import java.io.StringReader;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
Expand Down Expand Up @@ -80,6 +80,7 @@ class SessionUtilKeyPair {
SessionUtilKeyPair(
PrivateKey privateKey,
String privateKeyFile,
String privateKeyBase64,
String privateKeyFilePwd,
String accountName,
String userName)
Expand All @@ -100,17 +101,30 @@ class SessionUtilKeyPair {
}
}

// if there is both a file and a private key, there is a problem
// Ensure that we only received one of: privateKey, privateKeyFile, or privateKeyBase64
if (!Strings.isNullOrEmpty(privateKeyFile) && privateKey != null) {
ets marked this conversation as resolved.
Show resolved Hide resolved
throw new SFException(
ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
"Cannot have both private key value and private key file.");
"Cannot have both private key object and private key file.");
} else if (!Strings.isNullOrEmpty(privateKeyBase64) && privateKey != null) {
throw new SFException(
ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
"Cannot have both private key object and private key string value.");
} else if (!Strings.isNullOrEmpty(privateKeyBase64) && !Strings.isNullOrEmpty(privateKeyFile)) {
throw new SFException(
ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
"Cannot have both private key file and private key string value.");
} else {
// if privateKeyFile has a value and privateKey is null
this.privateKey =
Strings.isNullOrEmpty(privateKeyFile)
? privateKey
: extractPrivateKeyFromFile(privateKeyFile, privateKeyFilePwd);
if (!Strings.isNullOrEmpty(privateKeyBase64)) {
// privateKeyBase64 has a value and other options for passing private key are null
this.privateKey = extractPrivateKeyFromBase64(privateKeyBase64, privateKeyFilePwd);
} else {
// either extract from privateKeyFile or use the passed object
this.privateKey =
Strings.isNullOrEmpty(privateKeyFile)
? privateKey
: extractPrivateKeyFromFile(privateKeyFile, privateKeyFilePwd);
}
}
// construct public key from raw bytes
if (this.privateKey instanceof RSAPrivateCrtKey) {
Expand Down Expand Up @@ -148,33 +162,52 @@ private SecretKeyFactory getSecretKeyFactory(String algorithm) throws NoSuchAlgo

private PrivateKey extractPrivateKeyFromFile(String privateKeyFile, String privateKeyFilePwd)
throws SFException {

try {
Path privKeyPath = Paths.get(privateKeyFile);
FileUtil.logFileUsage(privKeyPath, "Extract private key from file", true);
byte[] bytes = Files.readAllBytes(privKeyPath);
return extractPrivateKeyFromBytes(bytes, privateKeyFilePwd);
} catch (IOException ie) {
logger.error("Could not read private key from file", ie);
throw new SFException(ie, ErrorCode.INVALID_PARAMETER_VALUE, ie.getCause());
}
}

private PrivateKey extractPrivateKeyFromBytes(byte[] privateKeyBytes, String privateKeyBytesPwd)
throws SFException {
if (isBouncyCastleProviderEnabled) {
try {
return extractPrivateKeyWithBouncyCastle(privateKeyFile, privateKeyFilePwd);
return extractPrivateKeyWithBouncyCastle(privateKeyBytes, privateKeyBytesPwd);
} catch (IOException | PKCSException | OperatorCreationException e) {
logger.error("Could not extract private key using Bouncy Castle provider", e);
throw new SFException(e, ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY, e.getCause());
}
} else {
try {
return extractPrivateKeyWithJdk(privateKeyFile, privateKeyFilePwd);
return extractPrivateKeyWithJdk(privateKeyBytes, privateKeyBytesPwd);
} catch (NoSuchAlgorithmException
| InvalidKeySpecException
| IOException
| IllegalArgumentException
| NullPointerException
| InvalidKeyException e) {
logger.error(
"Could not extract private key. Try setting the JVM argument: " + "-D{}" + "=TRUE",
"Could not extract private key using standard JDK. Try setting the JVM argument: "
+ "-D{}"
+ "=TRUE",
SecurityUtil.ENABLE_BOUNCYCASTLE_PROVIDER_JVM);
throw new SFException(
e,
ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
privateKeyFile + ": " + e.getMessage());
throw new SFException(e, ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY, e.getMessage());
}
}
}

private PrivateKey extractPrivateKeyFromBase64(String privateKeyBase64, String privateKeyBytesPwd)
throws SFException {
byte[] decodedKey = Base64.decodeBase64(privateKeyBase64);
return extractPrivateKeyFromBytes(decodedKey, privateKeyBytesPwd);
}

public String issueJwtToken() throws SFException {
JWTClaimsSet.Builder builder = new JWTClaimsSet.Builder();
String sub = String.format(SUBJECT_FMT, this.accountName, this.userName);
Expand Down Expand Up @@ -232,13 +265,12 @@ public static int getTimeout() {
}

private PrivateKey extractPrivateKeyWithBouncyCastle(
String privateKeyFile, String privateKeyFilePwd)
byte[] privateKeyBytes, String privateKeyFilePwd)
throws IOException, PKCSException, OperatorCreationException {
Path privKeyPath = Paths.get(privateKeyFile);
FileUtil.logFileUsage(
privKeyPath, "Extract private key from file using Bouncy Castle provider", true);

PrivateKeyInfo privateKeyInfo = null;
PEMParser pemParser = new PEMParser(new FileReader(privKeyPath.toFile()));
PEMParser pemParser =
new PEMParser(new StringReader(new String(privateKeyBytes, StandardCharsets.UTF_8)));
Object pemObject = pemParser.readObject();
if (pemObject instanceof PKCS8EncryptedPrivateKeyInfo) {
// Handle the case where the private key is encrypted.
Expand All @@ -264,11 +296,9 @@ private PrivateKey extractPrivateKeyWithBouncyCastle(
return converter.getPrivateKey(privateKeyInfo);
}

private PrivateKey extractPrivateKeyWithJdk(String privateKeyFile, String privateKeyFilePwd)
private PrivateKey extractPrivateKeyWithJdk(byte[] privateKeyFileBytes, String privateKeyFilePwd)
throws IOException, NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
Path privKeyPath = Paths.get(privateKeyFile);
FileUtil.logFileUsage(privKeyPath, "Extract private key from file using Jdk", true);
String privateKeyContent = new String(Files.readAllBytes(privKeyPath));
String privateKeyContent = new String(privateKeyFileBytes, StandardCharsets.UTF_8);
if (Strings.isNullOrEmpty(privateKeyFilePwd)) {
// unencrypted private key file
return generatePrivateKey(false, privateKeyContent, privateKeyFilePwd);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -232,7 +232,15 @@ public void setPrivateKeyFile(String location, String password) {
this.setAuthenticator(AUTHENTICATOR_SNOWFLAKE_JWT);
this.properties.put(SFSessionProperty.PRIVATE_KEY_FILE.getPropertyKey(), location);
if (!Strings.isNullOrEmpty(password)) {
this.properties.put(SFSessionProperty.PRIVATE_KEY_FILE_PWD.getPropertyKey(), password);
this.properties.put(SFSessionProperty.PRIVATE_KEY_PWD.getPropertyKey(), password);
}
}

public void setPrivateKeyBase64(String privateKeyBase64, String password) {
this.setAuthenticator(AUTHENTICATOR_SNOWFLAKE_JWT);
this.properties.put(SFSessionProperty.PRIVATE_KEY_BASE64.getPropertyKey(), privateKeyBase64);
if (!Strings.isNullOrEmpty(password)) {
this.properties.put(SFSessionProperty.PRIVATE_KEY_PWD.getPropertyKey(), password);
}
}

Expand Down
Loading
Loading