reduce calls to DetermineRoleFromLoginRequest from 3 to 1 for aws aut…

…h method (#22583)

* reduce calls to DetermineRoleFromLoginRequest from 3 to 1 for aws auth method

* change ordering of LoginCreateToken args

* replace another determineRoleFromLoginRequest function with role from context

* add changelog

* Check for role in context if not there make call to DeteremineRoleFromLoginRequest

* move context role check below nanmespace check

* Update changelog/22583.txt

Co-authored-by: Nick Cabatoff <ncabatoff@hashicorp.com>

* revert signature to same order

* make sure resp is last argument

* retrieve role from context closer to where role variable is needed

* remove failsafe for role in mfa login

* Update changelog/22583.txt

---------

Co-authored-by: Nick Cabatoff <ncabatoff@hashicorp.com>

changelog/22583.txt

@@ -0,0 +1,3 @@
+```release-note:bug
+core/quotas: Reduce overhead for role calculation when using cloud auth methods.
+```

http/util.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -60,11 +61,16 @@ func rateLimitQuotaWrapping(handler http.Handler, core *vault.Core) http.Handler
 		}
 		r.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
 
+		role := core.DetermineRoleFromLoginRequestFromBytes(mountPath, bodyBytes, r.Context())
+
+		// add an entry to the context to prevent recalculating request role unnecessarily
+		r = r.WithContext(context.WithValue(r.Context(), logical.CtxKeyRequestRole{}, role))
+
 		quotaResp, err := core.ApplyRateLimitQuota(r.Context(), &quotas.Request{
 			Type:          quotas.TypeRateLimit,
 			Path:          path,
 			MountPath:     mountPath,
-			Role:          core.DetermineRoleFromLoginRequestFromBytes(mountPath, bodyBytes, r.Context()),
+			Role:          role,
 			NamespacePath: ns.Path,
 			ClientAddress: parseRemoteIPAddress(r),
 		})

sdk/logical/request.go

@@ -392,3 +392,9 @@ type CtxKeyInFlightRequestID struct{}
 func (c CtxKeyInFlightRequestID) String() string {
 	return "in-flight-request-ID"
 }
+
+type CtxKeyRequestRole struct{}
+
+func (c CtxKeyRequestRole) String() string {
+	return "request-role"
+}

vault/login_mfa.go

@@ -756,12 +756,17 @@ func (c *Core) LoginMFACreateToken(ctx context.Context, reqPath string, cachedAu
 		return nil, fmt.Errorf("namespace not found: %w", err)
 	}
 
+	var role string
+	if reqRole := ctx.Value(logical.CtxKeyRequestRole{}); reqRole != nil {
+		role = reqRole.(string)
+	}
+
 	// The request successfully authenticated itself. Run the quota checks on
 	// the original login request path before creating the token.
 	quotaResp, quotaErr := c.applyLeaseCountQuota(ctx, &quotas.Request{
 		Path:          reqPath,
 		MountPath:     strings.TrimPrefix(mountPoint, ns.Path),
-		Role:          c.DetermineRoleFromLoginRequest(mountPoint, loginRequestData, ctx),
+		Role:          role,
 		NamespacePath: ns.Path,
 	})
 
@@ -781,7 +786,7 @@ func (c *Core) LoginMFACreateToken(ctx context.Context, reqPath string, cachedAu
 	// note that we don't need to handle the error for the following function right away.
 	// The function takes the response as in input variable and modify it. So, the returned
 	// arguments are resp and err.
-	leaseGenerated, resp, err := c.LoginCreateToken(ctx, ns, reqPath, mountPoint, resp, loginRequestData)
+	leaseGenerated, resp, err := c.LoginCreateToken(ctx, ns, reqPath, mountPoint, role, resp)
 
 	if quotaResp.Access != nil {
 		quotaAckErr := c.ackLeaseQuota(quotaResp.Access, leaseGenerated)

vault/request_handling.go

@@ -472,6 +472,10 @@ func (c *Core) switchedLockHandleRequest(httpCtx context.Context, req *logical.R
 	if ok {
 		ctx = context.WithValue(ctx, logical.CtxKeyInFlightRequestID{}, inFlightReqID)
 	}
+	requestRole, ok := httpCtx.Value(logical.CtxKeyRequestRole{}).(string)
+	if ok {
+		ctx = context.WithValue(ctx, logical.CtxKeyRequestRole{}, requestRole)
+	}
 	resp, err = c.handleCancelableRequest(ctx, req)
 	req.SetTokenEntry(nil)
 	cancel()
@@ -1198,7 +1202,13 @@ func (c *Core) handleRequest(ctx context.Context, req *logical.Request) (retResp
 					Path:        resp.Auth.CreationPath,
 					NamespaceID: ns.ID,
 				}
-				if err := c.expiration.RegisterAuth(ctx, registeredTokenEntry, resp.Auth, c.DetermineRoleFromLoginRequest(req.MountPoint, req.Data, ctx)); err != nil {
+
+				var role string
+				if reqRole := ctx.Value(logical.CtxKeyRequestRole{}); reqRole != nil {
+					role = reqRole.(string)
+				}
+
+				if err := c.expiration.RegisterAuth(ctx, registeredTokenEntry, resp.Auth, role); err != nil {
 					// Best-effort clean up on error, so we log the cleanup error as
 					// a warning but still return as internal error.
 					if err := c.tokenStore.revokeOrphan(ctx, resp.Auth.ClientToken); err != nil {
@@ -1392,12 +1402,18 @@ func (c *Core) handleLoginRequest(ctx context.Context, req *logical.Request) (re
 			return
 		}
 
+		// Check for request role
+		var role string
+		if reqRole := ctx.Value(logical.CtxKeyRequestRole{}); reqRole != nil {
+			role = reqRole.(string)
+		}
+
 		// The request successfully authenticated itself. Run the quota checks
 		// before creating lease.
 		quotaResp, quotaErr := c.applyLeaseCountQuota(ctx, &quotas.Request{
 			Path:          req.Path,
 			MountPath:     strings.TrimPrefix(req.MountPoint, ns.Path),
-			Role:          c.DetermineRoleFromLoginRequest(req.MountPoint, req.Data, ctx),
+			Role:          role,
 			NamespacePath: ns.Path,
 		})
 
@@ -1589,7 +1605,7 @@ func (c *Core) handleLoginRequest(ctx context.Context, req *logical.Request) (re
 		// Attach the display name, might be used by audit backends
 		req.DisplayName = auth.DisplayName
 
-		leaseGen, respTokenCreate, errCreateToken := c.LoginCreateToken(ctx, ns, req.Path, source, resp, req.Data)
+		leaseGen, respTokenCreate, errCreateToken := c.LoginCreateToken(ctx, ns, req.Path, source, role, resp)
 		leaseGenerated = leaseGen
 		if errCreateToken != nil {
 			return respTokenCreate, nil, errCreateToken
@@ -1620,9 +1636,8 @@ func (c *Core) handleLoginRequest(ctx context.Context, req *logical.Request) (re
 // LoginCreateToken creates a token as a result of a login request.
 // If MFA is enforced, mfa/validate endpoint calls this functions
 // after successful MFA validation to generate the token.
-func (c *Core) LoginCreateToken(ctx context.Context, ns *namespace.Namespace, reqPath, mountPoint string, resp *logical.Response, loginRequestData map[string]interface{}) (bool, *logical.Response, error) {
+func (c *Core) LoginCreateToken(ctx context.Context, ns *namespace.Namespace, reqPath, mountPoint, role string, resp *logical.Response) (bool, *logical.Response, error) {
 	auth := resp.Auth
-
 	source := strings.TrimPrefix(mountPoint, credentialRoutePrefix)
 	source = strings.ReplaceAll(source, "/", "-")
 
@@ -1682,7 +1697,7 @@ func (c *Core) LoginCreateToken(ctx context.Context, ns *namespace.Namespace, re
 	}
 
 	leaseGenerated := false
-	err = registerFunc(ctx, tokenTTL, reqPath, auth, c.DetermineRoleFromLoginRequest(mountPoint, loginRequestData, ctx))
+	err = registerFunc(ctx, tokenTTL, reqPath, auth, role)
 	switch {
 	case err == nil:
 		if auth.TokenType != logical.TokenTypeBatch {