DEBUG-2334 enforce probe type validity (#4013)

* DEBUG-2334 enforce probe type validity Dynamic instrumentation currently supports 4 probe types (log, metric, span and span decoration). Initial implementation of Ruby DI will only support log probes. To avoid installing the other probe types, potentially with incorrect results, add a validation of probe type and require that the probe is of the supported type (currently, just the log probe). * types --------- Co-authored-by: Oleg Pudeyev <[email protected]>
DataDog · Oct 23, 2024 · f8d7b5e · f8d7b5e
1 parent 38f6cc1
commit f8d7b5e
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 21 deletions.
diff --git a/lib/datadog/di/probe.rb b/lib/datadog/di/probe.rb
@@ -31,11 +31,17 @@ module DI
     #
     # @api private
     class Probe
+      KNOWN_TYPES = %i[log].freeze
+
       def initialize(id:, type:,
         file: nil, line_no: nil, type_name: nil, method_name: nil,
         template: nil, capture_snapshot: false, max_capture_depth: nil, rate_limit: nil)
         # Perform some sanity checks here to detect unexpected attribute
         # combinations, in order to not do them in subsequent code.
+        unless KNOWN_TYPES.include?(type)
+          raise ArgumentError, "Unknown probe type: #{type}"
+        end
+
         if line_no && method_name
           raise ArgumentError, "Probe contains both line number and method name: #{id}"
         end

diff --git a/lib/datadog/di/probe_builder.rb b/lib/datadog/di/probe_builder.rb
@@ -17,11 +17,17 @@ module DI
     #
     # @api private
     module ProbeBuilder
+      PROBE_TYPES = {
+        'LOG_PROBE' => :log,
+      }.freeze
+
       module_function def build_from_remote_config(config)
         # The validations here are not yet comprehensive.
+        type = config.fetch('type')
+        type_symbol = PROBE_TYPES[type] or raise ArgumentError, "Unrecognized probe type: #{type}"
         Probe.new(
           id: config.fetch("id"),
-          type: config.fetch("type"),
+          type: type_symbol,
           file: config["where"]&.[]("sourceFile"),
           # Sometimes lines are sometimes received as an array of nil
           # for some reason.

diff --git a/sig/datadog/di/probe.rbs b/sig/datadog/di/probe.rbs
@@ -1,9 +1,11 @@
 module Datadog
   module DI
     class Probe
+      KNOWN_TYPES: Array[Symbol]
+
       @id: String
 
-      @type: String
+      @type: Symbol
 
       @file: String?
 
@@ -19,12 +21,12 @@ module Datadog
 
       @rate_limiter: Datadog::Core::RateLimiter
 
-      def initialize: (id: String, type: String, ?file: String?, ?line_no: Integer?, ?type_name: String?, ?method_name: String?, ?template: String?, ?capture_snapshot: bool,
+      def initialize: (id: String, type: Symbol, ?file: String?, ?line_no: Integer?, ?type_name: String?, ?method_name: String?, ?template: String?, ?capture_snapshot: bool,
 	?max_capture_depth: Integer, ?rate_limit: Integer) -> void
 
       attr_reader id: String
 
-      attr_reader type: String
+      attr_reader type: Symbol
 
       attr_reader file: String?
 

diff --git a/sig/datadog/di/probe_builder.rbs b/sig/datadog/di/probe_builder.rbs
@@ -1,6 +1,8 @@
 module Datadog
   module DI
     module ProbeBuilder
+      PROBE_TYPES: { "LOG_PROBE" => :log }
+
       def self?.build_from_remote_config: (Hash[untyped,untyped] config) -> Probe
     end
   end

diff --git a/spec/datadog/di/probe_builder_spec.rb b/spec/datadog/di/probe_builder_spec.rb
@@ -31,7 +31,7 @@
 
       it "creates line probe with corresponding values" do
         expect(probe.id).to eq "3ecfd456-2d7c-4359-a51f-d4cc44141ffe"
-        expect(probe.type).to eq "LOG_PROBE"
+        expect(probe.type).to eq :log
         expect(probe.file).to eq "aaa.rb"
         expect(probe.line_no).to eq 4321
         expect(probe.type_name).to be nil

diff --git a/spec/datadog/di/probe_spec.rb b/spec/datadog/di/probe_spec.rb
@@ -6,13 +6,13 @@
 
   shared_context "method probe" do
     let(:probe) do
-      described_class.new(id: "42", type: "foo", type_name: "Foo", method_name: "bar")
+      described_class.new(id: "42", type: :log, type_name: 'Foo', method_name: "bar")
     end
   end
 
   shared_context "line probe" do
     let(:probe) do
-      described_class.new(id: "42", type: "foo", file: "foo.rb", line_no: 4)
+      described_class.new(id: "42", type: :log, file: "foo.rb", line_no: 4)
     end
   end
 
@@ -23,7 +23,7 @@
       it "creates an instance" do
         expect(probe).to be_a(described_class)
         expect(probe.id).to eq "42"
-        expect(probe.type).to eq "foo"
+        expect(probe.type).to eq :log
         expect(probe.type_name).to eq "Foo"
         expect(probe.method_name).to eq "bar"
         expect(probe.file).to be nil
@@ -37,17 +37,31 @@
       it "creates an instance" do
         expect(probe).to be_a(described_class)
         expect(probe.id).to eq "42"
-        expect(probe.type).to eq "foo"
+        expect(probe.type).to eq :log
         expect(probe.type_name).to be nil
         expect(probe.method_name).to be nil
         expect(probe.file).to eq "foo.rb"
         expect(probe.line_no).to eq 4
       end
     end
 
+    context 'unsupported type' do
+      let(:probe) do
+        # LOG_PROBE is a valid type in RC probe specification but not
+        # as an argument to Probe constructor.
+        described_class.new(id: '42', type: 'LOG_PROBE', file: 'x', line_no: 1)
+      end
+
+      it 'raises ArgumentError' do
+        expect do
+          probe
+        end.to raise_error(ArgumentError, /Unknown probe type/)
+      end
+    end
+
     context "neither method nor line" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo")
+        described_class.new(id: "42", type: :log)
       end
 
       it "raises ArgumentError" do
@@ -59,7 +73,7 @@
 
     context "both method and line" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo",
+        described_class.new(id: "42", type: :log,
           type_name: "foo", method_name: "bar", file: "baz", line_no: 4)
       end
 
@@ -74,7 +88,7 @@
   describe "#line?" do
     context "line probe" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo", file: "bar.rb", line_no: 5)
+        described_class.new(id: "42", type: :log, file: "bar.rb", line_no: 5)
       end
 
       it "is true" do
@@ -84,7 +98,7 @@
 
     context "method probe" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar")
+        described_class.new(id: "42", type: :log, type_name: "FooClass", method_name: "bar")
       end
 
       it "is false" do
@@ -94,7 +108,7 @@
 
     context "method probe with file name" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar", file: "quux.rb")
+        described_class.new(id: "42", type: :log, type_name: "FooClass", method_name: "bar", file: "quux.rb")
       end
 
       it "is false" do
@@ -106,7 +120,7 @@
   describe "#method?" do
     context "line probe" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo", file: "bar.rb", line_no: 5)
+        described_class.new(id: "42", type: :log, file: "bar.rb", line_no: 5)
       end
 
       it "is false" do
@@ -116,7 +130,7 @@
 
     context "method probe" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar")
+        described_class.new(id: "42", type: :log, type_name: "FooClass", method_name: "bar")
       end
 
       it "is true" do
@@ -126,7 +140,7 @@
 
     context "method probe with file name" do
       let(:probe) do
-        described_class.new(id: "42", type: "foo", type_name: "FooClass", method_name: "bar", file: "quux.rb")
+        described_class.new(id: "42", type: :log, type_name: "FooClass", method_name: "bar", file: "quux.rb")
       end
 
       it "is true" do
@@ -137,15 +151,15 @@
 
   describe "#line_no" do
     context "one line number" do
-      let(:probe) { described_class.new(id: "x", type: "log", line_no: 5) }
+      let(:probe) { described_class.new(id: "x", type: :log, line_no: 5) }
 
       it "returns the line number" do
         expect(probe.line_no).to eq 5
       end
     end
 
     context "nil line number" do
-      let(:probe) { described_class.new(id: "id", type: "LOG", type_name: "x", method_name: "y", line_no: nil) }
+      let(:probe) { described_class.new(id: "id", type: :log, type_name: "x", method_name: "y", line_no: nil) }
 
       it "returns nil" do
         expect(probe.line_no).to be nil
@@ -155,15 +169,15 @@
 
   describe "#line_no!" do
     context "one line number" do
-      let(:probe) { described_class.new(id: "x", type: "log", line_no: 5) }
+      let(:probe) { described_class.new(id: "x", type: :log, line_no: 5) }
 
       it "returns the line number" do
         expect(probe.line_no!).to eq 5
       end
     end
 
     context "nil line number" do
-      let(:probe) { described_class.new(id: "id", type: "LOG", type_name: "x", method_name: "y", line_no: nil) }
+      let(:probe) { described_class.new(id: "id", type: :log, type_name: "x", method_name: "y", line_no: nil) }
 
       it "raises MissingLineNumber" do
         expect do