From 2d7bf1c6b0307b420081db163ef0605367ececc6 Mon Sep 17 00:00:00 2001
From: Sean Tan <seant@splunk.com>
Date: Thu, 5 Aug 2021 12:39:00 -0700
Subject: [PATCH] Add roleArn to AwsS3IntegrationConfig

Signed-off-by: Sean Tan <seant@splunk.com>
---
 app-config.yaml                               |  1 +
 .../src/reading/AwsS3UrlReader.ts             | 63 +++++++++++++------
 packages/integration/config.d.ts              |  6 ++
 packages/integration/src/awsS3/config.ts      |  8 ++-
 4 files changed, 57 insertions(+), 21 deletions(-)

diff --git a/app-config.yaml b/app-config.yaml
index 9f365becef..a4146a811a 100644
--- a/app-config.yaml
+++ b/app-config.yaml
@@ -167,6 +167,7 @@ integrations:
     - host: amazonaws.com
       accessKeyId: ${AWS_ACCESS_KEY_ID}
       secretAccessKey: ${AWS_SECRET_ACCESS_KEY}
+      #roleArn:
 
 catalog:
   rules:
diff --git a/packages/backend-common/src/reading/AwsS3UrlReader.ts b/packages/backend-common/src/reading/AwsS3UrlReader.ts
index fdeb827f66..7e714a9589 100644
--- a/packages/backend-common/src/reading/AwsS3UrlReader.ts
+++ b/packages/backend-common/src/reading/AwsS3UrlReader.ts
@@ -15,6 +15,7 @@
  */
 
 import aws, { Credentials, S3 } from 'aws-sdk';
+import { CredentialsOptions } from 'aws-sdk/lib/credentials';
 import {
   ReaderFactory,
   ReadTreeResponse,
@@ -25,6 +26,7 @@ import {
 } from './types';
 import getRawBody from 'raw-body';
 import { AwsS3Integration, ScmIntegrations } from '@backstage/integration';
+import { constant } from 'lodash';
 
 const parseURL = (
   url: string,
@@ -60,29 +62,15 @@ const parseURL = (
 };
 
 export class AwsS3UrlReader implements UrlReader {
-  static factory: ReaderFactory = ({ config, logger }) => {
+  static factory: ReaderFactory = ({ config }) => {
     const integrations = ScmIntegrations.fromConfig(config);
 
     return integrations.awsS3.list().map(integration => {
-      let s3: S3;
-      if (
-        !integration.config.accessKeyId ||
-        !integration.config.secretAccessKey
-      ) {
-        logger.debug(
-          'integrations.awsS3 not found in app config. AWS S3 integration will use default AWS credentials if set in environment.',
-        );
-        s3 = new S3({});
-      } else {
-        const creds = new Credentials({
-          accessKeyId: integration.config.accessKeyId,
-          secretAccessKey: integration.config.secretAccessKey,
-        });
-        s3 = new S3({
-          apiVersion: '2006-03-01',
-          credentials: creds,
-        });
-      }
+      const creds = AwsS3UrlReader.buildCredentials(integration);
+      const s3 = new S3({
+        apiVersion: '2006-03-01',
+        credentials: creds,
+      });
       const reader = new AwsS3UrlReader(integration, s3);
       const predicate = (url: URL) =>
         url.host.endsWith(integration.config.host);
@@ -95,6 +83,41 @@ export class AwsS3UrlReader implements UrlReader {
     private readonly s3: S3,
   ) {}
 
+  private static buildCredentials(
+    integration?: AwsS3Integration,
+  ): Credentials | CredentialsOptions | undefined {
+    /*
+    Credentials is an optional config. 
+    If missing, the default ways of authenticating AWS SDK will be used.
+    */
+    if (!integration) {
+      return undefined;
+    }
+
+    const accessKeyId = integration.config.accessKeyId;
+    const secretAccessKey = integration.config.secretAccessKey;
+    let explicitCredentials: Credentials | undefined;
+    if (accessKeyId && secretAccessKey) {
+      explicitCredentials = new Credentials({
+        accessKeyId,
+        secretAccessKey,
+      });
+    }
+
+    const roleArn = integration.config.roleArn;
+    if (roleArn) {
+      return new aws.ChainableTemporaryCredentials({
+        masterCredentials: explicitCredentials,
+        params: {
+          RoleSessionName: 'backstage-aws-s3-url-reader',
+          RoleArn: roleArn,
+        },
+      });
+    }
+
+    return explicitCredentials;
+  }
+
   async read(url: string): Promise<Buffer> {
     const response = await this.readUrl(url);
     return response.buffer();
diff --git a/packages/integration/config.d.ts b/packages/integration/config.d.ts
index a4e07fdd2c..fc6a094aa4 100644
--- a/packages/integration/config.d.ts
+++ b/packages/integration/config.d.ts
@@ -183,6 +183,12 @@ export interface Config {
        * @visibility secret
        */
       secretAccessKey?: string;
+
+      /**
+       * ARN of the role to be assumed
+       * @visibility backend
+       */
+      roleArn?: string;
     }>;
   };
 }
diff --git a/packages/integration/src/awsS3/config.ts b/packages/integration/src/awsS3/config.ts
index e92d1162e2..8900425e57 100644
--- a/packages/integration/src/awsS3/config.ts
+++ b/packages/integration/src/awsS3/config.ts
@@ -40,6 +40,11 @@ export type AwsS3IntegrationConfig = {
    * secretAccessKey
    */
   secretAccessKey?: string;
+
+  /**
+   * roleArn
+   */
+  roleArn?: string;
 };
 
 /**
@@ -54,6 +59,7 @@ export function readAwsS3IntegrationConfig(
   const host = config.getOptionalString('host') ?? AMAZON_AWS_HOST;
   const accessKeyId = config.getOptionalString('accessKeyId');
   const secretAccessKey = config.getOptionalString('secretAccessKey');
+  const roleArn = config.getOptionalString('roleArn');
 
   if (!isValidHost(host)) {
     throw new Error(
@@ -61,7 +67,7 @@ export function readAwsS3IntegrationConfig(
     );
   }
 
-  return { host, accessKeyId, secretAccessKey };
+  return { host, accessKeyId, secretAccessKey, roleArn };
 }
 
 export function readAwsS3IntegrationConfigs(