From ac743df5fb3b69f3f07d6a1cdd8290a75bc6c26b Mon Sep 17 00:00:00 2001 From: Jon Gallant <2163001+jongio@users.noreply.github.com> Date: Sun, 19 Jul 2026 12:50:14 -0700 Subject: [PATCH] feat: add traceparent flag Closes #301 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- cli/src/internal/cmd/root.go | 4 + cli/src/internal/cmd/root_test.go | 20 +++++ cli/src/internal/config/config.go | 1 + cli/src/internal/service/service.go | 107 ++++++++++++++++++++++- cli/src/internal/service/service_test.go | 52 +++++++++++ 5 files changed, 182 insertions(+), 2 deletions(-) diff --git a/cli/src/internal/cmd/root.go b/cli/src/internal/cmd/root.go index 35e9e2a..19aef52 100644 --- a/cli/src/internal/cmd/root.go +++ b/cli/src/internal/cmd/root.go @@ -27,6 +27,7 @@ var ( noAuth bool apiVersion string clientRequestID string + traceparent string urlParams []string headers []string headerFile string @@ -192,6 +193,8 @@ Examples: rootCmd.PersistentFlags().StringVar(&clientRequestID, "client-request-id", "", "Set the x-ms-client-request-id header for Azure request correlation. Pass the flag without a value to generate a random ID.") // Passing --client-request-id without a value generates a fresh ID for this invocation. rootCmd.PersistentFlags().Lookup("client-request-id").NoOptDefVal = uuid.NewString() + rootCmd.PersistentFlags().StringVar(&traceparent, "traceparent", "", "Set the W3C traceparent header. Pass the flag without a value to generate one.") + rootCmd.PersistentFlags().Lookup("traceparent").NoOptDefVal = service.TraceparentAutoValue rootCmd.PersistentFlags().StringArrayVar(&urlParams, "url-param", []string{}, "Set or append a URL query parameter (repeatable, format: key=value)") rootCmd.PersistentFlags().StringArrayVarP(&headers, "header", "H", []string{}, "Custom headers (repeatable, format: Key:Value)") rootCmd.PersistentFlags().StringVar(&headerFile, "header-file", "", "Read headers from a file (one Key: Value per line; blank lines and # comments ignored). -H overrides on conflict.") @@ -267,6 +270,7 @@ func snapshotConfig() config.Config { NoAuth: noAuth, APIVersion: apiVersion, ClientRequestID: clientRequestID, + Traceparent: traceparent, URLParams: urlParams, Headers: headers, HeaderFile: headerFile, diff --git a/cli/src/internal/cmd/root_test.go b/cli/src/internal/cmd/root_test.go index a1268be..25843e0 100644 --- a/cli/src/internal/cmd/root_test.go +++ b/cli/src/internal/cmd/root_test.go @@ -12,6 +12,7 @@ import ( "time" "github.com/jongio/azd-rest/src/internal/config" + "github.com/jongio/azd-rest/src/internal/service" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -31,6 +32,7 @@ func resetGlobalFlags() { noAuth = false apiVersion = "" clientRequestID = "" + traceparent = "" urlParams = []string{} headers = []string{} headerFile = "" @@ -94,6 +96,17 @@ func TestNewRootCmd_SilentFlag(t *testing.T) { assert.Empty(t, flag.Shorthand, "--silent should have no short alias") } +func TestNewRootCmd_TraceparentFlag(t *testing.T) { + resetGlobalFlags() + cmd := NewRootCmd() + + flag := cmd.PersistentFlags().Lookup("traceparent") + require.NotNil(t, flag, "--traceparent persistent flag should be registered") + assert.Empty(t, flag.DefValue, "--traceparent should default to empty") + assert.Equal(t, service.TraceparentAutoValue, flag.NoOptDefVal) + assert.Empty(t, flag.Shorthand, "--traceparent should have no short alias") +} + func TestSnapshotConfig_Silent(t *testing.T) { resetGlobalFlags() silent = true @@ -101,6 +114,13 @@ func TestSnapshotConfig_Silent(t *testing.T) { assert.True(t, cfg.Silent, "snapshotConfig should carry the silent flag") } +func TestSnapshotConfig_Traceparent(t *testing.T) { + resetGlobalFlags() + traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + cfg := snapshotConfig() + assert.Equal(t, traceparent, cfg.Traceparent, "snapshotConfig should carry the traceparent flag") +} + func TestBuildRequestOptions_Headers(t *testing.T) { resetGlobalFlags() headers = []string{"X-Custom: value1", "Authorization: Bearer token"} diff --git a/cli/src/internal/config/config.go b/cli/src/internal/config/config.go index 0e8015c..51fce04 100644 --- a/cli/src/internal/config/config.go +++ b/cli/src/internal/config/config.go @@ -13,6 +13,7 @@ type Config struct { NoAuth bool APIVersion string ClientRequestID string + Traceparent string URLParams []string Headers []string HeaderFile string diff --git a/cli/src/internal/service/service.go b/cli/src/internal/service/service.go index c12646c..ab30df4 100644 --- a/cli/src/internal/service/service.go +++ b/cli/src/internal/service/service.go @@ -7,6 +7,8 @@ import ( "bufio" "bytes" "context" + "crypto/rand" + "encoding/hex" "encoding/json" "errors" "fmt" @@ -23,8 +25,14 @@ import ( "github.com/jongio/azd-rest/src/internal/config" ) -// clientRequestIDHeader is the Azure correlation header set by --client-request-id. -const clientRequestIDHeader = "x-ms-client-request-id" +const ( + // TraceparentAutoValue tells BuildRequestOptions to generate a fresh W3C traceparent header. + TraceparentAutoValue = "generate" + + // clientRequestIDHeader is the Azure correlation header set by --client-request-id. + clientRequestIDHeader = "x-ms-client-request-id" + traceparentHeader = "traceparent" +) // TokenProviderFactory creates a TokenProvider. Abstracting this allows tests // to inject mocks without touching real Azure credentials. @@ -286,6 +294,15 @@ func (s *RequestService) BuildRequestOptions(cfg config.Config, method, url stri opts.Headers[clientRequestIDHeader] = cfg.ClientRequestID } + // The --traceparent flag is authoritative and overrides a matching -H header. + if cfg.Traceparent != "" { + value, traceErr := prepareTraceparentHeader(cfg.Traceparent) + if traceErr != nil { + return opts, nil, traceErr + } + opts.Headers[traceparentHeader] = value + } + // Form fields (#202): build an application/x-www-form-urlencoded body from // repeatable --form-field flags. This is mutually exclusive with a raw body. if len(cfg.FormFields) > 0 { @@ -469,6 +486,92 @@ func (s *RequestService) Execute(ctx context.Context, cfg config.Config, method, return nil } +func prepareTraceparentHeader(value string) (string, error) { + if value == TraceparentAutoValue { + return generateTraceparent() + } + + normalized := strings.ToLower(strings.TrimSpace(value)) + if err := validateTraceparent(normalized); err != nil { + return "", err + } + return normalized, nil +} + +func generateTraceparent() (string, error) { + traceID, err := randomNonZeroHex(16) + if err != nil { + return "", err + } + parentID, err := randomNonZeroHex(8) + if err != nil { + return "", err + } + return fmt.Sprintf("00-%s-%s-01", traceID, parentID), nil +} + +func randomNonZeroHex(length int) (string, error) { + buf := make([]byte, length) + for { + if _, err := rand.Read(buf); err != nil { + return "", fmt.Errorf("failed to generate traceparent: %w", err) + } + if !allZeroBytes(buf) { + return hex.EncodeToString(buf), nil + } + } +} + +func validateTraceparent(value string) error { + parts := strings.Split(value, "-") + if len(parts) != 4 { + return fmt.Errorf("invalid traceparent %q: expected version, trace ID, parent ID, and flags", value) + } + if parts[0] != "00" { + return fmt.Errorf("invalid traceparent %q: only version 00 is supported", value) + } + if !isLowerHex(parts[1], 32) || allZeroHex(parts[1]) { + return fmt.Errorf("invalid traceparent %q: trace ID must be 32 non-zero lowercase hex characters", value) + } + if !isLowerHex(parts[2], 16) || allZeroHex(parts[2]) { + return fmt.Errorf("invalid traceparent %q: parent ID must be 16 non-zero lowercase hex characters", value) + } + if !isLowerHex(parts[3], 2) { + return fmt.Errorf("invalid traceparent %q: trace flags must be 2 lowercase hex characters", value) + } + return nil +} + +func isLowerHex(value string, length int) bool { + if len(value) != length { + return false + } + for _, ch := range value { + if (ch < '0' || ch > '9') && (ch < 'a' || ch > 'f') { + return false + } + } + return true +} + +func allZeroHex(value string) bool { + for _, ch := range value { + if ch != '0' { + return false + } + } + return true +} + +func allZeroBytes(value []byte) bool { + for _, b := range value { + if b != 0 { + return false + } + } + return true +} + // writeResponseOutput renders the response body to stdout or --output-file, // choosing the raw path for binary content and the formatter path otherwise. func (s *RequestService) writeResponseOutput(cfg config.Config, resp *client.Response) error { diff --git a/cli/src/internal/service/service_test.go b/cli/src/internal/service/service_test.go index ce3633d..a67bb79 100644 --- a/cli/src/internal/service/service_test.go +++ b/cli/src/internal/service/service_test.go @@ -68,6 +68,58 @@ func baseTestConfig(t *testing.T) config.Config { return cfg } +func TestBuildRequestOptions_TraceparentSetsHeader(t *testing.T) { + svc := newTestService() + cfg := baseTestConfig(t) + cfg.Headers = []string{"traceparent: 00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00"} + cfg.Traceparent = "00-4BF92F3577B34DA6A3CE929D0E0E4736-00F067AA0BA902B7-01" + + opts, cleanup, err := svc.BuildRequestOptions(cfg, "GET", "https://example.com") + if cleanup != nil { + cleanup() + } + require.NoError(t, err) + assert.Equal(t, "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", opts.Headers[traceparentHeader]) +} + +func TestBuildRequestOptions_TraceparentGeneratesHeader(t *testing.T) { + svc := newTestService() + cfg := baseTestConfig(t) + cfg.Traceparent = TraceparentAutoValue + + opts, cleanup, err := svc.BuildRequestOptions(cfg, "GET", "https://example.com") + if cleanup != nil { + cleanup() + } + require.NoError(t, err) + + got := opts.Headers[traceparentHeader] + require.NotEmpty(t, got) + require.NoError(t, validateTraceparent(got)) +} + +func TestBuildRequestOptions_TraceparentRejectsInvalidBeforeTokenProvider(t *testing.T) { + tokenProviderCalls := 0 + svc := NewRequestService( + func() (client.TokenProvider, error) { + tokenProviderCalls++ + return nil, nil + }, + DefaultHTTPClientFactory, + ) + + cfg := config.Defaults() + cfg.Traceparent = "not-a-traceparent" + + _, cleanup, err := svc.BuildRequestOptions(cfg, "GET", "https://management.azure.com/subscriptions?api-version=2021-04-01") + if cleanup != nil { + cleanup() + } + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid traceparent") + assert.Zero(t, tokenProviderCalls, "traceparent validation should run before token creation") +} + func TestBuildResponseHeaderBlock_StatusAndSortedHeaders(t *testing.T) { resp := &client.Response{ Status: "200 OK",