From 22b1c5a8c84a0048b9d687e711ef4f7dccbdc7e7 Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Tue, 23 Sep 2025 08:48:18 -0700 Subject: [PATCH 01/10] Debug confidant low cache hit --- kmsauth/__init__.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/kmsauth/__init__.py b/kmsauth/__init__.py index 451b976..bcd2fac 100644 --- a/kmsauth/__init__.py +++ b/kmsauth/__init__.py @@ -169,14 +169,16 @@ def _format_auth_key(self, keys): ) def _get_key_arn(self, key): + logging.debug(f'Getting key ARN for {key}') if key.startswith('arn:aws:kms:'): self.KEY_METADATA[key] = { 'KeyMetadata': {'Arn': key} } if key not in self.KEY_METADATA: - self.KEY_METADATA[key] = self.kms_client.describe_key( - KeyId='{0}'.format(key) - ) + with self.stats.timer('kms_describe_key'): + self.KEY_METADATA[key] = self.kms_client.describe_key( + KeyId='{0}'.format(key) + ) return self.KEY_METADATA[key]['KeyMetadata']['Arn'] def _get_key_alias_from_cache(self, key_arn): @@ -246,6 +248,9 @@ def decrypt_token(self, username, token): raise TokenValidationError('Unacceptable token version.') if self.stats: self.stats.incr('token_version_{0}'.format(version)) + self.stats.incr(f'cache_key.from.{_from}') + self.stats.incr(f'cache_key.to.{self.to_auth_context}') + self.stats.incr(f'cache_key.user_type.{user_type}') try: token_key = '{0}{1}{2}{3}'.format( hashlib.sha256(ensure_bytes(token)).hexdigest(), @@ -256,6 +261,11 @@ def decrypt_token(self, username, token): except Exception: raise TokenValidationError('Authentication error.') if token_key not in self.TOKENS: + self.stats.incr('token_cache.miss') + self.stats.gauge('token_cache.size_at_miss', len(self.TOKENS)) + if len(self.TOKENS) >= self.token_cache_size: + self.stats.incr('token_cache.eviction') + try: token = base64.b64decode(token) # Ensure normal context fields override whatever is in @@ -313,6 +323,7 @@ def decrypt_token(self, username, token): 'Authentication error. General error.' ) else: + self.stats.incr('token_cache.hit') ret = self.TOKENS[token_key] now = datetime.datetime.utcnow() try: @@ -342,6 +353,8 @@ def decrypt_token(self, username, token): raise TokenValidationError( 'Authentication error. Invalid time validity for token.' ) + self.stats.incr('token_cache.set') + self.stats.gauge('token_cache.size_at_set', len(self.TOKENS)) self.TOKENS[token_key] = ret return self.TOKENS[token_key] From e40704bcdc106299838404eb1ac077e626c7952e Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Tue, 23 Sep 2025 13:29:06 -0700 Subject: [PATCH 02/10] Delete python version 3.7.x from test --- .github/workflows/pull_request.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index 94f73b1..521b2dd 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -17,7 +17,7 @@ jobs: runs-on: ubuntu-22.04 strategy: matrix: - python-version: ['3.7.x', '3.8.x', '3.9.x', '3.10.x', '3.11.x'] + python-version: ['3.8.x', '3.9.x', '3.10.x', '3.11.x'] steps: - name: Checkout uses: actions/checkout@v1 From 7b72b0fc124353cec295fe58129a1d915ca9edbc Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Tue, 23 Sep 2025 13:32:21 -0700 Subject: [PATCH 03/10] Delete python version 3.8.x from test --- .github/workflows/pull_request.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index 521b2dd..91b4fb7 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -5,10 +5,10 @@ jobs: steps: - name: Checkout uses: actions/checkout@v1 - - name: Setup python 3.8 + - name: Setup python 3.10 uses: actions/setup-python@v1 with: - python-version: 3.8 + python-version: 3.10 - name: Install pre-commit run: pip install pre-commit - name: Run pre-commit @@ -17,7 +17,7 @@ jobs: runs-on: ubuntu-22.04 strategy: matrix: - python-version: ['3.8.x', '3.9.x', '3.10.x', '3.11.x'] + python-version: ['3.9.x', '3.10.x', '3.11.x'] steps: - name: Checkout uses: actions/checkout@v1 From b54ef315ea081a46193881fc30d7f50f88b84e2c Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Tue, 23 Sep 2025 13:37:32 -0700 Subject: [PATCH 04/10] Update python version --- .github/workflows/pull_request.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index 91b4fb7..d217124 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -4,11 +4,11 @@ jobs: runs-on: ubuntu-22.04 steps: - name: Checkout - uses: actions/checkout@v1 + uses: actions/checkout@v4 - name: Setup python 3.10 - uses: actions/setup-python@v1 + uses: actions/setup-python@v4 with: - python-version: 3.10 + python-version: "3.10" - name: Install pre-commit run: pip install pre-commit - name: Run pre-commit @@ -20,9 +20,9 @@ jobs: python-version: ['3.9.x', '3.10.x', '3.11.x'] steps: - name: Checkout - uses: actions/checkout@v1 + uses: actions/checkout@v4 - name: Setup python ${{ matrix.python-version }} - uses: actions/setup-python@v1 + uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - name: Install dependencies From 89df3e135a0eed53c4f813126c58b3f7dd856840 Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Tue, 23 Sep 2025 15:00:28 -0700 Subject: [PATCH 05/10] Fix stats error --- .github/workflows/push.yml | 8 ++++---- kmsauth/__init__.py | 26 +++++++++++++++++--------- 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 079342a..fbdb5a7 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -11,11 +11,11 @@ jobs: runs-on: ubuntu-22.04 steps: - name: Checkout - uses: actions/checkout@v1 - - name: Setup python 3.8 - uses: actions/setup-python@v1 + uses: actions/checkout@v4 + - name: Setup python 3.10 + uses: actions/setup-python@v4 with: - python-version: 3.8 + python-version: "3.10" - name: Add wheel dependency run: pip install wheel - name: Generate dist diff --git a/kmsauth/__init__.py b/kmsauth/__init__.py index bcd2fac..438449f 100644 --- a/kmsauth/__init__.py +++ b/kmsauth/__init__.py @@ -129,6 +129,7 @@ def __init__( else: self.extra_context = extra_context self.TOKENS = LRU(token_cache_size) + self.token_cache_size = token_cache_size self.KEY_METADATA = {} self.stats = stats self._validate() @@ -169,13 +170,17 @@ def _format_auth_key(self, keys): ) def _get_key_arn(self, key): - logging.debug(f'Getting key ARN for {key}') if key.startswith('arn:aws:kms:'): self.KEY_METADATA[key] = { 'KeyMetadata': {'Arn': key} } if key not in self.KEY_METADATA: - with self.stats.timer('kms_describe_key'): + if self.stats: + with self.stats.timer('kms_describe_key'): + self.KEY_METADATA[key] = self.kms_client.describe_key( + KeyId='{0}'.format(key) + ) + else: self.KEY_METADATA[key] = self.kms_client.describe_key( KeyId='{0}'.format(key) ) @@ -261,10 +266,11 @@ def decrypt_token(self, username, token): except Exception: raise TokenValidationError('Authentication error.') if token_key not in self.TOKENS: - self.stats.incr('token_cache.miss') - self.stats.gauge('token_cache.size_at_miss', len(self.TOKENS)) - if len(self.TOKENS) >= self.token_cache_size: - self.stats.incr('token_cache.eviction') + if self.stats: + self.stats.incr('token_cache.miss') + self.stats.gauge('token_cache.size_at_miss', len(self.TOKENS)) + if len(self.TOKENS) >= self.token_cache_size: + self.stats.incr('token_cache.eviction') try: token = base64.b64decode(token) @@ -323,7 +329,8 @@ def decrypt_token(self, username, token): 'Authentication error. General error.' ) else: - self.stats.incr('token_cache.hit') + if self.stats: + self.stats.incr('token_cache.hit') ret = self.TOKENS[token_key] now = datetime.datetime.utcnow() try: @@ -353,8 +360,9 @@ def decrypt_token(self, username, token): raise TokenValidationError( 'Authentication error. Invalid time validity for token.' ) - self.stats.incr('token_cache.set') - self.stats.gauge('token_cache.size_at_set', len(self.TOKENS)) + if self.stats: + self.stats.incr('token_cache.set') + self.stats.gauge('token_cache.size_at_set', len(self.TOKENS)) self.TOKENS[token_key] = ret return self.TOKENS[token_key] From 2f31039883ed285f0cdc739a523466dd506412ee Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Tue, 23 Sep 2025 16:03:39 -0700 Subject: [PATCH 06/10] Update the VERSION --- kmsauth/__init__.py | 33 ++++++++++++++++++++------------- setup.py | 2 +- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/kmsauth/__init__.py b/kmsauth/__init__.py index 438449f..68b157d 100644 --- a/kmsauth/__init__.py +++ b/kmsauth/__init__.py @@ -247,15 +247,16 @@ def decrypt_token(self, username, token): ''' Decrypt a token. ''' + time_start = datetime.datetime.utcnow() version, user_type, _from = self._parse_username(username) if (version > self.maximum_token_version or version < self.minimum_token_version): raise TokenValidationError('Unacceptable token version.') if self.stats: - self.stats.incr('token_version_{0}'.format(version)) - self.stats.incr(f'cache_key.from.{_from}') - self.stats.incr(f'cache_key.to.{self.to_auth_context}') - self.stats.incr(f'cache_key.user_type.{user_type}') + self.stats.incr('token_version_{version}') + self.stats.incr(f'cache_key_from_{_from}') + self.stats.incr(f'cache_key_to_{self.to_auth_context}') + self.stats.incr(f'cache_key_user_type_{user_type}') try: token_key = '{0}{1}{2}{3}'.format( hashlib.sha256(ensure_bytes(token)).hexdigest(), @@ -267,10 +268,10 @@ def decrypt_token(self, username, token): raise TokenValidationError('Authentication error.') if token_key not in self.TOKENS: if self.stats: - self.stats.incr('token_cache.miss') - self.stats.gauge('token_cache.size_at_miss', len(self.TOKENS)) + self.stats.incr('token_cache_miss') + self.stats.gauge('token_cache_size_at_miss', len(self.TOKENS)) if len(self.TOKENS) >= self.token_cache_size: - self.stats.incr('token_cache.eviction') + self.stats.incr('token_cache_eviction') try: token = base64.b64decode(token) @@ -330,9 +331,12 @@ def decrypt_token(self, username, token): ) else: if self.stats: - self.stats.incr('token_cache.hit') + self.stats.incr('token_cache_hit') ret = self.TOKENS[token_key] + now = datetime.datetime.utcnow() + if self.stats: + self.stats.timing('decrypt_token_validation_duration', (now - time_start).total_seconds() * 1000) # noqa: E501 try: not_before = datetime.datetime.strptime( ret['payload']['not_before'], @@ -344,14 +348,14 @@ def decrypt_token(self, username, token): ) except Exception: logging.exception( - 'Failed to get not_before and not_after from token payload.' + 'Failed to get not_before and not_after from token payload.' # noqa: E501 ) raise TokenValidationError( 'Authentication error. Missing validity.' ) delta = (not_after - not_before).seconds / 60 if delta > self.auth_token_max_lifetime: - logging.warning('Token used which exceeds max token lifetime.') + logging.warning('Token used which exceeds max token lifetime.') # noqa: E501 raise TokenValidationError( 'Authentication error. Token lifetime exceeded.' ) @@ -360,10 +364,13 @@ def decrypt_token(self, username, token): raise TokenValidationError( 'Authentication error. Invalid time validity for token.' ) - if self.stats: - self.stats.incr('token_cache.set') - self.stats.gauge('token_cache.size_at_set', len(self.TOKENS)) + self.TOKENS[token_key] = ret + duration = (datetime.datetime.utcnow() - now).total_seconds() * 1000 + if self.stats: + self.stats.timing('decrypt_token_duration_post_validation', duration) # noqa: E501 + self.stats.incr('token_cache_set') + self.stats.gauge('token_cache_size_at_set', len(self.TOKENS)) # noqa: E501 return self.TOKENS[token_key] diff --git a/setup.py b/setup.py index 2b6ce72..0fed2d2 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ from setuptools import setup, find_packages -VERSION = "0.6.3" +VERSION = "0.6.4.dev2" requirements = [ # Boto3 is the Amazon Web Services (AWS) Software Development Kit (SDK) From 40a7439e96068749fe5c06681e4671fc6e774756 Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Wed, 24 Sep 2025 23:04:57 -0700 Subject: [PATCH 07/10] Add more metric for cache hit debug --- kmsauth/__init__.py | 14 +++++++++++++- setup.py | 2 +- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/kmsauth/__init__.py b/kmsauth/__init__.py index 68b157d..371ef08 100644 --- a/kmsauth/__init__.py +++ b/kmsauth/__init__.py @@ -296,6 +296,7 @@ def decrypt_token(self, username, token): # Decrypt doesn't take KeyId as an argument. We need to verify # the correct key was used to do the decryption. # Annoyingly, the KeyId from the data is actually an arn. + key_validation_start = datetime.datetime.utcnow() key_arn = data['KeyId'] if user_type == 'service': if not self._valid_service_auth_key(key_arn): @@ -311,9 +312,16 @@ def decrypt_token(self, username, token): raise TokenValidationError( 'Authentication error. Unsupported user_type.' ) + if self.stats: + key_validation_duration = (datetime.datetime.utcnow() - key_validation_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('key_validation_duration', key_validation_duration) # noqa: E501 + json_start = datetime.datetime.utcnow() plaintext = data['Plaintext'] payload = json.loads(plaintext) key_alias = self._get_key_alias_from_cache(key_arn) + if self.stats: + json_duration = (datetime.datetime.utcnow() - json_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('json_processing_duration', json_duration) ret = {'payload': payload, 'key_alias': key_alias} except TokenValidationError: raise @@ -337,6 +345,7 @@ def decrypt_token(self, username, token): now = datetime.datetime.utcnow() if self.stats: self.stats.timing('decrypt_token_validation_duration', (now - time_start).total_seconds() * 1000) # noqa: E501 + time_validation_start = datetime.datetime.utcnow() try: not_before = datetime.datetime.strptime( ret['payload']['not_before'], @@ -364,9 +373,12 @@ def decrypt_token(self, username, token): raise TokenValidationError( 'Authentication error. Invalid time validity for token.' ) + if self.stats: + time_validation_duration = (datetime.datetime.utcnow() - time_validation_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('time_validation_duration', time_validation_duration) # noqa: E501 self.TOKENS[token_key] = ret - duration = (datetime.datetime.utcnow() - now).total_seconds() * 1000 + duration = (datetime.datetime.utcnow() - now).total_seconds() * 1000 # noqa: E501 if self.stats: self.stats.timing('decrypt_token_duration_post_validation', duration) # noqa: E501 self.stats.incr('token_cache_set') diff --git a/setup.py b/setup.py index 0fed2d2..90f0339 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ from setuptools import setup, find_packages -VERSION = "0.6.4.dev2" +VERSION = "0.6.4.dev3" requirements = [ # Boto3 is the Amazon Web Services (AWS) Software Development Kit (SDK) From bb7b70e6a937e5a24560c504b098603c8f1b2b11 Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Thu, 25 Sep 2025 06:00:58 -0700 Subject: [PATCH 08/10] Add more metrics --- kmsauth/__init__.py | 29 ++++++++++++++++++++++++++++- setup.py | 2 +- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/kmsauth/__init__.py b/kmsauth/__init__.py index 371ef08..5d7c45d 100644 --- a/kmsauth/__init__.py +++ b/kmsauth/__init__.py @@ -253,10 +253,14 @@ def decrypt_token(self, username, token): version < self.minimum_token_version): raise TokenValidationError('Unacceptable token version.') if self.stats: + parse_duration = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('username_parse_duration', parse_duration) # noqa: E501 self.stats.incr('token_version_{version}') self.stats.incr(f'cache_key_from_{_from}') self.stats.incr(f'cache_key_to_{self.to_auth_context}') self.stats.incr(f'cache_key_user_type_{user_type}') + + cache_key_gen_start = datetime.datetime.utcnow() try: token_key = '{0}{1}{2}{3}'.format( hashlib.sha256(ensure_bytes(token)).hexdigest(), @@ -266,7 +270,17 @@ def decrypt_token(self, username, token): ) except Exception: raise TokenValidationError('Authentication error.') - if token_key not in self.TOKENS: + if self.stats: + cache_key_duration = (datetime.datetime.utcnow() - cache_key_gen_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('cache_key_generation_duration', cache_key_duration) # noqa: E501 + + cache_lookup_start = datetime.datetime.utcnow() + cache_miss = token_key not in self.TOKENS + if self.stats: + cache_lookup_duration = (datetime.datetime.utcnow() - cache_lookup_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('cache_lookup_duration', cache_lookup_duration) # noqa: E501 + + if cache_miss: if self.stats: self.stats.incr('token_cache_miss') self.stats.gauge('token_cache_size_at_miss', len(self.TOKENS)) @@ -338,12 +352,20 @@ def decrypt_token(self, username, token): 'Authentication error. General error.' ) else: + cache_hit_start = datetime.datetime.utcnow() if self.stats: self.stats.incr('token_cache_hit') ret = self.TOKENS[token_key] + if self.stats: + cache_hit_duration = (datetime.datetime.utcnow() - cache_hit_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('cache_hit_lookup_duration', cache_hit_duration) # noqa: E501 now = datetime.datetime.utcnow() if self.stats: + # Total time from start to this point (before time validation) + pre_time_validation_duration = (now - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('pre_time_validation_duration', pre_time_validation_duration) # noqa: E501 + # Original total validation duration metric self.stats.timing('decrypt_token_validation_duration', (now - time_start).total_seconds() * 1000) # noqa: E501 time_validation_start = datetime.datetime.utcnow() try: @@ -377,7 +399,12 @@ def decrypt_token(self, username, token): time_validation_duration = (datetime.datetime.utcnow() - time_validation_start).total_seconds() * 1000 # noqa: E501 self.stats.timing('time_validation_duration', time_validation_duration) # noqa: E501 + cache_set_start = datetime.datetime.utcnow() self.TOKENS[token_key] = ret + if self.stats: + cache_set_duration = (datetime.datetime.utcnow() - cache_set_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('cache_set_duration', cache_set_duration) # noqa: E501 + duration = (datetime.datetime.utcnow() - now).total_seconds() * 1000 # noqa: E501 if self.stats: self.stats.timing('decrypt_token_duration_post_validation', duration) # noqa: E501 diff --git a/setup.py b/setup.py index 90f0339..ccf1911 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ from setuptools import setup, find_packages -VERSION = "0.6.4.dev3" +VERSION = "0.6.4.dev4" requirements = [ # Boto3 is the Amazon Web Services (AWS) Software Development Kit (SDK) From e5e02edbd54ce2b45db4f02dbac9c2350632238f Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Thu, 25 Sep 2025 07:51:19 -0700 Subject: [PATCH 09/10] Add more metrics --- kmsauth/__init__.py | 47 ++++++++++++++++++++++++++------------------- setup.py | 2 +- 2 files changed, 28 insertions(+), 21 deletions(-) diff --git a/kmsauth/__init__.py b/kmsauth/__init__.py index 5d7c45d..e048752 100644 --- a/kmsauth/__init__.py +++ b/kmsauth/__init__.py @@ -253,14 +253,15 @@ def decrypt_token(self, username, token): version < self.minimum_token_version): raise TokenValidationError('Unacceptable token version.') if self.stats: - parse_duration = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 - self.stats.timing('username_parse_duration', parse_duration) # noqa: E501 + # Checkpoint 1: After username parsing + checkpoint_1 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_1_after_parse', checkpoint_1) # noqa: E501 + self.stats.incr('token_version_{version}') self.stats.incr(f'cache_key_from_{_from}') self.stats.incr(f'cache_key_to_{self.to_auth_context}') self.stats.incr(f'cache_key_user_type_{user_type}') - cache_key_gen_start = datetime.datetime.utcnow() try: token_key = '{0}{1}{2}{3}'.format( hashlib.sha256(ensure_bytes(token)).hexdigest(), @@ -270,15 +271,18 @@ def decrypt_token(self, username, token): ) except Exception: raise TokenValidationError('Authentication error.') + if self.stats: - cache_key_duration = (datetime.datetime.utcnow() - cache_key_gen_start).total_seconds() * 1000 # noqa: E501 - self.stats.timing('cache_key_generation_duration', cache_key_duration) # noqa: E501 + # Checkpoint 2: After cache key generation + checkpoint_2 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_2_after_cache_key', checkpoint_2) # noqa: E501 - cache_lookup_start = datetime.datetime.utcnow() cache_miss = token_key not in self.TOKENS + if self.stats: - cache_lookup_duration = (datetime.datetime.utcnow() - cache_lookup_start).total_seconds() * 1000 # noqa: E501 - self.stats.timing('cache_lookup_duration', cache_lookup_duration) # noqa: E501 + # Checkpoint 3: After cache lookup + checkpoint_3 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_3_after_cache_lookup', checkpoint_3) # noqa: E501 if cache_miss: if self.stats: @@ -302,6 +306,9 @@ def decrypt_token(self, username, token): CiphertextBlob=token, EncryptionContext=context ) + # Checkpoint 4: After KMS decrypt + checkpoint_4 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_4_after_kms_decrypt', checkpoint_4) # noqa: E501 else: data = self.kms_client.decrypt( CiphertextBlob=token, @@ -310,7 +317,6 @@ def decrypt_token(self, username, token): # Decrypt doesn't take KeyId as an argument. We need to verify # the correct key was used to do the decryption. # Annoyingly, the KeyId from the data is actually an arn. - key_validation_start = datetime.datetime.utcnow() key_arn = data['KeyId'] if user_type == 'service': if not self._valid_service_auth_key(key_arn): @@ -327,15 +333,16 @@ def decrypt_token(self, username, token): 'Authentication error. Unsupported user_type.' ) if self.stats: - key_validation_duration = (datetime.datetime.utcnow() - key_validation_start).total_seconds() * 1000 # noqa: E501 - self.stats.timing('key_validation_duration', key_validation_duration) # noqa: E501 - json_start = datetime.datetime.utcnow() + # Checkpoint 5: After key validation + checkpoint_5 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_5_after_key_validation', checkpoint_5) # noqa: E501 plaintext = data['Plaintext'] payload = json.loads(plaintext) key_alias = self._get_key_alias_from_cache(key_arn) if self.stats: - json_duration = (datetime.datetime.utcnow() - json_start).total_seconds() * 1000 # noqa: E501 - self.stats.timing('json_processing_duration', json_duration) + # Checkpoint 6: After JSON processing + checkpoint_6 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_6_after_json_processing', checkpoint_6) # noqa: E501 ret = {'payload': payload, 'key_alias': key_alias} except TokenValidationError: raise @@ -352,13 +359,13 @@ def decrypt_token(self, username, token): 'Authentication error. General error.' ) else: - cache_hit_start = datetime.datetime.utcnow() if self.stats: self.stats.incr('token_cache_hit') ret = self.TOKENS[token_key] if self.stats: - cache_hit_duration = (datetime.datetime.utcnow() - cache_hit_start).total_seconds() * 1000 # noqa: E501 - self.stats.timing('cache_hit_lookup_duration', cache_hit_duration) # noqa: E501 + # Checkpoint 7: After cache hit + checkpoint_7 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_7_after_cache_hit', checkpoint_7) # noqa: E501 now = datetime.datetime.utcnow() if self.stats: @@ -367,7 +374,6 @@ def decrypt_token(self, username, token): self.stats.timing('pre_time_validation_duration', pre_time_validation_duration) # noqa: E501 # Original total validation duration metric self.stats.timing('decrypt_token_validation_duration', (now - time_start).total_seconds() * 1000) # noqa: E501 - time_validation_start = datetime.datetime.utcnow() try: not_before = datetime.datetime.strptime( ret['payload']['not_before'], @@ -396,8 +402,9 @@ def decrypt_token(self, username, token): 'Authentication error. Invalid time validity for token.' ) if self.stats: - time_validation_duration = (datetime.datetime.utcnow() - time_validation_start).total_seconds() * 1000 # noqa: E501 - self.stats.timing('time_validation_duration', time_validation_duration) # noqa: E501 + # Checkpoint 8: After time validation + checkpoint_8 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_8_after_time_validation', checkpoint_8) # noqa: E501 cache_set_start = datetime.datetime.utcnow() self.TOKENS[token_key] = ret diff --git a/setup.py b/setup.py index ccf1911..fc9ed98 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ from setuptools import setup, find_packages -VERSION = "0.6.4.dev4" +VERSION = "0.6.4.dev5" requirements = [ # Boto3 is the Amazon Web Services (AWS) Software Development Kit (SDK) From 68b25767c8f12322e89ff72a6e1dbc418428eb95 Mon Sep 17 00:00:00 2001 From: Wenchong Hu Date: Thu, 25 Sep 2025 09:33:55 -0700 Subject: [PATCH 10/10] More metrics --- kmsauth/__init__.py | 15 +++++++++++++++ setup.py | 2 +- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/kmsauth/__init__.py b/kmsauth/__init__.py index e048752..db929bc 100644 --- a/kmsauth/__init__.py +++ b/kmsauth/__init__.py @@ -253,6 +253,7 @@ def decrypt_token(self, username, token): version < self.minimum_token_version): raise TokenValidationError('Unacceptable token version.') if self.stats: + self.stats.incr('token_version_{0}'.format(version)) # Checkpoint 1: After username parsing checkpoint_1 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 self.stats.timing('checkpoint_1_after_parse', checkpoint_1) # noqa: E501 @@ -291,8 +292,17 @@ def decrypt_token(self, username, token): if len(self.TOKENS) >= self.token_cache_size: self.stats.incr('token_cache_eviction') + # Checkpoint 3.5: After stats calls in cache miss + checkpoint_3_5 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_3_5_after_cache_miss_stats', checkpoint_3_5) # noqa: E501 + try: token = base64.b64decode(token) + if self.stats: + # Checkpoint 3.7: After base64 decode + checkpoint_3_7 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_3_7_after_base64_decode', checkpoint_3_7) # noqa: E501 + # Ensure normal context fields override whatever is in # extra_context. context = copy.deepcopy(self.extra_context) @@ -300,6 +310,11 @@ def decrypt_token(self, username, token): context['from'] = _from if version > 1: context['user_type'] = user_type + + if self.stats: + # Checkpoint 3.9: After context setup + checkpoint_3_9 = (datetime.datetime.utcnow() - time_start).total_seconds() * 1000 # noqa: E501 + self.stats.timing('checkpoint_3_9_after_context_setup', checkpoint_3_9) # noqa: E501 if self.stats: with self.stats.timer('kms_decrypt_token'): data = self.kms_client.decrypt( diff --git a/setup.py b/setup.py index fc9ed98..82f1833 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ from setuptools import setup, find_packages -VERSION = "0.6.4.dev5" +VERSION = "0.6.4.dev7" requirements = [ # Boto3 is the Amazon Web Services (AWS) Software Development Kit (SDK)