AWS upgrade role entries (#7025)

* upgrade aws roles

* test upgrade aws roles

* Initialize aws credential backend at mount time

* add a TODO

* create end-to-end test for builtin/credential/aws

* fix bug in initializer

* improve comments

* add Initialize() to logical.Backend

* use Initialize() in Core.enableCredentialInternal()

* use InitializeRequest to call Initialize()

* improve unit testing for framework.Backend

* call logical.Backend.Initialize() from all of the places that it needs to be called.

* implement backend.proto changes for logical.Backend.Initialize()

* persist current role storage version when upgrading aws roles

* format comments correctly

* improve comments

* use postUnseal funcs to initialize backends

* simplify test suite

* improve test suite

* simplify logic in aws role upgrade

* simplify aws credential initialization logic

* simplify logic in aws role upgrade

* use the core's activeContext for initialization

* refactor builtin/plugin/Backend

* use a goroutine to upgrade the aws roles

* misc improvements and cleanup

* do not run AWS role upgrade on DR Secondary

* always call logical.Backend.Initialize() when loading a plugin.

* improve comments

* on standbys and DR secondaries we do not want to run any kind of upgrade logic

* fix awsVersion struct

* clarify aws version upgrade

* make the upgrade logic for aws auth more explicit

* aws upgrade is now called from a switch

* fix fallthrough bug

* simplify logic

* simplify logic

* rename things

* introduce currentAwsVersion const to track aws version

* improve comments

* rearrange things once more

* conglomerate things into one function

* stub out aws auth initialize e2e test

* improve aws auth initialize e2e test

* finish aws auth initialize e2e test

* tinker with aws auth initialize e2e test

* tinker with aws auth initialize e2e test

* tinker with aws auth initialize e2e test

* fix typo in test suite

* simplify logic a tad

* rearrange assignment

* Fix a few lifecycle related issues in #7025 (#7075)

* Fix panic when plugin fails to load

builtin/credential/aws/backend.go

@@ -35,7 +35,7 @@ type backend struct {
 	configMutex sync.RWMutex
 
 	// Lock to make changes to role entries
-	roleMutex sync.RWMutex
+	roleMutex sync.Mutex
 
 	// Lock to make changes to the blacklist entries
 	blacklistMutex sync.RWMutex
@@ -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) {
@@ -134,8 +138,10 @@ func Backend(conf *logical.BackendConfig) (*backend, error) {
 			pathIdentityWhitelist(b),
 			pathTidyIdentityWhitelist(b),
 		},
-		Invalidate:  b.invalidate,
-		BackendType: logical.TypeCredential,
+		Invalidate:     b.invalidate,
+		InitializeFunc: b.initialize,
+		BackendType:    logical.TypeCredential,
+		Clean:          b.cleanup,
 	}
 
 	return b, nil
@@ -205,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":

builtin/credential/aws/backend_e2e_test.go

@@ -0,0 +1,133 @@
+package awsauth
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	hclog "github.com/hashicorp/go-hclog"
+	"github.com/hashicorp/vault/api"
+	vaulthttp "github.com/hashicorp/vault/http"
+	"github.com/hashicorp/vault/sdk/helper/logging"
+	"github.com/hashicorp/vault/sdk/logical"
+	"github.com/hashicorp/vault/vault"
+)
+
+func TestBackend_E2E_Initialize(t *testing.T) {
+
+	ctx := context.Background()
+
+	// Set up the cluster.  This will trigger an Initialize(); we sleep briefly
+	// awaiting its completion.
+	cluster := setupAwsTestCluster(t, ctx)
+	defer cluster.Cleanup()
+	time.Sleep(time.Second)
+	core := cluster.Cores[0]
+
+	// Fetch the aws auth's path in storage.  This is a uuid that is different
+	// every time we run the test
+	authUuids, err := core.UnderlyingStorage.List(ctx, "auth/")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(authUuids) != 1 {
+		t.Fatalf("expected exactly one auth path")
+	}
+	awsPath := "auth/" + authUuids[0]
+
+	// Make sure that the upgrade happened, by fishing the 'config/version'
+	// entry out of storage.  We can't use core.Client.Logical().Read() to do
+	// this, because 'config/version' hasn't been exposed as a path.
+	// TODO: should we expose 'config/version' as a path?
+	version, err := core.UnderlyingStorage.Get(ctx, awsPath+"config/version")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if version == nil {
+		t.Fatalf("no config found")
+	}
+
+	// Nuke the version, so we can pretend that Initialize() has never been run
+	if err := core.UnderlyingStorage.Delete(ctx, awsPath+"config/version"); err != nil {
+		t.Fatal(err)
+	}
+	version, err = core.UnderlyingStorage.Get(ctx, awsPath+"config/version")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if version != nil {
+		t.Fatalf("version found")
+	}
+
+	// Create a role
+	data := map[string]interface{}{
+		"auth_type":       "ec2",
+		"policies":        "default",
+		"bound_subnet_id": "subnet-abcdef"}
+	if _, err := core.Client.Logical().Write("auth/aws/role/test-role", data); err != nil {
+		t.Fatal(err)
+	}
+	role, err := core.Client.Logical().Read("auth/aws/role/test-role")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if role == nil {
+		t.Fatalf("no role found")
+	}
+
+	// There should _still_ be no config version
+	version, err = core.UnderlyingStorage.Get(ctx, awsPath+"config/version")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if version != nil {
+		t.Fatalf("version found")
+	}
+
+	// Seal, and then Unseal. This will once again trigger an Initialize(),
+	// only this time there will be a role present during the upgrade.
+	core.Seal(t)
+	cluster.UnsealCores(t)
+	time.Sleep(time.Second)
+
+	// Now the config version should be there again
+	version, err = core.UnderlyingStorage.Get(ctx, awsPath+"config/version")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if version == nil {
+		t.Fatalf("no version found")
+	}
+}
+
+func setupAwsTestCluster(t *testing.T, ctx context.Context) *vault.TestCluster {
+
+	// create a cluster with the aws auth backend built-in
+	logger := logging.NewVaultLogger(hclog.Trace)
+	coreConfig := &vault.CoreConfig{
+		Logger: logger,
+		CredentialBackends: map[string]logical.Factory{
+			"aws": Factory,
+		},
+	}
+	cluster := vault.NewTestCluster(t, coreConfig, &vault.TestClusterOptions{
+		NumCores:    1,
+		HandlerFunc: vaulthttp.Handler,
+	})
+
+	cluster.Start()
+	if len(cluster.Cores) != 1 {
+		t.Fatalf("expected exactly one core")
+	}
+	core := cluster.Cores[0]
+	vault.TestWaitActive(t, core.Core)
+
+	// load the auth plugin
+	if err := core.Client.Sys().EnableAuthWithOptions("aws", &api.EnableAuthOptions{
+		Type: "aws",
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	return cluster
+}

builtin/credential/aws/path_role.go

@@ -319,6 +319,119 @@ func (b *backend) setRole(ctx context.Context, s logical.Storage, roleName strin
 	return nil
 }
 
+// initialize is used to initialize the AWS roles
+func (b *backend) initialize(ctx context.Context, req *logical.InitializationRequest) error {
+
+	// on standbys and DR secondaries we do not want to run any kind of upgrade logic
+	if b.System().ReplicationState().HasState(consts.ReplicationPerformanceStandby | consts.ReplicationDRSecondary) {
+		return nil
+	}
+
+	// Initialize only if we are either:
+	//   (1) A local mount.
+	//   (2) Are _NOT_ a replicated performance secondary
+	if b.System().LocalMount() || !b.System().ReplicationState().HasState(consts.ReplicationPerformanceSecondary) {
+
+		s := req.Storage
+
+		logger := b.Logger().Named("initialize")
+		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,
+			// so we could see some role requests block until the lock is
+			// released.  However we'd rather see those requests block (and
+			// potentially start timing out) than allow a non-upgraded role to
+			// be fetched.
+			b.roleMutex.Lock()
+			defer b.roleMutex.Unlock()
+
+			upgraded, err := b.upgrade(upgradeCtx, s)
+			if err != nil {
+				logger.Error("error running initialization", "error", err)
+				return
+			}
+			if upgraded {
+				logger.Info("an upgrade was performed during initialization")
+			}
+		}()
+
+	}
+
+	return nil
+}
+
+// awsVersion stores info about the the latest aws version that we have
+// upgraded to.
+type awsVersion struct {
+	Version int `json:"version"`
+}
+
+// currentAwsVersion stores the latest version that we have upgraded to.
+// Note that this is tracked independently from currentRoleStorageVersion.
+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
+	}
+	var version awsVersion
+	if entry != nil {
+		err = entry.DecodeJSON(&version)
+		if err != nil {
+			return false, err
+		}
+	}
+
+	upgraded := version.Version < currentAwsVersion
+	switch version.Version {
+	case 0:
+		// Read all the role names.
+		roleNames, err := s.List(ctx, "role/")
+		if err != nil {
+			return false, err
+		}
+
+		// Upgrade the roles as necessary.
+		for _, roleName := range roleNames {
+			// make sure the context hasn't been canceled
+			if ctx.Err() != nil {
+				return false, err
+			}
+			_, err := b.roleInternal(ctx, s, roleName)
+			if err != nil {
+				return false, err
+			}
+		}
+		fallthrough
+
+	case currentAwsVersion:
+		version.Version = currentAwsVersion
+
+	default:
+		return false, fmt.Errorf("unrecognized role version: %d", version.Version)
+	}
+
+	// save the current version
+	if upgraded {
+		entry, err = logical.StorageEntryJSON("config/version", &version)
+		if err != nil {
+			return false, err
+		}
+		err = s.Put(ctx, entry)
+		if err != nil {
+			return false, err
+		}
+	}
+
+	return upgraded, nil
+}
+
 // If needed, updates the role entry and returns a bool indicating if it was updated
 // (and thus needs to be persisted)
 func (b *backend) upgradeRole(ctx context.Context, s logical.Storage, roleEntry *awsRoleEntry) (bool, error) {