diff options
Diffstat (limited to 'internal/command/command.go')
| -rw-r--r-- | internal/command/command.go | 28 |
1 files changed, 27 insertions, 1 deletions
diff --git a/internal/command/command.go b/internal/command/command.go index 283b4a1..7e0617e 100644 --- a/internal/command/command.go +++ b/internal/command/command.go @@ -1,6 +1,8 @@ package command import ( + "context" + "gitlab.com/gitlab-org/gitlab-shell/internal/command/authorizedkeys" "gitlab.com/gitlab-org/gitlab-shell/internal/command/authorizedprincipals" "gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs" @@ -16,10 +18,13 @@ import ( "gitlab.com/gitlab-org/gitlab-shell/internal/command/uploadpack" "gitlab.com/gitlab-org/gitlab-shell/internal/config" "gitlab.com/gitlab-org/gitlab-shell/internal/executable" + "gitlab.com/gitlab-org/labkit/correlation" + "gitlab.com/gitlab-org/labkit/log" + "gitlab.com/gitlab-org/labkit/tracing" ) type Command interface { - Execute() error + Execute(ctx context.Context) error } func New(e *executable.Executable, arguments []string, config *config.Config, readWriter *readwriter.ReadWriter) (Command, error) { @@ -35,6 +40,27 @@ func New(e *executable.Executable, arguments []string, config *config.Config, re 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()) { + ctx, finished := tracing.ExtractFromEnv(context.Background()) + defer finished() + + 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) + } + } + + return ctx, finished +} + func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command { switch e.Name { case executable.GitlabShell: |