aikido-zen 1.5.0 → 1.6.0.beta.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2a9cc243a2ca18ad42f75596db9fb9558362371ee68646f2f75997b80dfd9d04
-  data.tar.gz: 7f9990113775ca615b9b50035a15779e5a9ceb777ad793c164dc6b69c8898732
+  metadata.gz: b0cf8df166776f46ad02f951fd01710135db8e90cbeede8c54c4ff643a6463d0
+  data.tar.gz: 60b5c0e61ce8627e8739441b92843a32f10d519f68afcd52b91aa172812b0207
 SHA512:
-  metadata.gz: 0f4711a8dacc3597ea2a143cb50e735f3234af4d8fcfa0f3d76bf993de51dc17a19c6967c338119cb0433d2c6a850a363a96a6326259316c9d2fc5d346c3809c
-  data.tar.gz: 1d85b479341bec9f2a4ad9b9b8bf12bfc71dbd98213a6020b8c03fde0ea3a6cea47c9e1a6c0f5895474c8e8de3cac823d655cf59396c16309640a45c976cb68b
+  metadata.gz: fcf428661f4688f886b9d4d86daa6d65fac6a348fbb6f222099af2d6723b70a27166ca8e74314afcba03c38fac30b7b3fccc6c39b0a290dbe7e95114db193a2a
+  data.tar.gz: 268766e8d50d75b1d218918fea7acac88bde561361bf176fb025509148781065bdf460e4446001373cc51478d6656f9f510457d15bf26ee33540110d3cd2ce25

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
@@ -80,16 +78,8 @@ module Aikido::Zen
       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.")
-        end
+        @api_stream.handle("config-updated") { |event| settings_updated(event) }
+        @api_stream.start!
       end
 
       poll_for_setting_updates
@@ -243,6 +233,7 @@ module Aikido::Zen
     def update_settings_from_runtime_config!(data)
       return unless @runtime_config_update_mutex.try_lock
       begin
+        Aikido::Zen.api_cache.runtime_config = data
         Aikido::Zen.runtime_settings.update_from_runtime_config_json(data)
       ensure
         @runtime_config_update_mutex.unlock
@@ -254,6 +245,7 @@ module Aikido::Zen
     def update_settings_from_runtime_firewall_lists!(data)
       return unless @runtime_firewall_lists_update_mutex.try_lock
       begin
+        Aikido::Zen.api_cache.runtime_firewall_lists = data
         Aikido::Zen.runtime_settings.update_from_runtime_firewall_lists_json(data)
       ensure
         @runtime_firewall_lists_update_mutex.unlock

data/lib/aikido/zen/api_cache.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+module Aikido::Zen
+  class APICache
+    attr_accessor :runtime_config
+    attr_accessor :runtime_firewall_lists
+  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

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
 

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

@@ -12,7 +12,7 @@ module Aikido::Zen
       end
 
       def self.from_json(data)
-        type = data[:type]
+        type = data["type"]
         subclass = @@registry[type]
         subclass.from_json(data)
       end
@@ -25,7 +25,7 @@ module Aikido::Zen
 
       def as_json
         {
-          type: @type
+          "type" => @type
         }
       end
 
@@ -72,7 +72,7 @@ module Aikido::Zen
         register "track_user_agent"
 
         def self.from_json(data)
-          new(data[:user_agent_keys])
+          new(data["user_agent_keys"])
         end
 
         def initialize(user_agent_keys)
@@ -82,7 +82,7 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            user_agent_keys: @user_agent_keys
+            "user_agent_keys" => @user_agent_keys
           })
         end
 
@@ -99,7 +99,7 @@ module Aikido::Zen
         register "track_ip_list"
 
         def self.from_json(data)
-          new(data[:ip_list_keys])
+          new(data["ip_list_keys"])
         end
 
         def initialize(ip_list_keys)
@@ -109,7 +109,7 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            ip_list_keys: @ip_list_keys
+            "ip_list_keys" => @ip_list_keys
           })
         end
 
@@ -127,7 +127,7 @@ module Aikido::Zen
 
         def self.from_json(data)
           new(
-            being_blocked: data[:being_blocked]
+            being_blocked: data["being_blocked"]
           )
         end
 
@@ -138,7 +138,7 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            being_blocked: @being_blocked
+            "being_blocked" => @being_blocked
           })
         end
 
@@ -156,9 +156,9 @@ module Aikido::Zen
 
         def self.from_json(data)
           new(
-            data[:sink_name],
-            data[:duration],
-            has_errors: data[:has_errors]
+            data["sink_name"],
+            data["duration"],
+            has_errors: data["has_errors"]
           )
         end
 
@@ -171,9 +171,9 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            sink_name: @sink_name,
-            duration: @duration,
-            has_errors: @has_errors
+            "sink_name" => @sink_name,
+            "duration" => @duration,
+            "has_errors" => @has_errors
           })
         end
 
@@ -191,8 +191,8 @@ module Aikido::Zen
 
         def self.from_json(data)
           new(
-            data[:sink_name],
-            being_blocked: data[:being_blocked]
+            data["sink_name"],
+            being_blocked: data["being_blocked"]
           )
         end
 
@@ -204,8 +204,8 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            sink_name: @sink_name,
-            being_blocked: @being_blocked
+            "sink_name" => @sink_name,
+            "being_blocked" => @being_blocked
           })
         end
 
@@ -222,7 +222,7 @@ module Aikido::Zen
         register "track_user"
 
         def self.from_json(data)
-          new(Aikido::Zen::Actor.from_json(data[:actor]))
+          new(Aikido::Zen::Actor.from_json(data["actor"]))
         end
 
         def initialize(actor)
@@ -232,7 +232,7 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            actor: @actor.as_json
+            "actor" => @actor.as_json
           })
         end
 
@@ -249,7 +249,7 @@ module Aikido::Zen
         register "track_outbound"
 
         def self.from_json(data)
-          new(OutboundConnection.from_json(data[:connection]))
+          new(OutboundConnection.from_json(data["connection"]))
         end
 
         def initialize(connection)
@@ -259,7 +259,7 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            connection: @connection.as_json
+            "connection" => @connection.as_json
           })
         end
 
@@ -277,8 +277,8 @@ module Aikido::Zen
 
         def self.from_json(data)
           new(
-            Route.from_json(data[:route]),
-            Request::Schema.from_json(data[:schema])
+            Route.from_json(data["route"]),
+            Request::Schema.from_json(data["schema"])
           )
         end
 
@@ -290,8 +290,8 @@ module Aikido::Zen
 
         def as_json
           super.update({
-            route: @route.as_json,
-            schema: @schema.as_json
+            "route" => @route.as_json,
+            "schema" => @schema.as_json
           })
         end
 

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

@@ -26,10 +26,10 @@ module Aikido::Zen
     def as_json
       @visits.map do |route, record|
         {
-          method: route.verb,
-          path: route.path,
-          hits: record.hits,
-          apispec: record.schema.as_json
+          "method" => route.verb,
+          "path" => route.path,
+          "hits" => record.hits,
+          "apispec" => record.schema.as_json
         }.compact
       end
     end