Transit: fix race in the key update api

 - The key update API would release the lock a little too early
   after it persisted the update so the reference could be updated
   when it was preparing the response to the caller across updates
   and/or key rotations
 - The storage updates were okay, just the response back to the caller
   of the update might see a mixture of different updates

Author: stevendpclark

builtin/logical/transit/path_keys.go

@@ -260,9 +260,10 @@ func (b *backend) pathPolicyWrite(ctx context.Context, req *logical.Request, d *
 	if p == nil {
 		return nil, fmt.Errorf("error generating key: returned policy was nil")
 	}
-	if b.System().CachingDisabled() {
-		p.Unlock()
+	if !b.System().CachingDisabled() {
+		p.Lock(true)
 	}
+	defer p.Unlock()
 
 	resp, err := b.formatKeyPolicy(p, nil)
 	if err != nil {

builtin/logical/transit/path_keys_config_test.go

@@ -10,14 +10,16 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
-	uuid "github.com/hashicorp/go-uuid"
+	"github.com/hashicorp/go-uuid"
 	"github.com/hashicorp/vault/api"
 	vaulthttp "github.com/hashicorp/vault/http"
 	"github.com/hashicorp/vault/sdk/logical"
 	"github.com/hashicorp/vault/vault"
+	"github.com/stretchr/testify/require"
 )
 
 func TestTransit_ConfigSettings(t *testing.T) {
@@ -406,3 +408,51 @@ func TestTransit_UpdateKeyConfigWithAutorotation(t *testing.T) {
 		})
 	}
 }
+
+// TestTransitLockRace expose race condition fixed in VAULT-23977
+func TestTransitLockRace(t *testing.T) {
+	b, s := createBackendWithStorage(t)
+
+	createReq := &logical.Request{
+		Operation: logical.UpdateOperation,
+		Path:      "keys/test",
+		Storage:   s,
+	}
+
+	_, err := b.HandleRequest(context.Background(), createReq)
+	require.NoError(t, err, "failed creating key")
+
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+	// Expose race between key rotation and key update api
+	go func() {
+		for i := 0; i < 100; i++ {
+			rotateKeyReq := &logical.Request{
+				Operation: logical.UpdateOperation,
+				Path:      "keys/test/rotate",
+				Storage:   s,
+			}
+
+			_, err := b.HandleRequest(context.Background(), rotateKeyReq)
+			require.NoError(t, err, "failed rotating key")
+		}
+
+		wg.Done()
+	}()
+
+	go func() {
+		for i := 0; i < 100; i++ {
+			updateKeyReq := &logical.Request{
+				Operation: logical.UpdateOperation,
+				Path:      "keys/test",
+				Storage:   s,
+			}
+
+			_, err := b.HandleRequest(context.Background(), updateKeyReq)
+			require.NoError(t, err, "failed updating key")
+		}
+		wg.Done()
+	}()
+
+	wg.Wait()
+}