aikido-zen 0.2.0-x86_64-mingw-64 → 1.0.1.beta.2-x86_64-mingw-64

data/lib/aikido/zen/sinks/net_http.rb

@@ -1,25 +1,32 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Net
       module HTTP
+        def self.load_sinks!
+          # In stdlib but not always required
+          require "net/http"
+
+          ::Net::HTTP.prepend(Net::HTTP::HTTPExtensions)
+        end
+
         SINK = Sinks.add("net-http", scanners: [
-          Aikido::Zen::Scanners::SSRFScanner,
-          Aikido::Zen::OutboundConnectionMonitor
+          Scanners::SSRFScanner,
+          OutboundConnectionMonitor
         ])
 
-        module Extensions
+        module Helpers
           # Maps a Net::HTTP connection to an Aikido OutboundConnection,
           # which our tooling expects.
           #
           # @param http [Net::HTTP]
           # @return [Aikido::Zen::OutboundConnection]
           def self.build_outbound(http)
-            Aikido::Zen::OutboundConnection.new(
+            OutboundConnection.new(
               host: http.address,
               port: http.port
             )
@@ -34,7 +41,7 @@ module Aikido::Zen
               path: req.path
             }))
 
-            Aikido::Zen::Scanners::SSRFScanner::Request.new(
+            Scanners::SSRFScanner::Request.new(
               verb: req.method,
               uri: uri,
               headers: req.to_hash,
@@ -43,33 +50,44 @@ module Aikido::Zen
           end
 
           def self.wrap_response(response)
-            Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            Scanners::SSRFScanner::Response.new(
               status: response.code.to_i,
               headers: response.to_hash,
               header_normalizer: ->(val) { Array(val).join(", ") }
             )
           end
 
-          def request(req, *)
-            wrapped_request = Extensions.wrap_request(req, self)
+          def self.scan(request, connection, operation)
+            SINK.scan(
+              request: request,
+              connection: connection,
+              operation: operation
+            )
+          end
+        end
+
+        module HTTPExtensions
+          extend Sinks::DSL
+
+          sink_around :request do |super_call, req|
+            wrapped_request = Helpers.wrap_request(req, self)
 
             # Store the request information so the DNS sinks can pick it up.
-            if (context = Aikido::Zen.current_context)
+            context = Aikido::Zen.current_context
+            if context
               prev_request = context["ssrf.request"]
               context["ssrf.request"] = wrapped_request
             end
 
-            SINK.scan(
-              connection: Extensions.build_outbound(self),
-              request: wrapped_request,
-              operation: "request"
-            )
+            connection = Helpers.build_outbound(self)
+
+            Helpers.scan(wrapped_request, connection, "request")
 
-            response = super
+            response = super_call.call
 
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+            Scanners::SSRFScanner.track_redirects(
               request: wrapped_request,
-              response: Extensions.wrap_response(response)
+              response: Helpers.wrap_response(response)
             )
 
             response
@@ -82,4 +100,4 @@ module Aikido::Zen
   end
 end
 
-::Net::HTTP.prepend(Aikido::Zen::Sinks::Net::HTTP::Extensions)
+Aikido::Zen::Sinks::Net::HTTP.load_sinks!

data/lib/aikido/zen/sinks/patron.rb

@@ -1,56 +1,75 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Patron
+      def self.load_sinks!
+        if Gem.loaded_specs["patron"]
+          require "patron"
+
+          ::Patron::Session.prepend(SessionExtensions)
+        end
+      end
+
       SINK = Sinks.add("patron", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
       ])
 
-      module Extensions
+      module Helpers
         def self.wrap_response(request, response)
           # In this case, automatic redirection happened inside libcurl.
           if response.url != request.url && !response.url.to_s.empty?
-            Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            Scanners::SSRFScanner::Response.new(
               status: 302, # We can't know what the actual status was, but we just need a 3XX
               headers: response.headers.merge("Location" => response.url)
             )
           else
-            Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            Scanners::SSRFScanner::Response.new(
               status: response.status,
               headers: response.headers
             )
           end
         end
 
-        def handle_request(request)
-          wrapped_request = Aikido::Zen::Scanners::SSRFScanner::Request.new(
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module SessionExtensions
+        extend Sinks::DSL
+
+        sink_around :handle_request do |super_call, request|
+          wrapped_request = Scanners::SSRFScanner::Request.new(
             verb: request.action,
             uri: URI(request.url),
             headers: request.headers
           )
 
           # Store the request information so the DNS sinks can pick it up.
-          if (context = Aikido::Zen.current_context)
+          context = Aikido::Zen.current_context
+          if context
             prev_request = context["ssrf.request"]
             context["ssrf.request"] = wrapped_request
           end
 
-          SINK.scan(
-            connection: Aikido::Zen::OutboundConnection.from_uri(URI(request.url)),
-            request: wrapped_request,
-            operation: "request"
-          )
+          connection = OutboundConnection.from_uri(URI(request.url))
+
+          Helpers.scan(wrapped_request, connection, "request")
 
-          response = super
+          response = super_call.call
 
-          Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+          Scanners::SSRFScanner.track_redirects(
             request: wrapped_request,
-            response: Extensions.wrap_response(request, response)
+            response: Helpers.wrap_response(request, response)
           )
 
           # When libcurl has follow_location set, it will handle redirections
@@ -62,18 +81,16 @@ module Aikido::Zen
           # stop the response from being exposed to the user. This downgrades
           # the SSRF into a blind SSRF, which is better than doing nothing.
           if request.url != response.url && !response.url.to_s.empty?
-            last_effective_request = Aikido::Zen::Scanners::SSRFScanner::Request.new(
+            last_effective_request = Scanners::SSRFScanner::Request.new(
               verb: request.action,
               uri: URI(response.url),
               headers: request.headers
             )
             context["ssrf.request"] = last_effective_request if context
 
-            SINK.scan(
-              connection: Aikido::Zen::OutboundConnection.from_uri(URI(response.url)),
-              request: last_effective_request,
-              operation: "request"
-            )
+            connection = OutboundConnection.from_uri(URI(response.url))
+
+            Helpers.scan(last_effective_request, connection, "request")
           end
 
           response
@@ -85,4 +102,4 @@ module Aikido::Zen
   end
 end
 
-::Patron::Session.prepend(Aikido::Zen::Sinks::Patron::Extensions)
+Aikido::Zen::Sinks::Patron.load_sinks!

data/lib/aikido/zen/sinks/pg.rb

@@ -1,51 +1,74 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
-
 module Aikido::Zen
   module Sinks
     module PG
+      def self.load_sinks!
+        if Gem.loaded_specs["pg"]
+          require "pg"
+
+          ::PG::Connection.prepend(PG::ConnectionExtensions)
+        end
+      end
+
       SINK = Sinks.add("pg", scanners: [Scanners::SQLInjectionScanner])
 
-      # For some reason, the ActiveRecord pg adapter does not wrap exceptions in
-      # StatementInvalid, which leads to inconsistent handling. This guarantees
-      # that all Zen errors are wrapped in a StatementInvalid, so documentation
-      # can be consistent.
-      WRAP_EXCEPTIONS = if defined?(ActiveRecord::StatementInvalid)
-        <<~RUBY
-          rescue Aikido::Zen::SQLInjectionError
-            raise ActiveRecord::StatementInvalid
-        RUBY
+      module Helpers
+        # For some reason, the ActiveRecord pg adaptor does not wrap exceptions
+        # in ActiveRecord::StatementInvalid, leading to inconsistent handling.
+        # This guarantees that Aikido::Zen::SQLInjectionErrors are wrapped in
+        # an ActiveRecord::StatementInvalid.
+        def self.safe(&block)
+          # Code coverage is disabled here because this ActiveRecord behavior is
+          # exercised in end-to-end tests, which are not covered by SimpleCov.
+          # :nocov:
+          if !defined?(ActiveRecord::StatementInvalid)
+            Sinks::DSL.safe(&block)
+          else
+            begin
+              Sinks::DSL.safe(&block)
+            rescue Aikido::Zen::SQLInjectionError => err
+              raise ActiveRecord::StatementInvalid, cause: err
+            end
+          end
+          # :nocov:
+        end
+
+        def self.scan(query, operation)
+          SINK.scan(
+            query: query,
+            dialect: :postgresql,
+            operation: operation
+          )
+        end
       end
 
-      module Extensions
+      module ConnectionExtensions
+        extend Sinks::DSL
+
         %i[
           send_query exec sync_exec async_exec
           send_query_params exec_params sync_exec_params async_exec_params
-        ].each do |method|
-          module_eval <<~RUBY, __FILE__, __LINE__ + 1
-            def #{method}(query, *)
-              SINK.scan(query: query, dialect: :postgresql, operation: :#{method})
-              super
-            #{WRAP_EXCEPTIONS}
+        ].each do |method_name|
+          presafe_sink_before method_name do |query|
+            Helpers.safe do
+              Helpers.scan(query, method_name)
             end
-          RUBY
+          end
         end
 
         %i[
           send_prepare prepare async_prepare sync_prepare
-        ].each do |method|
-          module_eval <<~RUBY, __FILE__, __LINE__ + 1
-            def #{method}(_, query, *)
-              SINK.scan(query: query, dialect: :postgresql, operation: :#{method})
-              super
-            #{WRAP_EXCEPTIONS}
+        ].each do |method_name|
+          presafe_sink_before method_name do |_, query|
+            Helpers.safe do
+              Helpers.scan(query, method_name)
             end
-          RUBY
+          end
         end
       end
     end
   end
 end
 
-::PG::Connection.prepend(Aikido::Zen::Sinks::PG::Extensions)
+Aikido::Zen::Sinks::PG.load_sinks!

data/lib/aikido/zen/sinks/resolv.rb

@@ -1,28 +1,28 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
 require_relative "../scanners/stored_ssrf_scanner"
 require_relative "../scanners/ssrf_scanner"
 
 module Aikido::Zen
   module Sinks
     module Resolv
+      def self.load_sinks!
+        # In stdlib but not always required
+        require "resolv"
+
+        ::Resolv.prepend(ResolvExtensions)
+      end
+
       SINK = Sinks.add("resolv", scanners: [
-        Aikido::Zen::Scanners::StoredSSRFScanner,
-        Aikido::Zen::Scanners::SSRFScanner
+        Scanners::StoredSSRFScanner,
+        Scanners::SSRFScanner
       ])
 
-      module Extensions
-        def each_address(name, &block)
-          addresses = []
-
-          super do |address|
-            addresses << address
-            yield address
-          end
-        ensure
-          if (context = Aikido::Zen.current_context)
-            context["dns.lookups"] ||= Aikido::Zen::Scanners::SSRF::DNSLookups.new
+      module Helpers
+        def self.scan(name, addresses, operation)
+          context = Aikido::Zen.current_context
+          if context
+            context["dns.lookups"] ||= Scanners::SSRF::DNSLookups.new
             context["dns.lookups"].add(name, addresses)
           end
 
@@ -30,12 +30,33 @@ module Aikido::Zen
             hostname: name,
             addresses: addresses,
             request: context && context["ssrf.request"],
-            operation: "lookup"
+            operation: operation
           )
         end
       end
+
+      module ResolvExtensions
+        def each_address(*args, **kwargs, &blk)
+          # each_address is defined "manually" because no sink method pattern
+          # is applicable.
+
+          name, = args
+
+          addresses = []
+          super(*args, **kwargs) do |address| # rubocop:disable Style/SuperArguments
+            addresses << address
+            blk.call(address)
+          end
+        ensure
+          # Ensure partial results are scanned.
+
+          Sinks::DSL.safe do
+            Helpers.scan(name, addresses, "lookup")
+          end
+        end
+      end
     end
   end
 end
 
-::Resolv.prepend(Aikido::Zen::Sinks::Resolv::Extensions)
+Aikido::Zen::Sinks::Resolv.load_sinks!

data/lib/aikido/zen/sinks/socket.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 require "socket"
+require_relative "../scanners/stored_ssrf_scanner"
+require_relative "../scanners/ssrf_scanner"
 
 module Aikido::Zen
   module Sinks
@@ -8,44 +10,71 @@ module Aikido::Zen
     # there's no way to access the internal DNS resolution that happens in C
     # when using the socket primitives.
     module Socket
+      def self.load_sinks!
+        ::IPSocket.singleton_class.prepend(Socket::IPSocketExtensions)
+      end
+
       SINK = Sinks.add("socket", scanners: [
-        Aikido::Zen::Scanners::StoredSSRFScanner,
-        Aikido::Zen::Scanners::SSRFScanner
+        Scanners::StoredSSRFScanner,
+        Scanners::SSRFScanner
       ])
 
-      module IPSocketExtensions
-        def self.scan_socket(hostname, socket)
+      module Helpers
+        def self.scan(hostname, socket, operation)
+          # We're patching IPSocket.open(..) method.
+          # The IPSocket class hierarchy is:
+          #             IPSocket
+          #            /        \
+          #       TCPSocket    UDPSocket
+          #       /       \
+          # TCPServer     SOCKSSocket
+          #
+          # Because we want to scan only HTTP requests, we skip in case the
+          # socket is not *exactly* an instance of TCPSocket — it's any
+          # of it subclasses
+          return unless socket.instance_of?(TCPSocket)
+
           # ["AF_INET", 80, "10.0.0.1", "10.0.0.1"]
-          addr_family, *, remote_address = socket.peeraddr
+          address_family, _port, _hostname, numeric_address = socket.peeraddr(:numeric)
 
           # We only care about IPv4 (AF_INET) or IPv6 (AF_INET6) sockets
           # This might be overcautious, since this is _IP_Socket, so you
           # would expect it's only used for IP connections?
-          return unless addr_family.start_with?("AF_INET")
 
-          if (context = Aikido::Zen.current_context)
-            context["dns.lookups"] ||= Aikido::Zen::Scanners::SSRF::DNSLookups.new
-            context["dns.lookups"].add(hostname, remote_address)
+          # Code coverage is disabled here because the then clause is a no-op,
+          # so there is nothing to cover.
+          # :nocov:
+          return unless address_family.start_with?("AF_INET")
+          # :nocov:
+
+          context = Aikido::Zen.current_context
+          if context
+            context["dns.lookups"] ||= Scanners::SSRF::DNSLookups.new
+            context["dns.lookups"].add(hostname, numeric_address)
           end
 
           SINK.scan(
             hostname: hostname,
-            addresses: [remote_address],
+            addresses: [numeric_address],
             request: context && context["ssrf.request"],
-            operation: "open"
+            operation: operation
           )
         end
+      end
 
-        def open(name, *)
-          socket = super
-
-          IPSocketExtensions.scan_socket(name, socket)
+      module IPSocketExtensions
+        extend Sinks::DSL
 
-          socket
+        sink_after :open do |socket, remote_host|
+          # Code coverage is disabled here because the tests are contrived and
+          # intentionally do not call open.
+          # :nocov:
+          Helpers.scan(remote_host, socket, "open")
+          # :nocov:
         end
       end
     end
   end
 end
 
-::IPSocket.singleton_class.prepend(Aikido::Zen::Sinks::Socket::IPSocketExtensions)
+Aikido::Zen::Sinks::Socket.load_sinks!

data/lib/aikido/zen/sinks/sqlite3.rb

@@ -1,30 +1,49 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
-
 module Aikido::Zen
   module Sinks
     module SQLite3
+      def self.load_sinks!
+        if Gem.loaded_specs["sqlite3"]
+          require "sqlite3"
+
+          ::SQLite3::Database.prepend(DatabaseExtensions)
+          ::SQLite3::Statement.prepend(StatementExtensions)
+        end
+      end
+
       SINK = Sinks.add("sqlite3", scanners: [Scanners::SQLInjectionScanner])
 
-      module DatabaseExt
-        def exec_batch(sql, *)
-          SINK.scan(query: sql, dialect: :sqlite, operation: "exec_batch")
+      module Helpers
+        def self.scan(query, operation)
+          SINK.scan(
+            query: query,
+            dialect: :sqlite,
+            operation: operation
+          )
+        end
+      end
+
+      module DatabaseExtensions
+        extend Sinks::DSL
+
+        private
 
-          super
+        # SQLite3::Database#exec_batch is an internal native private method.
+        sink_before :exec_batch do |sql|
+          Helpers.scan(sql, "exec_batch")
         end
       end
 
-      module StatementExt
-        def initialize(_, sql, *)
-          SINK.scan(query: sql, dialect: :sqlite, operation: "statement.execute")
+      module StatementExtensions
+        extend Sinks::DSL
 
-          super
+        sink_before :initialize do |_db, sql|
+          Helpers.scan(sql, "statement.execute")
         end
       end
     end
   end
 end
 
-::SQLite3::Database.prepend(Aikido::Zen::Sinks::SQLite3::DatabaseExt)
-::SQLite3::Statement.prepend(Aikido::Zen::Sinks::SQLite3::StatementExt)
+Aikido::Zen::Sinks::SQLite3.load_sinks!

data/lib/aikido/zen/sinks/trilogy.rb

@@ -1,21 +1,33 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
-
 module Aikido::Zen
   module Sinks
     module Trilogy
+      def self.load_sinks!
+        if Gem.loaded_specs["trilogy"]
+          require "trilogy"
+
+          ::Trilogy.prepend(TrilogyExtensions)
+        end
+      end
+
       SINK = Sinks.add("trilogy", scanners: [Scanners::SQLInjectionScanner])
 
-      module Extensions
-        def query(query, *)
-          SINK.scan(query: query, dialect: :mysql, operation: "query")
+      module Helpers
+        def self.scan(query, operation)
+          SINK.scan(query: query, dialect: :mysql, operation: operation)
+        end
+      end
+
+      module TrilogyExtensions
+        extend Sinks::DSL
 
-          super
+        sink_before :query do |query|
+          Helpers.scan(query, "query")
         end
       end
     end
   end
 end
 
-::Trilogy.prepend(Aikido::Zen::Sinks::Trilogy::Extensions)
+Aikido::Zen::Sinks::Trilogy.load_sinks!

data/lib/aikido/zen/sinks.rb

@@ -1,30 +1,39 @@
 # frozen_string_literal: true
 
-require_relative "sink"
+# Code coverage is disabled in this file because it is environment-specific and
+# not intended to be tested directly.
+# :nocov:
 
-require_relative "sinks/socket"
+require_relative "sink"
+require_relative "sinks_dsl"
 
 require_relative "sinks/action_controller" if defined?(::ActionController)
-require_relative "sinks/file" if defined?(::File)
 
 # Sadly, in ruby versions lower than 3.0, it's not possible to patch the
 # Kernel module because how the `prepend` method is applied
 # (https://stackoverflow.com/questions/78110397/prepend-kernel-module-function-globally#comment137713906_78112924)
-if RUBY_VERSION >= "3.0"
-  require_relative "sinks/kernel" if defined?(::Kernel)
-end
-require_relative "sinks/resolv" if defined?(::Resolv)
-require_relative "sinks/net_http" if defined?(::Net::HTTP)
-require_relative "sinks/http" if defined?(::HTTP)
-require_relative "sinks/httpx" if defined?(::HTTPX)
-require_relative "sinks/httpclient" if defined?(::HTTPClient)
-require_relative "sinks/excon" if defined?(::Excon)
-require_relative "sinks/curb" if defined?(::Curl)
-require_relative "sinks/patron" if defined?(::Patron)
+require_relative "sinks/kernel" if RUBY_VERSION >= "3.0"
+
+require_relative "sinks/file"
+require_relative "sinks/socket"
+require_relative "sinks/resolv"
+require_relative "sinks/net_http"
+
+# http.rb aims to support and is tested against Ruby 3.0+:
+# https://github.com/httprb/http?tab=readme-ov-file#supported-ruby-versions
+require_relative "sinks/http" if RUBY_VERSION >= "3.0"
+
+require_relative "sinks/httpx"
+require_relative "sinks/httpclient"
+require_relative "sinks/excon"
+require_relative "sinks/curb"
+require_relative "sinks/patron"
 require_relative "sinks/typhoeus" if defined?(::Typhoeus)
-require_relative "sinks/async_http" if defined?(::Async::HTTP)
-require_relative "sinks/em_http" if defined?(::EventMachine::HttpRequest)
-require_relative "sinks/mysql2" if defined?(::Mysql2)
-require_relative "sinks/pg" if defined?(::PG)
-require_relative "sinks/sqlite3" if defined?(::SQLite3)
-require_relative "sinks/trilogy" if defined?(::Trilogy)
+require_relative "sinks/async_http"
+require_relative "sinks/em_http"
+require_relative "sinks/mysql2"
+require_relative "sinks/pg"
+require_relative "sinks/sqlite3"
+require_relative "sinks/trilogy"
+
+# :nocov: