SNOW-1895884: Token cache refactor

src/main/java/net/snowflake/client/core/SFTrustManager.java

@@ -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++) {

src/main/java/net/snowflake/client/core/SecureStorageLinuxManager.java

@@ -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;
@@ -43,6 +44,7 @@ private SecureStorageLinuxManager() {
             .build();
     logger.debug(
         "Using temporary file: {} as a token cache storage", fileCacheManager.getCacheFilePath());
+    populateLocalCredCache();
   }
 
   private static class SecureStorageLinuxManagerHolder {
@@ -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<>());
+    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) {
-    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.");

src/main/java/net/snowflake/client/core/SecureStorageManager.java

@@ -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);
@@ -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 {

src/test/java/net/snowflake/client/core/SecureStorageManagerTest.java

@@ -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 {
@@ -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() {