Implement byterate-based flood protection

2 files changed, 58 insertions, 3 deletions

diff --git a/lib/rbot/ircbot.rb b/lib/rbot/ircbot.rb
index 7c95e52e..fb33b7f5 100644
--- a/lib/rbot/ircbot.rb
+++ b/lib/rbot/ircbot.rb
@@ -78,6 +78,7 @@ class IrcBot
   # create a new IrcBot with botclass +botclass+
   def initialize(botclass, params = {})
     # BotConfig for the core bot
+    # TODO should we split socket stuff into ircsocket, etc?
     BotConfig.register BotConfigStringValue.new('server.name',
       :default => "localhost", :requires_restart => true,
       :desc => "What server should the bot connect to?",
@@ -116,8 +117,12 @@ class IrcBot
       :on_change => Proc.new {|bot, v| bot.socket.sendq_delay = v })
     BotConfig.register BotConfigIntegerValue.new('server.sendq_burst',
       :default => 4, :validate => Proc.new{|v| v >= 0},
-      :desc => "(flood prevention) max lines to burst to the server before throttling. Most ircd's allow bursts of up 5 lines, with non-burst limits of 512 bytes/2 seconds",
+      :desc => "(flood prevention) max lines to burst to the server before throttling. Most ircd's allow bursts of up 5 lines",
       :on_change => Proc.new {|bot, v| bot.socket.sendq_burst = v })
+    BotConfig.register BotConfigStringValue.new('server.byterate',
+      :default => "400/2", :validate => Proc.new{|v| v.match(/\d+\/\d/)},
+      :desc => "(flood prevention) max bytes/seconds rate to send the server. Most ircd's have limits of 512 bytes/2 seconds",
+      :on_change => Proc.new {|bot, v| bot.socket.byterate = v })
     BotConfig.register BotConfigIntegerValue.new('server.ping_timeout',
       :default => 10, :validate => Proc.new{|v| v >= 0},
       :on_change => Proc.new {|bot, v| bot.start_server_pings},
diff --git a/lib/rbot/ircsocket.rb b/lib/rbot/ircsocket.rb
index 5e0c797b..e9fc884f 100644
--- a/lib/rbot/ircsocket.rb
+++ b/lib/rbot/ircsocket.rb
@@ -13,6 +13,15 @@ module Irc
     # total number of lines received from the irc server
     attr_reader :lines_received
 
+    # total number of bytes sent to the irc server
+    attr_reader :bytes_sent
+
+    # total number of bytes received from the irc server
+    attr_reader :bytes_received
+
+    # accumulator for the throttle
+    attr_reader :throttle_bytes
+
     # delay between lines sent
     attr_reader :sendq_delay
 
@@ -23,7 +32,7 @@ module Irc
     # port::   IRCd port
     # host::   optional local host to bind to (ruby 1.7+ required)
     # create a new IrcSocket
-    def initialize(server, port, host, sendq_delay=2, sendq_burst=4)
+    def initialize(server, port, host, sendq_delay=2, sendq_burst=4, brt="400/2")
       @timer = Timer::Timer.new
       @timer.add(0.2) do
         spool
@@ -41,12 +50,29 @@ module Irc
         @sendq_delay = 2
       end
       @last_send = Time.new - @sendq_delay
+      @last_throttle = Time.new
       @burst = 0
       if sendq_burst
         @sendq_burst = sendq_burst.to_i
       else
         @sendq_burst = 4
       end
+      @bytes_per = 400
+      @seconds_per = 2
+      @throttle_bytes = 0
+      setbyterate(brt)
+    end
+
+    def setbyterate(brt)
+      if brt.match(/(\d+)\/(\d)/)
+        @bytes_per = $1.to_i
+        @seconds_per = $2.to_i
+        debug "Byterate now #{byterate}"
+        return true
+      else
+        debug "Couldn't set byterate #{brt}"
+        return false
+      end
     end
 
     def connected?
@@ -96,6 +122,16 @@ module Irc
       end
     end
 
+    def byterate
+      return "#{@bytes_per}/#{@seconds_per}"
+    end
+
+    def byterate=(newrate)
+      @qmutex.synchronize do
+        setbyterate(newrate)
+      end
+    end
+
     # used to send lines to the remote IRCd
     # message: IRC message to send
     def puts(message)
@@ -142,6 +178,11 @@ module Irc
         return
       end
       now = Time.new
+      if @throttle_bytes > 0
+        @throttle_bytes -= (now - @last_throttle)*@bytes_per/@seconds_per
+        @throttle_bytes = 0 if @throttle_bytes < 0
+        @last_throttle = now
+      end
       if (now >= (@last_send + @sendq_delay))
         # reset burst counter after @sendq_delay has passed
         @burst = 0
@@ -155,7 +196,15 @@ module Irc
         debug "(can send #{@sendq_burst - @burst} lines, there are #{@sendq.length} to send)"
         (@sendq_burst - @burst).times do
           break if @sendq.empty?
-          puts_critical(@sendq.shift)
+          mess = @sendq[0]
+          if @throttle_bytes == 0 or mess.length+@throttle_bytes < @bytes_per
+            puts_critical(@sendq.shift)
+          else
+            debug "(flood protection: breaking at message of length #{mess.length})"
+            debug "(Throttle bytes: #{@throttle_bytes})"
+            debug "(Byterate: #{byterate})"
+            break
+          end
         end
       end
       if @sendq.empty?
@@ -201,6 +250,7 @@ module Irc
       @last_send = Time.new
       @lines_sent += 1
       @burst += 1
+      @throttle_bytes += message.length
     end
 
   end