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

@@ -5,10 +5,11 @@
       
       ##### 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
 
       ##### EventMachine
       def initialize(options)
         @debug = options[:debug] || false
         @servers = {}
@@ -16,10 +17,15 @@
 
       def receive_data(data)
         debug [:connection, data]
         processed = @on_data.call(data)
 
+        return if processed == :async or processed.nil?
+        relay_to_servers(processed)
+      end
+
+      def relay_to_servers(processed)
         if processed.is_a? Array
           data, servers = *processed
 
           # guard for "unbound" servers
           servers = servers.collect {|s| @servers[s]}.compact
@@ -30,24 +36,33 @@
 
         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|
           c.name = name
           c.plexer = self
+          c.proxy_incoming_to(self, 10240) if opts[:relay_server]
         end
+        self.proxy_incoming_to(srv, 10240) if opts[:relay_client]
 
         @servers[name] = srv
       end
 
       #
+      # [ip, port] of the connected client
+      #
+      def peer
+        @peer ||= Socket.unpack_sockaddr_in(get_peername).reverse
+      end
+
+      #
       # relay data from backend server to client
       #
       def relay_from_backend(name, data)
         debug [:relay_from_backend, name, data]
 
@@ -67,24 +82,27 @@
       def unbind_backend(name)
         debug [:unbind_backend, name]
         @servers[name] = nil
 
         # if all connections are terminated downstream, then notify client
-        close_connection if @servers.values.compact.size.zero?
+        close_connection_after_writing if @servers.values.compact.size.zero?
         
         if @on_finish
           @on_finish.call(name)
-          @on_finish.call(:done) if @servers.values.compact.size.zero?
+
+          # not sure if this is required
+          # @on_finish.call(:done) if @servers.values.compact.size.zero?
         end
       end
 
       private
       
       def debug(*data)
-        return unless @debug
-        require 'pp'
-        pp data
-        puts
+        if @debug
+          require 'pp'
+          pp data
+          puts
+        end
       end
     end
   end
 end