br,lightning: S3 URL support assuming role (#36893)

close #36891

DEPS.bzl

@@ -2754,8 +2754,8 @@ def go_deps():
         name = "com_github_pingcap_kvproto",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/kvproto",
-        sum = "h1:PAXtUVMJnyQQS8t9GzihIFmh6FBXu0JziWbIVknLniA=",
-        version = "v0.0.0-20220711062932-08b02befd813",
+        sum = "h1:4UQdx1acoUrQD0Q5Etz1ABd31duzSgp3XwEnb/cvV9I=",
+        version = "v0.0.0-20220804022843-f006036b1277",
     )
     go_repository(
         name = "com_github_pingcap_log",

br/pkg/storage/BUILD.bazel

@@ -29,6 +29,7 @@ go_library(
         "@com_github_aws_aws_sdk_go//aws/awserr",
         "@com_github_aws_aws_sdk_go//aws/client",
         "@com_github_aws_aws_sdk_go//aws/credentials",
+        "@com_github_aws_aws_sdk_go//aws/credentials/stscreds",
         "@com_github_aws_aws_sdk_go//aws/request",
         "@com_github_aws_aws_sdk_go//aws/session",
         "@com_github_aws_aws_sdk_go//service/s3",

br/pkg/storage/parse_test.go

@@ -3,6 +3,7 @@
 package storage
 
 import (
+	"fmt"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -78,6 +79,24 @@ func TestCreateStorage(t *testing.T) {
 	require.Equal(t, "nREY/7Dt+PaIbYKrKlEEMMF/ExCiJEX=XMLPUANw", s3.SecretAccessKey)
 	require.True(t, s3.ForcePathStyle)
 
+	// parse role ARN and external ID
+	testRoleARN := "arn:aws:iam::888888888888:role/my-role"
+	testExternalID := "abcd1234"
+	s, err = ParseBackend(
+		fmt.Sprintf(
+			"s3://bucket5/prefix/path?role-arn=%s&external-id=%s",
+			url.QueryEscape(testRoleARN),
+			url.QueryEscape(testExternalID),
+		), nil,
+	)
+	require.NoError(t, err)
+	s3 = s.GetS3()
+	require.NotNil(t, s3)
+	require.Equal(t, "bucket5", s3.Bucket)
+	require.Equal(t, "prefix/path", s3.Prefix)
+	require.Equal(t, testRoleARN, s3.RoleArn)
+	require.Equal(t, testExternalID, s3.ExternalId)
+
 	gcsOpt := &BackendOptions{
 		GCS: GCSBackendOptions{
 			Endpoint: "https://gcs.example.com/",

br/pkg/storage/s3.go

@@ -20,6 +20,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/client"
 	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/s3"
@@ -42,6 +43,8 @@ const (
 	s3SseKmsKeyIDOption  = "s3.sse-kms-key-id"
 	s3ACLOption          = "s3.acl"
 	s3ProviderOption     = "s3.provider"
+	s3RoleARNOption      = "s3.role-arn"
+	s3ExternalIDOption   = "s3.external-id"
 	notFound             = "NotFound"
 	// number of retries to make of operations.
 	maxRetries = 7
@@ -129,6 +132,8 @@ type S3BackendOptions struct {
 	Provider              string `json:"provider" toml:"provider"`
 	ForcePathStyle        bool   `json:"force-path-style" toml:"force-path-style"`
 	UseAccelerateEndpoint bool   `json:"use-accelerate-endpoint" toml:"use-accelerate-endpoint"`
+	RoleARN               string `json:"role-arn" toml:"role-arn"`
+	ExternalID            string `json:"external-id" toml:"external-id"`
 }
 
 // Apply apply s3 options on backuppb.S3.
@@ -168,6 +173,8 @@ func (options *S3BackendOptions) Apply(s3 *backuppb.S3) error {
 	s3.AccessKey = options.AccessKey
 	s3.SecretAccessKey = options.SecretAccessKey
 	s3.ForcePathStyle = options.ForcePathStyle
+	s3.RoleArn = options.RoleARN
+	s3.ExternalId = options.ExternalID
 	return nil
 }
 
@@ -183,6 +190,8 @@ func defineS3Flags(flags *pflag.FlagSet) {
 		"Leave empty to use S3 owned key.")
 	flags.String(s3ACLOption, "", "(experimental) Set the S3 canned ACLs, e.g. authenticated-read")
 	flags.String(s3ProviderOption, "", "(experimental) Set the S3 provider, e.g. aws, alibaba, ceph")
+	flags.String(s3RoleARNOption, "", "(experimental) Set the ARN of the IAM role to assume when accessing AWS S3")
+	flags.String(s3ExternalIDOption, "", "(experimental) Set the external ID when assuming the role to access AWS S3")
 }
 
 // parseFromFlags parse S3BackendOptions from command line flags.
@@ -218,6 +227,14 @@ func (options *S3BackendOptions) parseFromFlags(flags *pflag.FlagSet) error {
 	if err != nil {
 		return errors.Trace(err)
 	}
+	options.RoleARN, err = flags.GetString(s3RoleARNOption)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	options.ExternalID, err = flags.GetString(s3ExternalIDOption)
+	if err != nil {
+		return errors.Trace(err)
+	}
 	return nil
 }
 
@@ -273,7 +290,8 @@ func createOssRAMCred() (*credentials.Credentials, error) {
 func newS3Storage(backend *backuppb.S3, opts *ExternalStorageOptions) (obj *S3Storage, errRet error) {
 	qs := *backend
 	awsConfig := aws.NewConfig().
-		WithS3ForcePathStyle(qs.ForcePathStyle)
+		WithS3ForcePathStyle(qs.ForcePathStyle).
+		WithCredentialsChainVerboseErrors(true)
 	if qs.Region == "" {
 		awsConfig.WithRegion(defaultRegion)
 	} else {
@@ -317,7 +335,19 @@ func newS3Storage(backend *backuppb.S3, opts *ExternalStorageOptions) (obj *S3St
 		}
 	}
 
-	c := s3.New(ses)
+	s3CliConfigs := []*aws.Config{}
+	// if role ARN and external ID are provided, try to get the credential using this way
+	if len(qs.RoleArn) > 0 {
+		creds := stscreds.NewCredentials(ses, qs.RoleArn, func(p *stscreds.AssumeRoleProvider) {
+			if len(qs.ExternalId) > 0 {
+				p.ExternalID = &qs.ExternalId
+			}
+		})
+		s3CliConfigs = append(s3CliConfigs,
+			aws.NewConfig().WithCredentials(creds),
+		)
+	}
+	c := s3.New(ses, s3CliConfigs...)
 	// s3manager.GetBucketRegionWithClient will set credential anonymous, which works with s3.
 	// we need reassign credential to be compatible with minio authentication.
 	confCred := ses.Config.Credentials
@@ -340,8 +370,8 @@ func newS3Storage(backend *backuppb.S3, opts *ExternalStorageOptions) (obj *S3St
 		qs.Region = region
 		backend.Region = region
 		if region != defaultRegion {
-			awsConfig.WithRegion(region)
-			c = s3.New(ses, awsConfig)
+			s3CliConfigs = append(s3CliConfigs, aws.NewConfig().WithRegion(region))
+			c = s3.New(ses, s3CliConfigs...)
 		}
 	}
 	log.Info("succeed to get bucket region from s3", zap.String("bucket region", region))

go.mod

@@ -46,7 +46,7 @@ require (
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20220423142525-ae43b7f4e5c3
 	github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059
-	github.com/pingcap/kvproto v0.0.0-20220711062932-08b02befd813
+	github.com/pingcap/kvproto v0.0.0-20220804022843-f006036b1277
 	github.com/pingcap/log v1.1.0
 	github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4
 	github.com/pingcap/tidb/parser v0.0.0-20211011031125-9b13dc409c5e

go.sum

@@ -737,8 +737,8 @@ github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
 github.com/pingcap/kvproto v0.0.0-20220510035547-0e2f26c0a46a/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
-github.com/pingcap/kvproto v0.0.0-20220711062932-08b02befd813 h1:PAXtUVMJnyQQS8t9GzihIFmh6FBXu0JziWbIVknLniA=
-github.com/pingcap/kvproto v0.0.0-20220711062932-08b02befd813/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
+github.com/pingcap/kvproto v0.0.0-20220804022843-f006036b1277 h1:4UQdx1acoUrQD0Q5Etz1ABd31duzSgp3XwEnb/cvV9I=
+github.com/pingcap/kvproto v0.0.0-20220804022843-f006036b1277/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
 github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=

store/mockstore/unistore/tikv/server.go

@@ -44,6 +44,9 @@ var _ tikvpb.TikvServer = new(Server)
 
 // Server implements the tikvpb.TikvServer interface.
 type Server struct {
+	// After updating the kvproto, some methods of TikvServer are not implemented.
+	// Construct `Server` based on `UnimplementedTikvServer`, in order to compile successfully
+	tikvpb.UnimplementedTikvServer
 	mvccStore     *MVCCStore
 	regionManager RegionManager
 	innerServer   InnerServer