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

Feat: aws assume role setup

Author: VisargD

src/handlers/handlerUtils.ts

@@ -269,6 +269,7 @@ export async function tryPostProxy(
     : (providerOption.urlToFetch as string);
 
   const headers = await apiConfig.headers({
+    c,
     providerOptions: providerOption,
     fn,
     transformedRequestBody: params,
@@ -520,6 +521,7 @@ export async function tryPost(
   const url = `${baseUrl}${endpoint}`;
 
   const headers = await apiConfig.headers({
+    c,
     providerOptions: providerOption,
     fn,
     transformedRequestBody,
@@ -1037,6 +1039,9 @@ export function constructConfigFromRequestHeaders(
     awsSecretAccessKey: requestHeaders[`x-${POWERED_BY}-aws-secret-access-key`],
     awsSessionToken: requestHeaders[`x-${POWERED_BY}-aws-session-token`],
     awsRegion: requestHeaders[`x-${POWERED_BY}-aws-region`],
+    awsRoleArn: requestHeaders[`x-${POWERED_BY}-aws-role-arn`],
+    awsAuthType: requestHeaders[`x-${POWERED_BY}-aws-auth-type`],
+    awsExternalId: requestHeaders[`x-${POWERED_BY}-aws-external-id`],
   };
 
   const workersAiConfig = {

src/providers/bedrock/api.ts

@@ -1,12 +1,13 @@
 import { GatewayError } from '../../errors/GatewayError';
 import { ProviderAPIConfig } from '../types';
 import { bedrockInvokeModels } from './constants';
-import { generateAWSHeaders } from './utils';
+import { generateAWSHeaders, getAssumedRoleCredentials } from './utils';
 
 const BedrockAPIConfig: ProviderAPIConfig = {
   getBaseURL: ({ providerOptions }) =>
     `https://bedrock-runtime.${providerOptions.awsRegion || 'us-east-1'}.amazonaws.com`,
   headers: async ({
+    c,
     providerOptions,
     transformedRequestBody,
     transformedRequestUrl,
@@ -15,6 +16,19 @@ const BedrockAPIConfig: ProviderAPIConfig = {
       'content-type': 'application/json',
     };
 
+    if (providerOptions.awsAuthType === 'assumedRole') {
+      const { accessKeyId, secretAccessKey, sessionToken } =
+        (await getAssumedRoleCredentials(
+          c,
+          providerOptions.awsRoleArn || '',
+          providerOptions.awsExternalId || '',
+          providerOptions.awsRegion || ''
+        )) || {};
+      providerOptions.awsAccessKeyId = accessKeyId;
+      providerOptions.awsSecretAccessKey = secretAccessKey;
+      providerOptions.awsSessionToken = sessionToken;
+    }
+
     return generateAWSHeaders(
       transformedRequestBody,
       headers,

src/providers/bedrock/utils.ts

@@ -6,6 +6,8 @@ import {
   BedrockChatCompletionsParams,
   BedrockConverseCohereChatCompletionsParams,
 } from './chatComplete';
+import { Context } from 'hono';
+import { env } from 'hono/adapter';
 
 export const generateAWSHeaders = async (
   body: Record<string, any>,
@@ -154,3 +156,98 @@ export const transformAI21AdditionalModelRequestFields = (
   }
   return additionalModelRequestFields;
 };
+
+export async function getAssumedRoleCredentials(
+  c: Context,
+  awsRoleArn: string,
+  awsExternalId: string,
+  awsRegion: 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';
+  const service = 'sts';
+  const hostname = `sts.${region}.amazonaws.com`;
+  const signer = new SignatureV4({
+    service,
+    region,
+    credentials: {
+      accessKeyId,
+      secretAccessKey,
+    },
+    sha256: Sha256,
+  });
+  const url = `https://${hostname}?Action=AssumeRole&Version=2011-06-15&RoleArn=${awsRoleArn}&ExternalId=${awsExternalId}&RoleSessionName=random`;
+  const urlObj = new URL(url);
+  const requestHeaders = { host: hostname };
+  const options = {
+    method: 'GET',
+    path: urlObj.pathname,
+    protocol: urlObj.protocol,
+    hostname: urlObj.hostname,
+    headers: requestHeaders,
+    query: Object.fromEntries(urlObj.searchParams),
+  };
+  const { headers } = await signer.sign(options);
+
+  let credentials: any;
+  try {
+    const response = await fetch(url, {
+      method: 'GET',
+      headers: headers,
+    });
+
+    if (!response.ok) {
+      const resp = await response.text();
+      console.error({ message: resp });
+      throw new Error(`HTTP error! status: ${response.status}`);
+    }
+
+    const xmlData = await response.text();
+    credentials = parseXml(xmlData);
+    if (putInCacheWithValue) {
+      putInCacheWithValue(env(c), cacheKey, credentials, 60); //1 minute
+    }
+  } catch (error) {
+    console.error({ message: `Error assuming role:, ${error}` });
+  }
+
+  return credentials;
+}
+
+function parseXml(xml: string) {
+  // Simple XML parser for this specific use case
+  const getTagContent = (tag: string) => {
+    const regex = new RegExp(`<${tag}>(.*?)</${tag}>`, 's');
+    const match = xml.match(regex);
+    return match ? match[1] : null;
+  };
+
+  const credentials = getTagContent('Credentials');
+  if (!credentials) {
+    throw new Error('Failed to parse Credentials from XML response');
+  }
+
+  return {
+    accessKeyId: getTagContent('AccessKeyId'),
+    secretAccessKey: getTagContent('SecretAccessKey'),
+    sessionToken: getTagContent('SessionToken'),
+    expiration: getTagContent('Expiration'),
+  };
+}

src/providers/types.ts

@@ -1,3 +1,4 @@
+import { Context } from 'hono';
 import { Message, Options, Params } from '../types/requestBody';
 
 /**
@@ -35,6 +36,7 @@ export interface ProviderConfig {
 export interface ProviderAPIConfig {
   /** A function to generate the headers for the API request. */
   headers: (args: {
+    c: Context;
     providerOptions: Options;
     fn: string;
     transformedRequestBody: Record<string, any>;

src/types/requestBody.ts

@@ -78,6 +78,9 @@ export interface Options {
   awsAccessKeyId?: string;
   awsSessionToken?: string;
   awsRegion?: string;
+  awsAuthType?: string;
+  awsRoleArn?: string;
+  awsExternalId?: string;
 
   /** Stability AI specific */
   stabilityClientId?: string;