Encapsulate inferred route reading and writing

Author: y9v

lib/datadog/appsec/api_security/route_extractor.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require_relative '../../tracing/contrib/rack/route_from_path_inference'
+require_relative '../../tracing/contrib/rack/route_inference'
 
 module Datadog
   module AppSec
@@ -68,9 +68,7 @@ def self.route_pattern(request)
             #       to generic request path
             (pattern || request.path).delete_suffix(RAILS_FORMAT_SUFFIX)
           else
-            Tracing::Contrib::Rack::RouteFromPathInference.infer(request.path).tap do |inferred_route|
-              request.env[Tracing::Contrib::Rack::Ext::DATADOG_INFERRED_ROUTE] = inferred_route
-            end
+            Tracing::Contrib::Rack::RouteInference.read_or_infer(request.env)
           end
         end
       end

lib/datadog/tracing/contrib/rack/ext.rb

@@ -12,7 +12,6 @@ module Ext
           ENV_ANALYTICS_ENABLED = 'DD_TRACE_RACK_ANALYTICS_ENABLED'
           ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_RACK_ANALYTICS_SAMPLE_RATE'
           RACK_ENV_REQUEST_SPAN = 'datadog.rack_request_span'
-          DATADOG_INFERRED_ROUTE = 'datadog.inferred_route'
           SPAN_HTTP_PROXY_REQUEST = 'http.proxy.request'
           SPAN_HTTP_PROXY_QUEUE = 'http.proxy.queue'
           SPAN_REQUEST = 'rack.request'

lib/datadog/tracing/contrib/rack/middlewares.rb

@@ -13,7 +13,7 @@
 require_relative 'header_collection'
 require_relative 'header_tagging'
 require_relative 'request_queue'
-require_relative 'route_from_path_inference'
+require_relative 'route_inference'
 require_relative 'trace_proxy_middleware'
 
 module Datadog
@@ -235,7 +235,7 @@ def set_route_and_endpoint_tags(trace:, request_span:, status:, env:)
               # we infer the route from the full request path.
               if last_script_name == '' && env['SCRIPT_NAME'] != '' &&
                   !Datadog.configuration.tracing.resource_renaming.always_simplified_endpoint &&
-                  (inferred_route = RouteFromPathInference.infer(env['PATH_INFO']))
+                  (inferred_route = RouteInference.infer(env['PATH_INFO']))
                 set_endpoint_tag(request_span, last_route + inferred_route)
               end
 
@@ -250,16 +250,9 @@ def set_route_and_endpoint_tags(trace:, request_span:, status:, env:)
 
             if Datadog.configuration.tracing.resource_renaming.always_simplified_endpoint ||
                 request_span.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE).nil?
-              # when http.route tag wasn't set or resource_renaming.always_simplified_endpoint is set to true,
-              # we will try to infer the path for the full request path,
-              # which is SCRIPT_NAME + PATH_INFO for mounted rack applications
-              full_path = env['SCRIPT_NAME'].to_s + env['PATH_INFO'].to_s
-
-              # Inferred route calculation might have been done already for AppSec API Security sampling,
-              # in such case the calculation result is stored in request env
-              inferred_route = env[Tracing::Contrib::Rack::Ext::DATADOG_INFERRED_ROUTE] || RouteFromPathInference.infer(full_path)
-
-              set_endpoint_tag(request_span, inferred_route) if inferred_route
+              if (inferred_route = RouteInference.read_or_infer(env))
+                set_endpoint_tag(request_span, inferred_route) if inferred_route
+              end
             elsif !request_span.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ENDPOINT)
               set_endpoint_tag(request_span, request_span.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE))
             end

lib/datadog/tracing/contrib/rack/route_from_path_inference.rb renamed to lib/datadog/tracing/contrib/rack/route_inference.rb

@@ -6,7 +6,7 @@ module Contrib
       module Rack
         # This module provides logic for inferring HTTP route pattern
         # from an HTTP path.
-        module RouteFromPathInference
+        module RouteInference
           MAX_NUMBER_OF_SEGMENTS = 8
 
           INT_PARAM_REGEX = /\A[0-9]+\z/.freeze
@@ -15,8 +15,15 @@ module RouteFromPathInference
           HEX_ID_PARAM_REGEX = /\A(?=.*\d)[A-Fa-f0-9._-]{6,}\z/.freeze
           STRING_PARAM_REGEX = /\A.{20,}|.*[%&'()*+,:=@].*\z/.freeze
 
+          DATADOG_INFERRED_ROUTE_ENV_KEY = 'datadog.inferred_route'
+
           module_function
 
+          def read_or_infer(request_env)
+            request_env[DATADOG_INFERRED_ROUTE_ENV_KEY] ||=
+              infer(request_env['SCRIPT_NAME'].to_s + request_env['PATH_INFO'].to_s)
+          end
+
           def infer(path)
             segments = path.delete_prefix('/').split('/', MAX_NUMBER_OF_SEGMENTS + 1).first(MAX_NUMBER_OF_SEGMENTS)
 

sig/datadog/tracing/contrib/rack/ext.rbs

@@ -11,8 +11,6 @@ module Datadog
 
           RACK_ENV_REQUEST_SPAN: "datadog.rack_request_span"
 
-          DATADOG_INFERRED_ROUTE: ::String
-
           SPAN_HTTP_PROXY_REQUEST: "http.proxy.request"
 
           SPAN_HTTP_PROXY_QUEUE: "http.proxy.queue"

sig/datadog/tracing/contrib/rack/route_from_path_inference.rbs renamed to sig/datadog/tracing/contrib/rack/route_inference.rbs

@@ -2,7 +2,7 @@ module Datadog
   module Tracing
     module Contrib
       module Rack
-        module RouteFromPathInference
+        module RouteInference
           MAX_NUMBER_OF_SEGMENTS: ::Integer
 
           INT_PARAM_REGEX: ::Regexp
@@ -11,6 +11,10 @@ module Datadog
           HEX_ID_PARAM_REGEX: ::Regexp
           STRING_PARAM_REGEX: ::Regexp
 
+          DATADOG_INFERRED_ROUTE_ENV_KEY: ::String
+
+          def self.read_or_infer: (Hash[::String, untyped] request_env) -> ::String?
+
           def self.infer: (::String path) -> ::String?
         end
       end

spec/datadog/tracing/contrib/rack/route_from_path_inference_spec.rb renamed to spec/datadog/tracing/contrib/rack/route_inference_spec.rb

@@ -1,8 +1,46 @@
 require 'datadog/tracing/contrib/support/spec_helper'
 
-require 'datadog/tracing/contrib/rack/route_from_path_inference'
+require 'datadog/tracing/contrib/rack/route_inference'
+
+RSpec.describe Datadog::Tracing::Contrib::Rack::RouteInference do
+  describe '.read_or_infer' do
+    context 'when inferred route was not yet persisted in request env' do
+      let(:env) do
+        {
+          'SCRIPT_NAME' => '/api',
+          'PATH_INFO' => '/users/1'
+        }
+      end
+
+      it 'returns inferred route' do
+        expect(described_class.read_or_infer(env)).to eq('/api/users/{param:int}')
+      end
+
+      it 'persists inferred route in request env' do
+        expect { described_class.read_or_infer(env) }.to change { env['datadog.inferred_route'] }
+          .from(nil).to('/api/users/{param:int}')
+      end
+    end
+
+    context 'when inferred route already persisted in request env' do
+      let(:env) do
+        {
+          'SCRIPT_NAME' => '/api',
+          'PATH_INFO' => '/users/1',
+          'datadog.inferred_route' => '/some_route'
+        }
+      end
+
+      it 'returns persisted inferred route' do
+        expect(described_class.read_or_infer(env)).to eq('/some_route')
+      end
+
+      it 'does not change inferred route value in request env' do
+        expect { described_class.read_or_infer(env) }.not_to change { env['datadog.inferred_route'] }
+      end
+    end
+  end
 
-RSpec.describe Datadog::Tracing::Contrib::Rack::RouteFromPathInference do
   describe '.infer' do
     it 'works with an empty path' do
       expect(described_class.infer('/')).to eq('/')