From c902c804430d78451c672f4145dc278c5aa370f7 Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Mon, 24 Oct 2022 16:23:44 -0400
Subject: [PATCH] keyring: refactor to hold locks for less time (#15026)

Follow-up from https://github.com/hashicorp/nomad/pull/14987/files#r1003611644

We don't need to hold the lock when querying the state store, so move the
read-lock to the interior of the `activeKeySet` function.
---
 nomad/encrypter.go | 25 +++++++++----------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/nomad/encrypter.go b/nomad/encrypter.go
index f13dc84be..fb67bf8a1 100644
--- a/nomad/encrypter.go
+++ b/nomad/encrypter.go
@@ -105,10 +105,8 @@ func (e *Encrypter) loadKeystore() error {
 // root key, and returns the cipher text (including the nonce), and
 // the key ID used to encrypt it
 func (e *Encrypter) Encrypt(cleartext []byte) ([]byte, string, error) {
-	e.lock.RLock()
-	defer e.lock.RUnlock()
 
-	keyset, err := e.activeKeySetLocked()
+	keyset, err := e.activeKeySet()
 	if err != nil {
 		return nil, "", err
 	}
@@ -156,17 +154,10 @@ const keyIDHeader = "kid"
 // encoded JWT with both the claim and its signature
 func (e *Encrypter) SignClaims(claim *structs.IdentityClaims) (string, error) {
 
-	getActiveKeyset := func() (*keyset, error) {
-		e.lock.RLock()
-		defer e.lock.RUnlock()
-		keyset, err := e.activeKeySetLocked()
-		return keyset, err
-	}
-
 	// If a key is rotated immediately following a leader election, plans that
 	// are in-flight may get signed before the new leader has the key. Allow for
 	// a short timeout-and-retry to avoid rejecting plans
-	keyset, err := getActiveKeyset()
+	keyset, err := e.activeKeySet()
 	if err != nil {
 		ctx, cancel := context.WithTimeout(e.srv.shutdownCtx, 5*time.Second)
 		defer cancel()
@@ -176,7 +167,7 @@ func (e *Encrypter) SignClaims(claim *structs.IdentityClaims) (string, error) {
 				return "", err
 			default:
 				time.Sleep(50 * time.Millisecond)
-				keyset, err = getActiveKeyset()
+				keyset, err = e.activeKeySet()
 				if keyset != nil {
 					break
 				}
@@ -198,8 +189,6 @@ func (e *Encrypter) SignClaims(claim *structs.IdentityClaims) (string, error) {
 // VerifyClaim accepts a previously-signed encoded claim and validates
 // it before returning the claim
 func (e *Encrypter) VerifyClaim(tokenString string) (*structs.IdentityClaims, error) {
-	e.lock.RLock()
-	defer e.lock.RUnlock()
 
 	token, err := jwt.ParseWithClaims(tokenString, &structs.IdentityClaims{}, func(token *jwt.Token) (interface{}, error) {
 		if _, ok := token.Method.(*jwt.SigningMethodEd25519); !ok {
@@ -210,6 +199,9 @@ func (e *Encrypter) VerifyClaim(tokenString string) (*structs.IdentityClaims, er
 			return nil, fmt.Errorf("missing key ID header")
 		}
 		keyID := raw.(string)
+
+		e.lock.RLock()
+		defer e.lock.RUnlock()
 		keyset, err := e.keysetByIDLocked(keyID)
 		if err != nil {
 			return nil, err
@@ -292,7 +284,7 @@ func (e *Encrypter) GetKey(keyID string) ([]byte, error) {
 // activeKeySetLocked returns the keyset that belongs to the key marked as
 // active in the state store (so that it's consistent with raft). The
 // called must read-lock the keyring
-func (e *Encrypter) activeKeySetLocked() (*keyset, error) {
+func (e *Encrypter) activeKeySet() (*keyset, error) {
 	store := e.srv.fsm.State()
 	keyMeta, err := store.GetActiveRootKeyMeta(nil)
 	if err != nil {
@@ -301,7 +293,8 @@ func (e *Encrypter) activeKeySetLocked() (*keyset, error) {
 	if keyMeta == nil {
 		return nil, fmt.Errorf("keyring has not been initialized yet")
 	}
-
+	e.lock.RLock()
+	defer e.lock.RUnlock()
 	return e.keysetByIDLocked(keyMeta.KeyID)
 }