Merge pull request #714 from Portkey-AI/feat/aws-assume-role-setup

chore: update aws bedrock assume role logic to allow 2 step process f…

Author: VisargD

src/providers/bedrock/api.ts

@@ -1,3 +1,4 @@
+import { env } from 'hono/adapter';
 import { GatewayError } from '../../errors/GatewayError';
 import { ProviderAPIConfig } from '../types';
 import { bedrockInvokeModels } from './constants';
@@ -17,16 +18,38 @@ const BedrockAPIConfig: ProviderAPIConfig = {
     };
 
     if (providerOptions.awsAuthType === 'assumedRole') {
-      const { accessKeyId, secretAccessKey, sessionToken } =
-        (await getAssumedRoleCredentials(
+      try {
+        // Assume the role in the source account
+        const sourceRoleCredentials = await getAssumedRoleCredentials(
           c,
-          providerOptions.awsRoleArn || '',
-          providerOptions.awsExternalId || '',
+          env(c).AWS_ASSUME_ROLE_SOURCE_ARN, // Role ARN in the source account
+          env(c).AWS_ASSUME_ROLE_SOURCE_EXTERNAL_ID || '', // External ID for source role (if needed)
           providerOptions.awsRegion || ''
-        )) || {};
-      providerOptions.awsAccessKeyId = accessKeyId;
-      providerOptions.awsSecretAccessKey = secretAccessKey;
-      providerOptions.awsSessionToken = sessionToken;
+        );
+
+        if (!sourceRoleCredentials) {
+          throw new Error('Server Error while assuming internal role');
+        }
+
+        // Assume role in destination account using temporary creds obtained in first step
+        const { accessKeyId, secretAccessKey, sessionToken } =
+          (await getAssumedRoleCredentials(
+            c,
+            providerOptions.awsRoleArn || '',
+            providerOptions.awsExternalId || '',
+            providerOptions.awsRegion || '',
+            {
+              accessKeyId: sourceRoleCredentials.accessKeyId,
+              secretAccessKey: sourceRoleCredentials.secretAccessKey,
+              sessionToken: sourceRoleCredentials.sessionToken,
+            }
+          )) || {};
+        providerOptions.awsAccessKeyId = accessKeyId;
+        providerOptions.awsSecretAccessKey = secretAccessKey;
+        providerOptions.awsSessionToken = sessionToken;
+      } catch (e) {
+        throw new GatewayError('Error while assuming bedrock role');
+      }
     }
 
     return generateAWSHeaders(

src/providers/bedrock/utils.ts

@@ -161,27 +161,43 @@ export async function getAssumedRoleCredentials(
   c: Context,
   awsRoleArn: string,
   awsExternalId: string,
-  awsRegion: string
+  awsRegion: string,
+  creds?: {
+    accessKeyId: string;
+    secretAccessKey: string;
+    sessionToken?: string;
+  }
 ) {
   const cacheKey = `${awsRoleArn}/${awsExternalId}/${awsRegion}`;
   const getFromCacheByKey = c.get('getFromCacheByKey');
   const putInCacheWithValue = c.get('putInCacheWithValue');
 
-  const {
-    AWS_ASSUME_ROLE_ACCESS_KEY_ID,
-    AWS_ASSUME_ROLE_SECRET_ACCESS_KEY,
-    AWS_ASSUME_ROLE_REGION,
-  } = env(c);
   const resp = getFromCacheByKey
     ? await getFromCacheByKey(env(c), cacheKey)
     : null;
   if (resp) {
     return resp;
   }
-  // Long-term credentials to assume role, static values from ENV
-  const accessKeyId: string = AWS_ASSUME_ROLE_ACCESS_KEY_ID || '';
-  const secretAccessKey: string = AWS_ASSUME_ROLE_SECRET_ACCESS_KEY || '';
-  const region = awsRegion || AWS_ASSUME_ROLE_REGION || 'us-east-1';
+
+  // Determine which credentials to use
+  let accessKeyId: string;
+  let secretAccessKey: string;
+  let sessionToken: string | undefined;
+
+  if (creds) {
+    // Use provided credentials
+    accessKeyId = creds.accessKeyId;
+    secretAccessKey = creds.secretAccessKey;
+    sessionToken = creds.sessionToken;
+  } else {
+    // Use environment credentials
+    const { AWS_ASSUME_ROLE_ACCESS_KEY_ID, AWS_ASSUME_ROLE_SECRET_ACCESS_KEY } =
+      env(c);
+    accessKeyId = AWS_ASSUME_ROLE_ACCESS_KEY_ID || '';
+    secretAccessKey = AWS_ASSUME_ROLE_SECRET_ACCESS_KEY || '';
+  }
+
+  const region = awsRegion || 'us-east-1';
   const service = 'sts';
   const hostname = `sts.${region}.amazonaws.com`;
   const signer = new SignatureV4({
@@ -190,10 +206,13 @@ export async function getAssumedRoleCredentials(
     credentials: {
       accessKeyId,
       secretAccessKey,
+      sessionToken,
     },
     sha256: Sha256,
   });
-  const url = `https://${hostname}?Action=AssumeRole&Version=2011-06-15&RoleArn=${awsRoleArn}&ExternalId=${awsExternalId}&RoleSessionName=random`;
+  const date = new Date();
+  const sessionName = `${date.getFullYear()}${date.getMonth()}${date.getDay()}`;
+  const url = `https://${hostname}?Action=AssumeRole&Version=2011-06-15&RoleArn=${awsRoleArn}&RoleSessionName=${sessionName}${awsExternalId ? `&ExternalId=${awsExternalId}` : ''}`;
   const urlObj = new URL(url);
   const requestHeaders = { host: hostname };
   const options = {
@@ -227,7 +246,6 @@ export async function getAssumedRoleCredentials(
   } catch (error) {
     console.error({ message: `Error assuming role:, ${error}` });
   }
-
   return credentials;
 }