Addressing Yang's comments

go.mod

@@ -15,7 +15,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/shield v1.27.3
 	github.com/aws/aws-sdk-go-v2/service/wafregional v1.23.3
 	github.com/aws/aws-sdk-go-v2/service/wafv2 v1.51.4
-	github.com/aws/smithy-go v1.20.4
+	github.com/aws/smithy-go v1.21.0
 	github.com/evanphx/json-patch v5.7.0+incompatible
 	github.com/gavv/httpexpect/v2 v2.9.0
 	github.com/go-logr/logr v1.4.1

go.sum

@@ -78,8 +78,8 @@ github.com/aws/aws-sdk-go-v2/service/wafregional v1.23.3 h1:7dr6En0/6KRFoz8VmnYk
 github.com/aws/aws-sdk-go-v2/service/wafregional v1.23.3/go.mod h1:24TtlRsv4LKAE3VnRJQhpatr8cpX0yj8NSzg8/lxOCw=
 github.com/aws/aws-sdk-go-v2/service/wafv2 v1.51.4 h1:1khBA5uryBRJoCb4G2iR5RT06BkfPEjjDCHAiRb8P3Q=
 github.com/aws/aws-sdk-go-v2/service/wafv2 v1.51.4/go.mod h1:QpFImaPGKNwa+MiZ+oo6LbV1PVQBapc0CnrAMRScoxM=
-github.com/aws/smithy-go v1.20.4 h1:2HK1zBdPgRbjFOHlfeQZfpC4r72MOb9bZkiFwggKO+4=
-github.com/aws/smithy-go v1.20.4/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
+github.com/aws/smithy-go v1.21.0 h1:H7L8dtDRk0P1Qm6y0ji7MCYMQObJ5R9CRpyPhRUkLYA=
+github.com/aws/smithy-go v1.21.0/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=

pkg/annotations/parser.go

@@ -116,7 +116,7 @@ func (p *suffixAnnotationParser) ParseInt32Annotation(annotation string, value *
 	if !exists {
 		return false, nil
 	}
-	i, err := strconv.ParseInt(raw, 10, 64)
+	i, err := strconv.ParseInt(raw, 10, 32)
 	if err != nil {
 		return true, errors.Wrapf(err, "failed to parse int32 annotation, %v: %v", matchedKey, raw)
 	}

pkg/aws/cloud.go

@@ -23,6 +23,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	amerrors "k8s.io/apimachinery/pkg/util/errors"
+	epresolver "sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/services"
 )
 
@@ -77,12 +78,12 @@ func NewCloud(cfg CloudConfig, metricsRegisterer prometheus.Registerer, logger l
 	} else {
 		ec2IMDSEndpointMode = imds.EndpointModeStateIPv4
 	}
-
+	endpointsResolver := epresolver.NewResolver(cfg.AWSEndpoints)
 	ec2MetadataCfg, err := config.LoadDefaultConfig(context.TODO(),
 		config.WithRetryMaxAttempts(cfg.MaxRetries),
 		config.WithEC2IMDSEndpointMode(ec2IMDSEndpointMode),
 	)
-	ec2Metadata := services.NewEC2Metadata(ec2MetadataCfg)
+	ec2Metadata := services.NewEC2Metadata(ec2MetadataCfg, endpointsResolver)
 
 	if len(cfg.Region) == 0 {
 		region := os.Getenv("AWS_DEFAULT_REGION")
@@ -125,15 +126,10 @@ func NewCloud(cfg CloudConfig, metricsRegisterer prometheus.Registerer, logger l
 		if err != nil {
 			return nil, errors.Wrapf(err, "failed to initialize sdk metrics collector")
 		}
-		awsConfig.APIOptions = append(awsConfig.APIOptions, func(stack *smithymiddleware.Stack) error {
-			return metrics.WithSDKCallMetricCollector(metricsCollector)(stack)
-		})
-		awsConfig.APIOptions = append(awsConfig.APIOptions, func(stack *smithymiddleware.Stack) error {
-			return metrics.WithSDKRequestMetricCollector(metricsCollector)(stack)
-		})
+		awsConfig.APIOptions = metrics.WithSDKMetricCollector(metricsCollector, awsConfig.APIOptions)
 	}
 
-	ec2Service := services.NewEC2(awsConfig)
+	ec2Service := services.NewEC2(awsConfig, endpointsResolver)
 
 	vpcID, err := getVpcID(cfg, ec2Service, ec2Metadata, logger)
 	if err != nil {
@@ -143,12 +139,12 @@ func NewCloud(cfg CloudConfig, metricsRegisterer prometheus.Registerer, logger l
 	return &defaultCloud{
 		cfg:         cfg,
 		ec2:         ec2Service,
-		elbv2:       services.NewELBV2(awsConfig),
-		acm:         services.NewACM(awsConfig),
-		wafv2:       services.NewWAFv2(awsConfig),
-		wafRegional: services.NewWAFRegional(awsConfig, cfg.Region),
-		shield:      services.NewShield(awsConfig), //done
-		rgt:         services.NewRGT(awsConfig),
+		elbv2:       services.NewELBV2(awsConfig, endpointsResolver),
+		acm:         services.NewACM(awsConfig, endpointsResolver),
+		wafv2:       services.NewWAFv2(awsConfig, endpointsResolver),
+		wafRegional: services.NewWAFRegional(awsConfig, endpointsResolver, cfg.Region),
+		shield:      services.NewShield(awsConfig, endpointsResolver), //done
+		rgt:         services.NewRGT(awsConfig, endpointsResolver),
 	}, nil
 }
 

pkg/aws/endpoints/resolver.go

@@ -0,0 +1,24 @@
+package endpoints
+
+import "github.com/aws/aws-sdk-go-v2/aws"
+
+func NewResolver(configuration map[string]string) *Resolver {
+	return &Resolver{
+		configuration: configuration,
+	}
+}
+
+// Resolver is an AWS endpoints.Resolver that allows to customize AWS API endpoints.
+// It can be configured using the following format "${AWSServiceID}=${URL}"
+// e.g. "EC2=https://ec2.domain.com,Elastic Load Balancing v2=https://elbv2.domain.com"
+type Resolver struct {
+	configuration map[string]string
+}
+
+func (c *Resolver) EndpointFor(serviceId string) *string {
+	customEndpoint := c.configuration[serviceId]
+	if len(customEndpoint) != 0 {
+		return aws.String(customEndpoint)
+	}
+	return nil
+}

pkg/aws/endpoints/resolver_test.go

@@ -0,0 +1,53 @@
+package endpoints
+
+import (
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
+	"github.com/aws/aws-sdk-go-v2/service/wafv2"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestAWSEndpointResolver_EndpointFor(t *testing.T) {
+	configuration := map[string]string{
+		ec2.ServiceID:                    "https://ec2.domain.com",
+		elasticloadbalancingv2.ServiceID: "https://elbv2.domain.com",
+	}
+	c := &Resolver{
+		configuration: configuration,
+	}
+
+	type args struct {
+		val string
+	}
+
+	tests := []struct {
+		name    string
+		args    args
+		want    *string
+		wantErr error
+	}{
+		{
+			name: "when custom endpoint is configured",
+			args: args{
+				val: ec2.ServiceID,
+			},
+			want: aws.String(configuration[ec2.ServiceID]),
+		},
+		{
+			name: "when custom endpoint is unconfigured",
+			args: args{
+				val: wafv2.ServiceID,
+			},
+			want: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			res := c.EndpointFor(tt.args.val)
+			assert.Equal(t, tt.want, res)
+		})
+	}
+}

pkg/aws/metrics/collector.go

@@ -32,6 +32,18 @@ func NewCollector(registerer prometheus.Registerer) (*collector, error) {
 	}, nil
 }
 
+/*
+WithSDKMetricCollector is a function that collects prometheus metrics for the AWS SDK Go v2 API calls ad requests
+*/
+func WithSDKMetricCollector(c *collector, apiOptions []func(*smithymiddleware.Stack) error) []func(*smithymiddleware.Stack) error {
+	apiOptions = append(apiOptions, func(stack *smithymiddleware.Stack) error {
+		return WithSDKCallMetricCollector(c)(stack)
+	}, func(stack *smithymiddleware.Stack) error {
+		return WithSDKRequestMetricCollector(c)(stack)
+	})
+	return apiOptions
+}
+
 /*
 WithSDKCallMetricCollector is a middleware for the AWS SDK Go v2 that collects and reports metrics on API calls.
 The call metrics are collected after the call is completed

pkg/aws/services/acm.go

@@ -5,6 +5,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/acm"
 	"github.com/aws/aws-sdk-go-v2/service/acm/types"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 )
 
 type ACM interface {
@@ -14,9 +15,14 @@ type ACM interface {
 }
 
 // NewACM constructs new ACM implementation.
-func NewACM(cfg aws.Config) ACM {
+func NewACM(cfg aws.Config, endpointsResolver *endpoints.Resolver) ACM {
+	customEndpoint := endpointsResolver.EndpointFor(acm.ServiceID)
 	return &acmClient{
-		acmClient: acm.NewFromConfig(cfg),
+		acmClient: acm.NewFromConfig(cfg, func(o *acm.Options) {
+			if customEndpoint != nil {
+				o.BaseEndpoint = customEndpoint
+			}
+		}),
 	}
 }
 

pkg/aws/services/ec2.go

@@ -5,6 +5,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 )
 
 type EC2 interface {
@@ -36,9 +37,14 @@ type EC2 interface {
 }
 
 // NewEC2 constructs new EC2 implementation.
-func NewEC2(cfg aws.Config) EC2 {
+func NewEC2(cfg aws.Config, endpointsResolver *endpoints.Resolver) EC2 {
+	customEndpoint := endpointsResolver.EndpointFor(ec2.ServiceID)
 	return &ec2Client{
-		ec2Client: ec2.NewFromConfig(cfg),
+		ec2Client: ec2.NewFromConfig(cfg, func(o *ec2.Options) {
+			if customEndpoint != nil {
+				o.BaseEndpoint = customEndpoint
+			}
+		}),
 	}
 }
 

pkg/aws/services/ec2_metadata.go

@@ -6,6 +6,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"io"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 )
 
 type EC2Metadata interface {
@@ -14,9 +15,14 @@ type EC2Metadata interface {
 }
 
 // NewEC2Metadata constructs new EC2Metadata implementation.
-func NewEC2Metadata(cfg aws.Config) EC2Metadata {
+func NewEC2Metadata(cfg aws.Config, endpointsResolver *endpoints.Resolver) EC2Metadata {
+	customEndpoint := endpointsResolver.EndpointFor(imds.ServiceID)
 	return &ec2metadataClient{
-		ec2metadataClient: imds.NewFromConfig(cfg),
+		ec2metadataClient: imds.NewFromConfig(cfg, func(o *imds.Options) {
+			if customEndpoint != nil {
+				o.Endpoint = aws.ToString(customEndpoint)
+			}
+		}),
 	}
 }
 

pkg/aws/services/elbv2.go

@@ -5,6 +5,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
 	"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 	"time"
 )
 
@@ -58,9 +59,13 @@ type ELBV2 interface {
 	AddListenerCertificatesWithContext(ctx context.Context, input *elasticloadbalancingv2.AddListenerCertificatesInput) (*elasticloadbalancingv2.AddListenerCertificatesOutput, error)
 }
 
-// NewELBV2 constructs new ELBV2 implementation.
-func NewELBV2(cfg aws.Config) ELBV2 {
-	client := elasticloadbalancingv2.NewFromConfig(cfg)
+func NewELBV2(cfg aws.Config, endpointsResolver *endpoints.Resolver) ELBV2 {
+	customEndpoint := endpointsResolver.EndpointFor(elasticloadbalancingv2.ServiceID)
+	client := elasticloadbalancingv2.NewFromConfig(cfg, func(o *elasticloadbalancingv2.Options) {
+		if customEndpoint != nil {
+			o.BaseEndpoint = customEndpoint
+		}
+	})
 	return &elbv2Client{elbv2Client: client}
 }
 

pkg/aws/services/rgt.go

@@ -5,6 +5,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/resourcegroupstaggingapi"
 	rgttypes "github.com/aws/aws-sdk-go-v2/service/resourcegroupstaggingapi/types"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 )
 
 const (
@@ -17,8 +18,13 @@ type RGT interface {
 }
 
 // NewRGT constructs new RGT implementation.
-func NewRGT(cfg aws.Config) RGT {
-	client := resourcegroupstaggingapi.NewFromConfig(cfg)
+func NewRGT(cfg aws.Config, endpointsResolver *endpoints.Resolver) RGT {
+	customEndpoint := endpointsResolver.EndpointFor(resourcegroupstaggingapi.ServiceID)
+	client := resourcegroupstaggingapi.NewFromConfig(cfg, func(o *resourcegroupstaggingapi.Options) {
+		if customEndpoint != nil {
+			o.BaseEndpoint = customEndpoint
+		}
+	})
 	return &rgtClient{rgtClient: client}
 }
 

pkg/aws/services/shield.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	shieldsdk "github.com/aws/aws-sdk-go-v2/service/shield"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 )
 
 type Shield interface {
@@ -14,10 +15,12 @@ type Shield interface {
 }
 
 // NewShield constructs new Shield implementation.
-func NewShield(cfg aws.Config) Shield {
+func NewShield(cfg aws.Config, endpointsResolver *endpoints.Resolver) Shield {
+	customEndpoint := endpointsResolver.EndpointFor(shieldsdk.ServiceID)
 	// shield is only available as a global API in us-east-1.
 	client := shieldsdk.NewFromConfig(cfg, func(o *shieldsdk.Options) {
 		o.Region = "us-east-1"
+		o.BaseEndpoint = customEndpoint
 	})
 	return &shieldClient{shieldClient: client}
 }

pkg/aws/services/wafregional.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/wafregional"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 )
 
 type WAFRegional interface {
@@ -15,10 +16,12 @@ type WAFRegional interface {
 }
 
 // NewWAFRegional constructs new WAFRegional implementation.
-func NewWAFRegional(cfg aws.Config, region string) WAFRegional {
+func NewWAFRegional(cfg aws.Config, endpointsResolver *endpoints.Resolver, region string) WAFRegional {
+	customEndpoint := endpointsResolver.EndpointFor(wafregional.ServiceID)
 	return &wafRegionalClient{
 		wafRegionalClient: wafregional.NewFromConfig(cfg, func(o *wafregional.Options) {
 			o.Region = region
+			o.BaseEndpoint = customEndpoint
 		}),
 		region: region,
 	}

pkg/aws/services/wafv2.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/wafv2"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/endpoints"
 )
 
 type WAFv2 interface {
@@ -13,8 +14,13 @@ type WAFv2 interface {
 }
 
 // NewWAFv2 constructs new WAFv2 implementation.
-func NewWAFv2(cfg aws.Config) WAFv2 {
-	client := wafv2.NewFromConfig(cfg)
+func NewWAFv2(cfg aws.Config, endpointsResolver *endpoints.Resolver) WAFv2 {
+	customEndpoint := endpointsResolver.EndpointFor(wafv2.ServiceID)
+	client := wafv2.NewFromConfig(cfg, func(o *wafv2.Options) {
+		if customEndpoint != nil {
+			o.BaseEndpoint = customEndpoint
+		}
+	})
 	return &wafv2Client{wafv2Client: client}
 }