fix: per-profile file locking to prevent concurrent duplicate checkouts

Author: Turtdle

src/pybritive/britive_cli.py

@@ -24,6 +24,7 @@
 from . import __version__
 from .helpers import cloud_credential_printer as printer
 from .helpers.cache import Cache
+from .helpers.checkout_lock import CheckoutLock
 from .helpers.config import ConfigManager
 from .helpers.credentials import EncryptedFileCredentialManager, FileCredentialManager
 from .helpers.split import profile_split
@@ -792,6 +793,17 @@ def _checkout(
                 }
             raise e
 
+    def _check_cache(self, passphrase: Optional[str], profile_name: str, mode: str) -> Optional[dict]:
+        credentials = Cache(passphrase=passphrase).get_credentials(profile_name=profile_name, mode=mode)
+        if credentials:
+            expiration_timestamp_str = jmespath.search(
+                expression=self.cachable_modes[mode]['expiration_jmespath'], data=credentials
+            ).replace('Z', '')
+            expires = datetime.fromisoformat(expiration_timestamp_str)
+            if datetime.utcnow() < expires:
+                return credentials
+        return None
+
     @staticmethod
     def _should_check_force_renew(app, force_renew, console):
         return app in ['AWS', 'AWS Standalone'] and force_renew and not console
@@ -903,25 +915,14 @@ def _access_checkout(
         self._validate_justification(justification)
 
         if mode in self.cachable_modes:
-            self.silent = True  # CANNOT output anything other than the expected JSON
-            # we need to check the cache for the credentials first and then check to see if they are expired
-            # if not simply return those credentials, if they are expired, continue to do an actual checkout
+            self.silent = True
             app_type = self.cachable_modes[mode]['app_type']
-            credentials = Cache(passphrase=passphrase).get_credentials(profile_name=alias or profile, mode=mode)
+            credentials = self._check_cache(passphrase, alias or profile, mode)
             if credentials:
-                expiration_timestamp_str = jmespath.search(
-                    expression=self.cachable_modes[mode]['expiration_jmespath'], data=credentials
-                ).replace('Z', '')
-                expires = datetime.fromisoformat(expiration_timestamp_str)
-                now = datetime.utcnow()
-                if now >= expires:  # check to ensure the credentials are still valid, if not, set to None and get new
-                    credentials = None
-                else:
-                    cached_credentials_found = True
+                cached_credentials_found = True
 
         parts = self._split_profile_into_parts(profile)
 
-        # create this params once so we can use it multiple places
         params = {
             'app_name': parts['app'],
             'blocktime': blocktime,
@@ -936,30 +937,45 @@ def _access_checkout(
             'ticket_type': ticket_type,
         }
 
-        if not cached_credentials_found:  # nothing found in cache, cache is expired, or not a cachable mode
-            response = self._checkout(**params)
-            app_type = self._get_app_type(response['appContainerId'])
-            credentials = response['credentials']
-            console_fallback = response.get('console-fallback')
+        if not cached_credentials_found:
+            if mode in self.cachable_modes and self.config.checkout_lock_enabled():
+                with CheckoutLock(profile_key=alias or profile, mode=mode):
+                    credentials = self._check_cache(passphrase, alias or profile, mode)
+                    if credentials:
+                        cached_credentials_found = True
+                    else:
+                        response = self._checkout(**params)
+                        app_type = self._get_app_type(response['appContainerId'])
+                        credentials = response['credentials']
+                        console_fallback = response.get('console-fallback')
+                        Cache(passphrase=passphrase).save_credentials(
+                            profile_name=alias or profile, credentials=credentials, mode=mode
+                        )
+            else:
+                response = self._checkout(**params)
+                app_type = self._get_app_type(response['appContainerId'])
+                credentials = response['credentials']
+                console_fallback = response.get('console-fallback')
+                if mode in self.cachable_modes:
+                    Cache(passphrase=passphrase).save_credentials(
+                        profile_name=alias or profile, credentials=credentials, mode=mode
+                    )
 
-        # this handles the --force-renew flag
-        # lets check to see if we should checkin this profile first and check it out again
         if self._should_check_force_renew(app_type, force_renew, console):
             expiration = datetime.fromisoformat(credentials['expirationTime'].replace('Z', ''))
             now = datetime.utcnow()
             diff = (expiration - now).total_seconds() / 60.0
-            if diff < force_renew:  # time to checkin the profile so we can refresh creds
+            if diff < force_renew:
                 self.print('checking in the profile to get renewed credentials....standby')
                 self.checkin(profile=profile, console=console)
                 response = self._checkout(**params)
-                cached_credentials_found = False  # need to write new creds to cache
                 credentials = response['credentials']
                 console_fallback = response.get('console-fallback')
+                if mode in self.cachable_modes:
+                    Cache(passphrase=passphrase).save_credentials(
+                        profile_name=alias or profile, credentials=credentials, mode=mode
+                    )
 
-        if mode in self.cachable_modes and not cached_credentials_found:
-            Cache(passphrase=passphrase).save_credentials(
-                profile_name=alias or profile, credentials=credentials, mode=mode
-            )
         return app_type, console_fallback, credentials, k8s_processor
 
     def checkout(

src/pybritive/helpers/checkout_lock.py

@@ -0,0 +1,90 @@
+import hashlib
+import os
+import time
+from pathlib import Path
+from types import TracebackType
+from typing import Optional, Type
+
+
+class CheckoutLockTimeout(Exception):
+    pass
+
+
+class _WouldBlock(Exception):
+    pass
+
+
+try:
+    import fcntl
+
+    def _lock_fd(fd: int) -> None:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except (OSError, IOError):
+            raise _WouldBlock()
+
+    def _unlock_fd(fd: int) -> None:
+        fcntl.flock(fd, fcntl.LOCK_UN)
+
+except ImportError:
+    import msvcrt
+
+    def _lock_fd(fd: int) -> None:
+        try:
+            msvcrt.locking(fd, msvcrt.LK_NBLCK, 1)
+        except (OSError, IOError):
+            raise _WouldBlock()
+
+    def _unlock_fd(fd: int) -> None:
+        try:
+            msvcrt.locking(fd, msvcrt.LK_UNLCK, 1)
+        except (OSError, IOError):
+            pass
+
+
+class CheckoutLock:
+    def __init__(self, profile_key: str, mode: str, timeout: float = 120.0, poll_interval: float = 0.1) -> None:
+        self.timeout: float = timeout
+        self.poll_interval: float = poll_interval
+        self._fd: Optional[int] = None
+
+        home = os.getenv('PYBRITIVE_HOME_DIR', str(Path.home()))
+        lock_dir = Path(home) / '.britive' / 'locks'
+        lock_dir.mkdir(parents=True, exist_ok=True)
+
+        lock_name = hashlib.sha256(f'{mode}:{profile_key}'.lower().encode('utf-8')).hexdigest()[:16]
+        self.lock_path: str = str(lock_dir / f'{lock_name}.lock')
+
+    def acquire(self) -> None:
+        self._fd = os.open(self.lock_path, os.O_CREAT | os.O_RDWR)
+        deadline = time.monotonic() + self.timeout
+        while True:
+            try:
+                _lock_fd(self._fd)
+                return
+            except _WouldBlock:
+                if time.monotonic() >= deadline:
+                    os.close(self._fd)
+                    self._fd = None
+                    raise CheckoutLockTimeout(
+                        f'Timed out after {self.timeout}s waiting for checkout lock'
+                    )
+                time.sleep(self.poll_interval)
+
+    def release(self) -> None:
+        if self._fd is not None:
+            try:
+                _unlock_fd(self._fd)
+            finally:
+                os.close(self._fd)
+                self._fd = None
+
+    def __enter__(self) -> 'CheckoutLock':
+        self.acquire()
+        return self
+
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> bool:
+        self.release()
+        return False

src/pybritive/helpers/config.py

@@ -43,6 +43,7 @@ def coalesce(*arg):
     'auto_refresh_kube_config',
     'auto_refresh_profile_cache',
     'ca_bundle',
+    'checkout_lock',
     'credential_backend',
     'default_tenant',
     'output_format',
@@ -281,6 +282,9 @@ def validate_global(self, section, fields):
             if field.replace('-', '_') == 'auto_refresh_kube_config' and value not in ['true', 'false']:
                 error = f'Invalid {section} field {field} value {value} provided. Invalid value choice.'
                 self.validation_error_messages.append(error)
+            if field.replace('-', '_') == 'checkout_lock' and value not in ['true', 'false']:
+                error = f'Invalid {section} field {field} value {value} provided. Invalid value choice.'
+                self.validation_error_messages.append(error)
             if field == 'default_tenant':
                 tenant_aliases_from_sections = [extract_tenant(t) for t in self.config if t.startswith('tenant-')]
                 if value not in tenant_aliases_from_sections:
@@ -343,3 +347,10 @@ def auto_refresh_kube_config(self):
             'auto_refresh_kube_config', self.config.get('global', {}).get('auto-refresh-kube-config', 'false')
         )
         return value == 'true'
+
+    def checkout_lock_enabled(self) -> bool:
+        self.load()
+        value = self.config.get('global', {}).get(
+            'checkout_lock', self.config.get('global', {}).get('checkout-lock', 'false')
+        )
+        return value == 'true'