sanitize key tests

Author: axelsrz

libraries/botbuilder-azure/botbuilder/azure/__init__.py

@@ -6,8 +6,9 @@
 # --------------------------------------------------------------------------
 
 from .about import __version__
-from .cosmosdb_storage import CosmosDbStorage, CosmosDbConfig
+from .cosmosdb_storage import CosmosDbStorage, CosmosDbConfig, CosmosDbKeyEscape
 
 __all__ = ['CosmosDbStorage',
            'CosmosDbConfig',
+           'CosmosDbKeyEscape',
            '__version__']

libraries/botbuilder-azure/botbuilder/azure/cosmosdb_storage.py

@@ -6,7 +6,7 @@
 
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
-
+from hashlib import sha256
 from typing import Dict, List
 from threading import Semaphore
 import json
@@ -42,6 +42,42 @@ def __init__(self, endpoint: str = None, masterkey: str = None, database: str =
         self.container_creation_options = container_creation_options or kwargs.get('container_creation_options')
 
 
+class CosmosDbKeyEscape:
+
+    @staticmethod
+    def sanitize_key(key) -> str:
+        """Return the sanitized key.
+
+        Replace characters that are not allowed in keys in Cosmos.
+
+        :param key:
+        :return str:
+        """
+        # forbidden characters
+        bad_chars = ['\\', '?', '/', '#', '\t', '\n', '\r', '*']
+        # replace those with with '*' and the
+        # Unicode code point of the character and return the new string
+        key = ''.join(
+            map(
+                lambda x: '*' + str(ord(x)) if x in bad_chars else x, key
+            )
+        )
+
+        return CosmosDbKeyEscape.truncate_key(key)
+
+    @staticmethod
+    def truncate_key(key: str) -> str:
+        MAX_KEY_LEN = 255
+
+        if len(key) > MAX_KEY_LEN:
+            aux_hash = sha256(key.encode('utf-8'))
+            aux_hex = aux_hash.hexdigest()
+
+            key = key[0:MAX_KEY_LEN - len(aux_hex)] + aux_hex
+
+        return key
+
+
 class CosmosDbStorage(Storage):
     """The class for CosmosDB middleware for the Azure Bot Framework."""
 
@@ -77,7 +113,7 @@ async def read(self, keys: List[str]) -> Dict[str, object]:
             if keys:
                 # create the parameters object
                 parameters = [
-                    {'name': f'@id{i}', 'value': f'{self.__sanitize_key(key)}'}
+                    {'name': f'@id{i}', 'value': f'{CosmosDbKeyEscape.sanitize_key(key)}'}
                     for i, key in enumerate(keys)
                 ]
                 # get the names of the params
@@ -125,7 +161,7 @@ async def write(self, changes: Dict[str, StoreItem]):
                 # store the e_tag
                 e_tag = change.e_tag
                 # create the new document
-                doc = {'id': self.__sanitize_key(key),
+                doc = {'id': CosmosDbKeyEscape.sanitize_key(key),
                        'realId': key,
                        'document': self.__create_dict(change)
                        }
@@ -141,7 +177,7 @@ async def write(self, changes: Dict[str, StoreItem]):
                     access_condition = {'type': 'IfMatch', 'condition': e_tag}
                     self.client.ReplaceItem(
                         document_link=self.__item_link(
-                            self.__sanitize_key(key)),
+                            CosmosDbKeyEscape.sanitize_key(key)),
                         new_document=doc,
                         options={'accessCondition': access_condition}
                     )
@@ -169,7 +205,7 @@ async def delete(self, keys: List[str]):
             # call the function for each key
             for k in keys:
                 self.client.DeleteItem(
-                    document_link=self.__item_link(self.__sanitize_key(k)),
+                    document_link=self.__item_link(CosmosDbKeyEscape.sanitize_key(k)),
                     options=options
                 )
                 # print(res)
@@ -209,24 +245,6 @@ def __create_dict(self, si: StoreItem) -> Dict:
         return ({attr: getattr(si, attr)
                  for attr in non_magic_attr})
 
-    def __sanitize_key(self, key) -> str:
-        """Return the sanitized key.
-
-        Replace characters that are not allowed in keys in Cosmos.
-
-        :param key:
-        :return str:
-        """
-        # forbidden characters
-        bad_chars = ['\\', '?', '/', '#', '\t', '\n', '\r']
-        # replace those with with '*' and the
-        # Unicode code point of the character and return the new string
-        return ''.join(
-            map(
-                lambda x: '*' + str(ord(x)) if x in bad_chars else x, key
-            )
-        )
-
     def __item_link(self, id) -> str:
         """Return the item link of a item in the container.
 

libraries/botbuilder-azure/tests/test_key_validation.py

@@ -0,0 +1,86 @@
+from unittest import TestCase
+from botbuilder.azure import CosmosDbKeyEscape
+
+
+class TestKeyValidation(TestCase):
+    def test_should_not_change_a_valid_key(self):
+        valid_key = 'Abc12345'
+        sanitized_key = CosmosDbKeyEscape.sanitize_key(valid_key)
+        assert valid_key == sanitized_key, f'{valid_key} should be equal to {sanitized_key}'
+
+    def test_should_escape_illegal_characters_case_1(self):
+        # Ascii code of "?" is "3f"
+        sanitized_key = CosmosDbKeyEscape.sanitize_key('?test?')
+        assert sanitized_key == '*63test*63'
+
+    def test_should_escape_illegal_characters_case_2(self):
+        # Ascii code of "/" is "2f"
+        sanitized_key = CosmosDbKeyEscape.sanitize_key('/test/')
+        assert sanitized_key == '*47test*47'
+
+    def test_should_escape_illegal_characters_case_3(self):
+        # Ascii code of "\" is "5c"
+        sanitized_key = CosmosDbKeyEscape.sanitize_key('\\test\\')
+        assert sanitized_key == '*92test*92'
+
+    def test_should_escape_illegal_characters_case_4(self):
+        # Ascii code of "#" is "23"
+        sanitized_key = CosmosDbKeyEscape.sanitize_key('#test#')
+        assert sanitized_key == '*35test*35'
+
+    def test_should_escape_illegal_characters_case_5(self):
+        # Ascii code of "*" is "2a".
+        sanitized_key = CosmosDbKeyEscape.sanitize_key('*test*')
+        assert sanitized_key == '*42test*42'
+
+    def test_should_escape_illegal_characters_compound_key(self):
+        # Check a compound key
+        compoundsanitized_key = CosmosDbKeyEscape.sanitize_key('?#/')
+        assert compoundsanitized_key, '*3f*23*2f'
+
+    def test_should_handle_possible_collisions(self):
+        valid_key1 = '*2atest*2a'
+        valid_key2 = '*test*'
+
+        escaped1 = CosmosDbKeyEscape.sanitize_key(valid_key1)
+        escaped2 = CosmosDbKeyEscape.sanitize_key(valid_key2)
+
+        assert escaped1 != escaped2, f'{escaped1} should be different that {escaped2}'
+
+    def test_should_truncate_longer_keys(self):
+        # create an extra long key
+        # limit is 255
+        long_key = 'x' * 300
+        fixed = CosmosDbKeyEscape.sanitize_key(long_key)
+
+        assert len(fixed) <= 255, 'long key was not properly truncated'
+
+    def test_should_not_truncate_short_key(self):
+        # create a short key
+        short_key = 'x' * 16
+        fixed2 = CosmosDbKeyEscape.sanitize_key(short_key)
+
+        assert len(fixed2) == 16, 'short key was truncated improperly'
+
+    def test_should_create_sufficiently_different_truncated_keys_of_similar_origin(self):
+        # create 2 very similar extra long key where the difference will definitely be trimmed off by truncate function
+        long_key = 'x' * 300 + "1"
+        long_key2 = 'x' * 300 + "2"
+
+        fixed = CosmosDbKeyEscape.sanitize_key(long_key)
+        fixed2 = CosmosDbKeyEscape.sanitize_key(long_key2)
+
+        assert len(fixed) != fixed2, 'key truncation failed to create unique key'
+
+    def test_should_properly_truncate_keys_with_special_chars(self):
+        # create a short key
+        long_key = '*' * 300
+        fixed = CosmosDbKeyEscape.sanitize_key(long_key)
+
+        assert len(fixed) <= 255, 'long key with special char was truncated improperly'
+
+        # create a short key
+        short_key = '#' * 16
+        fixed2 = CosmosDbKeyEscape.sanitize_key(short_key)
+
+        assert len(fixed2) <= 255, 'short key with special char was truncated improperly'