From c8bd1e4185ebb68fd8a7082da86bd01135ebf41e Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <156273877+p-datadog@users.noreply.github.com>
Date: Thu, 7 Nov 2024 16:54:19 -0500
Subject: [PATCH] DEBUG-2334 repair DI transport (#4087)

* DEBUG-2334 repair DI transport

"input" endpoint requires a JSON payload, not a
multipart form post.

A better test is included to verify this functionality.
---
 lib/datadog/di/transport.rb       | 25 ++++++++++++++++---------
 sig/datadog/di/transport.rbs      |  2 +-
 spec/datadog/di/transport_spec.rb | 31 +++++++++++++++++++++++++------
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/lib/datadog/di/transport.rb b/lib/datadog/di/transport.rb
index 81b420073b5..b6cf3ed1ddf 100644
--- a/lib/datadog/di/transport.rb
+++ b/lib/datadog/di/transport.rb
@@ -35,11 +35,16 @@ def send_diagnostics(payload)
           StringIO.new(JSON.dump(payload)), 'application/json', 'event.json'
         )
         payload = {'event' => event_payload}
-        send_request('Probe status submission', DIAGNOSTICS_PATH, payload)
+        # Core transport unconditionally specifies headers to underlying
+        # Net::HTTP client, ends up passing 'nil' as headers if none are
+        # specified by us, which then causes Net::HTTP to die with an exception.
+        send_request('Probe status submission',
+          path: DIAGNOSTICS_PATH, form: payload, headers: {})
       end
 
       def send_input(payload)
-        send_request('Probe snapshot submission', INPUT_PATH, payload,
+        send_request('Probe snapshot submission',
+          path: INPUT_PATH, body: payload.to_s,
           headers: {'content-type' => 'application/json'},)
       end
 
@@ -47,19 +52,21 @@ def send_input(payload)
 
       attr_reader :client
 
-      def send_request(desc, path, payload, headers: {})
+      def send_request(desc, **options)
         # steep:ignore:start
-        env = OpenStruct.new(
-          path: path,
-          form: payload,
-          headers: headers,
-        )
+        env = OpenStruct.new(**options)
         # steep:ignore:end
         response = client.post(env)
         unless response.ok?
           raise Error::AgentCommunicationError, "#{desc} failed: #{response.code}: #{response.payload}"
         end
-      rescue IOError, SystemCallError => exc
+      # Datadog::Core::Transport does not perform any exception mapping,
+      # therefore we could have any exception here from failure to parse
+      # agent URI for example.
+      # If we ever implement retries for network errors, we should distinguish
+      # actual network errors from non-network errors that are raised by
+      # transport code.
+      rescue => exc
         raise Error::AgentCommunicationError, "#{desc} failed: #{exc.class}: #{exc}"
       end
     end
diff --git a/sig/datadog/di/transport.rbs b/sig/datadog/di/transport.rbs
index cfd6b1c07fa..e206f3cec3d 100644
--- a/sig/datadog/di/transport.rbs
+++ b/sig/datadog/di/transport.rbs
@@ -17,7 +17,7 @@ module Datadog
 
       attr_reader client: untyped
 
-      def send_request: (String desc, String path, Hash[untyped,untyped] payload, ?headers: ::Hash[untyped, untyped]) -> void
+      def send_request: (String desc, path: String, ?body: String, ?form: Hash[untyped,untyped], ?headers: ::Hash[untyped, untyped]) -> void
     end
   end
 end
diff --git a/spec/datadog/di/transport_spec.rb b/spec/datadog/di/transport_spec.rb
index 45631bc1afb..302cc3ecbe3 100644
--- a/spec/datadog/di/transport_spec.rb
+++ b/spec/datadog/di/transport_spec.rb
@@ -65,14 +65,33 @@
     end
 
     describe '.send_input' do
-      let(:payload) do
-        {}
+      context 'empty payload' do
+        let(:payload) do
+          {}
+        end
+
+        it 'does not raise exceptions' do
+          expect do
+            client.send_input(payload)
+          end.not_to raise_exception
+        end
       end
 
-      it 'does not raise exceptions' do
-        expect do
-          client.send_input(payload)
-        end.not_to raise_exception
+      context 'partial DI payload' do
+        let(:payload) do
+          {
+            service: 'rspec',
+            "debugger.snapshot": {
+              id: '1234',
+            },
+          }
+        end
+
+        it 'does not raise exceptions' do
+          expect do
+            client.send_input(payload)
+          end.not_to raise_exception
+        end
       end
     end