From d4a046820d8ba7f81b6a2c52aabd0ac6aab7abd2 Mon Sep 17 00:00:00 2001
From: Mike Palmiotto <mike.palmiotto@hashicorp.com>
Date: Wed, 1 May 2024 15:45:34 -0400
Subject: [PATCH] AOP: Add NeverDrop request priority (enterprise) (#26745)

This PR introduces the CE plumbing for a new high WritePriority, meant
to bypass rejection from the AOP write controller. We attach this
priority to any request on a sudo path, such that administrators can
still perform necessary operations during an overload.
---
 http/priority/priority.go | 11 ++++++-----
 sdk/logical/auth.go       |  4 ++++
 vault/request_handling.go | 13 +++++++++++++
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/http/priority/priority.go b/http/priority/priority.go
index 31bcbebeae..b273155482 100644
--- a/http/priority/priority.go
+++ b/http/priority/priority.go
@@ -33,15 +33,16 @@ const (
 
 	// StandardHTTP is the default AOPWritePriority for HTTP requests.
 	StandardHTTP AOPWritePriority = 128
+
+	// NeverDrop is used to mark a request such that it will never be rejected.
+	// This is currently used as an administrative priority used for requests on
+	// paths which require sudo capabilities.
+	NeverDrop AOPWritePriority = 255
 )
 
 // String returns the string representation of the AOPWritePriority.
 func (p AOPWritePriority) String() string {
-	switch p {
-	case AlwaysDrop:
-		return strconv.FormatUint(uint64(p), 8)
-	}
-	return ""
+	return strconv.FormatUint(uint64(p), 8)
 }
 
 // StringToAOPWritePriority converts a string to an AOPWritePriority.
diff --git a/sdk/logical/auth.go b/sdk/logical/auth.go
index 83d9daca12..7cf7992346 100644
--- a/sdk/logical/auth.go
+++ b/sdk/logical/auth.go
@@ -114,6 +114,10 @@ type Auth struct {
 
 	// EntityCreated is set to true if an entity is created as part of a login request
 	EntityCreated bool `json:"entity_created"`
+
+	// HTTPRequestPriority contains potential information about the request
+	// priority based on required path capabilities
+	HTTPRequestPriority *uint8 `json:"http_request_priority"`
 }
 
 func (a *Auth) GoString() string {
diff --git a/vault/request_handling.go b/vault/request_handling.go
index d447aff781..d6b6f8cf97 100644
--- a/vault/request_handling.go
+++ b/vault/request_handling.go
@@ -484,6 +484,12 @@ func (c *Core) CheckToken(ctx context.Context, req *logical.Request, unauth bool
 		RootPrivsRequired: rootPath,
 	})
 
+	// Assign the sudo path priority if the request is issued against a sudo path.
+	if rootPath {
+		pri := uint8(priority.NeverDrop)
+		auth.HTTPRequestPriority = &pri
+	}
+
 	auth.PolicyResults = &logical.PolicyResults{
 		Allowed: authResults.Allowed,
 	}
@@ -1014,6 +1020,13 @@ func (c *Core) handleRequest(ctx context.Context, req *logical.Request) (retResp
 		return nil, nil, ctErr
 	}
 
+	// See if the call to CheckToken set any request priority. We push the
+	// processing down into CheckToken so we only have to do a router lookup
+	// once.
+	if auth != nil && auth.HTTPRequestPriority != nil {
+		ctx = context.WithValue(ctx, logical.CtxKeyInFlightRequestPriority{}, *auth.HTTPRequestPriority)
+	}
+
 	// Updating in-flight request data with client/entity ID
 	inFlightReqID, ok := ctx.Value(logical.CtxKeyInFlightRequestID{}).(string)
 	if ok && req.ClientID != "" {