fixup for Add wsARN for SNS Config

Author: b-wu26

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/KimMachineGun/automemlimit v0.7.5
 	github.com/alecthomas/kingpin/v2 v2.4.0
 	github.com/alecthomas/units v0.0.0-20240927000941-0f3dac36c52b
+	github.com/aws/aws-sdk-go v1.55.8
 	github.com/aws/aws-sdk-go-v2 v1.41.1
 	github.com/aws/aws-sdk-go-v2/config v1.32.7
 	github.com/aws/aws-sdk-go-v2/credentials v1.19.7

go.sum

@@ -77,6 +77,8 @@ github.com/armon/go-metrics v0.4.1 h1:hR91U9KYmb6bLBYLQjyM+3j+rcd/UhE+G78SFnF8gJ
 github.com/armon/go-metrics v0.4.1/go.mod h1:E6amYzXo6aW1tqzoZGT755KkbgrJsSdpwZ+3JqfkOG4=
 github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
 github.com/armon/go-radix v1.0.0/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
+github.com/aws/aws-sdk-go v1.55.8 h1:JRmEUbU52aJQZ2AjX4q4Wu7t4uZjOu71uyNmaWlUkJQ=
+github.com/aws/aws-sdk-go v1.55.8/go.mod h1:ZkViS9AqA6otK+JBBNH2++sx1sgxrPKcSzPPvQkUtXk=
 github.com/aws/aws-sdk-go-v2 v1.41.1 h1:ABlyEARCDLN034NhxlRUSZr4l71mh+T5KAeGh6cerhU=
 github.com/aws/aws-sdk-go-v2 v1.41.1/go.mod h1:MayyLB8y+buD9hZqkCW3kX1AKq07Y5pXxtgB+rRFhz0=
 github.com/aws/aws-sdk-go-v2/config v1.32.7 h1:vxUyWGUwmkQ2g19n7JY/9YL8MfAIl7bTesIUykECXmY=

notify/sns/aws_round_tripper.go

@@ -0,0 +1,34 @@
+package sns
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/aws/aws-sdk-go/aws/arn"
+	"github.com/prometheus/alertmanager/config"
+)
+
+type confusedDeputyRoundTripper struct {
+	workspaceArn arn.ARN
+	rt           http.RoundTripper
+}
+
+func (rt *confusedDeputyRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	req.Header.Set("x-amz-delegation-source-account", rt.workspaceArn.AccountID)
+	req.Header.Set("x-amz-delegation-source-arn", rt.workspaceArn.String())
+	return rt.rt.RoundTrip(req)
+}
+
+// newConfusedDeputyRoundTripper adds confused deputy headers
+func newConfusedDeputyRoundTripper(c *config.SNSConfig, rt http.RoundTripper) (http.RoundTripper, error) {
+	if c.WorkspaceArn == "" {
+		return rt, nil
+	}
+
+	arn, err := arn.Parse(c.WorkspaceArn)
+
+	if err != nil {
+		return nil, fmt.Errorf("%s is not a valid arn", c.WorkspaceArn)
+	}
+	return &confusedDeputyRoundTripper{arn, rt}, nil
+}

notify/sns/aws_round_tripper_test.go

@@ -0,0 +1,97 @@
+package sns
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/prometheus/alertmanager/config"
+	commoncfg "github.com/prometheus/common/config"
+)
+
+func TestRoundTripperWithArnNotConfigured(t *testing.T) {
+	var testCases = []struct {
+		name                 string
+		snsConfig            config.SNSConfig
+		expectedHeaders      map[string]string
+		deniedHeaders        []string
+		expectedErrorMessage string
+	}{
+		{
+			name: "Workspace invalid Arn configured",
+			snsConfig: config.SNSConfig{
+				WorkspaceArn: "arn:--Invalid",
+			},
+			expectedHeaders:      map[string]string{},
+			deniedHeaders:        []string{},
+			expectedErrorMessage: "arn:--Invalid is not a valid arn",
+		},
+		{
+			name:            "Workspace Arn not configured",
+			snsConfig:       config.SNSConfig{},
+			expectedHeaders: map[string]string{},
+			deniedHeaders: []string{
+				"x-amz-source-account",
+				"x-amz-source-arn",
+				"x-amz-delegation-source-arn",
+				"x-amz-delegation-source-account",
+			},
+		},
+		{
+			name: "Workspace Arn configured",
+			snsConfig: config.SNSConfig{
+				WorkspaceArn: "arn:aws:aps:us-west-2:948363459592:workspace/ws-de4908b6-950e-4c4c-9e49-ec68169bc4c7",
+			},
+			expectedHeaders: map[string]string{
+				"x-amz-delegation-source-account": "948363459592",
+				"x-amz-delegation-source-arn":     "arn:aws:aps:us-west-2:948363459592:workspace/ws-de4908b6-950e-4c4c-9e49-ec68169bc4c7",
+			},
+			deniedHeaders: []string{},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			testServer := newTestServer(func(w http.ResponseWriter, r *http.Request) {
+				for _, name := range tc.deniedHeaders {
+					if _, ok := r.Header[name]; ok {
+						t.Fatalf("Header %s should not be set", name)
+					}
+				}
+
+				for key, value := range tc.expectedHeaders {
+					if r.Header.Get(key) != value {
+						t.Fatalf("The received Headers (%s) does not contain all expected headers (%s).", r.Header, tc.expectedHeaders)
+						return
+					}
+				}
+			})
+
+			defer testServer.Close()
+
+			client, err := commoncfg.NewClientFromConfig(commoncfg.HTTPClientConfig{}, "test")
+
+			if err != nil && err.Error() != tc.expectedErrorMessage {
+				t.Fatal(err.Error())
+			}
+
+			client.Transport, err = newConfusedDeputyRoundTripper(&tc.snsConfig, client.Transport)
+
+			if err != nil && err.Error() != tc.expectedErrorMessage {
+				t.Fatal(err.Error())
+			}
+
+			_, err = client.Get(testServer.URL)
+
+			if err != nil && err.Error() != tc.expectedErrorMessage {
+				t.Fatal(err.Error())
+			}
+		})
+	}
+}
+
+func newTestServer(handler func(w http.ResponseWriter, r *http.Request)) *httptest.Server {
+	testServer := httptest.NewUnstartedServer(http.HandlerFunc(handler))
+	testServer.Start()
+	return testServer
+}

notify/sns/sns.go

@@ -55,6 +55,14 @@ func New(c *config.SNSConfig, t *template.Template, l *slog.Logger, httpOpts ...
 	if err != nil {
 		return nil, err
 	}
+
+	// Custom AWS Round Tripper
+	client.Transport, err = newConfusedDeputyRoundTripper(c, client.Transport)
+
+	if err != nil {
+		return nil, err
+	}
+
 	return &Notifier{
 		conf:    c,
 		tmpl:    t,