feat: cleanup plugin & support different auth modes for plugin based …

plugins/bedrock/index.ts

@@ -1,18 +1,26 @@
-import { PluginHandler } from '../types';
+import { getRuntimeKey } from 'hono/adapter';
+import { PluginHandler, PluginOptions } from '../types';
 import {
   getCurrentContentPart,
   HttpError,
   setCurrentContentPart,
 } from '../utils';
 import { BedrockAccessKeyCreds, BedrockBody, BedrockParameters } from './type';
-import { bedrockPost, getAssumedRoleCredentials, redactPii } from './util';
+import {
+  bedrockPost,
+  getAssumedRoleCredentials,
+  getRegionFromEnv,
+  redactPii,
+} from './util';
 
 const REQUIRED_CREDENTIAL_KEYS = [
   'awsAccessKeyId',
   'awsSecretAccessKey',
   'awsRegion',
 ];
 
+const runtime = getRuntimeKey();
+
 export const validateCreds = (
   credentials?: BedrockParameters['credentials']
 ) => {
@@ -22,58 +30,83 @@ export const validateCreds = (
 };
 
 export const handleCredentials = async (
-  options: Record<string, any>,
-  credentials: BedrockParameters['credentials'] | null
+  credentials: BedrockParameters['credentials'] | null,
+  options?: PluginOptions
 ) => {
   const finalCredentials = {} as BedrockAccessKeyCreds;
   if (credentials?.awsAuthType === 'assumedRole') {
-    try {
-      // Assume the role in the source account
-      const sourceRoleCredentials = await getAssumedRoleCredentials(
-        options.getFromCacheByKey,
-        options.putInCacheWithValue,
-        options.env,
-        options.env.AWS_ASSUME_ROLE_SOURCE_ARN, // Role ARN in the source account
-        options.env.AWS_ASSUME_ROLE_SOURCE_EXTERNAL_ID || '', // External ID for source role (if needed)
-        credentials.awsRegion || ''
-      );
+    if (runtime === 'workerd') {
+      try {
+        // Assume the role in the source account
+        const sourceRoleCredentials = await getAssumedRoleCredentials(
+          options?.env.AWS_ASSUME_ROLE_SOURCE_ARN, // Role ARN in the source account
+          options?.env.AWS_ASSUME_ROLE_SOURCE_EXTERNAL_ID || '', // External ID for source role (if needed)
+          credentials.awsRegion || '',
+          undefined,
+          options
+        );
 
-      if (!sourceRoleCredentials) {
-        throw new Error('Failed to assume internal role');
-      }
+        if (!sourceRoleCredentials) {
+          throw new Error('Failed to assume internal role');
+        }
 
-      // Assume role in destination account using temporary creds obtained in first step
-      const destinationCredentials =
+        // Assume role in destination account using temporary creds obtained in first step
+        const destinationCredentials =
+          (await getAssumedRoleCredentials(
+            credentials.awsRoleArn || '',
+            credentials.awsExternalId || '',
+            credentials.awsRegion || '',
+            {
+              accessKeyId: sourceRoleCredentials.accessKeyId,
+              secretAccessKey: sourceRoleCredentials.secretAccessKey,
+              sessionToken: sourceRoleCredentials.sessionToken,
+            },
+            options
+          )) || {};
+        if (!destinationCredentials) {
+          throw new Error('Failed to assume destination role');
+        }
+        finalCredentials.awsAccessKeyId = destinationCredentials.accessKeyId;
+        finalCredentials.awsSecretAccessKey =
+          destinationCredentials.secretAccessKey;
+        finalCredentials.awsSessionToken = destinationCredentials.sessionToken;
+        finalCredentials.awsRegion = credentials.awsRegion || '';
+      } catch {
+        throw new Error('Error while assuming role');
+      }
+    }
+    if (runtime === 'node') {
+      const { accessKeyId, secretAccessKey, sessionToken } =
         (await getAssumedRoleCredentials(
-          options.getFromCacheByKey,
-          options.putInCacheWithValue,
-          options.env,
           credentials.awsRoleArn || '',
           credentials.awsExternalId || '',
-          credentials.awsRegion || '',
-          {
-            accessKeyId: sourceRoleCredentials.accessKeyId,
-            secretAccessKey: sourceRoleCredentials.secretAccessKey,
-            sessionToken: sourceRoleCredentials.sessionToken,
-          }
+          credentials.awsRegion || ''
         )) || {};
-      if (!destinationCredentials) {
-        throw new Error('Failed to assume destination role');
-      }
-      finalCredentials.awsAccessKeyId = destinationCredentials.accessKeyId;
-      finalCredentials.awsSecretAccessKey =
-        destinationCredentials.secretAccessKey;
-      finalCredentials.awsSessionToken = destinationCredentials.sessionToken;
-      finalCredentials.awsRegion = credentials.awsRegion || '';
-    } catch {
-      throw new Error('Error while assuming role');
+
+      finalCredentials.awsAccessKeyId = accessKeyId;
+      finalCredentials.awsSecretAccessKey = secretAccessKey;
+      finalCredentials.awsSessionToken = sessionToken;
+      finalCredentials.awsRegion = credentials?.awsRegion;
     }
+  } else if (credentials?.awsAuthType === 'serviceRole' && runtime === 'node') {
+    const { accessKeyId, secretAccessKey, sessionToken, awsRegion } =
+      await getAssumedRoleCredentials();
+
+    finalCredentials.awsAccessKeyId = accessKeyId;
+    finalCredentials.awsSecretAccessKey = secretAccessKey;
+    finalCredentials.awsSessionToken = sessionToken;
+    finalCredentials.awsRegion = awsRegion;
   } else {
     finalCredentials.awsAccessKeyId = credentials?.awsAccessKeyId || '';
     finalCredentials.awsSecretAccessKey = credentials?.awsSecretAccessKey || '';
     finalCredentials.awsSessionToken = credentials?.awsSessionToken || '';
     finalCredentials.awsRegion = credentials?.awsRegion || '';
   }
+
+  if (!finalCredentials.awsRegion) {
+    finalCredentials.awsRegion = getRegionFromEnv(options);
+  }
+
   return finalCredentials;
 };
 
@@ -102,10 +135,7 @@ export const pluginHandler: PluginHandler<
 
   try {
     const credentials = parameters.credentials || null;
-    const finalCredentials = await handleCredentials(
-      options as Record<string, any>,
-      credentials
-    );
+    const finalCredentials = await handleCredentials(credentials, options);
     const validate = validateCreds(finalCredentials);
 
     const guardrailVersion = parameters.guardrailVersion;
@@ -152,7 +182,7 @@ export const pluginHandler: PluginHandler<
     const interventionData = results.find(
       (result) => result && result.action === 'GUARDRAIL_INTERVENED'
     );
-    if (interventionData) {
+    if (interventionData?.action) {
       verdict = false;
     }
 
@@ -168,9 +198,8 @@ export const pluginHandler: PluginHandler<
       if (result.assessments[0]?.wordPolicy) {
         flaggedCategories.add('wordFilter');
       }
-
       if (hasTriggeredPII) {
-        continue;
+        break;
       }
 
       const sensitiveInfo = result.assessments[0]?.sensitiveInformationPolicy;

plugins/bedrock/type.ts

@@ -7,8 +7,11 @@ export type BedrockAccessKeyCreds = {
 };
 
 export type BedrockAssumedRoleCreds = {
-  awsAuthType: 'assumedRole';
+  awsAuthType: 'assumedRole' | 'serviceRole';
   awsRoleArn: string;
+  awsAccessKeyId: string;
+  awsSecretAccessKey: string;
+  awsSessionToken?: string;
   awsExternalId: string;
   awsRegion: string;
 };