From 6253db46441d015aeca001d76ec88d92e363995b Mon Sep 17 00:00:00 2001
From: "Alexander Nicholson 4584443+DragonStuff@users.noreply.github.com"
 <4584443+DragonStuff@users.noreply.github.com>
Date: Wed, 20 Nov 2024 10:58:15 +0900
Subject: [PATCH] feature: add IRSA and instance role support

---
 src/services/storage.ts | 130 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 118 insertions(+), 12 deletions(-)

diff --git a/src/services/storage.ts b/src/services/storage.ts
index 11eec9d..2550b61 100644
--- a/src/services/storage.ts
+++ b/src/services/storage.ts
@@ -3,30 +3,99 @@ import { S3, type ListObjectsV2Request } from "https://deno.land/x/aws_api@v0.8.
 import { ApiFactory } from "https://deno.land/x/aws_api@v0.8.1/client/mod.ts";
 import { getSignedUrl } from "https://deno.land/x/aws_s3_presign@2.2.1/mod.ts";
 
+interface Credentials {
+  awsAccessKeyId: string;
+  awsSecretKey: string;
+  sessionToken?: string;
+}
+
 export class StorageService {
   private client: S3;
+  private credentials: Credentials;
 
-  constructor(
-    private bucket: string,
-    private region: string,
-    private accessKeyId: string,
-    private secretAccessKey: string,
-  ) {
+  constructor(private bucket: string, private region: string) {
+    this.credentials = this.resolveCredentials();
     const factory = new ApiFactory({
       region: this.region,
-      credentials: {
-        awsAccessKeyId: this.accessKeyId,
-        awsSecretKey: this.secretAccessKey,
-      },
+      credentials: this.credentials,
     });
     
     this.client = new S3(factory);
   }
 
+  private async fetchInstanceProfileCredentials(): Promise<Credentials | null> {
+    try {
+      const metadataUrl = "http://169.254.169.254/latest/meta-data/iam/security-credentials/";
+      const roleName = await (await fetch(metadataUrl)).text();
+      const credentials = await (await fetch(`${metadataUrl}${roleName}`)).json();
+      
+      return {
+        awsAccessKeyId: credentials.AccessKeyId,
+        awsSecretKey: credentials.SecretAccessKey,
+        sessionToken: credentials.Token,
+      };
+    } catch {
+      return null;
+    }
+  }
+
+  private async fetchIRSACredentials(): Promise<Credentials | null> {
+    const roleArn = Deno.env.get("AWS_ROLE_ARN");
+    const tokenFile = Deno.env.get("AWS_WEB_IDENTITY_TOKEN_FILE");
+
+    if (!roleArn || !tokenFile) {
+      return null;
+    }
+
+    try {
+      const token = await Deno.readTextFile(tokenFile);
+      const sts = new WebIdentityCredentials(roleArn, token);
+      const creds = await sts.getCredentials();
+      
+      return {
+        awsAccessKeyId: creds.accessKeyId,
+        awsSecretKey: creds.secretAccessKey,
+        sessionToken: creds.sessionToken,
+      };
+    } catch {
+      return null;
+    }
+  }
+
+  private getStaticCredentials(): Credentials | null {
+    const accessKeyId = Deno.env.get("AWS_ACCESS_KEY_ID");
+    const secretKey = Deno.env.get("AWS_SECRET_ACCESS_KEY");
+
+    if (!accessKeyId || !secretKey) {
+      return null;
+    }
+
+    return {
+      awsAccessKeyId: accessKeyId,
+      awsSecretKey: secretKey,
+    };
+  }
+
+  private async resolveCredentials(): Promise<Credentials> {
+    // Try IRSA first
+    const irsaCreds = await this.fetchIRSACredentials();
+    if (irsaCreds) return irsaCreds;
+
+    // Then try static credentials
+    const staticCreds = this.getStaticCredentials();
+    if (staticCreds) return staticCreds;
+
+    // Finally try instance profile
+    const instanceCreds = await this.fetchInstanceProfileCredentials();
+    if (instanceCreds) return instanceCreds;
+
+    throw new Error("No valid AWS credentials found");
+  }
+
   async getSignedUrl(key: string, expiresIn = 3600): Promise<string> {
     return await getSignedUrl({
-      accessKeyId: this.accessKeyId,
-      secretAccessKey: this.secretAccessKey,
+      accessKeyId: this.credentials.awsAccessKeyId,
+      secretAccessKey: this.credentials.awsSecretKey,
       bucket: this.bucket,
       key,
       region: this.region,
@@ -83,4 +152,41 @@ export class StorageService {
       return undefined;
     }
   }
+}
+
+class WebIdentityCredentials {
+  constructor(
+    private roleArn: string,
+    private token: string,
+  ) {}
+
+  async getCredentials(): Promise<{ accessKeyId: string; secretAccessKey: string; sessionToken: string }> {
+    const params = new URLSearchParams({
+      Version: "2011-06-15",
+      Action: "AssumeRoleWithWebIdentity",
+      RoleArn: this.roleArn,
+      RoleSessionName: `deno-session-${Date.now()}`,
+      WebIdentityToken: this.token,
+    });
+
+    const response = await fetch(`https://sts.amazonaws.com?${params}`, {
+      method: "GET",
+    });
+
+    if (!response.ok) {
+      throw new Error(`Failed to assume role: ${await response.text()}`);
+    }
+
+    const xml = await response.text();
+    const result = new DOMParser().parseFromString(xml, "text/xml");
+    
+    const credentials = result.querySelector("Credentials");
+    if (!credentials) throw new Error("No credentials in response");
+
+    return {
+      accessKeyId: credentials.querySelector("AccessKeyId")?.textContent ?? "",
+      secretAccessKey: credentials.querySelector("SecretAccessKey")?.textContent ?? "",
+      sessionToken: credentials.querySelector("SessionToken")?.textContent ?? "",
+    };
+  }
 }