aikido-zen 1.5.0-aarch64-linux → 1.6.0-aarch64-linux

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fcc4f05fee5eb4c2b6bc4b75313c8d1a8948282b2e1ea3e4ef0e735055e55f42
-  data.tar.gz: 692ace187aae4b18bce5d1aefaa10c7aeacbfd28016645792556ae0aafe2c8dc
+  metadata.gz: eba16aa05cb3f3506e487a3dfd7c03a1d4bbc186a397053aaf43562fad1ed4f4
+  data.tar.gz: b9815eab35c5f245eba4587c64bf4bfda422c1bcfea1d1d21eee747057723e41
 SHA512:
-  metadata.gz: a32eea6aec3687a23e54c6cc3e1365b4e4bd9e2cdcf93d39acf46289ad6063ca0671061ed1aada13ab18267459af66d0dc5660f08745c0f290f705086cdb214e
-  data.tar.gz: a145a0fcd6e137f4ca276219e7a235811da1bb3aed09ea7e0a123ab7e69ae09b5df657a0050de99e30165040906db9b948979a24e0a736a61e992caad33fed53
+  metadata.gz: 9c99f6276ce2d33f42d3b9ec7f5430bb647cd72cb146e583066d97d918cc0e59b21d82999dada088197627be1226579cbea675f7750c609346a2d7487f39b3e5
+  data.tar.gz: ad2e3595f4fee8181ba2083f05a2a15181cf62db472f47c112a3480dd6575846cf50d7e9c3ac420a5f3ddbd48bda9b98cff815de50a12a35e427ebbce7e68427

data/benchmarks/ipc/benchmark.rb

@@ -0,0 +1,149 @@
+#!/usr/bin/env ruby
+
+require "benchmark"
+require "securerandom"
+require "aikido-zen"
+
+class BasicIO
+  def initialize(data, timeout: 5)
+    @data = data
+    @size = data.bytesize
+    @timeout = timeout
+  end
+
+  def read(socket)
+    buffer = String.new(encoding: Encoding::BINARY)
+
+    while buffer.bytesize < @size
+      case chunk = socket.read_nonblock(@size - buffer.bytesize, exception: false)
+      when :wait_readable
+        raise Errno::ETIMEDOUT, "read timed out" unless IO.select([socket], nil, nil, @timeout)
+      when nil
+        raise EOFError
+      else
+        buffer << chunk
+      end
+    end
+
+    buffer
+  end
+
+  def write(socket)
+    written = 0
+
+    while written < @data.bytesize
+      case n = socket.write_nonblock(@data.byteslice(written..), exception: false)
+      when :wait_writable
+        raise Errno::ETIMEDOUT, "write timed out" unless IO.select(nil, [socket], nil, @timeout)
+      else
+        written += n
+      end
+    end
+  end
+end
+
+class TimedIO
+  include Aikido::Zen::IPC::TimedIO
+
+  def initialize(data, timeout: 5)
+    @data = data
+    @size = data.bytesize
+    @timeout = timeout
+  end
+
+  def read(socket)
+    read_with_deadline(socket, @size, Process.clock_gettime(Process::CLOCK_MONOTONIC) + @timeout)
+  end
+
+  def write(socket)
+    write_with_deadline(socket, @data, Process.clock_gettime(Process::CLOCK_MONOTONIC) + @timeout)
+  end
+end
+
+class FramedIO
+  include Aikido::Zen::IPC::FramedIO
+
+  def initialize(data, timeout: 5)
+    @data = data
+    @timeout = timeout
+    @buffer = String.new(encoding: Encoding::BINARY)
+  end
+
+  def read(socket)
+    read_coalesced_frame_with_timeout(socket, @buffer, nil, @timeout)
+  end
+
+  def write(socket)
+    write_coalesced_frame_with_timeout(socket, @data, nil, @timeout)
+  end
+end
+
+def run(duration, io)
+  secret = SecureRandom.bytes(32)
+
+  server = Aikido::Zen::IPC::Server.start(secret) do |socket|
+    loop do
+      io.read(socket)
+    end
+  rescue EOFError, Errno::ECONNRESET, Errno::EPIPE
+    # disconnected
+  rescue IOError
+    # client stopped
+  end
+
+  client = Aikido::Zen::IPC::Client.new(secret, server.host, server.port)
+
+  sent = 0
+
+  deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + duration
+
+  result = Benchmark.measure do
+    while Process.clock_gettime(Process::CLOCK_MONOTONIC) < deadline
+      io.write(client.socket)
+
+      sent += 1
+    end
+  end
+
+  [sent, result.real]
+ensure
+  client.close
+  server.stop
+end
+
+def report(label, sent, real, data_size, baseline: nil)
+  rate = sent / real
+  mib_per_sec = rate * data_size / (1024 * 1024)
+
+  puts label
+  puts "#{rate.round} messages/sec, #{mib_per_sec.round(2)} MiB/sec"
+
+  if baseline
+    overhead = 100 - (100.0 * rate / baseline)
+    puts "#{overhead.round(1)}% fewer messages/sec than baseline"
+  end
+end
+
+if __FILE__ == $0
+  duration = (ENV["DURATION"] || 5).to_i
+  data_size = (ENV["DATA_SIZE"] || 1024).to_i
+
+  data = SecureRandom.bytes(data_size)
+
+  basic_io = BasicIO.new(data)
+  basic_io_sent, basic_io_real = run(duration, basic_io)
+  report("baseline", basic_io_sent, basic_io_real, data_size)
+  basic_io_rate = basic_io_sent / basic_io_real
+
+  puts
+
+  timed_io = TimedIO.new(data)
+  timed_io_sent, timed_io_real = run(duration, timed_io)
+  report("TimedIO", timed_io_sent, timed_io_real, data_size, baseline: basic_io_rate)
+
+  puts
+
+  framed_io = FramedIO.new(data)
+  framed_io_sent, framed_io_real = run(duration, framed_io)
+  report("FramedIO (coalesced)", framed_io_sent, framed_io_real, data_size, baseline: basic_io_rate)
+end

data/benchmarks/rpc/benchmark.rb

@@ -0,0 +1,57 @@
+#!/usr/bin/env ruby
+
+require "benchmark"
+require "securerandom"
+require "logger"
+require "aikido-zen"
+
+def run(duration, concurrency)
+  secret = SecureRandom.bytes(32)
+  logger = Logger.new(File::NULL)
+
+  server = Aikido::Zen::RPC::Server.start(secret, logger: logger) do |server|
+    server.handle("echo") { |respond, value| respond.call(value) }
+  end
+
+  client = Aikido::Zen::RPC::Client.start(secret, server.host, server.port, logger: logger)
+
+  calls = Concurrent::AtomicFixnum.new(0)
+
+  deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + duration
+
+  result = Benchmark.measure do
+    threads = concurrency.times.map do
+      Thread.new do
+        while Process.clock_gettime(Process::CLOCK_MONOTONIC) < deadline
+          client.invoke("echo", 5.0, "hello")
+
+          calls.increment
+        end
+      end
+    end
+
+    threads.each(&:join)
+  end
+
+  [calls.value, result.real]
+ensure
+  client.stop
+  server.stop
+end
+
+def report(label, total_calls, real)
+  rate = total_calls / real
+  latency_ms = real / total_calls * 1000
+
+  puts label
+  puts "#{rate.round} calls/sec, #{latency_ms.round(3)} ms/call"
+end
+
+if __FILE__ == $0
+  duration = (ENV["DURATION"] || 5).to_i
+  concurrency = (ENV["CONCURRENCY"] || 1).to_i
+
+  calls, real = run(duration, concurrency)
+
+  report("RPC echo (concurrency: #{concurrency})", calls, real)
+end

data/docs/rails.md

@@ -27,6 +27,28 @@ Bundler.require(*Rails.groups)
 That's it! Zen will start to run inside your app when it starts getting
 requests.
 
+## Running with multiple worker processes (Puma)
+
+Zen can share rate limiting state, runtime settings, and runtime stats
+across all of an app's worker processes. This only works if Rails is booted
+once and then forked into workers, which on Puma requires
+[`preload_app!`][puma-preload]:
+
+```ruby
+# config/puma.rb
+preload_app!
+```
+
+Without `preload_app!`, each worker boots Rails independently, so Zen
+tracks rate limits, settings, and stats separately per worker instead
+of sharing them.
+
+> [!NOTE]
+> This only matters for multi-process setups; it doesn't apply to
+> single-process or threaded-only deployments.
+
+[puma-preload]: https://github.com/puma/puma#cluster-mode
+
 ## Rate limiting and user blocking
 
 If you want to add the rate limiting feature to your app, modify your code like this:

data/lib/aikido/zen/actor.rb

@@ -40,11 +40,11 @@ module Aikido::Zen
   class Actor
     def self.from_json(data)
       new(
-        id: data[:id],
-        name: data[:name],
-        ip: data[:lastIpAddress],
-        first_seen_at: Time.at(data[:firstSeenAt] / 1000),
-        last_seen_at: Time.at(data[:lastSeenAt] / 1000)
+        id: data["id"],
+        name: data["name"],
+        ip: data["lastIpAddress"],
+        first_seen_at: Time.at(data["firstSeenAt"] / 1000),
+        last_seen_at: Time.at(data["lastSeenAt"] / 1000)
       )
     end
 
@@ -135,11 +135,11 @@ module Aikido::Zen
 
     def as_json
       {
-        id: id,
-        name: name,
-        lastIpAddress: ip,
-        firstSeenAt: first_seen_at.to_i * 1000,
-        lastSeenAt: last_seen_at.to_i * 1000
+        "id" => id,
+        "name" => name,
+        "lastIpAddress" => ip,
+        "firstSeenAt" => first_seen_at.to_i * 1000,
+        "lastSeenAt" => last_seen_at.to_i * 1000
       }.compact
     end
   end

data/lib/aikido/zen/agent.rb

@@ -19,14 +19,12 @@ module Aikido::Zen
     def initialize(
       config: Aikido::Zen.config,
       collector: Aikido::Zen.collector,
-      detached_agent: Aikido::Zen.detached_agent,
       worker: Aikido::Zen::Worker.new(config: config),
       api_client: Aikido::Zen::APIClient.new(config: config),
       api_stream: Aikido::Zen::APIStream.new(config: config)
     )
       @config = config
       @collector = collector
-      @detached_agent = detached_agent
       @worker = worker
       @api_client = api_client
       @api_stream = api_stream
@@ -61,35 +59,37 @@ module Aikido::Zen
         return
       end
 
+      # Code coverage is disabled here because the block is only called when
+      # the process exits.
+      # :nocov:
       at_exit { stop! if started? }
+      # :nocov:
 
       report(Events::Started.new(time: @started_at)) do |response|
-        if update_settings_from_runtime_config!(response)
+        if update_settings_from_runtime_config!(response, reason: "after start")
           updated_settings!
-          @config.logger.info("Updated runtime settings")
+        # :nocov:
+        else
+          # empty
         end
+        # :nocov:
       rescue => err
         @config.logger.error(err.message)
       end
 
       begin
-        update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists)
-        @config.logger.info("Updated runtime firewall list")
+        update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists, reason: "after start")
       rescue => err
         @config.logger.error(err.message)
       end
 
       if @config.realtime_settings_updates_enabled?
-        if @api_stream.can_connect?
-          @api_stream.handle("config-updated") { |event| settings_updated(event) }
-          @api_stream.start!
-
-          # Use the realtime setting updates endpoint when polling to check
-          # whether settings should be fetched.
-          @api_client.should_fetch_settings_endpoint = @config.realtime_settings_updates_endpoint
-        else
-          @config.logger.warn("Can't reach #{Aikido::Zen.config.realtime_settings_updates_endpoint}, make sure it's in your outbound firewall allowlist. Realtime config updates won't be available, switched to polling.")
+        @api_stream.handle("config-updated") do |event|
+          @config.logger.debug("Received server-sent event: config-updated")
+          settings_updated(event)
         end
+
+        @api_stream.start!
       end
 
       poll_for_setting_updates
@@ -177,13 +177,15 @@ module Aikido::Zen
 
       heartbeat = @collector.flush
       report(heartbeat) do |response|
-        if update_settings_from_runtime_config!(response)
+        if update_settings_from_runtime_config!(response, reason: "after heartbeat")
           updated_settings!
-          @config.logger.info("Updated runtime settings after heartbeat")
 
-          update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists)
-          @config.logger.info("Updated runtime firewall list after heartbeat")
+          update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists, reason: "after heartbeat")
+        # :nocov:
+        else
+          # empty
         end
+        # :nocov:
       end
     end
 
@@ -197,13 +199,15 @@ module Aikido::Zen
     def poll_for_setting_updates
       @worker.every(@config.polling_interval) do
         if @api_client.should_fetch_settings?
-          if update_settings_from_runtime_config!(@api_client.fetch_runtime_config)
+          if update_settings_from_runtime_config!(@api_client.fetch_runtime_config, reason: "after polling")
             updated_settings!
-            @config.logger.info("Updated runtime settings after polling")
+          # :nocov:
+          else
+            # empty
           end
+          # :nocov:
 
-          update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists)
-          @config.logger.info("Updated runtime firewall list after polling")
+          update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists, reason: "after polling")
         end
       end
     end
@@ -214,13 +218,15 @@ module Aikido::Zen
       updated_at = Time.at(event[:data]["configUpdatedAt"].to_i)
 
       if should_fetch_settings?(updated_at)
-        if update_settings_from_runtime_config!(@api_client.fetch_runtime_config)
+        if update_settings_from_runtime_config!(@api_client.fetch_runtime_config, reason: "after server-sent event")
           updated_settings!
-          @config.logger.info("Updated runtime settings after server-side event")
 
-          update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists)
-          @config.logger.info("Updated runtime firewall list after server-side event")
+          update_settings_from_runtime_firewall_lists!(@api_client.fetch_runtime_firewall_lists, reason: "after server-sent event")
+        # :nocov:
+        else
+          # empty
         end
+        # :nocov:
       end
     end
 
@@ -239,22 +245,44 @@ module Aikido::Zen
     end
 
     # @param data [Hash]
-    # @return [Boolean]
-    def update_settings_from_runtime_config!(data)
+    # @param reason [String] context appended to the log message when logged
+    # @return [Boolean] whether the settings were updated; or nil if another
+    #   thread is already applying an update.
+    def update_settings_from_runtime_config!(data, reason:)
+      # :nocov:
       return unless @runtime_config_update_mutex.try_lock
+      # :nocov:
       begin
-        Aikido::Zen.runtime_settings.update_from_runtime_config_json(data)
+        return false unless Aikido::Zen.api_cache.update_runtime_config(data)
+
+        if Aikido::Zen.runtime_settings.update_from_runtime_config_json(data)
+          @config.logger.info("Updated runtime settings #{reason}")
+          true
+        else
+          false
+        end
       ensure
         @runtime_config_update_mutex.unlock
       end
     end
 
     # @param data [Hash]
-    # @return [Boolean]
-    def update_settings_from_runtime_firewall_lists!(data)
+    # @param reason [String] context appended to the log message when logged
+    # @return [Boolean] whether the settings were updated; or nil if another
+    #   thread is already applying an update.
+    def update_settings_from_runtime_firewall_lists!(data, reason:)
+      # :nocov:
       return unless @runtime_firewall_lists_update_mutex.try_lock
+      # :nocov:
       begin
-        Aikido::Zen.runtime_settings.update_from_runtime_firewall_lists_json(data)
+        return false unless Aikido::Zen.api_cache.update_runtime_firewall_lists(data)
+
+        if Aikido::Zen.runtime_settings.update_from_runtime_firewall_lists_json(data)
+          @config.logger.info("Updated runtime firewall list #{reason}")
+          true
+        else
+          false
+        end
       ensure
         @runtime_firewall_lists_update_mutex.unlock
       end

data/lib/aikido/zen/api_cache.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  class APICache
+    attr_reader :runtime_config
+    attr_reader :runtime_firewall_lists
+    attr_reader :runtime_config_generation
+    attr_reader :runtime_firewall_lists_generation
+
+    def initialize
+      @runtime_config_generation = 0
+      @runtime_firewall_lists_generation = 0
+    end
+
+    def runtime_config=(value)
+      update_runtime_config(value)
+    end
+
+    def runtime_firewall_lists=(value)
+      update_runtime_firewall_lists(value)
+    end
+
+    # @param value [Hash, nil]
+    # @return [Boolean] whether the value actually changed
+    def update_runtime_config(value)
+      return false if value == @runtime_config
+
+      @runtime_config_generation += 1
+      @runtime_config = value
+
+      true
+    end
+
+    # @param value [Hash, nil]
+    # @return [Boolean] whether the value actually changed
+    def update_runtime_firewall_lists(value)
+      return false if value == @runtime_firewall_lists
+
+      @runtime_firewall_lists_generation += 1
+      @runtime_firewall_lists = value
+
+      true
+    end
+
+    # @param known_generation [Integer, nil]
+    # @return [Array(Object, Integer), nil] the current value and generation,
+    #   or nil if the known generation is already current.
+    def config_if_changed(known_generation)
+      return nil if known_generation == runtime_config_generation
+      [runtime_config, runtime_config_generation]
+    end
+
+    # @param known_generation [Integer, nil]
+    # @return [Array(Object, Integer), nil] the current value and generation,
+    #   or nil if the known generation is already current.
+    def firewall_lists_if_changed(known_generation)
+      return nil if known_generation == runtime_firewall_lists_generation
+      [runtime_firewall_lists, runtime_firewall_lists_generation]
+    end
+  end
+end

data/lib/aikido/zen/api_stream.rb

@@ -27,7 +27,7 @@ module Aikido::Zen
       @thread = nil
       @http = nil
 
-      endpoint = @config.realtime_settings_updates_endpoint
+      endpoint = @config.realtime_endpoint
 
       @host = endpoint.host
       @port = endpoint.port
@@ -37,31 +37,6 @@ module Aikido::Zen
       @handlers = Concurrent::Array.new
     end
 
-    # @return [Boolean] whether we could connect to the realtime endpoint
-    def can_connect?
-      http = Net::HTTP.new(@host, @port)
-      http.use_ssl = @use_ssl
-      http.open_timeout = 5
-      http.write_timeout = 5
-      http.read_timeout = 5
-      http.max_retries = 0
-
-      request = Net::HTTP::Get.new("/config")
-      request["Authorization"] = @token
-
-      begin
-        http.request(request)
-
-        return true
-      rescue Timeout::Error, SocketError, IOError, SystemCallError, OpenSSL::OpenSSLError => err
-        @config.logger.debug("Error probing realtime endpoint: #{err.class}: #{err.message}")
-      rescue => err
-        @config.logger.error("Error probing realtime endpoint: #{err.class}: #{err.message}")
-      end
-
-      false
-    end
-
     def running?
       @running.true?
     end
@@ -141,12 +116,18 @@ module Aikido::Zen
       http.open_timeout = @open_timeout
       http.write_timeout = @write_timeout
       http.read_timeout = @read_timeout
+      # Working around Net::HTTP cleverness; retries GET requests once by default.
       http.max_retries = 0
 
       request = Net::HTTP::Get.new("/api/runtime/stream")
       request["Authorization"] = @token
       request["Accept"] = "text/event-stream"
       request["Cache-Control"] = "no-cache"
+      request["X-Agent-Platform"] = "ruby"
+      request["X-Agent-Version"] = Aikido::Zen::VERSION
+      # Working around Net::HTTP cleverness; auto-negotiates gzip and buffers
+      # decompressed output internally.
+      request["Accept-Encoding"] = "identity"
 
       @config.logger.debug("API stream connecting")
       http.start

data/lib/aikido/zen/attack.rb

@@ -44,11 +44,11 @@ module Aikido::Zen
 
     def as_json
       {
-        kind: kind,
-        blocked: blocked?,
-        metadata: metadata,
-        operation: @operation,
-        stack: @stack
+        "kind" => kind,
+        "blocked" => blocked?,
+        "metadata" => metadata,
+        "operation" => @operation,
+        "stack" => @stack
       }.compact.merge(input.as_json)
     end
 
@@ -70,7 +70,7 @@ module Aikido::Zen
 
       def metadata
         {
-          filename: filepath
+          "filename" => filepath
         }
       end
 
@@ -107,7 +107,7 @@ module Aikido::Zen
 
       def metadata
         {
-          command: @command
+          "command" => @command
         }
       end
 
@@ -139,9 +139,9 @@ module Aikido::Zen
 
       def metadata
         {
-          sql: @query,
-          dialect: @dialect.name,
-          failedToTokenize: @failed_to_tokenize || nil
+          "sql" => @query,
+          "dialect" => @dialect.name,
+          "failedToTokenize" => @failed_to_tokenize || nil
         }.compact
       end
 
@@ -174,8 +174,8 @@ module Aikido::Zen
 
       def metadata
         {
-          hostname: @request.uri.hostname,
-          port: @request.uri.port.to_s
+          "hostname" => @request.uri.hostname,
+          "port" => @request.uri.port.to_s
         }
       end
     end
@@ -212,8 +212,8 @@ module Aikido::Zen
 
       def metadata
         {
-          hostname: @hostname,
-          privateIP: @address
+          "hostname" => @hostname,
+          "privateIP" => @address
         }
       end
     end

data/lib/aikido/zen/attack_wave.rb

@@ -69,9 +69,9 @@ module Aikido::Zen
 
       def as_json
         {
-          ipAddress: @ip_address,
-          userAgent: @user_agent,
-          source: @source
+          "ipAddress" => @ip_address,
+          "userAgent" => @user_agent,
+          "source" => @source
         }.compact
       end
 
@@ -101,10 +101,10 @@ module Aikido::Zen
 
       def as_json
         {
-          metadata: {
-            samples: @samples.as_json.to_json # The API only accepts string values in metadata
+          "metadata" => {
+            "samples" => @samples.as_json.to_json # The API only accepts string values in metadata
           },
-          user: @user.as_json
+          "user" => @user.as_json
         }.compact
       end
 
@@ -130,8 +130,8 @@ module Aikido::Zen
 
       def as_json
         {
-          method: @verb.as_json,
-          url: @path.as_json
+          "method" => @verb.as_json,
+          "url" => @path.as_json
         }.compact
       end