aikido-zen 1.0.2.beta.5-x86_64-linux → 1.0.2.beta.6-x86_64-linux

data/lib/aikido/zen/collector/stats.rb

@@ -35,10 +35,9 @@ module Aikido::Zen
     # Track the timestamp we start tracking this series of stats.
     #
     # @param at [Time]
-    # @return [self]
+    # @return [void]
     def start(at = Time.now.utc)
       @started_at = at
-      self
     end
 
     # Sets the end time for these stats block, freezes it to avoid any more
@@ -47,7 +46,7 @@ module Aikido::Zen
     #
     # @param at [Time] the time at which we're resetting, which is set as the
     #   ending time for the returned copy.
-    # @return [self]
+    # @return [void]
     def flush(at: Time.now.utc)
       # Make sure the timing stats are compressed before copying, since we
       # need these compressed when we serialize this for the API.
@@ -56,31 +55,29 @@ module Aikido::Zen
       freeze
     end
 
-    # @return [self]
+    # @return [void]
     def add_request
       @requests += 1
-      self
     end
 
-    # @param scan [Aikido::Zen::Scan]
-    # @return [self]
-    def add_scan(scan)
-      stats = @sinks[scan.sink.name]
+    # @param sink_name [String] the name of the sink
+    # @param duration [Float] the length the scan in seconds
+    # @param has_errors [Boolean] whether errors occurred during the scan
+    # @return [void]
+    def add_scan(sink_name, duration, has_errors:)
+      stats = @sinks[sink_name]
       stats.scans += 1
-      stats.errors += 1 if scan.errors?
-      stats.add_timing(scan.duration)
-      self
+      stats.errors += 1 if has_errors
+      stats.add_timing(duration)
     end
 
-    # @param attack [Aikido::Zen::Attack]
-    # @param being_blocked [Boolean] whether the Agent blocked the
-    #   request where this Attack happened or not.
-    # @return [self]
-    def add_attack(attack, being_blocked:)
-      stats = @sinks[attack.sink.name]
+    # @param sink_name [String] the name of the sink
+    # @param being_blocked [Boolean] whether the Agent blocked the request
+    # @return [void]
+    def add_attack(sink_name, being_blocked:)
+      stats = @sinks[sink_name]
       stats.attacks += 1
       stats.blocked_attacks += 1 if being_blocked
-      self
     end
 
     def as_json

data/lib/aikido/zen/collector/users.rb

@@ -11,9 +11,11 @@ module Aikido::Zen
       super(config.max_users_tracked)
     end
 
+    # @param actor [Aikido::Zen::Actor]
+    # @return [void]
     def add(actor)
       if key?(actor.id)
-        self[actor.id].update
+        self[actor.id] |= actor
       else
         self[actor.id] = actor
       end

data/lib/aikido/zen/collector.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "collector/event"
+
 module Aikido::Zen
   # Handles collecting all the runtime statistics to report back to the Aikido
   # servers.
@@ -7,14 +9,44 @@ module Aikido::Zen
     def initialize(config: Aikido::Zen.config)
       @config = config
 
+      @events = Queue.new
+
       @stats = Concurrent::AtomicReference.new(Stats.new(@config))
       @users = Concurrent::AtomicReference.new(Users.new(@config))
       @hosts = Concurrent::AtomicReference.new(Hosts.new(@config))
       @routes = Concurrent::AtomicReference.new(Routes.new(@config))
-      @heartbeats = Queue.new
+
       @middleware_installed = Concurrent::AtomicBoolean.new
     end
 
+    # Add an event, to be handled in the collector in current process or sent
+    # to a collector in another process.
+    #
+    # @param event [Aikido::Zen::Collector::Event] the event to add
+    # @return [void]
+    def add_event(event)
+      @events << event
+    end
+
+    # @return [Boolean] whether the collector has any events
+    def has_events?
+      !@events.empty?
+    end
+
+    # Flush all events.
+    #
+    # @return [Array<Aikido::Zen::Collector::Event>]
+    def flush_events
+      Array.new(@events.size) { @events.pop }
+    end
+
+    # Handle the events in the queue.
+    #
+    # @return [void]
+    def handle
+      flush_events.each { |event| event.handle(self) }
+    end
+
     # Flush all the stats into a Heartbeat event that can be reported back to
     # the Aikido servers.
     #
@@ -22,6 +54,8 @@ module Aikido::Zen
     #   of the new stats collection period. Defaults to now.
     # @return [Aikido::Zen::Events::Heartbeat]
     def flush(at: Time.now.utc)
+      handle
+
       stats = @stats.get_and_set(Stats.new(@config))
       users = @users.get_and_set(Users.new(@config))
       hosts = @hosts.get_and_set(Hosts.new(@config))
@@ -30,21 +64,11 @@ module Aikido::Zen
       start(at: at)
       stats = stats.flush(at: at)
 
-      Events::Heartbeat.new(
+      Aikido::Zen::Events::Heartbeat.new(
         stats: stats, users: users, hosts: hosts, routes: routes, middleware_installed: middleware_installed?
       )
     end
 
-    # Put heartbeats coming from child processes into the internal queue.
-    def push_heartbeat(heartbeat)
-      @heartbeats << heartbeat
-    end
-
-    # Drains into an array all the queued heartbeats
-    def flush_heartbeats
-      Array.new(@heartbeats.size) { @heartbeats.pop }
-    end
-
     # Sets the start time for this collection period.
     #
     # @param at [Time] defaults to now.
@@ -55,16 +79,13 @@ module Aikido::Zen
 
     # Track stats about the requests
     #
-    # @param request [Aikido::Zen::Request]
     # @return [void]
-    def track_request(*)
-      synchronize(@stats) { |stats| stats.add_request }
+    def track_request
+      add_event(Events::TrackRequest.new)
     end
 
-    #  Record the visited endpoint, and if enabled, the API schema for this endpoint.
-    # @param request [Aikido::Zen::Request]
-    def track_route(request)
-      synchronize(@routes) { |routes| routes.add(request) if request.route }
+    def handle_track_request
+      synchronize(@stats) { |stats| stats.add_request }
     end
 
     # Track stats about a scan performed by one of our sinks.
@@ -72,7 +93,13 @@ module Aikido::Zen
     # @param scan [Aikido::Zen::Scan]
     # @return [void]
     def track_scan(scan)
-      synchronize(@stats) { |stats| stats.add_scan(scan) }
+      add_event(Events::TrackScan.new(scan.sink.name, scan.duration, has_errors: scan.errors?))
+    end
+
+    def handle_track_scan(sink_name, duration, has_errors:)
+      synchronize(@stats) do |stats|
+        stats.add_scan(sink_name, duration, has_errors: has_errors)
+      end
     end
 
     # Track stats about an attack detected by our scanners.
@@ -80,25 +107,49 @@ module Aikido::Zen
     # @param attack [Aikido::Zen::Attack]
     # @return [void]
     def track_attack(attack)
+      add_event(Events::TrackAttack.new(attack.sink.name, being_blocked: attack.blocked?))
+    end
+
+    def handle_track_attack(sink_name, being_blocked:)
       synchronize(@stats) do |stats|
-        stats.add_attack(attack, being_blocked: attack.blocked?)
+        stats.add_attack(sink_name, being_blocked: being_blocked)
       end
     end
 
+    # Track the user reported by the developer to be behind this request.
+    #
+    # @param actor [Aikido::Zen::Actor]
+    # @return [void]
+    def track_user(actor)
+      add_event(Events::TrackUser.new(actor))
+    end
+
+    def handle_track_user(actor)
+      synchronize(@users) { |users| users.add(actor) }
+    end
+
     # Track an HTTP connections to an external host.
     #
     # @param connection [Aikido::Zen::OutboundConnection]
     # @return [void]
     def track_outbound(connection)
+      add_event(Events::TrackOutbound.new(connection))
+    end
+
+    def handle_track_outbound(connection)
       synchronize(@hosts) { |hosts| hosts.add(connection) }
     end
 
-    # Track the user reported by the developer to be behind this request.
+    # Record the visited endpoint, and if enabled, the API schema for this endpoint.
     #
-    # @param actor [Aikido::Zen::Actor]
+    # @param request [Aikido::Zen::Request]
     # @return [void]
-    def track_user(actor)
-      synchronize(@users) { |users| users.add(actor) }
+    def track_route(request)
+      add_event(Events::TrackRoute.new(request.route, request.schema))
+    end
+
+    def handle_track_route(route, schema)
+      synchronize(@routes) { |routes| routes.add(route, schema) }
     end
 
     def middleware_installed!
@@ -106,26 +157,40 @@ module Aikido::Zen
     end
 
     # @api private
-    def routes
-      @routes.get
+    #
+    # @note Visible for testing.
+    def stats
+      handle
+      @stats.get
     end
 
     # @api private
+    #
+    # @note Visible for testing.
     def users
+      handle
       @users.get
     end
 
     # @api private
+    #
+    # @note Visible for testing.
     def hosts
+      handle
       @hosts.get
     end
 
     # @api private
-    def stats
-      @stats.get
+    #
+    # @note Visible for testing.
+    def routes
+      handle
+      @routes.get
     end
 
     # @api private
+    #
+    # @note Visible for testing.
     def middleware_installed?
       @middleware_installed.true?
     end

data/lib/aikido/zen/config.rb

@@ -8,12 +8,6 @@ require_relative "context"
 
 module Aikido::Zen
   class Config
-    # @api private
-    # @return [Boolean] whether Aikido should protect.
-    def protect?
-      !api_token.nil? || blocking_mode? || debugging?
-    end
-
     # @return [Boolean] whether Aikido should be turned completely off (no
     #   intercepting calls to protect the app, no agent process running, no
     #   middleware installed). Defaults to false (so, enabled). Can be set
@@ -46,9 +40,9 @@ module Aikido::Zen
     #   settings changes. Defaults to evey 60 seconds.
     attr_accessor :polling_interval
 
-    # @return [Integer] the amount in seconds to wait before sending an initial
-    #   heartbeat event when the server reports no stats have been sent yet.
-    attr_accessor :initial_heartbeat_delay
+    # @return [Array<Integer>] the delays in seconds to wait before sending
+    #   each initial heartbeat event.
+    attr_accessor :initial_heartbeat_delays
 
     # @return [#call] Callable that can be passed an Object and returns a String
     #   of JSON. Defaults to the standard library's JSON.dump method.
@@ -167,7 +161,7 @@ module Aikido::Zen
       self.realtime_endpoint = ENV.fetch("AIKIDO_REALTIME_ENDPOINT", DEFAULT_RUNTIME_BASE_URL)
       self.api_token = ENV.fetch("AIKIDO_TOKEN", nil)
       self.polling_interval = 60
-      self.initial_heartbeat_delay = 60
+      self.initial_heartbeat_delays = [30, 60 * 2]
       self.json_encoder = DEFAULT_JSON_ENCODER
       self.json_decoder = DEFAULT_JSON_DECODER
       self.debugging = read_boolean_from_env(ENV.fetch("AIKIDO_DEBUG", false))

data/lib/aikido/zen/context.rb

@@ -92,17 +92,18 @@ module Aikido::Zen
 
     private
 
+    # @!visibility private
     def extract_payloads_from(data, source_type, prefix = nil)
       if data.respond_to?(:to_hash)
-        data.to_hash.flat_map { |name, val|
-          extract_payloads_from(val, source_type, [prefix, name].compact.join("."))
-        }
+        data.to_hash.flat_map do |key, value|
+          extract_payloads_from(value, source_type, [prefix, key].compact.join("."))
+        end
       elsif data.respond_to?(:to_ary)
-        data.to_ary.flat_map.with_index { |val, idx|
-          extract_payloads_from(val, source_type, [prefix, idx].compact.join("."))
-        }
+        data.to_ary.flat_map.with_index do |value, index|
+          extract_payloads_from(value, source_type, [prefix, index].compact.join("."))
+        end
       else
-        Payload.new(data, source_type, prefix.to_s)
+        [Payload.new(data, source_type, prefix.to_s)]
       end
     end
   end

data/lib/aikido/zen/detached_agent/agent.rb

@@ -14,53 +14,39 @@ module Aikido::Zen::DetachedAgent
   # parent process. We want to have the freshest data.
   #
   # It's possible to use `extend Forwardable` here for one-line forward calls to the
-  # @detached_agent_front object. Unfortunately, the methods to be called are
+  # @front_object object. Unfortunately, the methods to be called are
   # created at runtime by `DRbObject`, which leads to an ugly warning about
   # private methods after the delegator is bound.
   class Agent
     attr_reader :worker
 
     def initialize(
+      config: Aikido::Zen.config,
+      worker: Aikido::Zen::Worker.new(config: config),
       heartbeat_interval: 10,
       polling_interval: 10,
-      config: Aikido::Zen.config,
-      collector: Aikido::Zen.collector,
-      worker: Aikido::Zen::Worker.new(config: config)
+      collector: Aikido::Zen.collector
     )
       @config = config
+      @worker = worker
       @heartbeat_interval = heartbeat_interval
       @polling_interval = polling_interval
-      @worker = worker
+
       @collector = collector
-      @detached_agent_front = DRbObject.new_with_uri(config.detached_agent_socket_uri)
+
+      @front_object = DRbObject.new_with_uri(config.detached_agent_socket_uri)
+
       @has_forked = false
       schedule_tasks
     end
 
-    def send_heartbeat(at: Time.now.utc)
-      return unless @collector.stats.any?
-
-      heartbeat = @collector.flush(at: at)
-      @detached_agent_front.send_heartbeat_to_parent_process(heartbeat.as_json)
-    end
-
-    private def schedule_tasks
-      # For heartbeats is correct to send them from parent or child process. Otherwise, we'll lose
-      # stats made by the parent process.
-      @worker.every(@heartbeat_interval, run_now: false) { send_heartbeat }
-
-      # Runtime_settings fetch must happens only in the child processes, otherwise, due to
-      # we are updating the global runtime_settings, we could have an infinite recursion.
-      if @has_forked
-        @worker.every(@polling_interval) do
-          Aikido::Zen.runtime_settings = @detached_agent_front.updated_settings
-          @config.logger.debug "Updated runtime settings after polling from child process #{Process.pid}"
-        end
-      end
+    def send_collector_events
+      events_data = @collector.flush_events.map(&:as_json)
+      @front_object.send_collector_events(events_data)
     end
 
     def calculate_rate_limits(request)
-      @detached_agent_front.calculate_rate_limits(request.route, request.ip, request.actor.to_json)
+      @front_object.calculate_rate_limits(request.route.as_json, request.ip, request.actor.as_json)
     end
 
     # Every time a fork occurs (a new child process is created), we need to start
@@ -74,5 +60,20 @@ module Aikido::Zen::DetachedAgent
       @worker.restart
       schedule_tasks
     end
+
+    private
+
+    def schedule_tasks
+      @worker.every(@heartbeat_interval, run_now: false) { send_collector_events }
+
+      # Runtime_settings fetch must happens only in the child processes, otherwise, due to
+      # we are updating the global runtime_settings, we could have an infinite recursion.
+      if @has_forked
+        @worker.every(@polling_interval) do
+          Aikido::Zen.runtime_settings = @front_object.updated_settings
+          @config.logger.debug "Updated runtime settings after polling from child process #{Process.pid}"
+        end
+      end
+    end
   end
 end

data/lib/aikido/zen/detached_agent/front_object.rb

@@ -7,20 +7,23 @@ module Aikido::Zen::DetachedAgent
   class FrontObject
     def initialize(
       config: Aikido::Zen.config,
-      collector: Aikido::Zen.collector,
       runtime_settings: Aikido::Zen.runtime_settings,
+      collector: Aikido::Zen.collector,
       rate_limiter: Aikido::Zen::RateLimiter.new
     )
       @config = config
+      @runtime_settings = runtime_settings
       @collector = collector
       @rate_limiter = rate_limiter
-      @runtime_settings = runtime_settings
     end
 
     RequestKind = Struct.new(:route, :schema, :ip, :actor)
 
-    def send_heartbeat_to_parent_process(heartbeat)
-      @collector.push_heartbeat(heartbeat)
+    def send_collector_events(events_data)
+      events_data.each do |event_data|
+        event = Aikido::Zen::Collector::Event.from_json(event_data)
+        @collector.add_event(event)
+      end
     end
 
     # Method called by child processes to get an up-to-date version of the
@@ -29,8 +32,9 @@ module Aikido::Zen::DetachedAgent
       @runtime_settings
     end
 
-    def calculate_rate_limits(route, ip, actor_hash)
-      actor = Aikido::Zen::Actor(actor_hash) if actor_hash
+    def calculate_rate_limits(route_data, ip, actor_data)
+      actor = Aikido::Zen::Actor.from_json(actor_data)
+      route = Aikido::Zen::Route.from_json(route_data)
       @rate_limiter.calculate_rate_limits(RequestKind.new(route, nil, ip, actor))
     end
   end

data/lib/aikido/zen/internals.rb

@@ -60,11 +60,11 @@ module Aikido::Zen
       # @!method self.detect_sql_injection_native(query, input, dialect)
       # @param (see .detect_sql_injection)
       # @returns [Integer] 0 if no injection detected, 1 if an injection was
-      #   detected, or 2 if there was an internal error.
+      #   detected, 2 if there was an internal error, or 3 if SQL tokenization failed.
       # @raise [Aikido::Zen::InternalsError] if there's a problem loading or
       #   calling libzen.
       attach_function :detect_sql_injection_native, :detect_sql_injection,
-        [:string, :string, :int], :int
+        [:pointer, :size_t, :pointer, :size_t, :int], :int
     rescue LoadError, FFI::NotFoundError => err # rubocop:disable Lint/ShadowedException
       # :nocov:
 
@@ -90,14 +90,34 @@ module Aikido::Zen
       # @raise [Aikido::Zen::InternalsError] if there's a problem loading or
       #   calling libzen.
       def self.detect_sql_injection(query, input, dialect)
-        case detect_sql_injection_native(query, input, dialect)
+        query_bytes = encode_safely(query)
+        input_bytes = encode_safely(input)
+
+        query_ptr = FFI::MemoryPointer.new(:uint8, query_bytes.bytesize)
+        input_ptr = FFI::MemoryPointer.new(:uint8, input_bytes.bytesize)
+
+        query_ptr.put_bytes(0, query_bytes)
+        input_ptr.put_bytes(0, input_bytes)
+
+        case detect_sql_injection_native(query_ptr, query_bytes.bytesize, input_ptr, input_bytes.bytesize, dialect)
         when 0 then false
         when 1 then true
         when 2
           attempt = format("%s query %p with input %p", dialect, query, input)
           raise InternalsError.new(attempt, "calling detect_sql_injection in", libzen_name)
+        when 3
+          # SQL tokenization failed - return false (no injection detected)
+          false
         end
       end
     end
+
+    class << self
+      private
+
+      def encode_safely(string)
+        string.encode("UTF-8", invalid: :replace, undef: :replace)
+      end
+    end
   end
 end

data/lib/aikido/zen/middleware/fork_detector.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Aikido
+  module Zen
+    module Middleware
+      # This middleware is responsible for detecting when a process has forked
+      # (e.g., in a Puma or Unicorn worker) and resetting the state of the
+      # Aikido Zen agent. It should be inserted early in the middleware stack.
+      class ForkDetector
+        def initialize(app)
+          @app = app
+        end
+
+        def call(env)
+          # This is the single, reliable trigger point for the fork check.
+          Aikido::Zen.check_and_handle_fork
+
+          @app.call(env)
+        end
+      end
+    end
+  end
+end

data/lib/aikido/zen/middleware/request_tracker.rb

@@ -18,7 +18,7 @@ module Aikido::Zen
           route: request.route.path,
           http_method: request.request_method
         )
-          Aikido::Zen.track_request request
+          Aikido::Zen.track_request(request)
 
           if Aikido::Zen.config.collect_api_schema?
             Aikido::Zen.track_discovered_route(request)

data/lib/aikido/zen/outbound_connection.rb

@@ -3,6 +3,13 @@
 module Aikido::Zen
   # Simple data object to identify connections performed to outbound servers.
   class OutboundConnection
+    def self.from_json(data)
+      new(
+        host: data[:hostname],
+        port: data[:port]
+      )
+    end
+
     # Convenience factory to create connection descriptions out of URI objects.
     #
     # @param uri [URI]

data/lib/aikido/zen/rails_engine.rb

@@ -10,6 +10,8 @@ module Aikido::Zen
     end
 
     initializer "aikido.add_middleware" do |app|
+      app.middleware.insert_before 0, Aikido::Zen::Middleware::ForkDetector
+
       app.middleware.use Aikido::Zen::Middleware::SetContext
       app.middleware.use Aikido::Zen::Middleware::CheckAllowedAddresses
       # Request Tracker stats do not consider failed request or 40x, so the middleware
@@ -29,12 +31,6 @@ module Aikido::Zen
     end
 
     initializer "aikido.configuration" do |app|
-      # Allow the logger to be configured before checking if disabled? so we can
-      # let the user know that the agent is disabled.
-      logger = ::Rails.logger
-      logger = logger.tagged("aikido") if logger.respond_to?(:tagged)
-      app.config.zen.logger = logger
-
       app.config.zen.request_builder = Aikido::Zen::Context::RAILS_REQUEST_BUILDER
 
       # Plug Rails' JSON encoder/decoder, but only if the user hasn't changed

data/lib/aikido/zen/route.rb

@@ -5,6 +5,13 @@ module Aikido::Zen
   # framework to go from a given HTTP request to the code that handles said
   # request.
   class Route
+    def self.from_json(data)
+      new(
+        verb: data[:method],
+        path: data[:path]
+      )
+    end
+
     # @return [String] the HTTP verb used to request this route.
     attr_reader :verb
 

data/lib/aikido/zen/runtime_settings.rb

@@ -50,7 +50,7 @@ module Aikido::Zen
       last_updated_at = updated_at
 
       self.updated_at = Time.at(data["configUpdatedAt"].to_i / 1000)
-      self.heartbeat_interval = (data["heartbeatIntervalInMS"].to_i / 1000)
+      self.heartbeat_interval = data["heartbeatIntervalInMS"].to_i / 1000
       self.endpoints = RuntimeSettings::Endpoints.from_json(data["endpoints"])
       self.blocked_user_ids = data["blockedUserIds"]
       self.skip_protection_for_ips = RuntimeSettings::IPSet.from_json(data["allowedIPAddresses"])