aikido-zen 0.2.0-x86_64-linux → 1.0.1.beta.2-x86_64-linux

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

@@ -12,11 +12,11 @@ module Aikido::Zen
         def initialize(
           config: Aikido::Zen.config,
           settings: Aikido::Zen.runtime_settings,
-          rate_limiter: Aikido::Zen::RateLimiter.new
+          detached_agent: Aikido::Zen.detached_agent
         )
           @config = config
           @settings = settings
-          @rate_limiter = rate_limiter
+          @detached_agent = detached_agent
         end
 
         def block?(controller)
@@ -43,16 +43,21 @@ module Aikido::Zen
         end
 
         private def should_throttle?(request)
+          return false unless @settings.endpoints[request.route].rate_limiting.enabled?
           return false if @settings.skip_protection_for_ips.include?(request.ip)
 
-          @rate_limiter.throttle?(request)
+          result = @detached_agent.calculate_rate_limits(request)
+          return false unless result
+
+          request.env["aikido.rate_limiting"] = result
+          request.env["aikido.rate_limiting"].throttled?
         end
 
         # @param request [Aikido::Zen::Request]
         private def should_block_user?(request)
           return false if request.actor.nil?
 
-          @settings.blocked_user_ids&.include?(request.actor.id)
+          Aikido::Zen.runtime_settings.blocked_user_ids&.include?(request.actor.id)
         end
       end
 

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

@@ -1,26 +1,46 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Async
       module HTTP
+        def self.load_sinks!
+          if Gem.loaded_specs["async-http"]
+            require "async/http"
+
+            ::Async::HTTP::Client.prepend(Async::HTTP::ClientExtensions)
+          end
+        end
+
         SINK = Sinks.add("async-http", scanners: [
-          Aikido::Zen::Scanners::SSRFScanner,
-          Aikido::Zen::OutboundConnectionMonitor
+          Scanners::SSRFScanner,
+          OutboundConnectionMonitor
         ])
 
-        module Extensions
-          def call(request)
+        module Helpers
+          def self.scan(request, connection, operation)
+            SINK.scan(
+              request: request,
+              connection: connection,
+              operation: operation
+            )
+          end
+        end
+
+        module ClientExtensions
+          extend Sinks::DSL
+
+          sink_around :call do |super_call, request|
             uri = URI(format("%<scheme>s://%<authority>s%<path>s", {
               scheme: request.scheme || scheme,
               authority: request.authority || authority,
               path: request.path
             }))
 
-            wrapped_request = Aikido::Zen::Scanners::SSRFScanner::Request.new(
+            wrapped_request = Scanners::SSRFScanner::Request.new(
               verb: request.method,
               uri: uri,
               headers: request.headers.to_h,
@@ -28,22 +48,21 @@ module Aikido::Zen
             )
 
             # 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: wrapped_request,
-              operation: "request"
-            )
+            connection = OutboundConnection.from_uri(uri)
+
+            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: Aikido::Zen::Scanners::SSRFScanner::Response.new(
+              response: Scanners::SSRFScanner::Response.new(
                 status: response.status,
                 headers: response.headers.to_h,
                 header_normalizer: ->(value) { Array(value).join(", ") }
@@ -60,4 +79,4 @@ module Aikido::Zen
   end
 end
 
-::Async::HTTP::Client.prepend(Aikido::Zen::Sinks::Async::HTTP::Extensions)
+Aikido::Zen::Sinks::Async::HTTP.load_sinks!

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

@@ -1,19 +1,27 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Curl
+      def self.load_sinks!
+        if Gem.loaded_specs["curb"]
+          require "curb"
+
+          ::Curl::Easy.prepend(Curl::EasyExtensions)
+        end
+      end
+
       SINK = Sinks.add("curb", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
       ])
 
-      module Extensions
+      module Helpers
         def self.wrap_request(curl, url: curl.url)
-          Aikido::Zen::Scanners::SSRFScanner::Request.new(
+          Scanners::SSRFScanner::Request.new(
             verb: nil, # Curb hides this by directly setting an option in C
             uri: URI(url),
             headers: curl.headers
@@ -33,29 +41,40 @@ module Aikido::Zen
             status = curl.status.to_i
           end
 
-          Aikido::Zen::Scanners::SSRFScanner::Response.new(status: status, headers: headers)
+          Scanners::SSRFScanner::Response.new(status: status, headers: headers)
         end
 
-        def perform
-          wrapped_request = Extensions.wrap_request(self)
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module EasyExtensions
+        extend Sinks::DSL
+
+        sink_around :perform do |super_call|
+          wrapped_request = Helpers.wrap_request(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: Aikido::Zen::OutboundConnection.from_uri(URI(url)),
-            request: wrapped_request,
-            operation: "request"
-          )
+          connection = OutboundConnection.from_uri(URI(url))
 
-          response = super
+          Helpers.scan(wrapped_request, connection, "request")
 
-          Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+          response = super_call.call
+
+          Scanners::SSRFScanner.track_redirects(
             request: wrapped_request,
-            response: Extensions.wrap_response(self)
+            response: Helpers.wrap_response(self)
           )
 
           # When libcurl has follow_location set, it will handle redirections
@@ -67,14 +86,21 @@ 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 url != last_effective_url
-            last_effective_request = Extensions.wrap_request(self, url: last_effective_url)
-            context["ssrf.request"] = last_effective_request if context
-
-            SINK.scan(
-              connection: Aikido::Zen::OutboundConnection.from_uri(URI(last_effective_url)),
-              request: last_effective_request,
-              operation: "request"
-            )
+            last_effective_request = Helpers.wrap_request(self, url: last_effective_url)
+
+            # Code coverage is disabled here because the else clause is a no-op,
+            # so there is nothing to cover.
+            # :nocov:
+            if context
+              context["ssrf.request"] = last_effective_request
+            else
+              # empty
+            end
+            # :nocov:
+
+            connection = OutboundConnection.from_uri(URI(last_effective_url))
+
+            Helpers.scan(last_effective_request, connection, "request")
           end
 
           response
@@ -86,4 +112,4 @@ module Aikido::Zen
   end
 end
 
-::Curl::Easy.prepend(Aikido::Zen::Sinks::Curl::Extensions)
+Aikido::Zen::Sinks::Curl.load_sinks!

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

@@ -1,20 +1,45 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module EventMachine
       module HttpRequest
+        def self.load_sinks!
+          if Gem.loaded_specs["em-http-request"]
+            require "em-http-request"
+
+            ::EventMachine::HttpRequest.use(EventMachine::HttpRequest::Middleware)
+
+            # NOTE: We can't use middleware to intercept requests as we want to ensure any
+            # modifications to the request from user-supplied middleware are already applied
+            # before we scan the request.
+            ::EventMachine::HttpClient.prepend(EventMachine::HttpRequest::HttpClientExtensions)
+          end
+        end
+
         SINK = Sinks.add("em-http-request", scanners: [
-          Aikido::Zen::Scanners::SSRFScanner,
-          Aikido::Zen::OutboundConnectionMonitor
+          Scanners::SSRFScanner,
+          OutboundConnectionMonitor
         ])
 
-        module Extensions
-          def send_request(*)
-            wrapped_request = Aikido::Zen::Scanners::SSRFScanner::Request.new(
+        module Helpers
+          def self.scan(request, connection, operation)
+            SINK.scan(
+              request: request,
+              connection: connection,
+              operation: operation
+            )
+          end
+        end
+
+        module HttpClientExtensions
+          extend Sinks::DSL
+
+          sink_before :send_request do
+            wrapped_request = Scanners::SSRFScanner::Request.new(
               verb: req.method.to_s,
               uri: URI(req.uri),
               headers: req.headers
@@ -24,16 +49,12 @@ module Aikido::Zen
             context = Aikido::Zen.current_context
             context["ssrf.request"] = wrapped_request if context
 
-            SINK.scan(
-              connection: Aikido::Zen::OutboundConnection.new(
-                host: req.host,
-                port: req.port
-              ),
-              request: wrapped_request,
-              operation: "request"
+            connection = OutboundConnection.new(
+              host: req.host,
+              port: req.port
             )
 
-            super
+            Helpers.scan(wrapped_request, connection, "request")
           end
         end
 
@@ -43,13 +64,13 @@ module Aikido::Zen
             context = Aikido::Zen.current_context
             context["ssrf.request"] = nil if context
 
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
-              request: Aikido::Zen::Scanners::SSRFScanner::Request.new(
+            Scanners::SSRFScanner.track_redirects(
+              request: Scanners::SSRFScanner::Request.new(
                 verb: client.req.method,
                 uri: URI(client.req.uri),
                 headers: client.req.headers
               ),
-              response: Aikido::Zen::Scanners::SSRFScanner::Response.new(
+              response: Scanners::SSRFScanner::Response.new(
                 status: client.response_header.status,
                 headers: client.response_header.to_h
               )
@@ -61,11 +82,4 @@ module Aikido::Zen
   end
 end
 
-::EventMachine::HttpRequest
-  .use(Aikido::Zen::Sinks::EventMachine::HttpRequest::Middleware)
-
-# NOTE: We can't use middleware to intercept requests as we want to ensure any
-# modifications to the request from user-supplied middleware are already applied
-# before we scan the request.
-::EventMachine::HttpClient
-  .prepend(Aikido::Zen::Sinks::EventMachine::HttpRequest::Extensions)
+Aikido::Zen::Sinks::EventMachine::HttpRequest.load_sinks!

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

@@ -1,31 +1,26 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Excon
-      SINK = Sinks.add("excon", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
-      ])
+      def self.load_sinks!
+        if Gem.loaded_specs["excon"]
+          require "excon"
 
-      module Extensions
-        # Maps Excon request params to an Aikido OutboundConnection.
-        #
-        # @param connection [Hash<Symbol, Object>] the data set in the connection.
-        # @param request [Hash<Symbol, Object>] the data overrides sent for each
-        #   request.
-        #
-        # @return [Aikido::Zen::OutboundConnection]
-        def self.build_outbound(connection, request)
-          Aikido::Zen::OutboundConnection.new(
-            host: request.fetch(:hostname) { connection[:hostname] },
-            port: request.fetch(:port) { connection[:port] }
-          )
+          ::Excon::Connection.prepend(ConnectionExtensions)
+          ::Excon::Middleware::RedirectFollower.prepend(RedirectFollowerExtensions)
         end
+      end
+
+      SINK = Sinks.add("excon", scanners: [
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
+      ])
 
+      module Helpers
         def self.build_request(connection, request)
           uri = URI(format("%<scheme>s://%<host>s:%<port>i%<path>s", {
             scheme: request.fetch(:scheme) { connection[:scheme] },
@@ -35,45 +30,61 @@ module Aikido::Zen
           }))
           uri.query = request.fetch(:query) { connection[:query] }
 
-          Aikido::Zen::Scanners::SSRFScanner::Request.new(
+          Scanners::SSRFScanner::Request.new(
             verb: request.fetch(:method) { connection[:method] },
             uri: uri,
             headers: connection[:headers].to_h.merge(request[:headers].to_h)
           )
         end
 
-        def request(params = {}, *)
-          request = Extensions.build_request(@data, params)
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module ConnectionExtensions
+        extend Sinks::DSL
+
+        sink_around :request do |super_call, params = {}|
+          request = Helpers.build_request(@data, params)
 
           # 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"] = request
           end
 
-          SINK.scan(
-            connection: Aikido::Zen::OutboundConnection.from_uri(request.uri),
-            request: request,
-            operation: "request"
-          )
+          connection = OutboundConnection.from_uri(request.uri)
 
-          response = super
+          Helpers.scan(request, connection, "request")
 
-          Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+          response = super_call.call
+
+          Scanners::SSRFScanner.track_redirects(
             request: request,
-            response: Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            response: Scanners::SSRFScanner::Response.new(
               status: response.status,
               headers: response.headers.to_h
             )
           )
 
           response
-        rescue ::Excon::Error::Socket => err
-          # Excon wraps errors inside the lower level layer. This only happens
-          # to our scanning exceptions when a request is using RedirectFollower,
-          # so we unwrap them when it happens so host apps can handle errors
-          # consistently.
-          raise err.cause if err.cause.is_a?(Aikido::Zen::UnderAttackError)
+        rescue Sinks::DSL::PresafeError => err
+          outer_cause = err.cause
+          case outer_cause
+          when ::Excon::Error::Socket
+            inner_cause = outer_cause.cause
+            # Excon wraps errors inside the lower level layer. This only happens
+            # to our scanning exceptions when a request is using RedirectFollower,
+            # so we unwrap them when it happens so host apps can handle errors
+            # consistently.
+            raise inner_cause if inner_cause.is_a?(Aikido::Zen::UnderAttackError)
+          end
           raise
         ensure
           context["ssrf.request"] = prev_request if context
@@ -81,23 +92,30 @@ module Aikido::Zen
       end
 
       module RedirectFollowerExtensions
-        def response_call(data)
-          if (response = data[:response])
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
-              request: Extensions.build_request(data, {}),
-              response: Aikido::Zen::Scanners::SSRFScanner::Response.new(
+        extend Sinks::DSL
+
+        sink_before :response_call do |datum|
+          response = datum[:response]
+
+          # Code coverage is disabled here because the else clause is a no-op,
+          # so there is nothing to cover.
+          # :nocov:
+          if !response.nil?
+            Scanners::SSRFScanner.track_redirects(
+              request: Helpers.build_request(datum, {}),
+              response: Scanners::SSRFScanner::Response.new(
                 status: response[:status],
                 headers: response[:headers]
               )
             )
+          else
+            # empty
           end
-
-          super
+          # :nocov:
         end
       end
     end
   end
 end
 
-::Excon::Connection.prepend(Aikido::Zen::Sinks::Excon::Extensions)
-::Excon::Middleware::RedirectFollower.prepend(Aikido::Zen::Sinks::Excon::RedirectFollowerExtensions)
+Aikido::Zen::Sinks::Excon.load_sinks!