Add support for assuming AWS role via Azure VM managed identity

packages/openops/src/lib/aws/azure-aws-federation.ts

@@ -0,0 +1,117 @@
+import {
+  AssumeRoleCommand,
+  AssumeRoleWithWebIdentityCommand,
+  Credentials,
+  STSClient,
+} from '@aws-sdk/client-sts';
+import { logger, SharedSystemProp, system } from '@openops/server-shared';
+import { v4 as uuidv4 } from 'uuid';
+import { getAwsClient } from './get-client';
+
+let cachedCredentials: {
+  credentials: Credentials;
+  expiresAt: number;
+} | null = null;
+
+export function clearAzureFederationCache() {
+  cachedCredentials = null;
+}
+
+export async function assumeTargetRoleViaAzureFederation(
+  defaultRegion: string,
+  roleArn: string,
+  externalId?: string,
+  endpoint?: string | undefined | null,
+): Promise<Credentials | undefined> {
+  const sourceCredentials = await getAwsCredentialsFromAzureIdentity(
+    defaultRegion,
+  );
+
+  if (!sourceCredentials?.AccessKeyId || !sourceCredentials.SecretAccessKey) {
+    throw new Error('Failed to get AWS credentials from Azure identity');
+  }
+
+  const client = getAwsClient(
+    STSClient,
+    {
+      accessKeyId: sourceCredentials.AccessKeyId,
+      secretAccessKey: sourceCredentials.SecretAccessKey,
+      sessionToken: sourceCredentials.SessionToken,
+      endpoint,
+    },
+    defaultRegion,
+  );
+
+  const command = new AssumeRoleCommand({
+    RoleArn: roleArn,
+    ExternalId: externalId || undefined,
+    RoleSessionName: 'openops-' + uuidv4(),
+  });
+
+  const response = await client.send(command);
+
+  return response.Credentials;
+}
+
+export async function getAwsCredentialsFromAzureIdentity(
+  defaultRegion: string,
+): Promise<Credentials | undefined> {
+  const now = Date.now();
+  const buffer = 5 * 60 * 1000;
+
+  if (cachedCredentials && cachedCredentials.expiresAt > now + buffer) {
+    return cachedCredentials.credentials;
+  }
+
+  const webIdentityToken = await getAzureOidcTokenForAws();
+  const client = new STSClient({
+    region: defaultRegion,
+  });
+
+  const federationRoleArn = system.getOrThrow<string>(
+    SharedSystemProp.AWS_AZURE_FEDERATION_ROLE_ARN,
+  );
+
+  const command = new AssumeRoleWithWebIdentityCommand({
+    RoleArn: federationRoleArn,
+    RoleSessionName: 'openops-' + uuidv4(),
+    WebIdentityToken: webIdentityToken,
+  });
+
+  const response = await client.send(command);
+
+  if (response.Credentials) {
+    cachedCredentials = {
+      credentials: response.Credentials,
+      expiresAt: response.Credentials.Expiration
+        ? new Date(response.Credentials.Expiration).getTime()
+        : now + 3600 * 1000,
+    };
+  }
+
+  return response.Credentials;
+}
+
+async function getAzureOidcTokenForAws(): Promise<string> {
+  const resource = 'api://AzureADTokenExchange';
+
+  const url =
+    `http://169.254.169.254/metadata/identity/oauth2/token` +
+    `?api-version=2018-02-01` +
+    `&resource=${encodeURIComponent(resource)}`;
+
+  const response = await fetch(url, {
+    headers: {
+      Metadata: 'true',
+    },
+  });
+
+  if (!response.ok) {
+    logger.info('Failed to get Azure managed identity token.', response);
+    throw new Error('Failed to get Azure managed identity token.');
+  }
+
+  const data = (await response.json()) as { access_token: string };
+
+  return data.access_token;
+}

packages/openops/src/lib/aws/get-client.ts

@@ -1,26 +1,44 @@
 import { SharedSystemProp, system } from '@openops/server-shared';
 import { AwsCredentials } from './auth';
+import { getAwsCredentialsFromAzureIdentity } from './azure-aws-federation';
+
+type AwsClientConfig = {
+  region: string;
+  credentials?: AwsCredentials | (() => Promise<AwsCredentials>);
+  endpoint?: string;
+};
+
+type CachedAwsCredentials = AwsCredentials & {
+  expiration?: Date;
+};
+
+const azureCredentialCache = new Map<
+  string,
+  {
+    credentials: CachedAwsCredentials | null;
+    promise: Promise<CachedAwsCredentials> | null;
+  }
+>();
 
 export function getAwsClient<T>(
-  ClientConstructor: new (config: {
-    region: string;
-    credentials: AwsCredentials | undefined;
-    endpoint?: string;
-  }) => T,
+  ClientConstructor: new (config: AwsClientConfig) => T,
   credentials: AwsCredentials,
   region: string,
 ): T {
-  const config: any = { region };
+  const config: AwsClientConfig = {
+    region,
+  };
+
   if (credentials.accessKeyId) {
-    config.credentials = {
-      accessKeyId: credentials.accessKeyId,
-      secretAccessKey: credentials.secretAccessKey,
-      sessionToken: credentials.sessionToken,
-    };
+    config.credentials = createStaticCredentials(credentials);
   } else if (!system.getBoolean(SharedSystemProp.AWS_ENABLE_IMPLICIT_ROLE)) {
     throw new Error(
       'AWS credentials are required, please provide accessKeyId and secretAccessKey',
     );
+  } else if (
+    system.getBoolean(SharedSystemProp.AWS_USE_AZURE_MANAGED_IDENTITY)
+  ) {
+    config.credentials = createAzureManagedIdentityCredentialsProvider(region);
   }
 
   if (credentials.endpoint) {
@@ -29,3 +47,85 @@ export function getAwsClient<T>(
 
   return new ClientConstructor(config);
 }
+
+function createStaticCredentials(credentials: AwsCredentials): AwsCredentials {
+  return {
+    accessKeyId: credentials.accessKeyId,
+    secretAccessKey: credentials.secretAccessKey,
+    sessionToken: credentials.sessionToken,
+  };
+}
+
+function createAzureManagedIdentityCredentialsProvider(
+  region: string,
+): () => Promise<CachedAwsCredentials> {
+  const cache = getOrCreateAzureCredentialCache(region);
+
+  return async () => {
+    if (hasValidCredentials(cache.credentials)) {
+      return cache.credentials;
+    }
+
+    if (cache.promise) {
+      return cache.promise;
+    }
+
+    cache.promise = fetchAzureManagedIdentityCredentials(region);
+
+    try {
+      cache.credentials = await cache.promise;
+      return cache.credentials;
+    } finally {
+      cache.promise = null;
+    }
+  };
+}
+
+function getOrCreateAzureCredentialCache(region: string) {
+  let cache = azureCredentialCache.get(region);
+
+  if (!cache) {
+    cache = {
+      credentials: null,
+      promise: null,
+    };
+
+    azureCredentialCache.set(region, cache);
+  }
+
+  return cache;
+}
+
+function hasValidCredentials(
+  credentials: CachedAwsCredentials | null,
+): credentials is CachedAwsCredentials {
+  if (!credentials) {
+    return false;
+  }
+
+  if (!credentials.expiration) {
+    return true;
+  }
+
+  // Refresh 1 minute before expiration
+  return credentials.expiration.getTime() > Date.now() + 60_000;
+}
+
+async function fetchAzureManagedIdentityCredentials(
+  region: string,
+): Promise<CachedAwsCredentials> {
+  const stsCredentials = await getAwsCredentialsFromAzureIdentity(region);
+
+  if (!stsCredentials?.AccessKeyId || !stsCredentials?.SecretAccessKey) {
+    throw new Error(
+      'Failed to obtain AWS credentials from Azure managed identity',
+    );
+  }
+
+  return {
+    accessKeyId: stsCredentials.AccessKeyId,
+    secretAccessKey: stsCredentials.SecretAccessKey,
+    sessionToken: stsCredentials.SessionToken,
+    expiration: stsCredentials.Expiration,
+  };
+}

packages/openops/src/lib/aws/sts-common.ts

@@ -4,7 +4,9 @@ import {
   GetCallerIdentityCommand,
   STSClient,
 } from '@aws-sdk/client-sts';
+import { SharedSystemProp, system } from '@openops/server-shared';
 import { v4 as uuidv4 } from 'uuid';
+import { assumeTargetRoleViaAzureFederation } from './azure-aws-federation';
 import { getAwsClient } from './get-client';
 
 export async function getAccountId(
@@ -26,16 +28,31 @@ export async function assumeRole(
   externalId?: string,
   endpoint?: string | undefined | null,
 ): Promise<Credentials | undefined> {
+  if (
+    !accessKeyId &&
+    system.getBoolean(SharedSystemProp.AWS_ENABLE_IMPLICIT_ROLE) &&
+    system.getBoolean(SharedSystemProp.AWS_USE_AZURE_MANAGED_IDENTITY)
+  ) {
+    return assumeTargetRoleViaAzureFederation(
+      defaultRegion,
+      roleArn,
+      externalId,
+      endpoint,
+    );
+  }
+
   const client = getAwsClient(
     STSClient,
     { accessKeyId, secretAccessKey, endpoint },
     defaultRegion,
   );
+
   const command = new AssumeRoleCommand({
     RoleArn: roleArn,
     ExternalId: externalId || undefined,
     RoleSessionName: 'openops-' + uuidv4(),
   });
+
   const response = await client.send(command);
 
   return response.Credentials;