From c1470ebc216abaee0c597e56a3f697f075051e84 Mon Sep 17 00:00:00 2001
From: Marco Costa <marco.costa@datadoghq.com>
Date: Wed, 21 Dec 2022 16:35:23 -0800
Subject: [PATCH] Flexible header matching for HTTP propagator

---
 .../contrib/http/distributed/fetcher.rb       | 13 ++++++++---
 .../contrib/http/distributed/fetcher_spec.rb  | 23 +++++++++++++++----
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/lib/datadog/tracing/contrib/http/distributed/fetcher.rb b/lib/datadog/tracing/contrib/http/distributed/fetcher.rb
index b5c8936839..dc58f0491c 100644
--- a/lib/datadog/tracing/contrib/http/distributed/fetcher.rb
+++ b/lib/datadog/tracing/contrib/http/distributed/fetcher.rb
@@ -8,13 +8,20 @@ module Tracing
     module Contrib
       module HTTP
         module Distributed
-          # Retrieves Rack formatted headers from HTTP headers.
+          # Retrieves HTTP headers from carrier.
+          # Headers will also match if Rack-formatted:
+          # 'my-header' will match 'my-header' and 'HTTP_MY_HEADER'.
+          #
+          # In case both variants are present, the verbatim match will be used.
           class Fetcher < Tracing::Distributed::Fetcher
-            # TODO: Don't assume Rack format.
-            #       Make distributed tracing headers apathetic.
             # DEV: Should we try to parse both verbatim an Rack-formatted headers,
             # DEV: given Rack-formatted is the most common format in Ruby?
             def [](name)
+              # Try to fetch with the plain key
+              value = super(name)
+              return value if value && !value.empty?
+
+              # If not found, try the Rack-formatted key
               rack_header = "HTTP-#{name}"
               rack_header.upcase!
               rack_header.tr!('-'.freeze, '_'.freeze)
diff --git a/spec/datadog/tracing/contrib/http/distributed/fetcher_spec.rb b/spec/datadog/tracing/contrib/http/distributed/fetcher_spec.rb
index 3922c19b6b..07b74c2a02 100644
--- a/spec/datadog/tracing/contrib/http/distributed/fetcher_spec.rb
+++ b/spec/datadog/tracing/contrib/http/distributed/fetcher_spec.rb
@@ -29,14 +29,29 @@
       context 'that is not empty' do
         let(:env) { { 'HTTP_MY_KEY' => 'value' } }
         it { is_expected.to eq('value') }
+
+        context 'and a plain header' do
+          let(:env) { super().merge('my-key' => 'plain-match') }
+
+          it 'prefers the plain header match' do
+            is_expected.to eq('plain-match')
+          end
+        end
       end
     end
 
-    context 'with a header not Rack formatted' do
-      let(:key) { 'my-key' }
-      let(:env) { { key => 'value' } }
+    context 'with a plain header associated' do
+      let(:key) { 'rack.session' }
 
-      it { is_expected.to be_nil }
+      context 'that is empty' do
+        let(:env) { { key => '' } }
+        it { is_expected.to be_nil }
+      end
+
+      context 'that is not empty' do
+        let(:env) { { key => 'value' } }
+        it { is_expected.to eq('value') }
+      end
     end
   end
 end