lib/em-proxy/connection.rb in em-proxy-0.1.2 vs lib/em-proxy/connection.rb in em-proxy-0.1.3

@@ -1,25 +1,25 @@
 module EventMachine
   module ProxyServer
     class Connection < EventMachine::Connection
       attr_accessor :debug
-      
+
       ##### Proxy Methods
       def on_data(&blk); @on_data = blk; end
       def on_response(&blk); @on_response = blk; end
       def on_finish(&blk); @on_finish = blk; end
-      def on_connect(&blk); blk.call; end
+      def on_connect(&blk); @on_connect = blk; end
 
       ##### EventMachine
       def initialize(options)
         @debug = options[:debug] || false
         @servers = {}
       end
 
       def receive_data(data)
         debug [:connection, data]
-        processed = @on_data.call(data)
+        processed = @on_data.call(data) if @on_data
 
         return if processed == :async or processed.nil?
         relay_to_servers(processed)
       end
 
@@ -36,11 +36,11 @@
 
         servers.each do |s|
           s.send_data data unless data.nil?
         end
       end
-      
+
       #
       # initialize connections to backend servers
       #
       def server(name, opts)
         srv = EventMachine::connect(opts[:host], opts[:port], EventMachine::ProxyServer::Backend, @debug) do |c|
@@ -55,26 +55,32 @@
 
       #
       # [ip, port] of the connected client
       #
       def peer
-        @peer ||= Socket.unpack_sockaddr_in(get_peername).reverse
+        peername = get_peername
+        @peer ||= peername ? Socket.unpack_sockaddr_in(peername).reverse : nil
       end
 
       #
       # relay data from backend server to client
       #
       def relay_from_backend(name, data)
         debug [:relay_from_backend, name, data]
 
-        data = @on_response.call(name, data)
+        data = @on_response.call(name, data) if @on_response
         send_data data unless data.nil?
       end
 
+      def connected(name)
+        debug [:connected]
+        @on_connect.call(name) if @on_connect
+      end
+
       def unbind
         debug [:unbind, :connection]
-      
+
         # terminate any unfinished connections
         @servers.values.compact.each do |s|
           s.close_connection
         end
       end
@@ -83,20 +89,20 @@
         debug [:unbind_backend, name]
         @servers[name] = nil
 
         # if all connections are terminated downstream, then notify client
         close_connection_after_writing if @servers.values.compact.size.zero?
-        
+
         if @on_finish
           @on_finish.call(name)
 
           # not sure if this is required
           # @on_finish.call(:done) if @servers.values.compact.size.zero?
         end
       end
 
       private
-      
+
       def debug(*data)
         if @debug
           require 'pp'
           pp data
           puts