From 080780b4970a1c86540ac00b7ad8dcfc8701caf4 Mon Sep 17 00:00:00 2001
From: Mike Palmiotto <mike.palmiotto@hashicorp.com>
Date: Fri, 10 May 2024 13:04:01 -0400
Subject: [PATCH] Add a TraceID for forwarded request tracing (#26939)

---
 http/handler.go           |  7 ++++++-
 sdk/logical/request.go    | 17 +++++++++++++++++
 vault/request_handling.go |  5 +++++
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/http/handler.go b/http/handler.go
index 1023384d26..905838ba19 100644
--- a/http/handler.go
+++ b/http/handler.go
@@ -461,7 +461,12 @@ func wrapGenericHandler(core *vault.Core, h http.Handler, props *vault.HandlerPr
 		// The uuid for the request is going to be generated when a logical
 		// request is generated. But, here we generate one to be able to track
 		// in-flight requests, and use that to update the req data with clientID
-		inFlightReqID, err := uuid.GenerateUUID()
+		reqIDGen := props.RequestIDGenerator
+		if reqIDGen == nil {
+			// By default use a UUID
+			reqIDGen = uuid.GenerateUUID
+		}
+		inFlightReqID, err := reqIDGen()
 		if err != nil {
 			respondError(nw, http.StatusInternalServerError, fmt.Errorf("failed to generate an identifier for the in-flight request"))
 		}
diff --git a/sdk/logical/request.go b/sdk/logical/request.go
index 01ae7b948e..1c360d4dbc 100644
--- a/sdk/logical/request.go
+++ b/sdk/logical/request.go
@@ -484,6 +484,23 @@ func (c CtxKeyInFlightRequestPriority) String() string {
 	return "in-flight-request-priority"
 }
 
+// CtxKeyInFlightTraceID is used for passing a trace ID through request
+// forwarding. The CtxKeyInFlightRequestID created at the HTTP layer is
+// propagated on through any forwarded requests using this key.
+//
+// Note that this applies to replication service RPCs (including
+// ForwardingRequest from perf standbys or secondaries). The Forwarding RPC
+// service may propagate the context but the handling on the active node runs
+// back through the `http` package handler which builds a new context from HTTP
+// request properties and creates a fresh request ID. Forwarding RPC is used
+// exclusively in Community Edition but also in some special cases in Enterprise
+// such as when forwarding is forced by an HTTP header.
+type CtxKeyInFlightTraceID struct{}
+
+func (c CtxKeyInFlightTraceID) String() string {
+	return "in-flight-trace-ID"
+}
+
 type CtxKeyRequestRole struct{}
 
 func (c CtxKeyRequestRole) String() string {
diff --git a/vault/request_handling.go b/vault/request_handling.go
index f45a37b8a4..bae6460943 100644
--- a/vault/request_handling.go
+++ b/vault/request_handling.go
@@ -70,6 +70,11 @@ type HandlerProperties struct {
 	DisablePrintableCheck bool
 	RecoveryMode          bool
 	RecoveryToken         *uberAtomic.String
+
+	// RequestIDGenerator is primary used for testing purposes to allow tests to
+	// control the request IDs deterministically. In production code (i.e. if this
+	// is nil) the handler will generate UUIDs.
+	RequestIDGenerator func() (string, error)
 }
 
 // fetchEntityAndDerivedPolicies returns the entity object for the given entity