proxymachine 1.1.0 → 1.2.0

data/History.txt

@@ -1,3 +1,10 @@
+= 1.2.0 / 2010-02-09
+  * New Features
+    * Connection Errors and Timeouts
+    * Inactivity Timeouts
+  * Enhancements
+    * Better async retry logic
+
 = 1.1.0 / 2009-11-05
   * New Features
     * Add { :remote, :data, :reply } command [github.com/coderrr]

data/README.md

@@ -112,6 +112,48 @@ Valid return values
 `{ :close => true }` - Close the connection.  
 `{ :close => String }` - Close the connection after sending the String.  
 
+Connection Errors and Timeouts
+------------------------------
+
+It's possible to register a custom callback for handling connection
+errors. The callback is passed the remote when a connection is either
+rejected or a connection timeout occurs:
+
+    proxy do |data|
+      if data =~ /your thing/
+        { :remote => 'localhost:1234', :connect_timeout => 1.0 }
+      else
+        { :noop => true }
+      end
+    end
+
+    proxy_connect_error do |remote|
+      puts "error connecting to #{remote}"
+    end
+
+You must provide a `:connect_timeout` value in the `proxy` return value
+to enable connection timeouts. The `:connect_timeout` value is a float
+representing the number of seconds to wait before a connection is
+established. Hard connection rejections always trigger the callback, even
+when no `:connect_timeout` is provided.
+
+Inactivity Timeouts
+-------------------
+
+Inactivity timeouts work like connect timeouts but are triggered after
+the configured amount of time elapses without receiving the first byte
+of data from an already connected server:
+
+    proxy do |data|
+      { :remote => 'localhost:1234', :inactivity_timeout => 10.0 }
+    end
+
+    proxy_inactivity_error do |remote|
+      puts "#{remote} did not send any data for 10 seconds"
+    end
+
+If no `:inactivity_timeout` is provided, the `proxy_inactivity_error`
+callback is never triggered.
 
 Contribute
 ----------

data/VERSION.yml

@@ -1,4 +1,5 @@
 --- 
 :major: 1
-:minor: 1
+:minor: 2
 :patch: 0
+:build: 

data/lib/proxymachine.rb

@@ -70,12 +70,30 @@ class ProxyMachine
     end
   end
 
+  def self.set_connect_error_callback(&block)
+    @@connect_error_callback = block
+  end
+
+  def self.connect_error_callback
+    @@connect_error_callback
+  end
+
+  def self.set_inactivity_error_callback(&block)
+    @@inactivity_error_callback = block
+  end
+
+  def self.inactivity_error_callback
+    @@inactivity_error_callback
+  end
+
   def self.run(name, host, port)
     @@totalcounter = 0
     @@maxcounter = 0
     @@counter = 0
     @@name = name
     @@listen = "#{host}:#{port}"
+    @@connect_error_callback ||= proc { |remote| }
+    @@inactivity_error_callback ||= proc { |remote| }
     self.update_procline
     EM.epoll
 
@@ -107,4 +125,12 @@ module Kernel
   def proxy(&block)
     ProxyMachine.set_router(block)
   end
-end
+
+  def proxy_connect_error(&block)
+    ProxyMachine.set_connect_error_callback(&block)
+  end
+
+  def proxy_inactivity_error(&block)
+    ProxyMachine.set_inactivity_error_callback(&block)
+  end
+end

data/lib/proxymachine/client_connection.rb

@@ -10,7 +10,11 @@ class ProxyMachine
     def post_init
       LOGGER.info "Accepted #{peer}"
       @buffer = []
+      @remote = nil
       @tries = 0
+      @connected = false
+      @connect_timeout = nil
+      @inactivity_timeout = nil
       ProxyMachine.incr
     end
 
@@ -23,73 +27,95 @@ class ProxyMachine
     end
 
     def receive_data(data)
-      if !@server_side
+      if !@connected
         @buffer << data
-        ensure_server_side_connection
+        establish_remote_server if @remote.nil?
       end
     rescue => e
       close_connection
       LOGGER.info "#{e.class} - #{e.message}"
     end
 
-    def ensure_server_side_connection
-      @timer.cancel if @timer
-      unless @server_side
-        commands = ProxyMachine.router.call(@buffer.join)
-        LOGGER.info "#{peer} #{commands.inspect}"
-        close_connection unless commands.instance_of?(Hash)
-        if remote = commands[:remote]
-          m, host, port = *remote.match(/^(.+):(.+)$/)
-          if try_server_connect(host, port.to_i)
-            if data = commands[:data]
-              @buffer = [data]
-            end
-            if reply = commands[:reply]
-              send_data(reply)
-            end
-            send_and_clear_buffer
-          end
-        elsif close = commands[:close]
-          if close == true
-            close_connection
-          else
-            send_data(close)
-            close_connection_after_writing
-          end
-        elsif commands[:noop]
-          # do nothing
-        else
+    # Called when new data is available from the client but no remote
+    # server has been established. If a remote can be established, an
+    # attempt is made to connect and proxy to the remote server.
+    def establish_remote_server
+      fail "establish_remote_server called with remote established" if @remote
+      commands = ProxyMachine.router.call(@buffer.join)
+      LOGGER.info "#{peer} #{commands.inspect}"
+      close_connection unless commands.instance_of?(Hash)
+      if remote = commands[:remote]
+        m, host, port = *remote.match(/^(.+):(.+)$/)
+        @remote = [host, port]
+        if data = commands[:data]
+          @buffer = [data]
+        end
+        if reply = commands[:reply]
+          send_data(reply)
+        end
+        @connect_timeout = commands[:connect_timeout]
+        @inactivity_timeout = commands[:inactivity_timeout]
+        connect_to_server
+      elsif close = commands[:close]
+        if close == true
           close_connection
+        else
+          send_data(close)
+          close_connection_after_writing
         end
+      elsif commands[:noop]
+        # do nothing
+      else
+        close_connection
       end
     end
 
-    def try_server_connect(host, port)
+    # Connect to the remote server
+    def connect_to_server
+      fail "connect_server called without remote established" if @remote.nil?
+      host, port = @remote
+      LOGGER.info "Establishing new connection with #{host}:#{port}"
       @server_side = ServerConnection.request(host, port, self)
-      proxy_incoming_to(@server_side, 10240)
-      LOGGER.info "Successful connection to #{host}:#{port}."
-      true
-    rescue => e
+      @server_side.pending_connect_timeout = @connect_timeout
+      @server_side.comm_inactivity_timeout = @inactivity_timeout
+    end
+
+    # Called by the server side immediately after the server connection was
+    # successfully established. Send any buffer we've accumulated and start
+    # raw proxying.
+    def server_connection_success
+      LOGGER.info "Successful connection to #{@remote.join(':')}"
+      @connected = true
+      @buffer.each { |data| @server_side.send_data(data) }
+      proxy_incoming_to @server_side
+    end
+
+    # Called by the server side when a connection could not be established,
+    # either due to a hard connection failure or to a connection timeout.
+    # Leave the client connection open and retry the server connection up to
+    # 10 times.
+    def server_connection_failed
+      @server_side = nil
       if @tries < 10
         @tries += 1
-        LOGGER.info "Failed on server connect attempt #{@tries}. Trying again..."
-        @timer.cancel if @timer
-        @timer = EventMachine::Timer.new(0.1) do
-          self.ensure_server_side_connection
-        end
+        LOGGER.info "Retrying connection with #{@remote.join(':')} (##{@tries})"
+        EM.add_timer(0.1) { connect_to_server }
       else
-        LOGGER.info "Failed after ten connection attempts."
+        LOGGER.info "Connect #{@remote.join(':')} failed after ten attempts."
+        close_connection
+        ProxyMachine.connect_error_callback.call(@remote.join(':'))
       end
-      false
     end
 
-    def send_and_clear_buffer
-      if !@buffer.empty?
-        @buffer.each do |x|
-          @server_side.send_data(x)
-        end
-        @buffer = []
-      end
+    # Called by the server when an inactivity timeout is detected. The timeout
+    # argument is the configured inactivity timeout in seconds as a float; the
+    # elapsed argument is the amount of time that actually elapsed since
+    # connecting but not receiving any data.
+    def server_inactivity_timeout(timeout, elapsed)
+      LOGGER.info "Disconnecting #{@remote.join(':')} after #{elapsed}s of inactivity (> #{timeout.inspect})"
+      @server_side = nil
+      close_connection
+      ProxyMachine.inactivity_error_callback.call(@remote.join(':'))
     end
 
     def unbind

data/lib/proxymachine/server_connection.rb

@@ -6,14 +6,40 @@ class ProxyMachine
 
     def initialize(conn)
       @client_side = conn
+      @connected = false
+      @data_received = false
+      @timeout = nil
     end
 
-    def post_init
-      proxy_incoming_to(@client_side, 10240)
+    def receive_data(data)
+      fail "receive_data called after raw proxy enabled" if @data_received
+      @data_received = true
+      @client_side.send_data(data)
+      proxy_incoming_to @client_side
+    end
+
+    def connection_completed
+      @connected = Time.now
+      @timeout = comm_inactivity_timeout || 0.0
+      @client_side.server_connection_success
     end
 
     def unbind
-      @client_side.close_connection_after_writing
+      now = Time.now
+      if !@connected
+        @client_side.server_connection_failed
+      elsif !@data_received
+        if @timeout > 0.0 && (elapsed = now - @connected) >= @timeout
+          # EM aborted the connection due to an inactivity timeout
+          @client_side.server_inactivity_timeout(@timeout, elapsed)
+        else
+          # server disconnected soon after connecting without sending data
+          # treat this like a failed server connection
+          @client_side.server_connection_failed
+        end
+      else
+        @client_side.close_connection_after_writing
+      end
     end
   end
 end

data/proxymachine.gemspec

@@ -1,15 +1,15 @@
 # Generated by jeweler
-# DO NOT EDIT THIS FILE
-# Instead, edit Jeweler::Tasks in Rakefile, and run `rake gemspec`
+# DO NOT EDIT THIS FILE DIRECTLY
+# Instead, edit Jeweler::Tasks in rakefile, and run the gemspec command
 # -*- encoding: utf-8 -*-
 
 Gem::Specification.new do |s|
   s.name = %q{proxymachine}
-  s.version = "1.1.0"
+  s.version = "1.2.0"
 
   s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version=
   s.authors = ["Tom Preston-Werner"]
-  s.date = %q{2009-11-05}
+  s.date = %q{2010-02-09}
   s.default_executable = %q{proxymachine}
   s.email = %q{tom@mojombo.com}
   s.executables = ["proxymachine"]
@@ -64,3 +64,4 @@ Gem::Specification.new do |s|
     s.add_dependency(%q<eventmachine>, [">= 0.12.10"])
   end
 end
+

data/test/configs/simple.rb

@@ -15,7 +15,21 @@ proxy do |data|
     { :remote => "localhost:9980" }
   elsif data == 'g'
     { :remote => "localhost:9980", :data => 'g2', :reply => 'g3-' }
+  elsif data == 'connect reject'
+    { :remote => "localhost:9989" }
+  elsif data == 'inactivity'
+    { :remote => "localhost:9980", :data => 'sleep 3', :inactivity_timeout => 1 }
   else
     { :close => true }
   end
-end
+end
+
+ERROR_FILE = File.expand_path('../../proxy_error', __FILE__)
+
+proxy_connect_error do |remote|
+  File.open(ERROR_FILE, 'wb') { |fd| fd.write("connect error: #{remote}") }
+end
+
+proxy_inactivity_error do |remote|
+  File.open(ERROR_FILE, 'wb') { |fd| fd.write("activity error: #{remote}") }
+end

data/test/proxymachine_test.rb

@@ -8,6 +8,14 @@ def assert_proxy(host, port, send, recv)
 end
 
 class ProxymachineTest < Test::Unit::TestCase
+  def setup
+    @proxy_error_file = "#{File.dirname(__FILE__)}/proxy_error"
+  end
+
+  def teardown
+    File.unlink(@proxy_error_file) rescue nil
+  end
+
   should "handle simple routing" do
     assert_proxy('localhost', 9990, 'a', '9980:a')
     assert_proxy('localhost', 9990, 'b', '9981:b')
@@ -40,4 +48,24 @@ class ProxymachineTest < Test::Unit::TestCase
     assert_equal '9980:' + 'e' * 2048 + 'f', sock.read
     sock.close
   end
+
+  should "call proxy_connect_error when a connection is rejected" do
+    sock = TCPSocket.new('localhost', 9990)
+    sock.write('connect reject')
+    sock.flush
+    assert_equal "", sock.read
+    sock.close
+    assert_equal "connect error: localhost:9989", File.read(@proxy_error_file)
+  end
+
+  should "call proxy_inactivity_error when initial read times out" do
+    sock = TCPSocket.new('localhost', 9990)
+    sent = Time.now
+    sock.write('inactivity')
+    sock.flush
+    assert_equal "", sock.read
+    assert_operator Time.now - sent, :>=, 1.0
+    assert_equal "activity error: localhost:9980", File.read(@proxy_error_file)
+    sock.close
+  end
 end

data/test/test_helper.rb

@@ -17,6 +17,7 @@ module EventMachine
       end
 
       def receive_data(data)
+        sleep $1.to_f if data =~ /^sleep (.*)/
         send_data("#{@@port}:#{data}")
         close_connection_after_writing
       end
@@ -54,4 +55,4 @@ end
       EventMachine::Protocols::TestConnection.start('localhost', port)
     end
   end
-end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification 
 name: proxymachine
 version: !ruby/object:Gem::Version 
-  version: 1.1.0
+  version: 1.2.0
 platform: ruby
 authors: 
 - Tom Preston-Werner
@@ -9,7 +9,7 @@ autorequire:
 bindir: bin
 cert_chain: []
 
-date: 2009-11-05 00:00:00 -08:00
+date: 2010-02-09 00:00:00 -08:00
 default_executable: proxymachine
 dependencies: 
 - !ruby/object:Gem::Dependency