aikido-zen 1.0.1.beta.4-x86_64-linux-musl → 1.0.2.beta.1-x86_64-linux-musl

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

@@ -7,13 +7,6 @@ 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: [
           Scanners::SSRFScanner,
           OutboundConnectionMonitor
@@ -66,33 +59,38 @@ module Aikido::Zen
           end
         end
 
-        module HTTPExtensions
-          extend Sinks::DSL
+        def self.load_sinks!
+          # In stdlib but not always required
+          require "net/http"
 
-          sink_around :request do |super_call, req|
-            wrapped_request = Helpers.wrap_request(req, self)
+          ::Net::HTTP.class_eval do
+            extend Sinks::DSL
 
-            # Store the request information so the DNS sinks can pick it up.
-            context = Aikido::Zen.current_context
-            if context
-              prev_request = context["ssrf.request"]
-              context["ssrf.request"] = wrapped_request
-            end
+            sink_around :request do |original_call, req|
+              wrapped_request = Helpers.wrap_request(req, self)
 
-            connection = Helpers.build_outbound(self)
+              # Store the request information so the DNS sinks can pick it up.
+              context = Aikido::Zen.current_context
+              if context
+                prev_request = context["ssrf.request"]
+                context["ssrf.request"] = wrapped_request
+              end
 
-            Helpers.scan(wrapped_request, connection, "request")
+              connection = Helpers.build_outbound(self)
 
-            response = super_call.call
+              Helpers.scan(wrapped_request, connection, "request")
 
-            Scanners::SSRFScanner.track_redirects(
-              request: wrapped_request,
-              response: Helpers.wrap_response(response)
-            )
+              response = original_call.call
+
+              Scanners::SSRFScanner.track_redirects(
+                request: wrapped_request,
+                response: Helpers.wrap_response(response)
+              )
 
-            response
-          ensure
-            context["ssrf.request"] = prev_request if context
+              response
+            ensure
+              context["ssrf.request"] = prev_request if context
+            end
           end
         end
       end

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

@@ -6,14 +6,6 @@ require_relative "../outbound_connection_monitor"
 module Aikido::Zen
   module Sinks
     module Patron
-      def self.load_sinks!
-        if Aikido::Zen.satisfy "patron", ">= 0.6.4"
-          require "patron"
-
-          ::Patron::Session.prepend(SessionExtensions)
-        end
-      end
-
       SINK = Sinks.add("patron", scanners: [
         Scanners::SSRFScanner,
         OutboundConnectionMonitor
@@ -44,58 +36,64 @@ module Aikido::Zen
         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.
-          context = Aikido::Zen.current_context
-          if context
-            prev_request = context["ssrf.request"]
-            context["ssrf.request"] = wrapped_request
-          end
-
-          connection = OutboundConnection.from_uri(URI(request.url))
-
-          Helpers.scan(wrapped_request, connection, "request")
-
-          response = super_call.call
-
-          Scanners::SSRFScanner.track_redirects(
-            request: wrapped_request,
-            response: Helpers.wrap_response(request, response)
-          )
-
-          # When libcurl has follow_location set, it will handle redirections
-          # internally, and expose the response.url as the URI that was last
-          # requested in the redirect chain.
-          #
-          # In this case, we can't actually stop the request from happening, but
-          # we can scan again (now that we know another request happened), to
-          # 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 = Scanners::SSRFScanner::Request.new(
-              verb: request.action,
-              uri: URI(response.url),
-              headers: request.headers
-            )
-            context["ssrf.request"] = last_effective_request if context
-
-            connection = OutboundConnection.from_uri(URI(response.url))
+      def self.load_sinks!
+        if Aikido::Zen.satisfy "patron", ">= 0.6.4"
+          require "patron"
 
-            Helpers.scan(last_effective_request, connection, "request")
+          ::Patron::Session.class_eval do
+            extend Sinks::DSL
+
+            sink_around :handle_request do |original_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.
+              context = Aikido::Zen.current_context
+              if context
+                prev_request = context["ssrf.request"]
+                context["ssrf.request"] = wrapped_request
+              end
+
+              connection = OutboundConnection.from_uri(URI(request.url))
+
+              Helpers.scan(wrapped_request, connection, "request")
+
+              response = original_call.call
+
+              Scanners::SSRFScanner.track_redirects(
+                request: wrapped_request,
+                response: Helpers.wrap_response(request, response)
+              )
+
+              # When libcurl has follow_location set, it will handle redirections
+              # internally, and expose the response.url as the URI that was last
+              # requested in the redirect chain.
+              #
+              # In this case, we can't actually stop the request from happening, but
+              # we can scan again (now that we know another request happened), to
+              # 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 = Scanners::SSRFScanner::Request.new(
+                  verb: request.action,
+                  uri: URI(response.url),
+                  headers: request.headers
+                )
+                context["ssrf.request"] = last_effective_request if context
+
+                connection = OutboundConnection.from_uri(URI(response.url))
+
+                Helpers.scan(last_effective_request, connection, "request")
+              end
+
+              response
+            ensure
+              context["ssrf.request"] = prev_request if context
+            end
           end
-
-          response
-        ensure
-          context["ssrf.request"] = prev_request if context
         end
       end
     end

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

@@ -3,14 +3,6 @@
 module Aikido::Zen
   module Sinks
     module PG
-      def self.load_sinks!
-        if Aikido::Zen.satisfy "pg", ">= 1.0"
-          require "pg"
-
-          ::PG::Connection.prepend(PG::ConnectionExtensions)
-        end
-      end
-
       SINK = Sinks.add("pg", scanners: [Scanners::SQLInjectionScanner])
 
       module Helpers
@@ -43,26 +35,32 @@ module Aikido::Zen
         end
       end
 
-      module ConnectionExtensions
-        extend Sinks::DSL
+      def self.load_sinks!
+        if Aikido::Zen.satisfy "pg", ">= 1.0"
+          require "pg"
 
-        %i[
-          send_query exec sync_exec async_exec
-          send_query_params exec_params sync_exec_params async_exec_params
-        ].each do |method_name|
-          presafe_sink_before method_name do |query|
-            Helpers.safe do
-              Helpers.scan(query, method_name)
+          ::PG::Connection.class_eval do
+            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_name|
+              presafe_sink_before method_name do |query|
+                Helpers.safe do
+                  Helpers.scan(query, method_name)
+                end
+              end
             end
-          end
-        end
 
-        %i[
-          send_prepare prepare async_prepare sync_prepare
-        ].each do |method_name|
-          presafe_sink_before method_name do |_, query|
-            Helpers.safe do
-              Helpers.scan(query, method_name)
+            %i[
+              send_prepare prepare async_prepare sync_prepare
+            ].each do |method_name|
+              presafe_sink_before method_name do |_, query|
+                Helpers.safe do
+                  Helpers.scan(query, method_name)
+                end
+              end
             end
           end
         end

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

@@ -6,13 +6,6 @@ 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: [
         Scanners::StoredSSRFScanner,
         Scanners::SSRFScanner
@@ -35,23 +28,30 @@ module Aikido::Zen
         end
       end
 
-      module ResolvExtensions
-        def each_address(*args, **kwargs, &blk)
-          # each_address is defined "manually" because no sink method pattern
-          # is applicable.
+      def self.load_sinks!
+        # In stdlib but not always required
+        require "resolv"
 
-          name, = args
+        ::Resolv.class_eval do
+          alias_method :each_address__internal_for_aikido_zen, :each_address
 
-          addresses = []
-          super(*args, **kwargs) do |address| # rubocop:disable Style/SuperArguments
-            addresses << address
-            blk.call(address)
-          end
-        ensure
-          # Ensure partial results are scanned.
+          def each_address(*args, **kwargs, &blk)
+            # each_address is defined "manually" because no sink method pattern
+            # is applicable.
+
+            name, = args
+
+            addresses = []
+            each_address__internal_for_aikido_zen(*args, **kwargs) do |address|
+              addresses << address
+              blk.call(address)
+            end
+          ensure
+            # Ensure partial results are scanned.
 
-          Sinks::DSL.safe do
-            Helpers.scan(name, addresses, "lookup")
+            Sinks::DSL.safe do
+              Helpers.scan(name, addresses, "lookup")
+            end
           end
         end
       end

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

@@ -10,10 +10,6 @@ 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: [
         Scanners::StoredSSRFScanner,
         Scanners::SSRFScanner
@@ -62,15 +58,17 @@ module Aikido::Zen
         end
       end
 
-      module IPSocketExtensions
-        extend Sinks::DSL
+      def self.load_sinks!
+        ::IPSocket.singleton_class.class_eval do
+          extend Sinks::DSL
 
-        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:
+          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

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

@@ -3,15 +3,6 @@
 module Aikido::Zen
   module Sinks
     module SQLite3
-      def self.load_sinks!
-        if Aikido::Zen.satisfy "sqlite3", ">= 1.0"
-          require "sqlite3"
-
-          ::SQLite3::Database.prepend(DatabaseExtensions)
-          ::SQLite3::Statement.prepend(StatementExtensions)
-        end
-      end
-
       SINK = Sinks.add("sqlite3", scanners: [Scanners::SQLInjectionScanner])
 
       module Helpers
@@ -24,22 +15,28 @@ module Aikido::Zen
         end
       end
 
-      module DatabaseExtensions
-        extend Sinks::DSL
+      def self.load_sinks!
+        if Aikido::Zen.satisfy "sqlite3", ">= 1.0"
+          require "sqlite3"
 
-        private
+          ::SQLite3::Database.class_eval do
+            extend Sinks::DSL
 
-        # SQLite3::Database#exec_batch is an internal native private method.
-        sink_before :exec_batch do |sql|
-          Helpers.scan(sql, "exec_batch")
-        end
-      end
+            private
+
+            # SQLite3::Database#exec_batch is an internal native private method.
+            sink_before :exec_batch do |sql|
+              Helpers.scan(sql, "exec_batch")
+            end
+          end
 
-      module StatementExtensions
-        extend Sinks::DSL
+          ::SQLite3::Statement.class_eval do
+            extend Sinks::DSL
 
-        sink_before :initialize do |_db, sql|
-          Helpers.scan(sql, "statement.execute")
+            sink_before :initialize do |_db, sql|
+              Helpers.scan(sql, "statement.execute")
+            end
+          end
         end
       end
     end

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

@@ -3,14 +3,6 @@
 module Aikido::Zen
   module Sinks
     module Trilogy
-      def self.load_sinks!
-        if Aikido::Zen.satisfy "trilogy", ">= 2.0"
-          require "trilogy"
-
-          ::Trilogy.prepend(TrilogyExtensions)
-        end
-      end
-
       SINK = Sinks.add("trilogy", scanners: [Scanners::SQLInjectionScanner])
 
       module Helpers
@@ -19,11 +11,17 @@ module Aikido::Zen
         end
       end
 
-      module TrilogyExtensions
-        extend Sinks::DSL
+      def self.load_sinks!
+        if Aikido::Zen.satisfy "trilogy", ">= 2.0"
+          require "trilogy"
+
+          ::Trilogy.class_eval do
+            extend Sinks::DSL
 
-        sink_before :query do |query|
-          Helpers.scan(query, "query")
+            sink_before :query do |query|
+              Helpers.scan(query, "query")
+            end
+          end
         end
       end
     end

data/lib/aikido/zen/sinks.rb

@@ -9,10 +9,7 @@ require_relative "sinks_dsl"
 
 require_relative "sinks/action_controller" if defined?(::ActionController)
 
-# 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)
-require_relative "sinks/kernel" if RUBY_VERSION >= "3.0"
+require_relative "sinks/kernel"
 
 require_relative "sinks/file"
 require_relative "sinks/socket"

data/lib/aikido/zen/sinks_dsl.rb

@@ -109,10 +109,16 @@ module Aikido::Zen
       #
       # @return [void]
       def presafe_sink_before(method_name, &block)
+        raise ArgumentError, "block required" unless block
+
+        original = instance_method(method_name)
+
         define_method(method_name) do |*args, **kwargs, &blk|
           instance_exec(*args, **kwargs, &block)
-          super(*args, **kwargs, &blk)
+          original.bind_call(self, *args, **kwargs, &blk)
         end
+      rescue NameError
+        Aikido::Zen.config.logger.warn("cannot wrap method `#{method_name}' for class `#{self}'")
       end
 
       # Define a method `method_name` that safely executes the given block before
@@ -127,6 +133,8 @@ module Aikido::Zen
       #
       # @note the block is executed within `safe` to handle errors safely; the original method is executed outside of `safe` to preserve the original behavior
       def sink_before(method_name, &block)
+        raise ArgumentError, "block required" unless block
+
         presafe_sink_before(method_name) do |*args, **kwargs|
           DSL.safe do
             instance_exec(*args, **kwargs, &block)
@@ -145,11 +153,17 @@ module Aikido::Zen
       #
       # @return [void]
       def presafe_sink_after(method_name, &block)
+        raise ArgumentError, "block required" unless block
+
+        original = instance_method(method_name)
+
         define_method(method_name) do |*args, **kwargs, &blk|
-          result = super(*args, **kwargs, &blk)
+          result = original.bind_call(self, *args, **kwargs, &blk)
           instance_exec(result, *args, **kwargs, &block)
           result
         end
+      rescue NameError
+        Aikido::Zen.config.logger.warn("cannot wrap method `#{method_name}' for class `#{self}'")
       end
 
       # Define a method `method_name` that safely executes the given block after
@@ -165,6 +179,8 @@ module Aikido::Zen
       #
       # @note the block is executed within `safe` to handle errors safely; the original method is executed outside of `safe` to preserve the original behavior
       def sink_after(method_name, &block)
+        raise ArgumentError, "block required" unless block
+
         presafe_sink_after(method_name) do |result, *args, **kwargs|
           DSL.safe do
             instance_exec(result, *args, **kwargs, &block)
@@ -177,20 +193,26 @@ module Aikido::Zen
       #
       # @param method_name [Symbol, String] the name of the method to define
       # @yield the block to execute around the original method
-      # @yieldparam super_call [Proc] the proc that calls the original method
+      # @yieldparam original_call [Proc] the proc that calls the original method
       # @yieldparam args [Array] the positional arguments passed to the original method
       # @yieldparam kwargs [Hash] the keyword arguments passed to the original method
       #
       # @return [void]
       def presafe_sink_around(method_name, &block)
+        raise ArgumentError, "block required" unless block
+
+        original = instance_method(method_name)
+
         define_method(method_name) do |*args, **kwargs, &blk|
           result = nil
-          super_call = proc do
-            result = super(*args, **kwargs, &blk)
+          original_call = proc do
+            result = original.bind_call(self, *args, **kwargs, &blk)
           end
-          instance_exec(super_call, *args, **kwargs, &block)
+          instance_exec(original_call, *args, **kwargs, &block)
           result
         end
+      rescue NameError
+        Aikido::Zen.config.logger.warn("cannot wrap method `#{method_name}' for class `#{self}'")
       end
 
       # Define a method `method_name` that safely executes the given block around
@@ -198,27 +220,29 @@ module Aikido::Zen
       #
       # @param method_name [Symbol, String] the name of the method to define
       # @yield the block to execute around the original method
-      # @yieldparam super_call [Proc] the proc that calls the original method
+      # @yieldparam original_call [Proc] the proc that calls the original method
       # @yieldparam args [Array] the positional arguments passed to the original method
       # @yieldparam kwargs [Hash] the keyword arguments passed to the original method
       #
       # @return [void]
       #
       # @note the block is executed within `safe` to handle errors safely; the original method is executed within `presafe` to preserve the original behavior
-      # @note if the block does not call `super_call`, the original method is called automatically after the block is executed
+      # @note if the block does not call `original_call`, the original method is called automatically after the block is executed
       def sink_around(method_name, &block)
-        presafe_sink_around(method_name) do |presafe_super_call, *args, **kwargs|
-          super_called = false
-          super_call = proc do
-            super_called = true
+        raise ArgumentError, "block required" unless block
+
+        presafe_sink_around(method_name) do |presafe_original_call, *args, **kwargs|
+          original_called = false
+          original_call = proc do
+            original_called = true
             DSL.presafe do
-              presafe_super_call.call
+              presafe_original_call.call
             end
           end
           DSL.safe do
-            instance_exec(super_call, *args, **kwargs, &block)
+            instance_exec(original_call, *args, **kwargs, &block)
           end
-          presafe_super_call.call unless super_called
+          presafe_original_call.call unless original_called
         end
       end
     end

data/lib/aikido/zen/version.rb

@@ -2,7 +2,7 @@
 
 module Aikido
   module Zen
-    VERSION = "1.0.1.beta.4"
+    VERSION = "1.0.2.beta.1"
 
     # The version of libzen_internals that we build against.
     LIBZEN_VERSION = "0.1.39"

data/lib/aikido/zen.rb

@@ -39,7 +39,7 @@ module Aikido
       return unless config.protect?
 
       unless load_sources! && load_sinks!
-        config.logger.warn "Zen could not find any supported libraries or frameworks. Visit https://github.com/AikidoSec/firewall-ruby for more information."
+        config.logger.warn("Zen could not find any supported libraries or frameworks. Visit https://github.com/AikidoSec/firewall-ruby for more information.")
         return
       end
 
@@ -165,6 +165,11 @@ module Aikido
       end
     end
 
+    # Align with other Zen implementations, while keeping internal consistency.
+    class << self
+      alias_method :set_user, :track_user
+    end
+
     # Marks that the Zen middleware was installed properly
     # @return void
     def self.middleware_installed!
@@ -209,7 +214,7 @@ module Aikido
     end
 
     def self.detached_agent_server
-      @detached_agent_server ||= DetachedAgent::Server.start!
+      @detached_agent_server ||= DetachedAgent::Server.start
     end
 
     class << self

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: aikido-zen
 version: !ruby/object:Gem::Version
-  version: 1.0.1.beta.4
+  version: 1.0.2.beta.1
 platform: x86_64-linux-musl
 authors:
 - Aikido Security
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2025-08-11 00:00:00.000000000 Z
+date: 2025-08-27 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: concurrent-ruby