Merge pull request #9 from aws-samples/aksk

fix:support access key rotation

Author: YikaiHu

cmd/root.go

@@ -199,7 +199,6 @@ Supported types:
 		ctx := context.Background()
 
 		var job dth.Job
-
 		switch jobType {
 		case "Finder":
 			job = dth.NewFinder(ctx, cfg)
@@ -211,6 +210,7 @@ Supported types:
 			log.Fatalf("Unknown Job Type - %s. Type must be either Finder or Worker\n, please start again", jobType)
 
 		}
+		go dth.RunTicker(ctx, cfg)
 		job.Run(ctx)
 	},
 }

dth/client.go

@@ -59,6 +59,7 @@ type Client interface {
 // S3Client is an implementation of Client interface for Amazon S3
 type S3Client struct {
 	bucket, prefix, prefixList, region, sourceType string
+	isSrcClient                                    bool
 	client                                         *s3.Client
 }
 
@@ -85,6 +86,35 @@ func getEndpointURL(region, sourceType string) (url string) {
 	return url
 }
 
+func getClientCredentialsModifyFn(isSrc bool, src_cred, dst_cred *S3Credentials) func(o *s3.Options) {
+	if isSrc {
+		return func(o *s3.Options) {
+			if src_cred == nil {
+				log.Print("src_cred is nil, no modification")
+				return
+			}
+			if src_cred.noSignRequest {
+				o.Credentials = aws.AnonymousCredentials{}
+			}
+			if src_cred.accessKey != "" {
+				o.Credentials = aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(src_cred.accessKey, src_cred.secretKey, ""))
+			}
+		}
+	}
+	return func(o *s3.Options) {
+		if dst_cred == nil {
+			log.Print("dst_cred is nil, no modification")
+			return
+		}
+		if dst_cred.noSignRequest {
+			o.Credentials = aws.AnonymousCredentials{}
+		}
+		if dst_cred.accessKey != "" {
+			o.Credentials = aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(dst_cred.accessKey, dst_cred.secretKey, ""))
+		}
+	}
+}
+
 // NewS3Client creates a S3Client instance
 func NewS3Client(ctx context.Context, bucket, prefix, prefixList, endpoint, region, sourceType string, cred *S3Credentials) *S3Client {
 
@@ -147,7 +177,7 @@ func (c *S3Client) GetObject(ctx context.Context, key *string, size, start, chun
 		Range:  &bodyRange,
 	}
 
-	output, err := c.client.GetObject(ctx, input)
+	output, err := c.client.GetObject(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Unable to download %s with %d bytes start from %d - %s\n", *key, size, start, err.Error())
 		return nil, err
@@ -170,10 +200,10 @@ func (c *S3Client) GetObject(ctx context.Context, key *string, size, start, chun
 func (c *S3Client) listObjectFn(ctx context.Context, continuationToken, prefix, delimiter *string, maxKeys int32) (*s3.ListObjectsV2Output, error) {
 
 	input := &s3.ListObjectsV2Input{
-		Bucket:    &c.bucket,
-		Prefix:    prefix,
-		MaxKeys:   maxKeys,
-		Delimiter: delimiter,
+		Bucket:       &c.bucket,
+		Prefix:       prefix,
+		MaxKeys:      maxKeys,
+		Delimiter:    delimiter,
 		EncodingType: "url",
 	}
 
@@ -182,7 +212,7 @@ func (c *S3Client) listObjectFn(ctx context.Context, continuationToken, prefix,
 	}
 
 	// start := time.Now()
-	output, err := c.client.ListObjectsV2(ctx, input)
+	output, err := c.client.ListObjectsV2(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("Unable to list objects in /%s - %s\n", *prefix, err.Error())
 		return nil, err
@@ -305,7 +335,7 @@ func (c *S3Client) HeadObject(ctx context.Context, key *string) *Metadata {
 		Key:    key,
 	}
 
-	output, err := c.client.HeadObject(ctx, input)
+	output, err := c.client.HeadObject(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("Failed to head object for %s - %s\n", *key, err.Error())
 		return nil
@@ -324,7 +354,9 @@ func (c *S3Client) HeadObject(ctx context.Context, key *string) *Metadata {
 // ListSelectedPrefixes is a function to list prefixes from a customized list file.
 func (c *S3Client) ListSelectedPrefixes(ctx context.Context, key *string) (prefixes []*string) {
 
-	downloader := manager.NewDownloader(c.client)
+	downloader := manager.NewDownloader(c.client, func(d *manager.Downloader) {
+		d.ClientOptions = []func(o *s3.Options){getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED)}
+	})
 
 	getBuf := manager.NewWriteAtBuffer([]byte{})
 
@@ -391,7 +423,7 @@ func (c *S3Client) PutObject(ctx context.Context, key *string, body []byte, stor
 		input.Metadata = meta.Metadata
 	}
 
-	output, err := c.client.PutObject(ctx, input)
+	output, err := c.client.PutObject(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Got an error uploading file - %s\n", err.Error())
 		// return nil, err
@@ -413,7 +445,7 @@ func (c *S3Client) DeleteObject(ctx context.Context, key *string) (err error) {
 		Bucket: &c.bucket,
 		Key:    key,
 	}
-	_, err = c.client.DeleteObject(ctx, input)
+	_, err = c.client.DeleteObject(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Failed to delete object %s - %s\n", *key, err.Error())
 		return err
@@ -444,7 +476,7 @@ func (c *S3Client) CreateMultipartUpload(ctx context.Context, key, storageClass,
 		input.Metadata = meta.Metadata
 	}
 
-	output, err := c.client.CreateMultipartUpload(ctx, input)
+	output, err := c.client.CreateMultipartUpload(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Unable to create multipart upload for %s - %s\n", *key, err.Error())
 	} else {
@@ -475,7 +507,7 @@ func (c *S3Client) UploadPart(ctx context.Context, key *string, body []byte, upl
 		ContentMD5: &contentMD5,
 	}
 
-	output, err := c.client.UploadPart(ctx, input)
+	output, err := c.client.UploadPart(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Failed to upload part for %s - %s\n", *key, err.Error())
 		// return nil, err
@@ -513,7 +545,7 @@ func (c *S3Client) CompleteMultipartUpload(ctx context.Context, key, uploadID *s
 		MultipartUpload: &types.CompletedMultipartUpload{Parts: completedPart},
 	}
 
-	output, err := c.client.CompleteMultipartUpload(ctx, input)
+	output, err := c.client.CompleteMultipartUpload(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Unable to complete multipart upload for %s - %s\n", *key, err.Error())
 	} else {
@@ -540,7 +572,7 @@ func (c *S3Client) ListParts(ctx context.Context, key, uploadID *string) (parts
 	parts = make(map[int]*Part, 10000)
 
 	for {
-		output, err := c.client.ListParts(ctx, input)
+		output, err := c.client.ListParts(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 		if err != nil {
 			log.Printf("Failed to list parts for %s - %s\n", *key, err.Error())
 			break
@@ -574,7 +606,7 @@ func (c *S3Client) GetUploadID(ctx context.Context, key *string) (uploadID *stri
 		// MaxUploads: 1,
 	}
 
-	output, err := c.client.ListMultipartUploads(ctx, input)
+	output, err := c.client.ListMultipartUploads(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Failed to list multipart uploads - %s\n", err.Error())
 		return nil
@@ -600,7 +632,7 @@ func (c *S3Client) AbortMultipartUpload(ctx context.Context, key, uploadID *stri
 		Key:      key,
 		UploadId: uploadID,
 	}
-	_, err = c.client.AbortMultipartUpload(ctx, input)
+	_, err = c.client.AbortMultipartUpload(ctx, input, getClientCredentialsModifyFn(c.isSrcClient, SRC_CRED, DST_CRED))
 	if err != nil {
 		log.Printf("S3> Failed to abort multipart upload for %s - %s\n", *key, err.Error())
 		return err

dth/common.go

@@ -33,6 +33,13 @@ var (
 	MB int = 1 << 20
 )
 
+// Global variables which are updated automatically every 02:00 a.m.
+// When calling a S3 API, the credentials of the client will be updated according to these two variables
+var (
+	SRC_CRED *S3Credentials
+	DST_CRED *S3Credentials
+)
+
 // Source is an interface represents a type of cloud storage services
 // type Source interface {
 // 	GetEndpointURL()

dth/job.go

@@ -117,6 +117,13 @@ func NewFinder(ctx context.Context, cfg *JobConfig) (f *Finder) {
 	srcClient := NewS3Client(ctx, cfg.SrcBucket, cfg.SrcPrefix, cfg.SrcPrefixList, cfg.SrcEndpoint, cfg.SrcRegion, cfg.SrcType, srcCred)
 	desClient := NewS3Client(ctx, cfg.DestBucket, cfg.DestPrefix, "", "", cfg.DestRegion, "Amazon_S3", desCred)
 
+	if srcClient != nil {
+		srcClient.isSrcClient = true
+	}
+
+	SRC_CRED = srcCred
+	DST_CRED = desCred
+
 	f = &Finder{
 		srcClient: srcClient,
 		desClient: desClient,
@@ -126,6 +133,36 @@ func NewFinder(ctx context.Context, cfg *JobConfig) (f *Finder) {
 	return
 }
 
+func RunTicker(ctx context.Context, cfg *JobConfig) {
+	t := &JobTicker{}
+	t.updateTimer()
+	for {
+		select {
+		case <-ctx.Done():
+			log.Printf("timer has been stopped")
+			t.timer.Stop()
+		case <-t.timer.C:
+			log.Printf("ticker update S3 Credentials")
+			updateCreds(ctx, cfg)
+			t.updateTimer()
+		}
+	}
+}
+
+func updateCreds(ctx context.Context, cfg *JobConfig) {
+	if cfg == nil {
+		log.Print("nil config")
+		return
+	}
+	sm, err := NewSecretService(ctx)
+	if err != nil {
+		log.Printf("Warning - Unable to load credentials, use default setting - %s\n", err.Error())
+	}
+	SRC_CRED = getCredentials(ctx, cfg.SrcCredential, cfg.SrcInCurrentAccount, sm)
+	DST_CRED = getCredentials(ctx, cfg.DestCredential, cfg.DestInCurrentAccount, sm)
+	log.Print("src_cred and dst_cred has been updated")
+}
+
 // Run is main execution function for Finder.
 func (f *Finder) Run(ctx context.Context) {
 
@@ -422,6 +459,13 @@ func NewWorker(ctx context.Context, cfg *JobConfig) (w *Worker) {
 	srcClient := NewS3Client(ctx, cfg.SrcBucket, cfg.SrcPrefix, cfg.SrcPrefixList, cfg.SrcEndpoint, cfg.SrcRegion, cfg.SrcType, srcCred)
 	desClient := NewS3Client(ctx, cfg.DestBucket, cfg.DestPrefix, "", "", cfg.DestRegion, "Amazon_S3", desCred)
 
+	if srcClient != nil {
+		srcClient.isSrcClient = true
+	}
+
+	SRC_CRED = srcCred
+	DST_CRED = desCred
+
 	return &Worker{
 		srcClient: srcClient,
 		desClient: desClient,
@@ -600,6 +644,13 @@ func (w *Worker) processResult(ctx context.Context, obj *Object, rh *string, res
 	log.Printf("----->Transferred 1 object %s with status %s\n", obj.Key, res.status)
 	w.db.UpdateItem(ctx, &obj.Key, res)
 
+	if res.status == "ERROR" {
+		if strings.Contains(res.err.Error(), "403") {
+			log.Printf("Authentication failed, will update credentials")
+			updateCreds(ctx, w.cfg)
+		}
+	}
+
 	if res.status == "DONE" || res.status == "CANCEL" {
 		w.sqs.DeleteMessage(ctx, rh)
 	}

dth/ticker.go

@@ -0,0 +1,34 @@
+package dth
+
+import (
+	"log"
+	"time"
+)
+
+const INTERVAL_PERIOD time.Duration = 24 * time.Hour
+
+const HOUR_TO_TICK int = 02
+const MINUTE_TO_TICK int = 00
+const SECOND_TO_TICK int = 00
+
+type JobTicker struct {
+	timer *time.Timer
+}
+
+func (t *JobTicker) updateTimer() {
+	nextTick := time.Date(time.Now().Year(), time.Now().Month(),
+		time.Now().Day(), HOUR_TO_TICK, MINUTE_TO_TICK, SECOND_TO_TICK, 0, time.Local)
+	if !nextTick.After(time.Now()) {
+		nextTick = nextTick.Add(INTERVAL_PERIOD)
+	}
+	log.Printf("next tick is:%v", nextTick)
+	diff := nextTick.Sub(time.Now())
+	if diff < 0 {
+		diff = INTERVAL_PERIOD
+	}
+	if t.timer == nil {
+		t.timer = time.NewTimer(diff)
+	} else {
+		t.timer.Reset(diff)
+	}
+}

go.mod

@@ -15,4 +15,5 @@ require (
 	github.com/aws/smithy-go v1.7.0
 	github.com/spf13/cobra v1.1.3
 	github.com/spf13/viper v1.7.1
+	golang.org/x/sys v0.3.0 // indirect
 )

go.sum

@@ -314,6 +314,8 @@ golang.org/x/sys v0.0.0-20190507160741-ecd444e8653b/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20190606165138-5da285871e9c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0 h1:HyfiK1WMnHj5FXFXatD+Qs1A/xC2Run6RzeW1SyHxpc=
 golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.3.0 h1:w8ZOecv6NaNa/zC8944JTU3vz4u6Lagfk4RPQxv92NQ=
+golang.org/x/sys v0.3.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=