feat: Add support for correlationId

internal/daemon/controller/gateway.go

@@ -5,8 +5,11 @@ package controller
 
 import (
 	"context"
+	"fmt"
 	"math"
 	"net"
+	"net/http"
+	"strings"
 
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
@@ -19,8 +22,10 @@ import (
 	"github.com/hashicorp/boundary/internal/event"
 	"github.com/hashicorp/boundary/internal/kms"
 	"github.com/hashicorp/boundary/internal/types/subtypes"
+	"github.com/hashicorp/go-uuid"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/test/bufconn"
 )
 
@@ -52,6 +57,7 @@ func (noDelimiterStreamingMarshaler) Delimiter() []byte {
 
 func newGrpcGatewayMux() *runtime.ServeMux {
 	return runtime.NewServeMux(
+		runtime.WithMetadata(correlationIdAnnotator),
 		runtime.WithMarshalerOption(runtime.MIMEWildcard, &noDelimiterStreamingMarshaler{
 			&runtime.HTTPBodyMarshaler{
 				Marshaler: handlers.JSONMarshaler(),
@@ -62,6 +68,30 @@ func newGrpcGatewayMux() *runtime.ServeMux {
 	)
 }
 
+func correlationIdAnnotator(_ context.Context, req *http.Request) metadata.MD {
+	var correlationId string
+	for k, v := range req.Header {
+		if strings.ToLower(k) == globals.CorrelationIdKey {
+			correlationId = v[0]
+			break
+		}
+	}
+	if correlationId == "" {
+		var err error
+		correlationId, err = uuid.GenerateUUID()
+
+		// GenerateUUID should not return an error. If it does, panic since there is no
+		// err return path here.
+		if err != nil {
+			panic(fmt.Sprintf("failed to generate correlation id: %v", err))
+		}
+	}
+
+	return metadata.New(map[string]string{
+		globals.CorrelationIdKey: correlationId,
+	})
+}
+
 // newGrpcServerListener will create an in-memory listener for the gRPC server.
 func newGrpcServerListener() grpcServerListener {
 	buffer := globals.DefaultMaxRequestSize // seems like a reasonable size for the ring buffer, but then happily change the size if more info becomes available
@@ -116,6 +146,7 @@ func newGrpcServer(
 		grpc.UnaryInterceptor(
 			grpc_middleware.ChainUnaryServer(
 				unaryCtxInterceptor,                           // populated requestInfo from headers into the request ctx
+				correlationIdInterceptor(ctx),                 // populate correlationId from headers or generate random id
 				errorInterceptor(ctx),                         // convert domain and api errors into headers for the http proxy
 				aliasResolutionInterceptor(ctx, aliasRepoFn),  // Resolve ids when an alias is provided
 				subtypes.AttributeTransformerInterceptor(ctx), // convert to/from generic attributes from/to subtype specific attributes

internal/daemon/controller/gateway_test.go

@@ -4,15 +4,19 @@
 package controller
 
 import (
+	"context"
 	"fmt"
 	"math"
+	"net/http"
 	"testing"
 
 	"github.com/hashicorp/boundary/api/targets"
+	"github.com/hashicorp/boundary/globals"
 	_ "github.com/hashicorp/boundary/internal/daemon/controller/handlers/targets/tcp"
 	"github.com/hashicorp/boundary/internal/iam"
 	"github.com/hashicorp/boundary/internal/target"
 	"github.com/hashicorp/boundary/internal/target/tcp"
+	"github.com/hashicorp/go-uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -48,3 +52,43 @@ func Test_gatewayDialOptions(t *testing.T) {
 	assert.NotEmpty(res)
 	assert.Equal(targetCount, len(res.Items))
 }
+
+func Test_correlationIdAnnotator(t *testing.T) {
+	corId, err := uuid.GenerateUUID()
+	require.NoError(t, err)
+	req := &http.Request{
+		Header: map[string][]string{
+			globals.CorrelationIdKey: {corId},
+		},
+	}
+
+	md := correlationIdAnnotator(context.Background(), req)
+	require.NotNil(t, md)
+	corIds := md.Get(globals.CorrelationIdKey)
+	require.Len(t, corIds, 1)
+	assert.Equal(t, corId, corIds[0])
+
+	// Now see if we do not pass a correlation id it generates one
+	md = correlationIdAnnotator(context.Background(), &http.Request{})
+	require.NotNil(t, md)
+	corIds = md.Get(globals.CorrelationIdKey)
+	require.Len(t, corIds, 1)
+	assert.NotEqual(t, corId, corIds[0])
+	// Validate it parses as valid uuid
+	_, err = uuid.ParseUUID(corIds[0])
+	assert.NoError(t, err)
+
+	// Validate correlationIdAnnotator is case-insensitive
+	corId, err = uuid.GenerateUUID()
+	require.NoError(t, err)
+	req = &http.Request{
+		Header: map[string][]string{
+			"X-CorReLAtion-id": {corId},
+		},
+	}
+	md = correlationIdAnnotator(context.Background(), req)
+	require.NotNil(t, md)
+	corIds = md.Get(globals.CorrelationIdKey)
+	require.Len(t, corIds, 1)
+	assert.Equal(t, corId, corIds[0])
+}

internal/daemon/controller/handlers/targets/target_service.go

@@ -25,6 +25,7 @@ import (
 	"github.com/hashicorp/boundary/internal/daemon/controller/handlers"
 	"github.com/hashicorp/boundary/internal/db/timestamp"
 	"github.com/hashicorp/boundary/internal/errors"
+	"github.com/hashicorp/boundary/internal/event"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/api/services"
 	intglobals "github.com/hashicorp/boundary/internal/globals"
 	"github.com/hashicorp/boundary/internal/host"
@@ -817,6 +818,10 @@ func (s Service) AuthorizeSession(ctx context.Context, req *pbs.AuthorizeSession
 	if authResults.Error != nil {
 		return nil, authResults.Error
 	}
+	correlationId, ok := event.CorrelationIdFromContext(ctx)
+	if !ok {
+		return nil, stderrors.New("authorize session: missing correlation id")
+	}
 
 	if authResults.RoundTripValue == nil {
 		return nil, stderrors.New("authorize session: expected to get a target back from auth results")
@@ -1039,6 +1044,7 @@ func (s Service) AuthorizeSession(ctx context.Context, req *pbs.AuthorizeSession
 		IngressWorkerFilter: t.GetIngressWorkerFilter(),
 		DynamicCredentials:  dynCreds,
 		StaticCredentials:   staticCreds,
+		CorrelationId:       correlationId,
 	}
 	if protoWorker != nil {
 		sessionComposition.ProtocolWorkerId = protoWorker.GetPublicId()

internal/daemon/controller/handlers/targets/tcp/target_service_test.go

@@ -34,6 +34,7 @@ import (
 	"github.com/hashicorp/boundary/internal/daemon/controller/handlers/credentials"
 	"github.com/hashicorp/boundary/internal/daemon/controller/handlers/targets"
 	"github.com/hashicorp/boundary/internal/db"
+	"github.com/hashicorp/boundary/internal/event"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/api/services"
 	authpb "github.com/hashicorp/boundary/internal/gen/controller/auth"
 	hostplugin "github.com/hashicorp/boundary/internal/host/plugin"
@@ -56,6 +57,7 @@ import (
 	pb "github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/targets"
 	plgpb "github.com/hashicorp/boundary/sdk/pbs/plugin"
 	wrapping "github.com/hashicorp/go-kms-wrapping/v2"
+	"github.com/hashicorp/go-uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh/testdata"
@@ -3272,6 +3274,12 @@ func TestAuthorizeSession(t *testing.T) {
 	require.NoError(t, err)
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
+			// Generate correlation Id and add it to the context
+			corId, err := uuid.GenerateUUID()
+			require.NoError(t, err)
+			ctx, err = event.NewCorrelationIdContext(ctx, corId)
+			require.NoError(t, err)
+
 			tar := tcp.TestTarget(ctx, t, conn, proj.GetPublicId(), tc.name, target.WithDefaultPort(defaultPort))
 			apiTar, err := s.AddTargetHostSources(ctx, &pbs.AddTargetHostSourcesRequest{
 				Id:            tar.GetPublicId(),
@@ -3380,6 +3388,13 @@ func TestAuthorizeSession(t *testing.T) {
 
 			gotCred.Secret = nil
 
+			// CorrelationId is not part of the authSession resp, query the session directly
+			sessRepo, err := sessionRepoFn()
+			require.NoError(t, err)
+			gotSess, _, err := sessRepo.LookupSession(ctx, got.SessionId)
+			require.NoError(t, err)
+			require.Equal(t, corId, gotSess.CorrelationId)
+
 			got.AuthorizationToken, got.SessionId, got.CreatedTime = "", "", nil
 			assert.Empty(t, cmp.Diff(
 				got,
@@ -3951,6 +3966,12 @@ func TestAuthorizeSessionTypedCredentials(t *testing.T) {
 			// Tell our DB that there is a worker ready to serve the data
 			server.TestKmsWorker(t, conn, wrapper)
 
+			// Generate correlation Id and add it to the context
+			corId, err := uuid.GenerateUUID()
+			require.NoError(t, err)
+			ctx, err = event.NewCorrelationIdContext(ctx, corId)
+			require.NoError(t, err)
+
 			asRes, err := s.AuthorizeSession(ctx, &pbs.AuthorizeSessionRequest{
 				Id: tar.GetPublicId(),
 			})
@@ -4205,8 +4226,15 @@ func TestAuthorizeSession_Errors(t *testing.T) {
 		return tr.GetItem().GetVersion()
 	}
 
+	// Generate correlation Id and add it to the context
+	corId, err := uuid.GenerateUUID()
+	require.NoError(t, err)
+	ctxWithCor, err := event.NewCorrelationIdContext(ctx, corId)
+	require.NoError(t, err)
+
 	cases := []struct {
 		name            string
+		ctx             context.Context
 		setup           []func(target.Target) uint32
 		useTargetId     bool
 		wantErr         bool
@@ -4215,28 +4243,33 @@ func TestAuthorizeSession_Errors(t *testing.T) {
 		{
 			// This one must be run first since it relies on the DB not having any worker details
 			name:            "no worker",
+			ctx:             ctxWithCor,
 			setup:           []func(tcpTarget target.Target) uint32{hostExists, libraryExists},
 			useTargetId:     true,
 			wantErrContains: "No workers are available to handle this session",
 		},
 		{
 			name:        "success",
+			ctx:         ctxWithCor,
 			setup:       []func(tcpTarget target.Target) uint32{workerExists, hostExists, libraryExists},
 			useTargetId: true,
 		},
 		{
 			name:            "no target",
+			ctx:             ctxWithCor,
 			setup:           []func(tcpTarget target.Target) uint32{workerExists, hostExists, libraryExists},
 			useTargetId:     false,
 			wantErrContains: "Resource not found",
 		},
 		{
 			name:        "no host port",
+			ctx:         ctxWithCor,
 			setup:       []func(tcpTarget target.Target) uint32{workerExists, hostWithoutPort, libraryExists},
 			useTargetId: true,
 		},
 		{
 			name: "host port",
+			ctx:  ctxWithCor,
 			setup: []func(tcpTarget target.Target) uint32{
 				workerExists, func(tcpTarget target.Target) uint32 {
 					tcpTarget.SetAddress("127.0.0.1:22")
@@ -4252,25 +4285,36 @@ func TestAuthorizeSession_Errors(t *testing.T) {
 		},
 		{
 			name:            "no hosts",
+			ctx:             ctxWithCor,
 			setup:           []func(tcpTarget target.Target) uint32{workerExists, hostSetNoHostExists, libraryExists},
 			useTargetId:     true,
 			wantErrContains: "No host sources or address found for given target",
 		},
 		{
 			name:            "bad library configuration",
+			ctx:             ctxWithCor,
 			setup:           []func(tcpTarget target.Target) uint32{workerExists, hostExists, misConfiguredlibraryExists},
 			useTargetId:     true,
 			wantErrContains: "external system issue: error #3014: Error making API request",
 		},
 		{
 			name:            "expired token library",
+			ctx:             ctxWithCor,
 			setup:           []func(tcpTarget target.Target) uint32{workerExists, hostExists, expiredTokenLibrary},
 			useTargetId:     true,
 			wantErrContains: "vault.newClient: invalid configuration",
 		},
+		{
+			name:            "no correaltion id",
+			ctx:             ctx,
+			setup:           []func(tcpTarget target.Target) uint32{workerExists, hostExists, libraryExists},
+			useTargetId:     true,
+			wantErrContains: "authorize session: missing correlation id",
+		},
 	}
 	for i, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
+			ctx := tc.ctx
 			tar := tcp.TestTarget(ctx, t, conn, proj.GetPublicId(), fmt.Sprintf("test-%d", i), target.WithDefaultPort(22))
 
 			for _, fn := range tc.setup {

internal/daemon/controller/interceptor.go

@@ -11,6 +11,7 @@ import (
 	"runtime/debug"
 
 	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
+	"github.com/hashicorp/boundary/globals"
 	"github.com/hashicorp/boundary/internal/alias"
 	commonSrv "github.com/hashicorp/boundary/internal/daemon/common"
 	"github.com/hashicorp/boundary/internal/daemon/controller/auth"
@@ -23,6 +24,7 @@ import (
 	pberrors "github.com/hashicorp/boundary/internal/gen/errors"
 	"github.com/hashicorp/boundary/internal/kms"
 	"github.com/hashicorp/boundary/internal/requests"
+	"github.com/hashicorp/go-uuid"
 	"github.com/mr-tron/base58"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -278,6 +280,42 @@ func sharedRequestInterceptorLogic(
 	return interceptorCtx, err // not convinced we want to wrap every error and turn them into domain errors...
 }
 
+func correlationIdInterceptor(
+	_ context.Context,
+) grpc.UnaryServerInterceptor {
+	const op = "controller.correlationIdInterceptor"
+	return func(interceptorCtx context.Context, req any,
+		_ *grpc.UnaryServerInfo, handler grpc.UnaryHandler,
+	) (any, error) {
+		md, ok := metadata.FromIncomingContext(interceptorCtx)
+		if !ok {
+			return nil, errors.New(interceptorCtx, errors.Internal, op, "no metadata")
+		}
+
+		values := md.Get(globals.CorrelationIdKey)
+		if len(values) == 0 {
+			return nil, errors.New(interceptorCtx, errors.Internal, op, "missing correlation id metadata")
+		}
+		if len(values) > 1 {
+			return nil, errors.New(interceptorCtx, errors.Internal, op, fmt.Sprintf("expected 1 value for %s metadata and got %d", globals.CorrelationIdKey, len(values)))
+		}
+		correlationId := values[0]
+
+		// Validate the correlationId
+		if _, err := uuid.ParseUUID(correlationId); err != nil {
+			return nil, errors.Wrap(interceptorCtx, err, op, errors.WithMsg("failed to validated correlation id"))
+		}
+
+		interceptorCtx, err := event.NewCorrelationIdContext(interceptorCtx, correlationId)
+		if err != nil {
+			return nil, errors.Wrap(interceptorCtx, err, op, errors.WithCode(errors.Internal), errors.WithMsg("unable to create context with correlation id"))
+		}
+
+		// call the handler...
+		return handler(interceptorCtx, req)
+	}
+}
+
 func errorInterceptor(
 	_ context.Context,
 ) grpc.UnaryServerInterceptor {

internal/daemon/controller/interceptor_test.go

@@ -13,6 +13,7 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/hashicorp/boundary/globals"
 	"github.com/hashicorp/boundary/internal/alias"
 	talias "github.com/hashicorp/boundary/internal/alias/target"
 	"github.com/hashicorp/boundary/internal/authtoken"
@@ -32,6 +33,7 @@ import (
 	"github.com/hashicorp/boundary/internal/server"
 	"github.com/hashicorp/boundary/internal/target/tcp"
 	"github.com/hashicorp/go-hclog"
+	"github.com/hashicorp/go-uuid"
 	"github.com/mr-tron/base58"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -988,3 +990,79 @@ func (g *testGreeter) SayHello(ctx context.Context, req *interceptor.SayHelloReq
 		return &interceptor.SayHelloResponse{Message: "hello"}, nil
 	}
 }
+
+func Test_correlationIdInterceptor(t *testing.T) {
+	interceptor := correlationIdInterceptor(context.Background())
+	require.NotNil(t, interceptor)
+
+	corId, err := uuid.GenerateUUID()
+	require.NoError(t, err)
+
+	returnCtxHandler := func(ctx context.Context, req any) (any, error) {
+		return ctx, nil
+	}
+
+	cases := []struct {
+		name       string
+		ctx        context.Context
+		wantCorId  string
+		wantErr    bool
+		wantErrStr string
+	}{
+		{
+			name:       "no metadata",
+			ctx:        context.Background(),
+			wantErr:    true,
+			wantErrStr: "controller.correlationIdInterceptor: no metadata",
+		},
+		{
+			name: "no correlation id",
+			ctx: metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{
+				"not-correlation-id": "this is not a correlation id",
+			})),
+			wantErr:    true,
+			wantErrStr: "controller.correlationIdInterceptor: missing correlation id metadata",
+		},
+		{
+			name: "too many correlation ids",
+			ctx: metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{
+				globals.CorrelationIdKey: corId,
+				"x-Correlation-id":       corId, // metadata.New does a toLower so this is an easy way to add multiple of same key
+			})),
+			wantErr:    true,
+			wantErrStr: "controller.correlationIdInterceptor: expected 1 value for x-correlation-id metadata and got 2",
+		},
+		{
+			name: "invalid correlation id",
+			ctx: metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{
+				globals.CorrelationIdKey: "this is wrong",
+			})),
+			wantErr:    true,
+			wantErrStr: "controller.correlationIdInterceptor: failed to validated correlation id",
+		},
+		{
+			name: "valid correlation id",
+			ctx: metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{
+				globals.CorrelationIdKey: corId,
+			})),
+			wantCorId: corId,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			req := &pbs.GetAccountRequest{Id: "test"}
+
+			retCtx, err := interceptor(tc.ctx, req, nil, returnCtxHandler)
+			if tc.wantErr {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tc.wantErrStr)
+				return
+			}
+			assert.NoError(t, err)
+
+			corId, ok := event.CorrelationIdFromContext(retCtx.(context.Context))
+			require.True(t, ok)
+			assert.Equal(t, tc.wantCorId, corId)
+		})
+	}
+}