From 8cb13c072fc3b414f0922b4630a2c19644ad4985 Mon Sep 17 00:00:00 2001
From: Mark Wolfe <mark.wolfe@buildkite.com>
Date: Mon, 9 Sep 2024 14:54:24 +1000
Subject: [PATCH] Fix for region discovery issue with aws sdkv2 when running in
 ec2

---
 agent/awsv2.go                | 40 +++++++++++++++++++++++++++++++++++
 clicommand/agent_start.go     |  4 ++--
 clicommand/pipeline_upload.go |  3 +--
 clicommand/tool_sign.go       |  4 ++--
 go.mod                        |  2 +-
 5 files changed, 46 insertions(+), 7 deletions(-)
 create mode 100644 agent/awsv2.go

diff --git a/agent/awsv2.go b/agent/awsv2.go
new file mode 100644
index 0000000000..3f56546c3a
--- /dev/null
+++ b/agent/awsv2.go
@@ -0,0 +1,40 @@
+package agent
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
+)
+
+func GetAWSConfigV2(ctx context.Context, optFns ...func(*config.LoadOptions) error) (cfg aws.Config, err error) {
+	cfg, err = config.LoadDefaultConfig(ctx, optFns...)
+	if err != nil {
+		return cfg, fmt.Errorf("error loading default config: %w", err)
+	}
+
+	// local configuration resolved a region so we can return
+	if cfg.Region != "" {
+		return cfg, nil
+	}
+
+	// we need to fall back to the ec2 imds service to get the region
+	client := imds.NewFromConfig(cfg)
+
+	var regionResult *imds.GetRegionOutput
+	regionResult, err = client.GetRegion(ctx, &imds.GetRegionInput{})
+	if err != nil {
+		return cfg, fmt.Errorf("error getting region using imds: %w", err)
+	}
+
+	optFns = append(optFns, config.WithRegion(regionResult.Region))
+
+	cfg, err = config.LoadDefaultConfig(ctx, optFns...)
+	if err != nil {
+		return cfg, fmt.Errorf("error loading default config using imds region: %w", err)
+	}
+
+	return cfg, nil
+}
diff --git a/clicommand/agent_start.go b/clicommand/agent_start.go
index c432c2a9ba..29299617cb 100644
--- a/clicommand/agent_start.go
+++ b/clicommand/agent_start.go
@@ -904,12 +904,12 @@ var AgentStartCommand = cli.Command{
 
 			// this is currently loaded here to ensure it is ONLY loaded if the agent is using KMS for signing
 			// this will limit the possible impact of this new SDK on the rest of the agent users
-			awscfg, err := config.LoadDefaultConfig(
+			awscfg, err := agent.GetAWSConfigV2(
 				ctx,
 				config.WithClientLogMode(logMode),
 			)
 			if err != nil {
-				return fmt.Errorf("failed to load AWS config: %w", err)
+				return err
 			}
 
 			// assign a crypto signer which uses the KMS key to sign the pipeline
diff --git a/clicommand/pipeline_upload.go b/clicommand/pipeline_upload.go
index 7267b3fce9..dce6629e7e 100644
--- a/clicommand/pipeline_upload.go
+++ b/clicommand/pipeline_upload.go
@@ -14,7 +14,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/kms"
 	"github.com/buildkite/agent/v3/agent"
 	"github.com/buildkite/agent/v3/api"
@@ -290,7 +289,7 @@ var PipelineUploadCommand = cli.Command{
 
 		switch {
 		case cfg.SigningAWSKMSKey != "":
-			awscfg, err := config.LoadDefaultConfig(ctx)
+			awscfg, err := agent.GetAWSConfigV2(ctx)
 			if err != nil {
 				return fmt.Errorf("couldn't load AWS config: %w", err)
 			}
diff --git a/clicommand/tool_sign.go b/clicommand/tool_sign.go
index e675b7487d..27b744432d 100644
--- a/clicommand/tool_sign.go
+++ b/clicommand/tool_sign.go
@@ -9,8 +9,8 @@ import (
 	"os"
 	"strings"
 
-	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/kms"
+	"github.com/buildkite/agent/v3/agent"
 	"github.com/buildkite/agent/v3/internal/bkgql"
 	awssigner "github.com/buildkite/agent/v3/internal/cryptosigner/aws"
 	"github.com/buildkite/agent/v3/internal/stdin"
@@ -190,7 +190,7 @@ Signing a pipeline from a file:
 		switch {
 		case cfg.AWSKMSKeyID != "":
 			// load the AWS SDK V2 config
-			awscfg, err := config.LoadDefaultConfig(ctx)
+			awscfg, err := agent.GetAWSConfigV2(ctx)
 			if err != nil {
 				return fmt.Errorf("couldn't load AWS config: %w", err)
 			}
diff --git a/go.mod b/go.mod
index 20971524ea..da5d9bf462 100644
--- a/go.mod
+++ b/go.mod
@@ -12,6 +12,7 @@ require (
 	github.com/aws/aws-sdk-go v1.55.5
 	github.com/aws/aws-sdk-go-v2 v1.30.4
 	github.com/aws/aws-sdk-go-v2/config v1.27.31
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.12
 	github.com/aws/aws-sdk-go-v2/service/kms v1.35.5
 	github.com/brunoscheufler/aws-ecs-metadata-go v0.0.0-20220812150832-b6b31c6eeeaf
 	github.com/buildkite/bintest/v3 v3.3.0
@@ -77,7 +78,6 @@ require (
 	github.com/alexflint/go-scalar v1.0.0 // indirect
 	github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.30 // indirect
-	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.12 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.16 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.16 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1 // indirect