From c75e2bf60635542a4ea81a5453297f213b42b1ac Mon Sep 17 00:00:00 2001
From: Amy Tobey <tobert@gmail.com>
Date: Tue, 23 May 2023 11:42:55 -0400
Subject: [PATCH] add --force-trace-id and --force-span-id options (#173)
 (#210)

* add span options to status command

Useful to test the new --force-span-id and --force-trace-id options.

* add --force-trace-id and --force-span-id and test

As predicted, this was straightforward after #175 :)

* correct option names in comments

* match json key names to command line flags
---
 data_for_test.go         | 28 ++++++++++++++++++++++++++++
 otelcli/config.go        |  4 ++++
 otelcli/protobuf_span.go | 24 ++++++++++++++++++++++++
 otelcli/root.go          |  4 ++++
 otelcli/status.go        |  1 +
 5 files changed, 61 insertions(+)

diff --git a/data_for_test.go b/data_for_test.go
index 900c52b..d151225 100644
--- a/data_for_test.go
+++ b/data_for_test.go
@@ -846,4 +846,32 @@ var suites = []FixtureSuite{
 			},
 		},
 	},
+	// --force-trace-id and --force-span-id allow setting/forcing custom trace & span ids
+	{
+		{
+			Name: "forced trace & span ids",
+			Config: FixtureConfig{
+				CliArgs: []string{
+					"status",
+					"--endpoint", "{{endpoint}}",
+					"--force-trace-id", "00112233445566778899aabbccddeeff",
+					"--force-span-id", "beefcafefacedead",
+				},
+			},
+			Expect: Results{
+				Config: otelcli.DefaultConfig().WithEndpoint("{{endpoint}}"),
+				SpanData: map[string]string{
+					"trace_id": "00112233445566778899aabbccddeeff",
+					"span_id":  "beefcafefacedead",
+				},
+				SpanCount: 1,
+				Diagnostics: otelcli.Diagnostics{
+					NumArgs:           7,
+					IsRecording:       true,
+					DetectedLocalhost: true,
+					ParsedTimeoutMs:   1000,
+				},
+			},
+		},
+	},
 }
diff --git a/otelcli/config.go b/otelcli/config.go
index ae357e6..0047d93 100644
--- a/otelcli/config.go
+++ b/otelcli/config.go
@@ -33,6 +33,8 @@ func DefaultConfig() Config {
 		ServiceName:                  "otel-cli",
 		SpanName:                     "todo-generate-default-span-names",
 		Kind:                         "client",
+		ForceTraceId:                 "",
+		ForceSpanId:                  "",
 		Attributes:                   map[string]string{},
 		TraceparentCarrierFile:       "",
 		TraceparentIgnoreEnv:         false,
@@ -78,6 +80,8 @@ type Config struct {
 	Attributes        map[string]string `json:"span_attributes" env:"OTEL_CLI_ATTRIBUTES"`
 	StatusCode        string            `json:"span_status_code" env:"OTEL_CLI_STATUS_CODE"`
 	StatusDescription string            `json:"span_status_description" env:"OTEL_CLI_STATUS_DESCRIPTION"`
+	ForceSpanId       string            `json:"force_span_id" env:"OTEL_CLI_FORCE_SPAN_ID"`
+	ForceTraceId      string            `json:"force_trace_id" env:"OTEL_CLI_FORCE_TRACE_ID"`
 
 	TraceparentCarrierFile string `json:"traceparent_carrier_file" env:"OTEL_CLI_CARRIER_FILE"`
 	TraceparentIgnoreEnv   bool   `json:"traceparent_ignore_env" env:"OTEL_CLI_IGNORE_ENV"`
diff --git a/otelcli/protobuf_span.go b/otelcli/protobuf_span.go
index f3c6961..5a6c554 100644
--- a/otelcli/protobuf_span.go
+++ b/otelcli/protobuf_span.go
@@ -8,6 +8,7 @@ package otelcli
 
 import (
 	"crypto/rand"
+	"encoding/hex"
 	"time"
 
 	commonpb "go.opentelemetry.io/proto/otlp/common/v1"
@@ -81,11 +82,21 @@ func NewProtobufSpanWithConfig(c Config) tracepb.Span {
 			span.TraceId = tp.TraceId
 			span.ParentSpanId = tp.SpanId
 		}
+
 	} else {
 		span.TraceId = emptyTraceId
 		span.SpanId = emptySpanId
 	}
 
+	// --force-trace-id and --force-span-id let the user set their own trace & span ids
+	// these work in non-recording mode and will stomp trace id from the traceparent
+	if config.ForceTraceId != "" {
+		span.TraceId = parseHex(config.ForceTraceId, 16)
+	}
+	if config.ForceSpanId != "" {
+		span.SpanId = parseHex(config.ForceSpanId, 8)
+	}
+
 	SetSpanStatus(&span, c)
 
 	return span
@@ -131,6 +142,19 @@ func generateSpanId(c Config) []byte {
 	}
 }
 
+// parseHex parses hex into a []byte of length provided. Errors if the input is
+// not valid hex or the converted hex is not the right number of bytes.
+func parseHex(in string, expectedLen int) []byte {
+	out, err := hex.DecodeString(in)
+	if err != nil {
+		softFail("error parsing hex string %q: %s", in, err)
+	}
+	if len(out) != expectedLen {
+		softFail("hex string %q is the wrong length, expected %d bytes but got %d", in, expectedLen, len(out))
+	}
+	return out
+}
+
 // SpanKindIntToString takes an integer/constant protobuf span kind value
 // and returns the string representation used in otel-cli.
 func SpanKindIntToString(kind tracepb.Span_SpanKind) string {
diff --git a/otelcli/root.go b/otelcli/root.go
index 3b412b0..bed3013 100644
--- a/otelcli/root.go
+++ b/otelcli/root.go
@@ -109,6 +109,10 @@ func addSpanParams(cmd *cobra.Command) {
 	// --kind / -k
 	cmd.Flags().StringVarP(&config.Kind, "kind", "k", defaults.Kind, "set the trace kind, e.g. internal, server, client, producer, consumer")
 
+	// expert options: --force-trace-id, --force-span-id allow setting custom trace & span ids
+	cmd.Flags().StringVar(&config.ForceTraceId, "force-trace-id", defaults.ForceTraceId, "expert: force the trace id to be the one provided in hex")
+	cmd.Flags().StringVar(&config.ForceSpanId, "force-span-id", defaults.ForceSpanId, "expert: force the span id to be the one provided in hex")
+
 	addSpanStatusParams(cmd)
 }
 
diff --git a/otelcli/status.go b/otelcli/status.go
index 0557f51..a57d8c3 100644
--- a/otelcli/status.go
+++ b/otelcli/status.go
@@ -36,6 +36,7 @@ func init() {
 	rootCmd.AddCommand(statusCmd)
 	addCommonParams(statusCmd)
 	addClientParams(statusCmd)
+	addSpanParams(statusCmd)
 }
 
 func doStatus(cmd *cobra.Command, args []string) {