fix: Improve STS credential refresh and logging (#327)

Refactored credential refresh logic to handle expired STS credentials and avoid unnecessary AssumeRole calls. Reduced verbose logging and clarified cache handling for AssumeRole credentials.

Refactor credential refresh logic in IdentityClient

Simplifies and restructures the credential refresh flow to prioritize constructor/environment credentials, handle expired session tokens, and streamline VeFaaS IAM and AssumeRole fallback logic. Improves readability and maintainability by removing redundant checks and helper functions.

Author: loveyana

veadk/integrations/ve_identity/identity_client.py

@@ -56,67 +56,56 @@ def refresh_credentials(func):
     """
     import asyncio
 
+    def _try_get_vefaas_credentials():
+        """Attempt to retrieve credentials from VeFaaS IAM."""
+        try:
+            ve_iam_cred = get_credential_from_vefaas_iam()
+            return (
+                ve_iam_cred.access_key_id,
+                ve_iam_cred.secret_access_key,
+                ve_iam_cred.session_token,
+            )
+        except FileNotFoundError:
+            pass  # VeFaaS IAM file not found, ignore
+        except Exception as e:
+            logger.warning(f"Failed to retrieve credentials from VeFaaS IAM: {e}")
+        return None
+
     @wraps(func)
     def _refresh_creds(self: IdentityClient):
         """Helper to refresh credentials."""
-        # Try to get credentials from environment variables first
+        # Step 1: Get initial credentials from constructor or environment variables
         ak = self._initial_access_key or os.getenv("VOLCENGINE_ACCESS_KEY", "")
         sk = self._initial_secret_key or os.getenv("VOLCENGINE_SECRET_KEY", "")
         session_token = self._initial_session_token or os.getenv(
             "VOLCENGINE_SESSION_TOKEN", ""
         )
 
-        # Helper function to attempt VeFaaS IAM credential retrieval
-        def try_get_vefaas_credentials():
-            """Attempt to retrieve credentials from VeFaaS IAM."""
-            try:
-                logger.info("Attempting to fetch credentials from VeFaaS IAM...")
-                ve_iam_cred = get_credential_from_vefaas_iam()
-                return (
-                    ve_iam_cred.access_key_id,
-                    ve_iam_cred.secret_access_key,
-                    ve_iam_cred.session_token,
-                )
-            except FileNotFoundError as e:
-                logger.warning(f"VeFaaS IAM credentials not available: {e}")
-            except Exception as e:
-                logger.warning(f"Failed to retrieve credentials from VeFaaS IAM: {e}")
-            return None
+        # Step 2: Clear expired session_token
+        if self._is_sts_credential_expired():
+            logger.info("STS credentials expired, clearing...")
+            session_token = ""
 
-        # If no AK/SK, try to get from VeFaaS IAM
-        if not (ak and sk):
-            logger.info(
-                "Credentials not found in environment, attempting to fetch from VeFaaS IAM..."
-            )
-            credentials = try_get_vefaas_credentials()
-            if credentials:
+        # Step 3: Try VeFaaS IAM if no credentials or no session_token
+        # VeFaaS IAM provides complete credentials (ak, sk, session_token)
+        if not (ak and sk) or (ak and sk and not session_token):
+            if credentials := _try_get_vefaas_credentials():
                 ak, sk, session_token = credentials
 
-        # If we have AK/SK but no session token, try to get complete credentials
+        # Step 4: If still no session_token, try AssumeRole
         if ak and sk and not session_token:
-            # First attempt: try VeFaaS IAM
-            credentials = try_get_vefaas_credentials()
-            if credentials:
-                ak, sk, session_token = credentials
-
-            # Second attempt: if still no session token, try AssumeRole
-            if not session_token:
-                role_trn = self._get_iam_role_trn_from_vefaas_iam() or os.getenv(
-                    "RUNTIME_IAM_ROLE_TRN", ""
-                )
+            if role_trn := self._get_iam_role_trn_from_vefaas_iam() or os.getenv(
+                "RUNTIME_IAM_ROLE_TRN", ""
+            ):
+                try:
+                    sts_cred = self._assume_role(ak, sk, role_trn)
+                    ak = sts_cred.access_key_id
+                    sk = sts_cred.secret_access_key
+                    session_token = sts_cred.session_token
+                except Exception as e:
+                    logger.warning(f"Failed to assume role: {e}")
 
-                if role_trn:
-                    try:
-                        logger.info(f"Attempting AssumeRole with role: {role_trn}")
-                        sts_credentials = self._assume_role(ak, sk, role_trn)
-                        ak = sts_credentials.access_key_id
-                        sk = sts_credentials.secret_access_key
-                        session_token = sts_credentials.session_token
-                        logger.info("Successfully obtained credentials via AssumeRole")
-                    except Exception as e:
-                        logger.warning(f"Failed to assume role: {e}")
-
-        # Update configuration with the credentials
+        # Step 5: Update configuration with the credentials
         self._api_client.api_client.configuration.ak = ak
         self._api_client.api_client.configuration.sk = sk
         self._api_client.api_client.configuration.session_token = session_token
@@ -192,16 +181,9 @@ def __init__(
         self._sts_credential_expires_at: Optional[int] = None
 
     def _get_iam_role_trn_from_vefaas_iam(self) -> Optional[str]:
-        logger.info(
-            f"Try to get IAM Role TRN from VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH})."
-        )
-
         path = Path(VEFAAS_IAM_CRIDENTIAL_PATH)
 
         if not path.exists():
-            logger.error(
-                f"Get IAM Role TRN from IAM file failed, and VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH}) not exists. Please check your configuration."
-            )
             return None
 
         with open(VEFAAS_IAM_CRIDENTIAL_PATH, "r") as f:
@@ -233,6 +215,9 @@ def _assume_role(
     ) -> AssumeRoleCredential:
         """Execute AssumeRole to get STS temporary credentials.
 
+        This method performs the AssumeRole operation and caches the result.
+        Cache validation is handled by the caller (refresh_credentials decorator).
+
         Args:
             access_key: VolcEngine access key
             secret_key: VolcEngine secret key
@@ -244,16 +229,9 @@ def _assume_role(
         Raises:
             Exception: If AssumeRole fails
         """
-        # Check if the cached credentials are still valid
-        if (
-            self._cached_sts_credential is not None
-            and not self._is_sts_credential_expired()
-        ):
-            logger.info("Using cached STS credentials")
-            return self._cached_sts_credential
-
         logger.info(
-            "Cached STS credentials expired or not found, requesting new credentials..."
+            f"Requesting new STS credentials for role: {role_trn}, "
+            f"session: {settings.veidentity.role_session_name}"
         )
 
         # Create STS client configuration
@@ -272,11 +250,7 @@ def _assume_role(
             role_session_name=settings.veidentity.role_session_name,
         )
 
-        logger.info(
-            f"Executing AssumeRole for role: {role_trn}, "
-            f"session: {settings.veidentity.role_session_name}"
-        )
-
+        # Execute AssumeRole
         response: volcenginesdksts.AssumeRoleResponse = sts_client.assume_role(
             assume_role_request
         )
@@ -298,18 +272,19 @@ def _assume_role(
             expires_at_timestamp = calendar.timegm(dt.timetuple())
         except Exception as e:
             logger.warning(f"Failed to parse STS credential expiration time: {e}")
-            # Expires in 1 hour by default
+            # Default to 1 hour expiration
             import time
 
             expires_at_timestamp = int(time.time()) + 3600
 
+        # Create credential object
         sts_credential = AssumeRoleCredential(
             access_key_id=credentials.access_key_id,
             secret_access_key=credentials.secret_access_key,
             session_token=credentials.session_token,
         )
 
-        # Cached credentials and expiration time
+        # Cache credentials and expiration time
         self._cached_sts_credential = sts_credential
         self._sts_credential_expires_at = expires_at_timestamp