lib/rbot/ircsocket.rb

   1 require 'monitor'
   2
   3 class ::String
   4   # Calculate the penalty which will be assigned to this message
   5   # by the IRCd
   6   def irc_send_penalty
   7     # According to eggrdop, the initial penalty is
   8     penalty = 1 + self.size/100
   9     # on everything but UnderNET where it's
  10     # penalty = 2 + self.size/120
  11
  12     cmd, pars = self.split($;,2)
  13     debug "cmd: #{cmd}, pars: #{pars.inspect}"
  14     case cmd.to_sym
  15     when :KICK
  16       chan, nick, msg = pars.split
  17       chan = chan.split(',')
  18       nick = nick.split(',')
  19       penalty += nick.size
  20       penalty *= chan.size
  21     when :MODE
  22       chan, modes, argument = pars.split
  23       extra = 0
  24       if modes
  25         extra = 1
  26         if argument
  27           extra += modes.split(/\+|-/).size
  28         else
  29           extra += 3 * modes.split(/\+|-/).size
  30         end
  31       end
  32       if argument
  33         extra += 2 * argument.split.size
  34       end
  35       penalty += extra * chan.split.size
  36     when :TOPIC
  37       penalty += 1
  38       penalty += 2 unless pars.split.size < 2
  39     when :PRIVMSG, :NOTICE
  40       dests = pars.split($;,2).first
  41       penalty += dests.split(',').size
  42     when :WHO
  43       # I'm too lazy to implement this one correctly
  44       penalty += 5
  45     when :AWAY, :JOIN, :VERSION, :TIME, :TRACE, :WHOIS, :DNS
  46       penalty += 2
  47     when :INVITE, :NICK
  48       penalty += 3
  49     when :ISON
  50       penalty += 1
  51     else # Unknown messages
  52       penalty += 1
  53     end
  54     if penalty > 99
  55       debug "Wow, more than 99 secs of penalty!"
  56       penalty = 99
  57     end
  58     if penalty < 2
  59       debug "Wow, less than 2 secs of penalty!"
  60       penalty = 2
  61     end
  62     debug "penalty: #{penalty}"
  63     return penalty
  64   end
  65 end
  66
  67 module Irc
  68
  69   require 'socket'
  70   require 'thread'
  71
  72   class QueueRing
  73     # A QueueRing is implemented as an array with elements in the form
  74     # [chan, [message1, message2, ...]
  75     # Note that the channel +chan+ has no actual bearing with the channels
  76     # to which messages will be sent
  77
  78     def initialize
  79       @storage = Array.new
  80       @last_idx = -1
  81     end
  82
  83     def clear
  84       @storage.clear
  85       @last_idx = -1
  86     end
  87
  88     def length
  89       len = 0
  90       @storage.each {|c|
  91         len += c[1].size
  92       }
  93       return len
  94     end
  95     alias :size :length
  96
  97     def empty?
  98       @storage.empty?
  99     end
 100
 101     def push(mess, chan)
 102       cmess = @storage.assoc(chan)
 103       if cmess
 104         idx = @storage.index(cmess)
 105         cmess[1] << mess
 106         @storage[idx] = cmess
 107       else
 108         @storage << [chan, [mess]]
 109       end
 110     end
 111
 112     def next
 113       if empty?
 114         warning "trying to access empty ring"
 115         return nil
 116       end
 117       save_idx = @last_idx
 118       @last_idx = (@last_idx + 1) % @storage.size
 119       mess = @storage[@last_idx][1].first
 120       @last_idx = save_idx
 121       return mess
 122     end
 123
 124     def shift
 125       if empty?
 126         warning "trying to access empty ring"
 127         return nil
 128       end
 129       @last_idx = (@last_idx + 1) % @storage.size
 130       mess = @storage[@last_idx][1].shift
 131       @storage.delete(@storage[@last_idx]) if @storage[@last_idx][1] == []
 132       return mess
 133     end
 134
 135   end
 136
 137   class MessageQueue
 138
 139     def initialize
 140       # a MessageQueue is an array of QueueRings
 141       # rings have decreasing priority, so messages in ring 0
 142       # are more important than messages in ring 1, and so on
 143       @rings = Array.new(3) { |i|
 144         if i > 0
 145           QueueRing.new
 146         else
 147           # ring 0 is special in that if it's not empty, it will
 148           # be popped. IOW, ring 0 can starve the other rings
 149           # ring 0 is strictly FIFO and is therefore implemented
 150           # as an array
 151           Array.new
 152         end
 153       }
 154       # the other rings are satisfied round-robin
 155       @last_ring = 0
 156       self.extend(MonitorMixin)
 157       @non_empty = self.new_cond
 158     end
 159
 160     def clear
 161       self.synchronize do
 162         @rings.each { |r| r.clear }
 163         @last_ring = 0
 164       end
 165     end
 166
 167     def push(mess, chan=nil, cring=0)
 168       ring = cring
 169       self.synchronize do
 170         if ring == 0
 171           warning "message #{mess} at ring 0 has channel #{chan}: channel will be ignored" if !chan.nil?
 172           @rings[0] << mess
 173         else
 174           error "message #{mess} at ring #{ring} must have a channel" if chan.nil?
 175           @rings[ring].push mess, chan
 176         end
 177         @non_empty.signal
 178       end
 179     end
 180
 181     def shift(tmout = nil)
 182       self.synchronize do
 183         @non_empty.wait(tmout) if self.empty?
 184         return unsafe_shift
 185       end
 186     end
 187
 188     protected
 189
 190     def empty?
 191       !@rings.find { |r| !r.empty? }
 192     end
 193
 194     def length
 195       @rings.inject(0) { |s, r| s + r.size }
 196     end
 197     alias :size :length
 198
 199     def unsafe_shift
 200       if !@rings[0].empty?
 201         return @rings[0].shift
 202       end
 203       (@rings.size - 1).times do
 204         @last_ring = (@last_ring % (@rings.size - 1)) + 1
 205         return @rings[@last_ring].shift unless @rings[@last_ring].empty?
 206       end
 207       warning "trying to access an empty message queue"
 208       return nil
 209     end
 210
 211   end
 212
 213   # wrapped TCPSocket for communication with the server.
 214   # emulates a subset of TCPSocket functionality
 215   class Socket
 216
 217     MAX_IRC_SEND_PENALTY = 10
 218
 219     # total number of lines sent to the irc server
 220     attr_reader :lines_sent
 221
 222     # total number of lines received from the irc server
 223     attr_reader :lines_received
 224
 225     # total number of bytes sent to the irc server
 226     attr_reader :bytes_sent
 227
 228     # total number of bytes received from the irc server
 229     attr_reader :bytes_received
 230
 231     # accumulator for the throttle
 232     attr_reader :throttle_bytes
 233
 234     # an optional filter object. we call @filter.in(data) for
 235     # all incoming data and @filter.out(data) for all outgoing data
 236     attr_reader :filter
 237
 238     # normalized uri of the current server
 239     attr_reader :server_uri
 240
 241     # default trivial filter class
 242     class IdentityFilter
 243         def in(x)
 244             x
 245         end
 246
 247         def out(x)
 248             x
 249         end
 250     end
 251
 252     # set filter to identity, not to nil
 253     def filter=(f)
 254         @filter = f || IdentityFilter.new
 255     end
 256
 257     # server_list:: list of servers to connect to
 258     # host::   optional local host to bind to (ruby 1.7+ required)
 259     # create a new Irc::Socket
 260     def initialize(server_list, host, opts={})
 261       @server_list = server_list.dup
 262       @server_uri = nil
 263       @conn_count = 0
 264       @host = host
 265       @sock = nil
 266       @filter = IdentityFilter.new
 267       @spooler = false
 268       @lines_sent = 0
 269       @lines_received = 0
 270       if opts.kind_of?(Hash) and opts.key?(:ssl)
 271         @ssl = opts[:ssl]
 272       else
 273         @ssl = false
 274       end
 275     end
 276
 277     def connected?
 278       !@sock.nil?
 279     end
 280
 281     # open a TCP connection to the server
 282     def connect
 283       if connected?
 284         warning "reconnecting while connected"
 285         return
 286       end
 287       srv_uri = @server_list[@conn_count % @server_list.size].dup
 288       srv_uri = 'irc://' + srv_uri if !(srv_uri =~ /:\/\//)
 289       @conn_count += 1
 290       @server_uri = URI.parse(srv_uri)
 291       @server_uri.port = 6667 if !@server_uri.port
 292       debug "connection attempt \##{@conn_count} (#{@server_uri.host}:#{@server_uri.port})"
 293
 294       if(@host)
 295         begin
 296           sock=TCPSocket.new(@server_uri.host, @server_uri.port, @host)
 297         rescue ArgumentError => e
 298           error "Your version of ruby does not support binding to a "
 299           error "specific local address, please upgrade if you wish "
 300           error "to use HOST = foo"
 301           error "(this option has been disabled in order to continue)"
 302           sock=TCPSocket.new(@server_uri.host, @server_uri.port)
 303         end
 304       else
 305         sock=TCPSocket.new(@server_uri.host, @server_uri.port)
 306       end
 307       if(@ssl)
 308         require 'openssl'
 309         ssl_context = OpenSSL::SSL::SSLContext.new()
 310         ssl_context.verify_mode = OpenSSL::SSL::VERIFY_NONE
 311         sock = OpenSSL::SSL::SSLSocket.new(sock, ssl_context)
 312         sock.sync_close = true
 313         sock.connect
 314       end
 315       @sock = sock
 316       @last_send = Time.new
 317       @flood_send = Time.new
 318       @burst = 0
 319       @sock.extend(MonitorMixin)
 320       @sendq = MessageQueue.new
 321       @qthread = Thread.new { writer_loop }
 322     end
 323
 324     # used to send lines to the remote IRCd by skipping the queue
 325     # message: IRC message to send
 326     # it should only be used for stuff that *must not* be queued,
 327     # i.e. the initial PASS, NICK and USER command
 328     # or the final QUIT message
 329     def emergency_puts(message, penalty = false)
 330       @sock.synchronize do
 331         # debug "In puts - got @sock"
 332         puts_critical(message, penalty)
 333       end
 334     end
 335
 336     def handle_socket_error(string, e)
 337       error "#{string} failed: #{e.pretty_inspect}"
 338       # We assume that an error means that there are connection
 339       # problems and that we should reconnect, so we
 340       shutdown
 341       raise SocketError.new(e.inspect)
 342     end
 343
 344     # get the next line from the server (blocks)
 345     def gets
 346       if @sock.nil?
 347         warning "socket get attempted while closed"
 348         return nil
 349       end
 350       begin
 351         reply = @filter.in(@sock.gets)
 352         @lines_received += 1
 353         reply.strip! if reply
 354         debug "RECV: #{reply.inspect}"
 355         return reply
 356       rescue Exception => e
 357         handle_socket_error(:RECV, e)
 358       end
 359     end
 360
 361     def queue(msg, chan=nil, ring=0)
 362       @sendq.push msg, chan, ring
 363     end
 364
 365     def clearq
 366       @sendq.clear
 367     end
 368
 369     # flush the TCPSocket
 370     def flush
 371       @sock.flush
 372     end
 373
 374     # Wraps Kernel.select on the socket
 375     def select(timeout=nil)
 376       Kernel.select([@sock], nil, nil, timeout)
 377     end
 378
 379     # shutdown the connection to the server
 380     def shutdown(how=2)
 381       return unless connected?
 382       @qthread.kill
 383       @qthread = nil
 384       begin
 385         @sock.close
 386       rescue Exception => e
 387         error "error while shutting down: #{e.pretty_inspect}"
 388       end
 389       @sock = nil
 390       @sendq.clear
 391     end
 392
 393     private
 394
 395     def writer_loop
 396       loop do
 397         begin
 398           now = Time.now
 399           flood_delay = @flood_send - MAX_IRC_SEND_PENALTY - now
 400           delay = [flood_delay, 0].max
 401           if delay > 0
 402             debug "sleep(#{delay}) # (f: #{flood_delay})"
 403             sleep(delay)
 404           end
 405           msg = @sendq.shift
 406           debug "got #{msg.inspect} from queue, sending"
 407           emergency_puts(msg, true)
 408         rescue Exception => e
 409           error "Spooling failed: #{e.pretty_inspect}"
 410           debug e.backtrace.join("\n")
 411           raise e
 412         end
 413       end
 414     end
 415
 416     # same as puts, but expects to be called with a lock held on @sock
 417     def puts_critical(message, penalty=false)
 418       # debug "in puts_critical"
 419       begin
 420         debug "SEND: #{message.inspect}"
 421         if @sock.nil?
 422           error "SEND attempted on closed socket"
 423         else
 424           # we use Socket#syswrite() instead of Socket#puts() because
 425           # the latter is racy and can cause double message output in
 426           # some circumstances
 427           actual = @filter.out(message) + "\n"
 428           now = Time.new
 429           @sock.syswrite actual
 430           @last_send = now
 431           @flood_send = now if @flood_send < now
 432           @flood_send += message.irc_send_penalty if penalty
 433           @lines_sent += 1
 434         end
 435       rescue Exception => e
 436         handle_socket_error(:SEND, e)
 437       end
 438     end
 439
 440   end
 441
 442 end