Fix a few lifecycle related issues in #7025

hashicorp · Jul 5, 2019 · 0acbf33 · 0acbf33
1 parent 9d2275c
commit 0acbf33
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 3 deletions.
diff --git a/builtin/credential/aws/backend.go b/builtin/credential/aws/backend.go
@@ -81,6 +81,10 @@ type backend struct {
 	roleCache *cache.Cache
 
 	resolveArnToUniqueIDFunc func(context.Context, logical.Storage, string) (string, error)
+
+	// upgradeCancelFunc is used to cancel the context used in the upgrade
+	// function
+	upgradeCancelFunc context.CancelFunc
 }
 
 func Backend(conf *logical.BackendConfig) (*backend, error) {
@@ -137,6 +141,7 @@ func Backend(conf *logical.BackendConfig) (*backend, error) {
 		Invalidate:     b.invalidate,
 		InitializeFunc: b.initialize,
 		BackendType:    logical.TypeCredential,
+		Clean:          b.cleanup,
 	}
 
 	return b, nil
@@ -206,6 +211,12 @@ func (b *backend) periodicFunc(ctx context.Context, req *logical.Request) error
 	return nil
 }
 
+func (b *backend) cleanup(ctx context.Context) {
+	if b.upgradeCancelFunc != nil {
+		b.upgradeCancelFunc()
+	}
+}
+
 func (b *backend) invalidate(ctx context.Context, key string) {
 	switch {
 	case key == "config/client":

diff --git a/builtin/credential/aws/path_role.go b/builtin/credential/aws/path_role.go
@@ -335,7 +335,10 @@ func (b *backend) initialize(ctx context.Context, req *logical.InitializationReq
 		s := req.Storage
 
 		logger := b.Logger().Named("initialize")
-		logger.Info("starting initialization")
+		logger.Debug("starting initialization")
+
+		var upgradeCtx context.Context
+		upgradeCtx, b.upgradeCancelFunc = context.WithCancel(context.Background())
 
 		go func() {
 			// The vault will become unsealed while this goroutine is running,
@@ -346,7 +349,7 @@ func (b *backend) initialize(ctx context.Context, req *logical.InitializationReq
 			b.roleMutex.Lock()
 			defer b.roleMutex.Unlock()
 
-			upgraded, err := b.upgrade(ctx, s)
+			upgraded, err := b.upgrade(upgradeCtx, s)
 			if err != nil {
 				logger.Error("error running initialization", "error", err)
 				return
@@ -373,7 +376,6 @@ const currentAwsVersion = 1
 
 // upgrade does an upgrade, if necessary
 func (b *backend) upgrade(ctx context.Context, s logical.Storage) (bool, error) {
-
 	entry, err := s.Get(ctx, "config/version")
 	if err != nil {
 		return false, err

diff --git a/sdk/plugin/grpc_backend_client.go b/sdk/plugin/grpc_backend_client.go
@@ -7,6 +7,8 @@ import (
 	"sync/atomic"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	log "github.com/hashicorp/go-hclog"
 	plugin "github.com/hashicorp/go-plugin"
@@ -59,6 +61,14 @@ func (b *backendGRPCPluginClient) Initialize(ctx context.Context, _ *logical.Ini
 		if b.doneCtx.Err() != nil {
 			return ErrPluginShutdown
 		}
+
+		// If the plugin doesn't have Initialize implemented we should not fail
+		// the initalize call; otherwise this could halt startup of vault.
+		grpcStatus, ok := status.FromError(err)
+		if ok && grpcStatus.Code() == codes.Unimplemented {
+			return nil
+		}
+
 		return err
 	}
 	if reply.Err != nil {