From 501d3addcb50e5872c7b8e91de95212c8b5741a5 Mon Sep 17 00:00:00 2001
From: Boris Petersen <boris.petersen@idealo.de>
Date: Tue, 17 May 2022 13:32:22 +0200
Subject: [PATCH] Add ability to sign requests for all AWS services

This add the ability to utilize sigv4 signing for all AWS services not
just "aps". When the newly introduced property "service" is not set in
config it will default to "aps".

Signed-off-by: Boris Petersen <boris.petersen@idealo.de>
---
 sigv4/sigv4.go                         | 19 ++++++++++++-------
 sigv4/sigv4_config.go                  |  1 +
 sigv4/sigv4_config_test.go             |  7 +++++++
 sigv4/testdata/sigv4_good_service.yaml |  4 ++++
 4 files changed, 24 insertions(+), 7 deletions(-)
 create mode 100644 sigv4/testdata/sigv4_good_service.yaml

diff --git a/sigv4/sigv4.go b/sigv4/sigv4.go
index ae0f76e5..61269b52 100644
--- a/sigv4/sigv4.go
+++ b/sigv4/sigv4.go
@@ -37,9 +37,10 @@ var sigv4HeaderDenylist = []string{
 }
 
 type sigV4RoundTripper struct {
-	region string
-	next   http.RoundTripper
-	pool   sync.Pool
+	region  string
+	next    http.RoundTripper
+	pool    sync.Pool
+	service string
 
 	signer *signer.Signer
 }
@@ -88,11 +89,15 @@ func NewSigV4RoundTripper(cfg *SigV4Config, next http.RoundTripper) (http.RoundT
 	if cfg.RoleARN != "" {
 		signerCreds = stscreds.NewCredentials(sess, cfg.RoleARN)
 	}
+	if cfg.Service == "" {
+		cfg.Service = "aps"
+	}
 
 	rt := &sigV4RoundTripper{
-		region: cfg.Region,
-		next:   next,
-		signer: signer.NewSigner(signerCreds),
+		region:  cfg.Region,
+		next:    next,
+		signer:  signer.NewSigner(signerCreds),
+		service: cfg.Service,
 	}
 	rt.pool.New = rt.newBuf
 	return rt, nil
@@ -136,7 +141,7 @@ func (rt *sigV4RoundTripper) RoundTrip(req *http.Request) (*http.Response, error
 		signReq.Header.Del(header)
 	}
 
-	headers, err := rt.signer.Sign(signReq, seeker, "aps", rt.region, time.Now().UTC())
+	headers, err := rt.signer.Sign(signReq, seeker, rt.service, rt.region, time.Now().UTC())
 	if err != nil {
 		return nil, fmt.Errorf("failed to sign request: %w", err)
 	}
diff --git a/sigv4/sigv4_config.go b/sigv4/sigv4_config.go
index 83ef73d8..3afad8a5 100644
--- a/sigv4/sigv4_config.go
+++ b/sigv4/sigv4_config.go
@@ -29,6 +29,7 @@ type SigV4Config struct {
 	Profile            string        `yaml:"profile,omitempty"`
 	RoleARN            string        `yaml:"role_arn,omitempty"`
 	UseFIPSSTSEndpoint bool          `yaml:"use_fips_sts_endpoint,omitempty"`
+	Service   		   string        `yaml:"service,omitempty"`
 }
 
 func (c *SigV4Config) Validate() error {
diff --git a/sigv4/sigv4_config_test.go b/sigv4/sigv4_config_test.go
index f88340da..a8d2cfee 100644
--- a/sigv4/sigv4_config_test.go
+++ b/sigv4/sigv4_config_test.go
@@ -47,6 +47,13 @@ func TestGoodSigV4Configs(t *testing.T) {
 	}
 }
 
+func TestGoodSigV4ServiceConfigs(t *testing.T) {
+	filesToTest := []string{"testdata/sigv4_good_service.yaml", "testdata/sigv4_good_service.yaml"}
+	for _, filename := range filesToTest {
+		testGoodConfig(t, filename)
+	}
+}
+
 func TestBadSigV4Config(t *testing.T) {
 	filename := "testdata/sigv4_bad.yaml"
 	_, err := loadSigv4Config(filename)
diff --git a/sigv4/testdata/sigv4_good_service.yaml b/sigv4/testdata/sigv4_good_service.yaml
new file mode 100644
index 00000000..249d92fc
--- /dev/null
+++ b/sigv4/testdata/sigv4_good_service.yaml
@@ -0,0 +1,4 @@
+region: us-east-2
+profile: profile
+role_arn: blah:role/arn
+service: exectute-api