SNOW-69778: Internal change for future enhancement

Author: smtakeda

ocsp_snowflake.py

@@ -641,9 +641,9 @@ def clear_ssd_cache():
             SFSsd.SSD_CACHE = {}
 
     @staticmethod
-    def find_in_ssd_cache(sfc_endpoint):
-        if sfc_endpoint in SFSsd.SSD_CACHE:
-            return True, SFSsd.SSD_CACHE[sfc_endpoint]
+    def find_in_ssd_cache(account_name):
+        if account_name in SFSsd.SSD_CACHE:
+            return True, SFSsd.SSD_CACHE[account_name]
         return False, None
 
     @staticmethod
@@ -801,6 +801,20 @@ def is_cert_id_in_cache(self, cert_id, subject):
             self, cert_id, subject)
         return found, cache
 
+    def get_account_from_hostname(self, hostname):
+        """
+        Extract the account name part
+        from the hostname
+        :param hostname:
+        :return: account name
+        """
+        split_hname = hostname.split('.')
+        if "global" in split_hname:
+            acc_name = split_hname[0].split('-')[0]
+        else:
+            acc_name = split_hname[0]
+        return acc_name
+
     def validate_by_direct_connection(self, issuer, subject, hostname=None, do_retry=True):
         ssd_cache_status = False
         cache_status = False
@@ -809,7 +823,7 @@ def validate_by_direct_connection(self, issuer, subject, hostname=None, do_retry
 
         cert_id, req = self.create_ocsp_request(issuer, subject)
         if SnowflakeOCSP.SSD.ACTIVATE_SSD:
-            ssd_cache_status, ssd = SnowflakeOCSP.SSD.find_in_ssd_cache(hostname)
+            ssd_cache_status, ssd = SnowflakeOCSP.SSD.find_in_ssd_cache(self.get_account_from_hostname(hostname))
 
         if not ssd_cache_status:
             cache_status, ocsp_response = \
@@ -1173,8 +1187,8 @@ def read_directives():
             with codecs. open(host_specific_ssd, 'r', encoding='utf-8',
                               errors='ignore') as f:
                 ssd_json = json.load(f)
-                for hostname, ssd in ssd_json.items():
-                    SnowflakeOCSP.SSD.add_to_ssd_persistent_cache(hostname, ssd)
+                for account_name, ssd in ssd_json.items():
+                    SnowflakeOCSP.SSD.add_to_ssd_persistent_cache(account_name, ssd)
 
     def process_ocsp_bypass_directive(self, ssd_dir_enc, sfc_cert_id, sfc_endpoint):
         """
@@ -1211,11 +1225,22 @@ def process_ocsp_bypass_directive(self, ssd_dir_enc, sfc_cert_id, sfc_endpoint):
             return False
 
         # Check if the directive is generic (endpoint = *)
-        # or if it is meant for a specific endpoint
+        # or if it is meant for a specific account
         if jwt_ssd_decoded['sfcEndpoint'] != '*':
-                if sfc_endpoint != jwt_ssd_decoded['sfcEndpoint']:
-                    return False
+            """
+            In case there are multiple hostnames
+            associated with the same account,
+            (client failover, different region
+            same account, the sfc_endpoint field
+            would be expected to have a space separated
+            list of all the hostnames that can be
+            associated with the account in question.
+            """
+            split_string = jwt_ssd_decoded['sfcEndpoint'].split()
+            if sfc_endpoint in split_string:
                 return True
+            else:
+                return False
 
         ssd_cert_id_b64 = jwt_ssd_decoded['certId']
         ssd_cert_id = self.decode_cert_id_base64(ssd_cert_id_b64)

test/test_ocsp_ssd.py

@@ -111,12 +111,14 @@ def _create_host_spec_ocsp_bypass_ssd(ocsp, priv_key, hostname):
         exp_val = nbf_val+tdelta
         header = {'ssd_iss':'dep1'}
         payload = {}
-        payload.update({'sfcEndpoint': hostname})
+        hname_string = " ".join(hostname)
+        acc_name = ocsp.get_account_from_hostname(hostname[0])
+        payload.update({'sfcEndpoint': hname_string})
         payload.update({'certId': '*'})
         payload.update({'nbf': nbf_val})
         payload.update({'exp': exp_val})
         host_spec_jwt_token = jwt.encode(payload, priv_key, algorithm='RS512', headers=header)
-        host_spec_bypass_ssd = {hostname: host_spec_jwt_token.decode("utf-8")}
+        host_spec_bypass_ssd = {acc_name: host_spec_jwt_token.decode("utf-8")}
         json.dump(host_spec_bypass_ssd, jwt_host_spec_fp)
 
 
@@ -137,15 +139,16 @@ def test_host_spec_ocsp_bypass_ssd():
     ocsp = _setup_ssd_test(temp_ocsp_file_path)
     priv_key = _get_test_priv_key(1)
 
-    hostname = 'sfcsupport.us-east-1.snowflakecomputing.com'
+    hostname = ['sfcsupport.us-east-1.snowflakecomputing.com']
     try:
         _create_host_spec_ocsp_bypass_ssd(ocsp, priv_key, hostname)
     except Exception as ex:
         print("Exception occurred %s" %ex.message)
 
     ocsp.read_directives()
 
-    cache_status, cur_host_spec_token = ocsp.SSD.find_in_ssd_cache(hostname)
+    acc_name = ocsp.get_account_from_hostname(hostname[0])
+    cache_status, cur_host_spec_token = ocsp.SSD.find_in_ssd_cache(acc_name)
     assert((cur_host_spec_token is not None), "Failed to read host specific directive")
 
     try:
@@ -155,6 +158,84 @@ def test_host_spec_ocsp_bypass_ssd():
         print("Exception while processing SSD :"+ex)
 
 
+def test_host_spec_ocsp_bypass_updated_ssd():
+
+    """
+    Clean any skeletons of past tests
+    """
+    _teardown_ssd_test_setup()
+
+    """
+    Setup OCSP instance to use test keys
+    for authenticating SSD
+    """
+    tmp_dir = str(tempfile.gettempdir())
+    temp_ocsp_file_path = path.join(tmp_dir, "ocsp_cache_backup.json")
+    copy(OCSP_RESPONSE_CACHE_URI, temp_ocsp_file_path)
+    ocsp = _setup_ssd_test(temp_ocsp_file_path)
+    priv_key = _get_test_priv_key(1)
+
+    hostname = ['sfcsupport-test12345.global.us-east-1.snowflakecomputing.com',
+                'sfcsupport-test67890.global.us-east-1.snowflakecomputing.com',
+                'sfcsupport.us-east-1.snowflakecomputing.com',
+                'sfcsupport.us-east-2.snowflakecomputing.com']
+    try:
+        _create_host_spec_ocsp_bypass_ssd(ocsp, priv_key, hostname)
+    except Exception as ex:
+        print("Exception occurred %s" %ex.message)
+
+    ocsp.read_directives()
+
+    acc_name = ocsp.get_account_from_hostname(hostname[0])
+    cache_status, cur_host_spec_token = ocsp.SSD.find_in_ssd_cache(acc_name)
+    assert((cur_host_spec_token is not None), "Failed to read host specific directive")
+
+    try:
+        assert ocsp.process_ocsp_bypass_directive(cur_host_spec_token, '*', hostname[1]),\
+            "Failed to process host specific bypass ssd"
+    except Exception as ex:
+        print("Exception while processing SSD :"+ex)
+
+
+def test_invalid_host_spec_ocsp_bypass_updated_ssd():
+
+    """
+    Clean any skeletons of past tests
+    """
+    _teardown_ssd_test_setup()
+
+    """
+    Setup OCSP instance to use test keys
+    for authenticating SSD
+    """
+    tmp_dir = str(tempfile.gettempdir())
+    temp_ocsp_file_path = path.join(tmp_dir, "ocsp_cache_backup.json")
+    copy(OCSP_RESPONSE_CACHE_URI, temp_ocsp_file_path)
+    ocsp = _setup_ssd_test(temp_ocsp_file_path)
+    priv_key = _get_test_priv_key(1)
+
+    hostname = ['sfcsupport-test12345.global.us-east-1.snowflakecomputing.com',
+                'sfcsupport-test67890.global.us-east-1.snowflakecomputing.com',
+                'sfcsupport.us-east-1.snowflakecomputing.com',
+                'sfcsupport.us-east-2.snowflakecomputing.com']
+    try:
+        _create_host_spec_ocsp_bypass_ssd(ocsp, priv_key, hostname)
+    except Exception as ex:
+        print("Exception occurred %s" %ex.message)
+
+    ocsp.read_directives()
+
+    acc_name = ocsp.get_account_from_hostname(hostname[0])
+    cache_status, cur_host_spec_token = ocsp.SSD.find_in_ssd_cache(acc_name)
+    assert((cur_host_spec_token is not None), "Failed to read host specific directive")
+
+    try:
+        assert ocsp.process_ocsp_bypass_directive(cur_host_spec_token, '*', "sonytv.snowflakecomputing.com") is False,\
+            "SSD should not match hostname specified"
+    except Exception as ex:
+        print("Exception while processing SSD :"+ex)
+
+
 def _create_cert_spec_ocsp_bypass_token(priv_key, cid, validity_days=1):
 
     tdelta = timedelta(days=validity_days)