From 29bf4d226dbe39129a62a4e369c850c998cea7e9 Mon Sep 17 00:00:00 2001
From: Tony Hsu <tonyc.t.hsu@gmail.com>
Date: Mon, 17 Jul 2023 13:32:38 +0200
Subject: [PATCH] Enforce lower 64 bit coercion on #trace_id

---
 lib/datadog/tracing/correlation.rb       | 21 +++----
 spec/datadog/tracing/correlation_spec.rb | 70 ++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 12 deletions(-)

diff --git a/lib/datadog/tracing/correlation.rb b/lib/datadog/tracing/correlation.rb
index 653dd1f8c7d..7fd8edd520b 100644
--- a/lib/datadog/tracing/correlation.rb
+++ b/lib/datadog/tracing/correlation.rb
@@ -23,7 +23,6 @@ class Identifier
           :span_resource,
           :span_service,
           :span_type,
-          :trace_id,
           :trace_name,
           :trace_resource,
           :trace_service,
@@ -65,22 +64,20 @@ def to_log_format
             attributes << "#{LOG_ATTR_ENV}=#{env}" unless env.nil?
             attributes << "#{LOG_ATTR_SERVICE}=#{service}"
             attributes << "#{LOG_ATTR_VERSION}=#{version}" unless version.nil?
-            attributes << "#{LOG_ATTR_TRACE_ID}=#{logging_trace_id}"
+            attributes << "#{LOG_ATTR_TRACE_ID}=#{trace_id}"
             attributes << "#{LOG_ATTR_SPAN_ID}=#{span_id}"
             attributes.join(' ')
           end
         end
 
-        private
-
-        def logging_trace_id
-          @logging_trace_id ||=
-            if Datadog.configuration.tracing.trace_id_128_bit_logging_enabled &&
-                !Tracing::Utils::TraceId.to_high_order(@trace_id).zero?
-              Kernel.format('%032x', trace_id)
-            else
-              Tracing::Utils::TraceId.to_low_order(@trace_id)
-            end
+        # DEV-2.0: This public method was returning an Integer, but with 128 bit trace id it would return a String.
+        def trace_id
+          if Datadog.configuration.tracing.trace_id_128_bit_logging_enabled &&
+              !Tracing::Utils::TraceId.to_high_order(@trace_id).zero?
+            Kernel.format('%032x', @trace_id)
+          else
+            Tracing::Utils::TraceId.to_low_order(@trace_id)
+          end
         end
       end
 
diff --git a/spec/datadog/tracing/correlation_spec.rb b/spec/datadog/tracing/correlation_spec.rb
index caa3d814207..6145d1f37d5 100644
--- a/spec/datadog/tracing/correlation_spec.rb
+++ b/spec/datadog/tracing/correlation_spec.rb
@@ -424,5 +424,75 @@ def have_attribute(attribute)
         end
       end
     end
+
+    describe '#trace_id' do
+      context 'is defined' do
+        context 'when 128 bit trace id logging is not enabled' do
+          before do
+            allow(Datadog.configuration.tracing).to receive(:trace_id_128_bit_logging_enabled).and_return(false)
+          end
+
+          context 'when given 64 bit trace id' do
+            it 'returns to lower 64 bits of trace id' do
+              trace_id = 0xaaaaaaaaaaaaaaaa
+              expected_trace_id = 0xaaaaaaaaaaaaaaaa
+
+              identifier = described_class.new(trace_id: trace_id)
+
+              expect(identifier.trace_id).to eq(expected_trace_id)
+            end
+          end
+
+          context 'when given 128 bit trace id' do
+            it 'returns to lower 64 bits of trace id' do
+              trace_id = 0xaaaaaaaaaaaaaaaaffffffffffffffff
+              expected_trace_id = 0xffffffffffffffff
+
+              identifier = described_class.new(trace_id: trace_id)
+
+              expect(identifier.trace_id).to eq(expected_trace_id)
+            end
+          end
+        end
+
+        context 'when 128 bit trace id logging is enabled' do
+          before do
+            allow(Datadog.configuration.tracing).to receive(:trace_id_128_bit_logging_enabled).and_return(true)
+          end
+
+          context 'when given 64 bit trace id' do
+            it 'returns lower 64 bits of trace id' do
+              trace_id = 0xaaaaaaaaaaaaaaaa
+              expected_trace_id = 0xaaaaaaaaaaaaaaaa
+
+              identifier = described_class.new(trace_id: trace_id)
+
+              expect(identifier.trace_id).to eq(expected_trace_id)
+            end
+          end
+
+          context 'when given > 64 bit trace id' do
+            it 'returns the entire trace id in hex encoded and zero padded format' do
+              trace_id = 0x00ffffffffffffffaaaaaaaaaaaaaaaa
+
+              identifier = described_class.new(trace_id: trace_id)
+
+              expect(identifier.trace_id).to eq('00ffffffffffffffaaaaaaaaaaaaaaaa')
+            end
+          end
+        end
+
+        context 'when given > 64 bit trace id but high order is 0' do
+          it 'returns to lower 64 bits of trace id' do
+            trace_id = 0x00000000000000000aaaaaaaaaaaaaaaa
+            expected_trace_id = 0xaaaaaaaaaaaaaaaa
+
+            identifier = described_class.new(trace_id: trace_id)
+
+            expect(identifier.trace_id).to eq(expected_trace_id)
+          end
+        end
+      end
+    end
   end
 end