From ec89c771ce5e343ad6ee6cf47ae68fa7c1ceab3c Mon Sep 17 00:00:00 2001
From: Eric Beard <ericzbeard@gmail.com>
Date: Wed, 30 Oct 2024 14:57:25 -0700
Subject: [PATCH] Allow overriding the expected bucket owner

---
 internal/aws/s3/s3.go     | 27 +++++++++++++++++++++++----
 internal/cmd/rain/rain.go |  1 +
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/internal/aws/s3/s3.go b/internal/aws/s3/s3.go
index 65d7d129..b37d5995 100644
--- a/internal/aws/s3/s3.go
+++ b/internal/aws/s3/s3.go
@@ -30,9 +30,15 @@ import (
 
 const RAIN_BUCKET_SSM_KEY = "rain-bucket"
 
+// BucketName is set by the --s3-bucket param to deploy and pkg commands
 var BucketName = ""
+
+// BucketKeyPrefix is set by the --s3-prefix param to deploy and pkg commands
 var BucketKeyPrefix = ""
 
+// ExpectedBucketOwner is set by the --s3-owner param to deploy and pkg commands
+var ExpectedBucketOwner = ""
+
 func getClient() *s3.Client {
 	return s3.NewFromConfig(aws.Config())
 }
@@ -53,10 +59,23 @@ func BucketHasContents(bucketName string) (bool, error) {
 	return false, nil
 }
 
+func getAccountId() (string, error) {
+
+	accountId := ExpectedBucketOwner
+	var err error
+	if accountId == "" {
+		accountId, err = sts.GetAccountID()
+		if err != nil {
+			return "", err
+		}
+	}
+	return accountId, nil
+}
+
 // BucketExists checks whether the named bucket exists
 func BucketExists(bucketName string) (bool, error) {
 
-	accountId, err := sts.GetAccountID()
+	accountId, err := getAccountId()
 	if err != nil {
 		return false, err
 	}
@@ -167,7 +186,7 @@ func Upload(bucketName string, content []byte) (string, error) {
 
 	key := filepath.Join(BucketKeyPrefix, fmt.Sprintf("%x", sha256.Sum256(content)))
 
-	accountId, err := sts.GetAccountID()
+	accountId, err := getAccountId()
 	if err != nil {
 		return "", err
 	}
@@ -276,7 +295,7 @@ func RainBucket(forceCreation bool) string {
 // GetObject gets an object by key from an S3 bucket
 func GetObject(bucketName string, key string) ([]byte, error) {
 
-	accountId, err := sts.GetAccountID()
+	accountId, err := getAccountId()
 	if err != nil {
 		return nil, err
 	}
@@ -395,7 +414,7 @@ func PutObject(bucketName string, key string, body []byte) error {
 
 	config.Debugf("PutObject final mime type for %s: %s", key, contentType)
 
-	accountId, err := sts.GetAccountID()
+	accountId, err := getAccountId()
 	if err != nil {
 		return err
 	}
diff --git a/internal/cmd/rain/rain.go b/internal/cmd/rain/rain.go
index b2347b4c..cb9a46cb 100644
--- a/internal/cmd/rain/rain.go
+++ b/internal/cmd/rain/rain.go
@@ -76,6 +76,7 @@ func addCommand(label string, profileOptions, bucketOptions bool, c *cobra.Comma
 	if bucketOptions {
 		c.Flags().StringVar(&s3.BucketName, "s3-bucket", "", "Name of the S3 bucket that is used to upload assets")
 		c.Flags().StringVar(&s3.BucketKeyPrefix, "s3-prefix", "", "Prefix to add to objects uploaded to S3 bucket")
+		c.Flags().StringVar(&s3.ExpectedBucketOwner, "s3-owner", "", "The account where S3 assets are stored")
 	}
 
 	Cmd.AddCommand(c)