[Integration][AWS] | Fix InvalidToken Exceptions Due to Improper Token Refresh Calls

integrations/aws/.port/spec.yaml

@@ -42,6 +42,11 @@ configurations:
     require: false
     description: The number of concurrent accounts to scan. By default, it is set to 50.
     default: 50
+  - name: assumeRoleDuration
+    type: integer
+    require: false
+    description: The duration in seconds for which the credentials are valid. By default, it is set to 3600 seconds.
+    default: 3600
 deploymentMethodRequirements:
   - type: default
     configurations: ['awsAccessKeyId', 'awsSecretAccessKey']

integrations/aws/CHANGELOG.md

@@ -7,13 +7,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 <!-- towncrier release notes start -->
 
+## 0.2.64 (2024-11-28)
+
+
+### Bug Fixes
+
+- Fixed a bug where token refresh fails because its triggered while an active session was still using the old token.
+
+
 ## 0.2.63 (2024-11-25)
 
 
 ### Bug Fixes
 
 - Do not break delete entities when a region is not accessible
 
+
 ## 0.2.62 (2024-11-25)
 
 

integrations/aws/aws/session_manager.py

@@ -15,9 +15,6 @@ class AccountNotFoundError(OceanAbortException):
     pass
 
 
-ASSUME_ROLE_DURATION_SECONDS = 3600  # 1 hour
-
-
 class SessionManager:
     def __init__(self) -> None:
         """
@@ -98,7 +95,7 @@ async def _get_organization_session(self) -> aioboto3.Session | None:
                 organizations_client = await sts_client.assume_role(
                     RoleArn=organization_role_arn,
                     RoleSessionName="OceanOrgAssumeRoleSession",
-                    DurationSeconds=ASSUME_ROLE_DURATION_SECONDS,
+                    DurationSeconds=self._assume_role_duration_seconds(),
                 )
 
                 credentials = organizations_client["Credentials"]
@@ -121,6 +118,10 @@ async def _get_organization_session(self) -> aioboto3.Session | None:
     def _get_account_read_role_name(self) -> str:
         return ocean.integration_config.get("account_read_role_name", "")
 
+    @staticmethod
+    def _assume_role_duration_seconds() -> int:
+        return int(ocean.integration_config["assume_role_duration"])
+
     async def _update_available_access_credentials(self) -> None:
         logger.info("Updating AWS credentials")
         async with (
@@ -156,7 +157,7 @@ async def _assume_role_and_update_credentials(
             account_role = await sts_client.assume_role(
                 RoleArn=f'arn:aws:iam::{account["Id"]}:role/{self._get_account_read_role_name()}',
                 RoleSessionName="OceanMemberAssumeRoleSession",
-                DurationSeconds=ASSUME_ROLE_DURATION_SECONDS,
+                DurationSeconds=self._assume_role_duration_seconds(),
             )
             raw_credentials = account_role["Credentials"]
             credentials = AwsCredentials(

integrations/aws/main.py

@@ -89,6 +89,9 @@ async def resync_resources_for_account(
     aws_resource_config = typing.cast(AWSResourceConfig, event.resource_config)
 
     if is_global_resource(kind):
+        logger.info(
+            f"Handling global resource {kind} for account {credentials.account_id}"
+        )
         async for batch in _handle_global_resource_resync(
             kind, credentials, aws_resource_config
         ):
@@ -119,6 +122,7 @@ async def resync_all(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
         return
 
     await update_available_access_credentials()
+
     tasks = [
         semaphore_async_iterator(
             semaphore,
@@ -128,7 +132,6 @@ async def resync_all(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     ]
     if tasks:
         async for batch in stream_async_iterators_tasks(*tasks):
-            await update_available_access_credentials()
             yield batch
 
 
@@ -162,7 +165,6 @@ async def resync_elasticache(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     ]
     if tasks:
         async for batch in stream_async_iterators_tasks(*tasks):
-            await update_available_access_credentials()
             yield batch
 
 
@@ -190,7 +192,6 @@ async def resync_elv2_load_balancer(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
 
     if tasks:
         async for batch in stream_async_iterators_tasks(*tasks):
-            await update_available_access_credentials()
             yield batch
 
 
@@ -218,7 +219,6 @@ async def resync_acm(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
 
     if tasks:
         async for batch in stream_async_iterators_tasks(*tasks):
-            await update_available_access_credentials()
             yield batch
 
 
@@ -246,7 +246,6 @@ async def resync_ami(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     ]
     if tasks:
         async for batch in stream_async_iterators_tasks(*tasks):
-            await update_available_access_credentials()
             yield batch
 
 
@@ -274,7 +273,6 @@ async def resync_cloudformation(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
 
     if tasks:
         async for batch in stream_async_iterators_tasks(*tasks):
-            await update_available_access_credentials()
             yield batch
 
 

integrations/aws/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aws"
-version = "0.2.63"
+version = "0.2.64"
 description = "This integration will map all your resources in all the available accounts to your Port entities"
 authors = ["Shalev Avhar <[email protected]>", "Erik Zaadi <[email protected]>"]
 

integrations/aws/tests/utils/test_aws.py

@@ -23,8 +23,12 @@ async def _create_iterator_tasks(func: Any, count: int) -> List[Any]:
 
     @patch("utils.aws._session_manager.reset", new_callable=AsyncMock)
     @patch("utils.aws.lock", new_callable=AsyncMock)
+    @patch(
+        "port_ocean.context.ocean.PortOceanContext.integration_config",
+        return_value={"assume_role_duration": 3600},
+    )
     async def test_multiple_task_execution(
-        self, mock_lock: AsyncMock, mock_reset: AsyncMock
+        self, integration_config_mock: Any, mock_lock: AsyncMock, mock_reset: AsyncMock
     ) -> None:
         tasks: List[Any] = await self._create_iterator_tasks(
             self._run_update_access_iterator_result, 10
@@ -35,20 +39,22 @@ async def test_multiple_task_execution(
         # Assert that the reset method was awaited exactly once (i.e., no thundering herd)
         mock_reset.assert_awaited_once()
 
-        mock_lock.__aenter__.assert_awaited_once()
-        mock_lock.__aexit__.assert_awaited_once()
+        self.assertEqual(mock_lock.__aenter__.call_count, 10)
+        self.assertEqual(mock_lock.__aexit__.call_count, 10)
 
 
 class TestAwsSessions(unittest.IsolatedAsyncioTestCase):
     def setUp(self) -> None:
-        self.session_manager_mock: AsyncMock = patch(
-            "utils.aws._session_manager", autospec=SessionManager
-        ).start()
+        patcher = patch("utils.aws._session_manager", autospec=SessionManager)
+        self.session_manager_mock = patcher.start()
+        self.addCleanup(patcher.stop)
+
+        self.session_manager_mock._assume_role_duration_seconds.return_value = 3600
 
         self.credentials_mock: AsyncMock = AsyncMock(spec=AwsCredentials)
         self.session_mock: AsyncMock = AsyncMock(spec=Session)
 
-    def tearDown(self) -> None:
+    async def asyncTearDown(self) -> None:
         patch.stopall()
 
     async def test_get_sessions_with_custom_account_id(self) -> None:

integrations/aws/utils/aws.py

@@ -4,34 +4,33 @@
 from port_ocean.context.ocean import ocean
 from starlette.requests import Request
 
-from aws.session_manager import SessionManager, ASSUME_ROLE_DURATION_SECONDS
+from aws.session_manager import SessionManager
 from aws.aws_credentials import AwsCredentials
 
-from aiocache import cached, Cache  # type: ignore
+from aiocache import Cache  # type: ignore
 from asyncio import Lock
 
 from port_ocean.utils.async_iterators import stream_async_iterators_tasks
 
+
 _session_manager: SessionManager = SessionManager()
+lock = Lock()
+cache = Cache(Cache.MEMORY)
 
-CACHE_DURATION_SECONDS = (
-    0.80 * ASSUME_ROLE_DURATION_SECONDS
-)  # Refresh role credentials after exhausting 80% of the session duration
 
-lock = Lock()
+def _get_cache_duration_seconds() -> float:
+    return 0.50 * _session_manager._assume_role_duration_seconds()
 
 
-@cached(ttl=CACHE_DURATION_SECONDS, cache=Cache.MEMORY)
 async def update_available_access_credentials() -> bool:
-    """
-    Fetches the AWS account IDs that the current IAM role can access.
-    and saves them up to use as sessions
-
-    :return: List of AWS account IDs.
-    """
+    cache_key = "update_available_access_credentials"
     async with lock:
+        result = await cache.get(cache_key)
+        if result is not None:
+            return result
+
         await _session_manager.reset()
-        # makes this run once per DurationSeconds
+        await cache.set(cache_key, True, ttl=_get_cache_duration_seconds())
         return True
 
 
@@ -50,6 +49,7 @@ async def get_accounts() -> AsyncIterator[AwsCredentials]:
     Gets the AWS account IDs that the current IAM role can access.
     """
     await update_available_access_credentials()
+
     for credentials in _session_manager._aws_credentials:
         yield credentials
 
@@ -78,7 +78,6 @@ async def get_sessions(
     """
     Gets boto3 sessions for the AWS regions.
     """
-    await update_available_access_credentials()
 
     if custom_account_id:
         credentials = _session_manager.find_credentials_by_account_id(custom_account_id)

integrations/aws/utils/misc.py

@@ -6,6 +6,10 @@
 import asyncio
 
 
+def get_assume_role_duration_seconds() -> int:
+    return int(ocean.integration_config["assume_role_duration"])
+
+
 def get_semaphore() -> asyncio.BoundedSemaphore:
     max_concurrent_accounts: int = int(
         ocean.integration_config["maximum_concurrent_accounts"]