aikido-zen 0.1.1-arm64-darwin → 1.0.0.pre.beta.1-arm64-darwin

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

@@ -3,26 +3,34 @@
 module Aikido::Zen
   module Sinks
     module ActionController
-      # Implements the "middleware" for rate limiting in Rails apps, where we
-      # need to check at the end of the `before_action` chain, rather than in
-      # an actual Rack middleware, to allow for calls to Zen.track_user being
-      # made from before_actions in the host app, thus allowing rate-limiting
-      # by user ID rather than solely by IP.
-      class Throttler
+      # Implements the "middleware" for blocking requests (i.e.: rate-limiting or blocking
+      # user/bots) in Rails apps, where we need to check at the end of the `before_action`
+      # chain, rather than in an actual Rack middleware, to allow for calls to
+      # `Zen.track_user` being made from before_actions in the host app, thus allowing
+      # block/rate-limit by user ID rather than solely by IP.
+      class BlockRequestChecker
         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 throttle(controller)
+        def block?(controller)
           context = controller.request.env[Aikido::Zen::ENV_KEY]
           request = context.request
 
+          if should_block_user?(request)
+            status, headers, body = @config.blocked_responder.call(request, :user)
+            controller.headers.update(headers)
+            controller.render plain: Array(body).join, status: status
+
+            return true
+          end
+
           if should_throttle?(request)
             status, headers, body = @config.rate_limited_responder.call(request)
             controller.headers.update(headers)
@@ -35,14 +43,26 @@ 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?
+
+          Aikido::Zen.runtime_settings.blocked_user_ids&.include?(request.actor.id)
         end
       end
 
-      def self.throttler
-        @throttler ||= Aikido::Zen::Sinks::ActionController::Throttler.new
+      def self.block_request_checker
+        @block_request_checker ||= Aikido::Zen::Sinks::ActionController::BlockRequestChecker.new
       end
 
       module Extensions
@@ -50,10 +70,9 @@ module Aikido::Zen
           return super unless kind == :process_action
 
           super do
-            rate_limiter = Aikido::Zen::Sinks::ActionController.throttler
-            throttled = rate_limiter.throttle(self)
+            checker = Aikido::Zen::Sinks::ActionController.block_request_checker
 
-            yield if block_given? && !throttled
+            yield if block_given? && !checker.block?(self)
           end
         end
       end

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

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  module Sinks
+    module File
+      SINK = Sinks.add("File", scanners: [
+        Aikido::Zen::Scanners::PathTraversalScanner
+      ])
+
+      module Extensions
+        def self.scan_path(filepath, operation)
+          SINK.scan(
+            filepath: filepath,
+            operation: operation
+          )
+        end
+
+        # Module to extend only the initializer method of `File` (`File.new`)
+        module Initiliazer
+          def initialize(filename, *, **)
+            Extensions.scan_path(filename, "new")
+            super
+          end
+        end
+
+        def open(filename, *, **)
+          Extensions.scan_path(filename, "open")
+          super
+        end
+
+        def read(filename, *)
+          Extensions.scan_path(filename, "read")
+          super
+        end
+
+        def write(filename, *, **)
+          Extensions.scan_path(filename, "write")
+          super
+        end
+
+        def join(*)
+          joined = super
+          Extensions.scan_path(joined, "join")
+          joined
+        end
+
+        def chmod(mode, *paths)
+          paths.each { |path| Extensions.scan_path(path, "chmod") }
+          super
+        end
+
+        def chown(user, group, *paths)
+          paths.each { |path| Extensions.scan_path(path, "chown") }
+          super
+        end
+
+        def rename(from, to)
+          Extensions.scan_path(from, "rename")
+          Extensions.scan_path(to, "rename")
+          super
+        end
+
+        def symlink(from, to)
+          Extensions.scan_path(from, "symlink")
+          Extensions.scan_path(to, "symlink")
+          super
+        end
+
+        def truncate(file_name, *)
+          Extensions.scan_path(file_name, "truncate")
+          super
+        end
+
+        def unlink(*args)
+          args.each do |arg|
+            Extensions.scan_path(arg, "unlink")
+          end
+          super
+        end
+
+        def delete(*args)
+          args.each do |arg|
+            Extensions.scan_path(arg, "delete")
+          end
+          super
+        end
+
+        def utime(atime, mtime, *args)
+          args.each do |arg|
+            Extensions.scan_path(arg, "utime")
+          end
+          super
+        end
+
+        def expand_path(filename, *)
+          Extensions.scan_path(filename, "expand_path")
+          super
+        end
+
+        def realpath(filename, *)
+          Extensions.scan_path(filename, "realpath")
+          super
+        end
+
+        def realdirpath(filename, *)
+          Extensions.scan_path(filename, "realdirpath")
+          super
+        end
+      end
+    end
+  end
+end
+
+# Internally, Path Traversal's scanner logic uses `expand_path`, in order to avoid recursion issues we keep
+# a copy of the original method, only to be used internally.
+# It's important to keep this line before prepend the Extensions module, otherwise the alias will call
+# the extended method.
+::File.singleton_class.alias_method :expand_path__internal_for_aikido_zen, :expand_path
+::File.singleton_class.prepend(Aikido::Zen::Sinks::File::Extensions)
+::File.prepend Aikido::Zen::Sinks::File::Extensions::Initiliazer

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

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  module Sinks
+    module Kernel
+      SINK = Sinks.add("Kernel", scanners: [
+        Aikido::Zen::Scanners::ShellInjectionScanner
+      ])
+
+      module Extensions
+        # Checks if the user introduced input is trying to execute other commands
+        # using Shell Injection kind of attacks.
+        #
+        # @param command [String] the _full command_ that will be executed.
+        # @param context [Aikido::Zen::Context]
+        # @param sink [Aikido::Zen::Sink] the Sink that is running the scan.
+        # @param operation [Symbol, String] name of the method being scanned.
+        #
+        # @return [Aikido::Zen::Attacks::ShellInjectionAttack, nil] an Attack if any
+        # user input is detected as part of a Shell Injection Attack, or +nil+ if it's safe.
+        def self.scan_command(command, operation)
+          SINK.scan(
+            command: command,
+            operation: operation
+          )
+        end
+
+        # `system, spawn` functions can be invoked in several ways. For more details,
+        # see [the documentation](https://ruby-doc.org/3.4.1/Kernel.html#method-i-spawn)
+        #
+        # In our context, we care primarily about two common scenarios:
+        #   - one argument (String)
+        #       e.g.: system("ls"), system("echo something")
+        #   - two arguments (Hash, String)
+        #       e.g.: system({"foo" => "bar"}, "ls"), system({"foo" => "bar"}, "echo something")
+        #
+        # In all other cases, we do not protect against shell argument injections. Specifically:
+        #
+        # If a user input contains something like $(whoami) and is passed as part of the command
+        # arguments (e.g., user_input = "$(whoami)"):
+        #
+        #   system("echo", user_input)   This is safe because Ruby automatically escapes arguments
+        #                                passed to system/spawn in this form.
+        #
+        #   system("echo #{user_input}") This is not safe because Ruby interpolates the user_input
+        #                                into the command string, resulting in a potentially harmful
+        #                                command like `echo $(whoami)`.
+        def send_arg_to_scan(args, operation)
+          if args.size == 1 && args[0].is_a?(String)
+            Extensions.scan_command(args[0], operation)
+          end
+
+          if args.size == 2 && args[0].is_a?(Hash)
+            Extensions.scan_command(args[1], operation)
+          end
+        end
+
+        def system(*args, **)
+          send_arg_to_scan(args, "system")
+          super
+        end
+
+        def spawn(*args, **)
+          send_arg_to_scan(args, "spawn")
+          super
+        end
+      end
+    end
+  end
+end
+
+::Kernel.singleton_class.prepend Aikido::Zen::Sinks::Kernel::Extensions
+::Kernel.prepend Aikido::Zen::Sinks::Kernel::Extensions

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

@@ -15,6 +15,19 @@ module Aikido::Zen
 
       module IPSocketExtensions
         def self.scan_socket(hostname, socket)
+          # 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
 

data/lib/aikido/zen/sinks.rb

@@ -5,6 +5,14 @@ require_relative "sink"
 require_relative "sinks/socket"
 
 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)

data/lib/aikido/zen/system_info.rb

@@ -60,7 +60,7 @@ module Aikido::Zen
     end
 
     def os_version
-      Gem::Platform.local.version
+      Gem::Platform.local.version || "unknown"
     end
 
     def as_json

data/lib/aikido/zen/version.rb

@@ -2,9 +2,9 @@
 
 module Aikido
   module Zen
-    VERSION = "0.1.1"
+    VERSION = "1.0.0-beta.1"
 
     # The version of libzen_internals that we build against.
-    LIBZEN_VERSION = "0.1.31"
+    LIBZEN_VERSION = "0.1.37"
   end
 end

data/lib/aikido/zen/worker.rb

@@ -78,5 +78,10 @@ module Aikido::Zen
       @executor.shutdown
       @executor.wait_for_termination(30)
     end
+
+    def restart
+      shutdown
+      @executor = Concurrent::SingleThreadExecutor.new
+    end
   end
 end

data/lib/aikido/zen.rb

@@ -10,13 +10,16 @@ require_relative "zen/worker"
 require_relative "zen/agent"
 require_relative "zen/api_client"
 require_relative "zen/context"
+require_relative "zen/detached_agent"
+require_relative "zen/middleware/check_allowed_addresses"
+require_relative "zen/middleware/middleware"
+require_relative "zen/middleware/request_tracker"
 require_relative "zen/middleware/set_context"
 require_relative "zen/outbound_connection"
 require_relative "zen/outbound_connection_monitor"
 require_relative "zen/runtime_settings"
 require_relative "zen/rate_limiter"
 require_relative "zen/scanners"
-require_relative "zen/middleware/check_allowed_addresses"
 require_relative "zen/rails_engine" if defined?(::Rails)
 
 module Aikido
@@ -32,6 +35,10 @@ module Aikido
       @runtime_settings ||= RuntimeSettings.new
     end
 
+    def self.runtime_settings=(settings)
+      @runtime_settings = settings
+    end
+
     # Gets information about the current system configuration, which is sent to
     # the server along with any events.
     def self.system_info
@@ -41,9 +48,15 @@ module Aikido
     # Manages runtime metrics extracted from your app, which are uploaded to the
     # Aikido servers if configured to do so.
     def self.collector
+      check_and_handle_fork
       @collector ||= Collector.new
     end
 
+    def self.detached_agent
+      check_and_handle_fork
+      @detached_agent ||= DetachedAgent::Agent.new
+    end
+
     # Gets the current context object that holds all information about the
     # current request.
     #
@@ -66,8 +79,11 @@ module Aikido
     # @param request [Aikido::Zen::Request]
     # @return [void]
     def self.track_request(request)
-      autostart
-      collector.track_request(request)
+      collector.track_request
+    end
+
+    def self.track_discovered_route(request)
+      collector.track_route(request)
     end
 
     # Tracks a network connection made to an external service.
@@ -75,7 +91,6 @@ module Aikido
     # @param connection [Aikido::Zen::OutboundConnection]
     # @return [void]
     def self.track_outbound(connection)
-      autostart
       collector.track_outbound(connection)
     end
 
@@ -87,7 +102,6 @@ module Aikido
     # @raise [Aikido::Zen::UnderAttackError] if the scan detected an Attack
     #   and blocking_mode is enabled.
     def self.track_scan(scan)
-      autostart
       collector.track_scan(scan)
       agent.handle_attack(scan.attack) if scan.attack?
     end
@@ -100,7 +114,6 @@ module Aikido
       return if config.disabled?
 
       if (actor = Aikido::Zen::Actor(user))
-        autostart
         collector.track_user(actor)
         current_context.request.actor = actor if current_context
       else
@@ -113,6 +126,12 @@ module Aikido
       end
     end
 
+    # Marks that the Zen middleware was installed properly
+    # @return void
+    def self.middleware_installed!
+      collector.middleware_installed!
+    end
+
     # Load all sinks matching libraries loaded into memory. This method should
     # be called after all other dependencies have been loaded into memory (i.e.
     # at the end of the initialization process).
@@ -127,7 +146,8 @@ module Aikido
     # @!visibility private
     # Stop any background threads.
     def self.stop!
-      agent&.stop!
+      @agent&.stop!
+      @detached_agent_server&.stop!
     end
 
     # @!visibility private
@@ -136,8 +156,34 @@ module Aikido
       @agent ||= Agent.start
     end
 
+    def self.detached_agent_server
+      @detached_agent_server ||= DetachedAgent::Server.start!
+    end
+
     class << self
-      alias_method :autostart, :agent
+      # `agent` and `detached_agent` are started on the first method call.
+      # A mutex controls thread execution to prevent multiple attempts.
+      LOCK = Mutex.new
+
+      def start!
+        @pid = Process.pid
+        LOCK.synchronize do
+          agent
+          detached_agent_server
+        end
+      end
+
+      def check_and_handle_fork
+        if has_forked
+          @detached_agent&.handle_fork
+        end
+      end
+
+      def has_forked
+        pid_changed = Process.pid != @pid
+        @pid = Process.pid
+        pid_changed
+      end
     end
   end
 end

data/tasklib/bench.rake

@@ -2,8 +2,11 @@
 
 require "socket"
 require "timeout"
+require_relative "wrk"
 
 SERVER_PIDS = {}
+PORT_PROTECTED = 3001
+PORT_UNPROTECTED = 3002
 
 def stop_servers
   SERVER_PIDS.each { |_, pid| Process.kill("TERM", pid) }
@@ -11,12 +14,15 @@ def stop_servers
 end
 
 def boot_server(dir, port:, env: {})
+  env["RAILS_MIN_THREADS"] = NUMBER_OF_THREADS
+  env["RAILS_MAX_THREADS"] = NUMBER_OF_THREADS
   env["PORT"] = port.to_s
+  env["SECRET_KEY_BASE"] = rand(36**64).to_s(36)
 
   Dir.chdir(dir) do
     SERVER_PIDS[port] = Process.spawn(
       env,
-      "rails", "server", "--pid", "#{Dir.pwd}/tmp/pids/server.#{port}.pid",
+      "rails", "server", "--pid", "#{Dir.pwd}/tmp/pids/server.#{port}.pid", "-e", "production",
       out: "/dev/null"
     )
   rescue
@@ -48,8 +54,28 @@ end
 Pathname.glob("sample_apps/*").select(&:directory?).each do |dir|
   namespace :bench do
     namespace dir.basename.to_s do
-      desc "Run benchmarks for the #{dir.basename} sample app"
-      task run: [:boot_protected_app, :boot_unprotected_app] do
+      desc "Run WRK benchmarks for the #{dir.basename} sample app"
+      task wrk_run: [:boot_protected_app, :boot_unprotected_app] do
+        throughput_decrease_limit_perc = 25
+        if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new("3.0.0") && Gem::Version.new(RUBY_VERSION) < Gem::Version.new("3.1.0")
+          # add higher limit for ruby 3.0
+          throughput_decrease_limit_perc = 35
+        end
+
+        wait_for_servers
+        run_benchmark(
+          route_zen: "http://localhost:#{PORT_PROTECTED}/benchmark", # Application with Zen
+          route_no_zen: "http://localhost:#{PORT_UNPROTECTED}/benchmark", # Application without Zen
+          description: "An empty route (1ms simulated delay)",
+          throughput_decrease_limit_perc: throughput_decrease_limit_perc,
+          latency_increase_limit_ms: 200
+        )
+      ensure
+        stop_servers
+      end
+
+      desc "Run K6 benchmarks for the #{dir.basename} sample app"
+      task k6_run: [:boot_protected_app, :boot_unprotected_app] do
         wait_for_servers
         Dir.chdir("benchmarks") { sh "k6 run #{dir.basename}.js" }
       ensure
@@ -57,14 +83,12 @@ Pathname.glob("sample_apps/*").select(&:directory?).each do |dir|
       end
 
       task :boot_protected_app do
-        boot_server(dir, port: 3001)
+        boot_server(dir, port: PORT_PROTECTED)
       end
 
       task :boot_unprotected_app do
-        boot_server(dir, port: 3002, env: {"AIKIDO_DISABLE" => "true"})
+        boot_server(dir, port: PORT_UNPROTECTED, env: {"AIKIDO_DISABLED" => "true"})
       end
     end
-
-    task default: "#{dir.basename}:run"
   end
 end

data/tasklib/wrk.rb

@@ -0,0 +1,88 @@
+require "open3"
+require "time"
+
+NUMBER_OF_THREADS = ENV.fetch("BENCHMARK_NUMBER_OF_THREADS") { 12 }.to_s
+CONNECTIONS = ENV.fetch("BENCHMARK_WRK_CONNECTIONS") { 400 }
+
+def generate_wrk_command_for_url(url)
+  # Define the command with wrk included
+  "wrk --threads #{NUMBER_OF_THREADS} --connections #{CONNECTIONS} --duration 15s --timeout 5s --latency #{url}"
+end
+
+def cold_start(url)
+  10.times do
+    _, err, status = Open3.capture3("curl #{url}")
+
+    if status != 0
+      puts err
+      exit(-1)
+    end
+  end
+end
+
+def extract_requests_and_latency_tuple(out, err, status)
+  if status == 0
+    # Extracting requests/sec
+    requests_sec_match = out.match(/Requests\/sec:\s+([\d.]+)/)
+    requests_sec = requests_sec_match[1].to_f if requests_sec_match
+
+    # Extracting latency
+    latency_match = out.match(/Latency\s+([\d.]+)(ms|s)/)
+    latency = latency_match[1].to_f if latency_match
+    latency_unit = latency_match[2] if latency_match
+
+    if latency_unit == "s"
+      latency *= 1000
+    end
+
+    {requests_sec: requests_sec, latency: latency}
+  else
+    puts "Error occurred running benchmark command:"
+    puts err.strip
+    exit(1)
+  end
+end
+
+def run_benchmark(route_no_zen:, route_zen:, description:, throughput_decrease_limit_perc:, latency_increase_limit_ms:)
+  # Cold start
+  cold_start(route_no_zen)
+  cold_start(route_zen)
+
+  out, err, status = Open3.capture3(generate_wrk_command_for_url(route_zen))
+  puts <<~MSG
+    WRK OUTPUT
+    ================
+    FIREWALL ENABLED:
+        #{out}
+    ----------------
+  MSG
+  result_zen_enabled = extract_requests_and_latency_tuple(out, err, status)
+
+  out, err, status = Open3.capture3(generate_wrk_command_for_url(route_no_zen))
+  puts <<~MSG
+    FIREWALL DISABLED:
+        #{out}
+    ================
+  MSG
+  result_zen_disabled = extract_requests_and_latency_tuple(out, err, status)
+
+  # Check if the command was successful
+  if result_zen_enabled && result_zen_disabled
+    # Print the output, which should be the Requests/sec value
+    puts "[ZEN ENABLED ] Requests/sec: #{result_zen_enabled[:requests_sec]} | Latency in ms: #{result_zen_enabled[:latency]}"
+    puts "[ZEN DISABLED] Requests/sec: #{result_zen_disabled[:requests_sec]} | Latency in ms: #{result_zen_disabled[:latency]}"
+
+    latency_increase_ms = (result_zen_enabled[:latency] - result_zen_disabled[:latency]).round(2)
+    puts "-> Delta in ms: #{latency_increase_ms}ms after running load test on #{description}"
+
+    throughput_decrease_perc = ((result_zen_disabled[:requests_sec] - result_zen_enabled[:requests_sec]) / result_zen_disabled[:requests_sec] * 100).round
+    puts "-> #{throughput_decrease_perc}% decrease in throughput after running load test on #{description}\n"
+
+    if latency_increase_ms >= latency_increase_limit_ms
+      exit(1)
+    end
+    if throughput_decrease_perc >= throughput_decrease_limit_perc
+      exit(1)
+    end
+  end
+end