From 8511e5b2602de813d8d4df1e7d7304e2a642b323 Mon Sep 17 00:00:00 2001
From: Ketan Umare <kumare@lyft.com>
Date: Tue, 11 Aug 2020 10:25:47 -0700
Subject: [PATCH] Replacing need for IAM Wait (#111)

---
 flyteplugins/copilot/cmd/root.go | 46 ++++++++++++++++++++++++++++++++
 flyteplugins/copilot/go.mod      |  1 +
 2 files changed, 47 insertions(+)

diff --git a/flyteplugins/copilot/cmd/root.go b/flyteplugins/copilot/cmd/root.go
index fc21f6ab80..65e0e24c17 100644
--- a/flyteplugins/copilot/cmd/root.go
+++ b/flyteplugins/copilot/cmd/root.go
@@ -6,7 +6,12 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"time"
 
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
+	"github.com/aws/aws-sdk-go/aws/ec2metadata"
+	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/lyft/flyteidl/gen/pb-go/flyteidl/core"
 	"github.com/lyft/flytestdlib/config"
 	"github.com/lyft/flytestdlib/config/viper"
@@ -19,6 +24,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
+	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
@@ -66,6 +72,41 @@ func (r RootOptions) UploadError(ctx context.Context, code string, recvErr error
 	})
 }
 
+func PollUntilTimeout(ctx context.Context, pollInterval, timeout time.Duration, condition wait.ConditionFunc) error {
+	childCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+	return wait.PollUntil(pollInterval, condition, childCtx.Done())
+}
+
+func checkAWSCreds() error {
+	cfg, err := session.NewSession(&aws.Config{})
+	if err != nil {
+		return err
+	}
+
+	role := &ec2rolecreds.EC2RoleProvider{
+		Client: ec2metadata.New(cfg),
+	}
+	creds, err := role.Retrieve()
+	if err != nil {
+		return err
+	}
+	if creds.AccessKeyID == "" || creds.SecretAccessKey == "" || creds.SessionToken == "" {
+		return fmt.Errorf("invalid data in credential fetch")
+	}
+	return nil
+}
+
+func waitForAWSCreds(ctx context.Context, timeout time.Duration) error {
+	return PollUntilTimeout(ctx, time.Second*5, timeout, func() (bool, error) {
+		if err := checkAWSCreds(); err != nil {
+			logger.Errorf(ctx, "Failed to Get credentials.")
+			return false, nil
+		}
+		return true, nil
+	})
+}
+
 // NewCommand returns a new instance of the co-pilot root command
 func NewDataCommand() *cobra.Command {
 	rootOpts := &RootOptions{}
@@ -79,6 +120,11 @@ func NewDataCommand() *cobra.Command {
 			}
 			rootOpts.Scope = promutils.NewScope("flyte:data")
 			cfg := storage.GetConfig()
+			if cfg.Type == storage.TypeS3 {
+				if err := waitForAWSCreds(context.Background(), time.Minute*10); err != nil {
+					return err
+				}
+			}
 			store, err := storage.NewDataStore(cfg, rootOpts.Scope)
 			if err != nil {
 				return errors.Wrap(err, "failed to create datastore client")
diff --git a/flyteplugins/copilot/go.mod b/flyteplugins/copilot/go.mod
index dd9cb27f94..b789155dea 100644
--- a/flyteplugins/copilot/go.mod
+++ b/flyteplugins/copilot/go.mod
@@ -3,6 +3,7 @@ module github.com/lyft/flyteplugins/copilot
 go 1.13
 
 require (
+	github.com/aws/aws-sdk-go v1.28.11
 	github.com/fsnotify/fsnotify v1.4.9
 	github.com/gogo/protobuf v1.3.1
 	github.com/golang/protobuf v1.4.2