5 files changed, 39 insertions, 17 deletions
diff --git a/internal/command/command.go b/internal/command/command.go
index c2709a1..6696f0f 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -21,7 +21,6 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/executable"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 	"gitlab.com/gitlab-org/labkit/correlation"
-	"gitlab.com/gitlab-org/labkit/log"
 	"gitlab.com/gitlab-org/labkit/tracing"
 )
 
@@ -42,24 +41,40 @@ func New(e *executable.Executable, arguments []string, env sshenv.Env, config *c
 	return nil, disallowedcommand.Error
 }
 
-// ContextWithCorrelationID() will always return a background Context
-// with a correlation ID.  It will first attempt to extract the ID from
-// an environment variable. If is not available, a random one will be
-// generated.
-func ContextWithCorrelationID() (context.Context, func()) {
+// Setup() initializes tracing from the configuration file and generates a
+// background context from which all other contexts in the process should derive
+// from, as it has a service name and initial correlation ID set.
+func Setup(serviceName string, config *config.Config) (context.Context, func()) {
+	closer := tracing.Initialize(
+		tracing.WithServiceName(serviceName),
+
+		// For GitLab-Shell, we explicitly initialize tracing from a config file
+		// instead of the default environment variable (using GITLAB_TRACING)
+		// This decision was made owing to the difficulty in passing environment
+		// variables into GitLab-Shell processes.
+		//
+		// Processes are spawned as children of the SSH daemon, which tightly
+		// controls environment variables; doing this means we don't have to
+		// enable PermitUserEnvironment
+		//
+		// gitlab-sshd could use the standard GITLAB_TRACING envvar, but that
+		// would lead to inconsistencies between the two forms of operation
+		tracing.WithConnectionString(config.GitlabTracing),
+	)
+
 	ctx, finished := tracing.ExtractFromEnv(context.Background())
+	ctx = correlation.ContextWithClientName(ctx, serviceName)
 
 	correlationID := correlation.ExtractFromContext(ctx)
 	if correlationID == "" {
-		correlationID, err := correlation.RandomID()
-		if err != nil {
-			log.WithError(err).Warn("unable to generate correlation ID")
-		} else {
-			ctx = correlation.ContextWithCorrelation(ctx, correlationID)
-		}
+		correlationID := correlation.SafeRandomID()
+		ctx = correlation.ContextWithCorrelation(ctx, correlationID)
 	}
 
-	return ctx, finished
+	return ctx, func() {
+		finished()
+		closer.Close()
+	}
 }
 
 func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
diff --git a/internal/command/command_test.go b/internal/command/command_test.go
index a70ea84..3617d39 100644
--- a/internal/command/command_test.go
+++ b/internal/command/command_test.go
@@ -168,7 +168,7 @@ func TestFailingNew(t *testing.T) {
 	}
 }
 
-func TestContextWithCorrelationID(t *testing.T) {
+func TestSetup(t *testing.T) {
 	testCases := []struct {
 		name                  string
 		additionalEnv         map[string]string
@@ -190,16 +190,20 @@ func TestContextWithCorrelationID(t *testing.T) {
 			resetEnvironment := addAdditionalEnv(tc.additionalEnv)
 			defer resetEnvironment()
 
-			ctx, finished := ContextWithCorrelationID()
+			ctx, finished := Setup("foo", &config.Config{})
+			defer finished()
+
 			require.NotNil(t, ctx, "ctx is nil")
 			require.NotNil(t, finished, "finished is nil")
+
 			correlationID := correlation.ExtractFromContext(ctx)
 			require.NotEmpty(t, correlationID)
-
 			if tc.expectedCorrelationID != "" {
 				require.Equal(t, tc.expectedCorrelationID, correlationID)
 			}
-			defer finished()
+
+			clientName := correlation.ExtractClientNameFromContext(ctx)
+			require.Equal(t, "foo", clientName)
 		})
 	}
 }
diff --git a/internal/command/receivepack/gitalycall_test.go b/internal/command/receivepack/gitalycall_test.go
index 70f2b2e..12c603d 100644
--- a/internal/command/receivepack/gitalycall_test.go
+++ b/internal/command/receivepack/gitalycall_test.go
@@ -62,6 +62,7 @@ func TestReceivePack(t *testing.T) {
 
 		hook := testhelper.SetupLogger()
 		ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+		ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
 
 		err := cmd.Execute(ctx)
 		require.NoError(t, err)
diff --git a/internal/command/uploadarchive/gitalycall_test.go b/internal/command/uploadarchive/gitalycall_test.go
index b89b749..3ec1449 100644
--- a/internal/command/uploadarchive/gitalycall_test.go
+++ b/internal/command/uploadarchive/gitalycall_test.go
@@ -37,6 +37,7 @@ func TestUploadPack(t *testing.T) {
 
 	hook := testhelper.SetupLogger()
 	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
 
 	err := cmd.Execute(ctx)
 	require.NoError(t, err)
diff --git a/internal/command/uploadpack/gitalycall_test.go b/internal/command/uploadpack/gitalycall_test.go
index bedc6e6..6b4c6ba 100644
--- a/internal/command/uploadpack/gitalycall_test.go
+++ b/internal/command/uploadpack/gitalycall_test.go
@@ -37,6 +37,7 @@ func TestUploadPack(t *testing.T) {
 
 	hook := testhelper.SetupLogger()
 	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
 
 	err := cmd.Execute(ctx)
 	require.NoError(t, err)