From d1dfeeb1d17cddad10e78cce5da94cb1438fed81 Mon Sep 17 00:00:00 2001
From: Nick Cabatoff <ncabatoff@hashicorp.com>
Date: Tue, 22 Oct 2019 09:06:17 -0400
Subject: [PATCH] TestSysRekey_Verification would fail sometimes when
 recovery=true because when unsealing it wouldn't wait for core 0 to come up
 and become the active node. Much of our testing code assumes that core0 is
 the active node.

---
 .../external_tests/api/sys_rekey_ext_test.go  | 11 ++--
 vault/testing.go                              | 54 ++++++++-----------
 2 files changed, 25 insertions(+), 40 deletions(-)

diff --git a/vault/external_tests/api/sys_rekey_ext_test.go b/vault/external_tests/api/sys_rekey_ext_test.go
index c84d802dc3f8..34f5df74c964 100644
--- a/vault/external_tests/api/sys_rekey_ext_test.go
+++ b/vault/external_tests/api/sys_rekey_ext_test.go
@@ -139,12 +139,9 @@ func testSysRekey_Verification(t *testing.T, recovery bool, legacyShamir bool) {
 	// Sealing should clear state, so after this we should be able to perform
 	// the above again
 	cluster.EnsureCoresSealed(t)
-	if recovery {
-		cluster.UnsealWithStoredKeys(t)
-	} else {
-		cluster.UnsealCores(t)
+	if err := cluster.UnsealCoresWithError(recovery); err != nil {
+		t.Fatal(err)
 	}
-	vault.TestWaitActive(t, cluster.Cores[0].Core)
 	doRekeyInitialSteps()
 
 	doStartVerify := func() {
@@ -258,7 +255,7 @@ func testSysRekey_Verification(t *testing.T, recovery bool, legacyShamir bool) {
 		cluster.Start()
 		defer cluster.Cleanup()
 
-		if err := cluster.UnsealCoresWithError(); err == nil {
+		if err := cluster.UnsealCoresWithError(false); err == nil {
 			t.Fatal("expected error")
 		}
 
@@ -272,7 +269,7 @@ func testSysRekey_Verification(t *testing.T, recovery bool, legacyShamir bool) {
 			newKeyBytes = append(newKeyBytes, val)
 		}
 		cluster.BarrierKeys = newKeyBytes
-		if err := cluster.UnsealCoresWithError(); err != nil {
+		if err := cluster.UnsealCoresWithError(false); err != nil {
 			t.Fatal(err)
 		}
 	} else {
diff --git a/vault/testing.go b/vault/testing.go
index ba6e554dd54a..7f56e402bc11 100644
--- a/vault/testing.go
+++ b/vault/testing.go
@@ -829,19 +829,29 @@ func (c *TestCluster) Start() {
 // UnsealCores uses the cluster barrier keys to unseal the test cluster cores
 func (c *TestCluster) UnsealCores(t testing.T) {
 	t.Helper()
-	if err := c.UnsealCoresWithError(); err != nil {
+	if err := c.UnsealCoresWithError(false); err != nil {
 		t.Fatal(err)
 	}
 }
 
-func (c *TestCluster) UnsealCoresWithError() error {
-	numCores := len(c.Cores)
+func (c *TestCluster) UnsealCoresWithError(useStoredKeys bool) error {
+	unseal := func(core *Core) error {
+		for _, key := range c.BarrierKeys {
+			if _, err := core.Unseal(TestKeyCopy(key)); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+	if useStoredKeys {
+		unseal = func(core *Core) error {
+			return core.UnsealWithStoredKeys(context.Background())
+		}
+	}
 
 	// Unseal first core
-	for _, key := range c.BarrierKeys {
-		if _, err := c.Cores[0].Unseal(TestKeyCopy(key)); err != nil {
-			return fmt.Errorf("unseal core %d err: %s", 0, err)
-		}
+	if err := unseal(c.Cores[0].Core); err != nil {
+		return fmt.Errorf("unseal core %d err: %s", 0, err)
 	}
 
 	// Verify unsealed
@@ -854,11 +864,9 @@ func (c *TestCluster) UnsealCoresWithError() error {
 	}
 
 	// Unseal other cores
-	for i := 1; i < numCores; i++ {
-		for _, key := range c.BarrierKeys {
-			if _, err := c.Cores[i].Core.Unseal(TestKeyCopy(key)); err != nil {
-				return fmt.Errorf("unseal core %d err: %s", i, err)
-			}
+	for i := 1; i < len(c.Cores); i++ {
+		if err := unseal(c.Cores[i].Core); err != nil {
+			return fmt.Errorf("unseal core %d err: %s", i, err)
 		}
 	}
 
@@ -867,7 +875,7 @@ func (c *TestCluster) UnsealCoresWithError() error {
 
 	// Ensure cluster connection info is populated.
 	// Other cores should not come up as leaders.
-	for i := 1; i < numCores; i++ {
+	for i := 1; i < len(c.Cores); i++ {
 		isLeader, _, _, err := c.Cores[i].Leader()
 		if err != nil {
 			return err
@@ -989,26 +997,6 @@ func (c *TestCluster) ensureCoresSealed() error {
 	return nil
 }
 
-// UnsealWithStoredKeys uses stored keys to unseal the test cluster cores
-func (c *TestCluster) UnsealWithStoredKeys(t testing.T) error {
-	for _, core := range c.Cores {
-		if err := core.UnsealWithStoredKeys(context.Background()); err != nil {
-			return err
-		}
-		timeout := time.Now().Add(60 * time.Second)
-		for {
-			if time.Now().After(timeout) {
-				return fmt.Errorf("timeout waiting for core to unseal")
-			}
-			if !core.Sealed() {
-				break
-			}
-			time.Sleep(250 * time.Millisecond)
-		}
-	}
-	return nil
-}
-
 func SetReplicationFailureMode(core *TestClusterCore, mode uint32) {
 	atomic.StoreUint32(core.Core.replicationFailure, mode)
 }