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.

Sign up for GitHub

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-1895884: Token cache refactor #2044

Open
wants to merge 7 commits into
base: oauth-code-flow
Choose a base branch
from
+178 −152
Open

SNOW-1895884: Token cache refactor #2044

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c9fdff8
SNOW-1895884: Add new token cache keys
sfc-gh-dheyman Jan 24, 2025
1d0b235
Change cache json file structure
sfc-gh-dheyman Jan 29, 2025
197242f
Merge branch 'oauth-code-flow' into SNOW-1895884-token-cache-keys
sfc-gh-dheyman Jan 29, 2025
0e1f28f
Merge branch 'oauth-code-flow' of github.com:snowflakedb/snowflake-jd…
sfc-gh-dheyman Jan 29, 2025
4bee3f1
Merge branch 'SNOW-1895884-token-cache-keys' of github.com:snowflaked…
sfc-gh-dheyman Jan 29, 2025
8b7aa40
refactor convertTarget name
sfc-gh-dheyman Jan 30, 2025
6575f5d
Refactor cache locks
sfc-gh-dheyman Jan 30, 2025
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
2 changes: 1 addition & 1 deletion src/main/java/net/snowflake/client/core/SFTrustManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -544,7 +544,7 @@ private static void verifySignature(
* @param bytes a byte array
* @return a string in hexadecimal code
*/
private static String byteToHexString(byte[] bytes) {
static String byteToHexString(byte[] bytes) {
final char[] hexArray = "0123456789ABCDEF".toCharArray();
char[] hexChars = new char[bytes.length * 2];
for (int j = 0; j < bytes.length; j++) {
Expand Down
66 changes: 33 additions & 33 deletions src/main/java/net/snowflake/client/core/SecureStorageLinuxManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ public class SecureStorageLinuxManager implements SecureStorageManager {
private static final String CACHE_FILE_NAME = "temporary_credential.json";
private static final String CACHE_DIR_PROP = "net.snowflake.jdbc.temporaryCredentialCacheDir";
private static final String CACHE_DIR_ENV = "SF_TEMPORARY_CREDENTIAL_CACHE_DIR";
private static final String CACHE_FILE_TOKENS_OBJECT_NAME = "tokens";
private static final long CACHE_EXPIRATION_IN_SECONDS = 86400L;
private static final long CACHE_FILE_LOCK_EXPIRATION_IN_SECONDS = 60L;
private FileCacheManager fileCacheManager;
Expand All @@ -43,6 +44,7 @@ private SecureStorageLinuxManager() {
.build();
logger.debug(
"Using temporary file: {} as a token cache storage", fileCacheManager.getCacheFilePath());
populateLocalCredCache();
}

private static class SecureStorageLinuxManagerHolder {
Expand All @@ -53,62 +55,60 @@ public static SecureStorageLinuxManager getInstance() {
return SecureStorageLinuxManagerHolder.INSTANCE;
}

private ObjectNode localCacheToJson() {
ObjectNode res = mapper.createObjectNode();
for (Map.Entry<String, Map<String, String>> elem : localCredCache.entrySet()) {
String elemHost = elem.getKey();
Map<String, String> hostMap = elem.getValue();
ObjectNode hostNode = mapper.createObjectNode();
for (Map.Entry<String, String> elem0 : hostMap.entrySet()) {
hostNode.put(elem0.getKey(), elem0.getValue());
}
res.set(elemHost, hostNode);
}
return res;
}

public synchronized SecureStorageStatus setCredential(
String host, String user, String type, String token) {
if (Strings.isNullOrEmpty(token)) {
logger.warn("No token provided", false);
return SecureStorageStatus.SUCCESS;
}

localCredCache.computeIfAbsent(host.toUpperCase(), newMap -> new HashMap<>());

Map<String, String> hostMap = localCredCache.get(host.toUpperCase());
hostMap.put(SecureStorageManager.convertTarget(host, user, type), token);

localCredCache.computeIfAbsent(CACHE_FILE_TOKENS_OBJECT_NAME, tokensMap -> new HashMap<>());
Copy link
Contributor

Choose a reason for hiding this comment

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

We should read, update and write cache here under one file lock

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Updated logic

Map<String, String> tokensMap = localCredCache.get(CACHE_FILE_TOKENS_OBJECT_NAME);
tokensMap.put(SecureStorageManager.convertTarget(host, user, type), token);
fileCacheManager.writeCacheFile(localCacheToJson());
return SecureStorageStatus.SUCCESS;
}

public synchronized String getCredential(String host, String user, String type) {
JsonNode res = fileCacheManager.readCacheFile();
readJsonStoreCache(res);

Map<String, String> hostMap = localCredCache.get(host.toUpperCase());

if (hostMap == null) {
populateLocalCredCache();
Map<String, String> tokensMap = localCredCache.get(CACHE_FILE_TOKENS_OBJECT_NAME);
if (tokensMap == null) {
return null;
}

return hostMap.get(SecureStorageManager.convertTarget(host, user, type));
return tokensMap.get(SecureStorageManager.convertTarget(host, user, type));
}

/** May delete credentials which doesn't belong to this process */
public synchronized SecureStorageStatus deleteCredential(String host, String user, String type) {
Copy link
Contributor

Choose a reason for hiding this comment

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

We should read, delete entry and write cache here under one file lock

Copy link
Contributor Author

Choose a reason for hiding this comment

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

refactored

Map<String, String> hostMap = localCredCache.get(host.toUpperCase());
if (hostMap != null) {
hostMap.remove(SecureStorageManager.convertTarget(host, user, type));
if (hostMap.isEmpty()) {
localCredCache.remove(host.toUpperCase());
Map<String, String> tokensMap = localCredCache.get(CACHE_FILE_TOKENS_OBJECT_NAME);
if (tokensMap != null) {
tokensMap.remove(SecureStorageManager.convertTarget(host, user, type));
if (tokensMap.isEmpty()) {
localCredCache.remove(CACHE_FILE_TOKENS_OBJECT_NAME);
}
}
fileCacheManager.writeCacheFile(localCacheToJson());
return SecureStorageStatus.SUCCESS;
}

private ObjectNode localCacheToJson() {
ObjectNode res = mapper.createObjectNode();
for (Map.Entry<String, Map<String, String>> elem : localCredCache.entrySet()) {
String elemHost = elem.getKey();
Map<String, String> hostMap = elem.getValue();
ObjectNode hostNode = mapper.createObjectNode();
for (Map.Entry<String, String> elem0 : hostMap.entrySet()) {
hostNode.put(elem0.getKey(), elem0.getValue());
}
res.set(elemHost, hostNode);
}
return res;
}

private void populateLocalCredCache() {
JsonNode res = fileCacheManager.readCacheFile();
readJsonStoreCache(res);
}

private void readJsonStoreCache(JsonNode m) {
if (m == null || !m.getNodeType().equals(JsonNodeType.OBJECT)) {
logger.debug("Invalid cache file format.");
Expand Down
21 changes: 11 additions & 10 deletions src/main/java/net/snowflake/client/core/SecureStorageManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,14 @@

package net.snowflake.client.core;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

/**
* Interface for accessing Platform specific Local Secure Storage E.g. keychain on Mac credential
* manager on Windows
*/
interface SecureStorageManager {
String DRIVER_NAME = "SNOWFLAKE-JDBC-DRIVER";
int COLON_CHAR_LENGTH = 1;

SecureStorageStatus setCredential(String host, String user, String type, String token);
Expand All @@ -20,22 +22,21 @@ interface SecureStorageManager {

static String convertTarget(String host, String user, String type) {
StringBuilder target =
new StringBuilder(
host.length()
+ user.length()
+ DRIVER_NAME.length()
+ type.length()
+ 3 * COLON_CHAR_LENGTH);
new StringBuilder(host.length() + user.length() + type.length() + 3 * COLON_CHAR_LENGTH);

target.append(host.toUpperCase());
target.append(":");
target.append(user.toUpperCase());
target.append(":");
target.append(DRIVER_NAME);
target.append(":");
target.append(type.toUpperCase());

return target.toString();
try {
MessageDigest md = MessageDigest.getInstance("SHA-256");
byte[] hash = md.digest(target.toString().getBytes());
return SFTrustManager.byteToHexString(hash);
} catch (NoSuchAlgorithmException e) {
throw new RuntimeException(e);
}
}

enum SecureStorageStatus {
Expand Down
27 changes: 27 additions & 0 deletions src/test/java/net/snowflake/client/core/SecureStorageManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import net.snowflake.client.annotations.RunOnMac;
import net.snowflake.client.annotations.RunOnWindows;
import net.snowflake.client.annotations.RunOnWindowsOrMac;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

class MockAdvapi32Lib implements SecureStorageWindowsManager.Advapi32Lib {
Expand Down Expand Up @@ -224,6 +225,32 @@ public class SecureStorageManagerTest {
private static final String ID_TOKEN = "ID_TOKEN";
private static final String MFA_TOKEN = "MFATOKEN";

@Test
public void testConvertTarget() {
// hex values obtained using https://emn178.github.io/online-tools/sha256.html
String hashedKey =
SecureStorageManager.convertTarget(
host, user, CachedCredentialType.OAUTH_ACCESS_TOKEN.getValue());
Assertions.assertEquals(
"A7C7EBB89312E88552CD00664A0E20929801FACFBD682BF7C2363FB6EC8F914E", hashedKey);

hashedKey =
SecureStorageManager.convertTarget(
host, user, CachedCredentialType.OAUTH_REFRESH_TOKEN.getValue());
Assertions.assertEquals(
"DB37028833FA02B125FBD6DE8CE679C7E62E7D38FAC585E98060E00987F96772", hashedKey);

hashedKey =
SecureStorageManager.convertTarget(host, user, CachedCredentialType.ID_TOKEN.getValue());
Assertions.assertEquals(
"6AA3F783E07D1D2182DAB59442806E2433C55C2BD4D9240790FD5B4B91FD4FDB", hashedKey);

hashedKey =
SecureStorageManager.convertTarget(host, user, CachedCredentialType.MFA_TOKEN.getValue());
Assertions.assertEquals(
"9D10D4EFE45605D85993C6AC95334F1B63D36611B83615656EC7F277A947BF4B", hashedKey);
}

@Test
@RunOnWindowsOrMac
public void testLoadNativeLibrary() {
Expand Down