├── .document ├── .gitignore ├── Gemfile ├── History.txt ├── LICENSE ├── README.md ├── Rakefile ├── bin └── proxymachine ├── examples ├── git.rb ├── long.rb └── transparent.rb ├── lib ├── proxymachine.rb └── proxymachine │ ├── client_connection.rb │ └── server_connection.rb ├── proxymachine.gemspec └── test ├── configs └── simple.rb ├── proxymachine_test.rb └── test_helper.rb /.document: -------------------------------------------------------------------------------- 1 | README.rdoc 2 | lib/**/*.rb 3 | bin/* 4 | features/**/*.feature 5 | LICENSE 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.sw? 2 | .DS_Store 3 | coverage 4 | rdoc 5 | pkg 6 | Gemfile.lock 7 | -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | source "http://rubygems.org" 2 | 3 | gemspec 4 | -------------------------------------------------------------------------------- /History.txt: -------------------------------------------------------------------------------- 1 | = 1.2.4 / 2011-02-01 2 | * Bug fixes 3 | * Fix version number in procline 4 | * Unrequire rubygems. Some people consider this a bug. 5 | 6 | = 1.2.3 / 2010-02-23 7 | * Bug fixes 8 | * Never retry after connection is established and data is sent 9 | 10 | = 1.2.2 / 2010-02-19 11 | * Bug fixes 12 | * Bring back the buffer limit 13 | 14 | = 1.2.1 / 2010-02-11 15 | * Bug fixes 16 | * Don't count client closes as connection errors 17 | 18 | = 1.2.0 / 2010-02-09 19 | * New Features 20 | * Connection Errors and Timeouts 21 | * Inactivity Timeouts 22 | * Enhancements 23 | * Better async retry logic 24 | 25 | = 1.1.0 / 2009-11-05 26 | * New Features 27 | * Add { :remote, :data, :reply } command [github.com/coderrr] 28 | * Minor Changes 29 | * Namespace connection classes under ProxyMachine instead of EM [github.com/cmelbye] 30 | * Require socket [github.com/cmelbye] 31 | * Up EM dep to 0.12.10 32 | * Add SOCKS4 Proxy example [github.com/coderrr] 33 | 34 | = 1.0.0 / 2009-10-19 35 | * No changes. Production ready! 36 | 37 | = 0.2.8 / 2009-10-14 38 | * Minor changes 39 | * Always log proxy connection 40 | * Add version and total connection count to procline 41 | * Add max connection count to procline 42 | * Use Logger for logging 43 | 44 | = 0.2.7 / 2009-10-12 45 | * Minor changes 46 | * Use a 10k buffer to prevent memory growth due to slow clients 47 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2009 Tom Preston-Werner 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining 4 | a copy of this software and associated documentation files (the 5 | "Software"), to deal in the Software without restriction, including 6 | without limitation the rights to use, copy, modify, merge, publish, 7 | distribute, sublicense, and/or sell copies of the Software, and to 8 | permit persons to whom the Software is furnished to do so, subject to 9 | the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be 12 | included in all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ProxyMachine 2 | ============ 3 | 4 | By Tom Preston-Werner (tom@mojombo.com) 5 | 6 | 7 | Description 8 | ----------- 9 | 10 | ProxyMachine is a simple content aware (layer 7) TCP routing proxy built on 11 | EventMachine that lets you configure the routing logic in Ruby. 12 | 13 | If you need to proxy connections to different backend servers depending on the 14 | contents of the transmission, then ProxyMachine will make your life easy! 15 | 16 | The idea here is simple. For each client connection, start receiving data 17 | chunks and placing them into a buffer. Each time a new chunk arrives, send the 18 | buffer to a user specified block. The block's job is to parse the buffer to 19 | determine where the connection should be proxied. If the buffer contains 20 | enough data to make a determination, the block returns the address and port of 21 | the correct backend server. If not, it can choose to do nothing and wait for 22 | more data to arrive, close the connection, or close the connection after 23 | sending custom data. Once the block returns an address, a connection to the 24 | backend is made, the buffer is replayed to the backend, and the client and 25 | backend connections are hooked up to form a transparent proxy. This 26 | bidirectional proxy continues to exist until either the client or backend 27 | close the connection. 28 | 29 | ProxyMachine was developed for GitHub's federated architecture and is 30 | successfully used in production to proxy millions of requests every day. The 31 | performance and memory profile have both proven to be excellent. 32 | 33 | 34 | Installation 35 | ------------ 36 | 37 | $ gem install proxymachine -s http://gemcutter.org 38 | 39 | 40 | Running 41 | ------- 42 | 43 | Usage: 44 | proxymachine -c [-h ] [-p ] 45 | 46 | Options: 47 | -c, --config CONFIG Configuration file 48 | -h, --host HOST Hostname to bind. Default 0.0.0.0 49 | -p, --port PORT Port to listen on. Default 5432 50 | 51 | 52 | Signals 53 | ------- 54 | 55 | QUIT - Graceful shutdown. Stop accepting connections immediately and 56 | wait as long as necessary for all connections to close. 57 | 58 | TERM - Fast shutdown. Stop accepting connections immediately and wait 59 | up to 10 seconds for connections to close before forcing 60 | termination. 61 | 62 | INT - Same as TERM 63 | 64 | 65 | Example routing config file 66 | --------------------------- 67 | 68 | class GitRouter 69 | # Look at the routing table and return the correct address for +name+ 70 | # Returns ":" e.g. "ae8f31c.example.com:9418" 71 | def self.lookup(name) 72 | ... 73 | end 74 | end 75 | 76 | # Perform content-aware routing based on the stream data. Here, the 77 | # header information from the Git protocol is parsed to find the 78 | # username and a lookup routine is run on the name to find the correct 79 | # backend server. If no match can be made yet, do nothing with the 80 | # connection. 81 | proxy do |data| 82 | if data =~ %r{^....git-upload-pack /([\w\.\-]+)/[\w\.\-]+\000host=\w+\000} 83 | name = $1 84 | { :remote => GitRouter.lookup(name) } 85 | else 86 | { :noop => true } 87 | end 88 | end 89 | 90 | 91 | Example SOCKS4 Proxy in 7 Lines 92 | ------------------------------- 93 | 94 | proxy do |data| 95 | return if data.size < 9 96 | v, c, port, o1, o2, o3, o4, user = data.unpack("CCnC4a*") 97 | return { :close => "\0\x5b\0\0\0\0\0\0" } if v != 4 or c != 1 98 | return if ! idx = user.index("\0") 99 | { :remote => "#{[o1,o2,o3,o4]*'.'}:#{port}", 100 | :reply => "\0\x5a\0\0\0\0\0\0", 101 | :data => data[idx+9..-1] } 102 | end 103 | 104 | 105 | Valid return values 106 | ------------------- 107 | 108 | `{ :remote => String }` - String is the host:port of the backend server that will be proxied. 109 | `{ :remote => String, :data => String }` - Same as above, but send the given data instead. 110 | `{ :remote => String, :data => String, :reply => String}` - Same as above, but reply with given data back to the client 111 | `{ :noop => true }` - Do nothing. 112 | `{ :close => true }` - Close the connection. 113 | `{ :close => String }` - Close the connection after sending the String. 114 | 115 | Connection Errors and Timeouts 116 | ------------------------------ 117 | 118 | It's possible to register a custom callback for handling connection 119 | errors. The callback is passed the remote when a connection is either 120 | rejected or a connection timeout occurs: 121 | 122 | proxy do |data| 123 | if data =~ /your thing/ 124 | { :remote => 'localhost:1234', :connect_timeout => 1.0 } 125 | else 126 | { :noop => true } 127 | end 128 | end 129 | 130 | proxy_connect_error do |remote| 131 | puts "error connecting to #{remote}" 132 | end 133 | 134 | You must provide a `:connect_timeout` value in the `proxy` return value 135 | to enable connection timeouts. The `:connect_timeout` value is a float 136 | representing the number of seconds to wait before a connection is 137 | established. Hard connection rejections always trigger the callback, even 138 | when no `:connect_timeout` is provided. 139 | 140 | Inactivity Timeouts 141 | ------------------- 142 | 143 | Inactivity timeouts work like connect timeouts but are triggered after 144 | the configured amount of time elapses without receiving the first byte 145 | of data from an already connected server: 146 | 147 | proxy do |data| 148 | { :remote => 'localhost:1234', :inactivity_timeout => 10.0 } 149 | end 150 | 151 | proxy_inactivity_error do |remote| 152 | puts "#{remote} did not send any data for 10 seconds" 153 | end 154 | 155 | If no `:inactivity_timeout` is provided, the `proxy_inactivity_error` 156 | callback is never triggered. 157 | 158 | Contribute 159 | ---------- 160 | 161 | If you'd like to hack on ProxyMachine, start by forking my repo on GitHub: 162 | 163 | http://github.com/mojombo/proxymachine 164 | 165 | To get all of the dependencies, install the gem first. The best way to get 166 | your changes merged back into core is as follows: 167 | 168 | 1. Clone down your fork 169 | 1. Create a topic branch to contain your change 170 | 1. Hack away 171 | 1. Add tests and make sure everything still passes by running `rake` 172 | 1. If you are adding new functionality, document it in the README.md 173 | 1. Do not change the version number, I will do that on my end 174 | 1. If necessary, rebase your commits into logical chunks, without errors 175 | 1. Push the branch up to GitHub 176 | 1. Send me (mojombo) a pull request for your branch 177 | 178 | 179 | Copyright 180 | --------- 181 | 182 | Copyright (c) 2009 Tom Preston-Werner. See LICENSE for details. 183 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | require 'rubygems' 2 | require 'rake' 3 | require 'date' 4 | 5 | ############################################################################# 6 | # 7 | # Helper functions 8 | # 9 | ############################################################################# 10 | 11 | def name 12 | @name ||= Dir['*.gemspec'].first.split('.').first 13 | end 14 | 15 | def version 16 | line = File.read("lib/#{name}.rb")[/^\s*VERSION\s*=\s*.*/] 17 | line.match(/.*VERSION\s*=\s*['"](.*)['"]/)[1] 18 | end 19 | 20 | def date 21 | Date.today.to_s 22 | end 23 | 24 | def rubyforge_project 25 | name 26 | end 27 | 28 | def gemspec_file 29 | "#{name}.gemspec" 30 | end 31 | 32 | def gem_file 33 | "#{name}-#{version}.gem" 34 | end 35 | 36 | def replace_header(head, header_name) 37 | head.sub!(/(\.#{header_name}\s*= ').*'/) { "#{$1}#{send(header_name)}'"} 38 | end 39 | 40 | ############################################################################# 41 | # 42 | # Standard tasks 43 | # 44 | ############################################################################# 45 | 46 | task :default => :test 47 | 48 | require 'rake/testtask' 49 | Rake::TestTask.new(:test) do |test| 50 | test.libs << 'lib' << 'test' 51 | test.pattern = 'test/**/test_*.rb' 52 | test.verbose = true 53 | end 54 | 55 | desc "Generate RCov test coverage and open in your browser" 56 | task :coverage do 57 | require 'rcov' 58 | sh "rm -fr coverage" 59 | sh "rcov test/test_*.rb" 60 | sh "open coverage/index.html" 61 | end 62 | 63 | require 'rake/rdoctask' 64 | Rake::RDocTask.new do |rdoc| 65 | rdoc.rdoc_dir = 'rdoc' 66 | rdoc.title = "#{name} #{version}" 67 | rdoc.rdoc_files.include('README*') 68 | rdoc.rdoc_files.include('lib/**/*.rb') 69 | end 70 | 71 | desc "Open an irb session preloaded with this library" 72 | task :console do 73 | sh "irb -rubygems -r ./lib/#{name}.rb" 74 | end 75 | 76 | ############################################################################# 77 | # 78 | # Custom tasks (add your own tasks here) 79 | # 80 | ############################################################################# 81 | 82 | 83 | 84 | ############################################################################# 85 | # 86 | # Packaging tasks 87 | # 88 | ############################################################################# 89 | 90 | desc "Create tag v#{version} and build and push #{gem_file} to Rubygems" 91 | task :release => :build do 92 | unless `git branch` =~ /^\* master$/ 93 | puts "You must be on the master branch to release!" 94 | exit! 95 | end 96 | sh "git commit --allow-empty -a -m 'Release #{version}'" 97 | sh "git tag v#{version}" 98 | sh "git push origin master" 99 | sh "git push origin v#{version}" 100 | sh "gem push pkg/#{name}-#{version}.gem" 101 | end 102 | 103 | desc "Build #{gem_file} into the pkg directory" 104 | task :build => :gemspec do 105 | sh "mkdir -p pkg" 106 | sh "gem build #{gemspec_file}" 107 | sh "mv #{gem_file} pkg" 108 | end 109 | 110 | desc "Generate #{gemspec_file}" 111 | task :gemspec => :validate do 112 | # read spec file and split out manifest section 113 | spec = File.read(gemspec_file) 114 | head, manifest, tail = spec.split(" # = MANIFEST =\n") 115 | 116 | # replace name version and date 117 | replace_header(head, :name) 118 | replace_header(head, :version) 119 | replace_header(head, :date) 120 | #comment this out if your rubyforge_project has a different name 121 | replace_header(head, :rubyforge_project) 122 | 123 | # determine file list from git ls-files 124 | files = `git ls-files`. 125 | split("\n"). 126 | sort. 127 | reject { |file| file =~ /^\./ }. 128 | reject { |file| file =~ /^(rdoc|pkg)/ }. 129 | map { |file| " #{file}" }. 130 | join("\n") 131 | 132 | # piece file back together and write 133 | manifest = " s.files = %w[\n#{files}\n ]\n" 134 | spec = [head, manifest, tail].join(" # = MANIFEST =\n") 135 | File.open(gemspec_file, 'w') { |io| io.write(spec) } 136 | puts "Updated #{gemspec_file}" 137 | end 138 | 139 | desc "Validate #{gemspec_file}" 140 | task :validate do 141 | libfiles = Dir['lib/*'] - ["lib/#{name}.rb", "lib/#{name}"] 142 | unless libfiles.empty? 143 | puts "Directory `lib` should only contain a `#{name}.rb` file and `#{name}` dir." 144 | exit! 145 | end 146 | unless Dir['VERSION*'].empty? 147 | puts "A `VERSION` file at root level violates Gem best practices." 148 | exit! 149 | end 150 | end 151 | -------------------------------------------------------------------------------- /bin/proxymachine: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env ruby 2 | 3 | $:.unshift File.join(File.dirname(__FILE__), *%w[.. lib]) 4 | 5 | require 'optparse' 6 | require 'proxymachine' 7 | 8 | begin 9 | options = {:host => "0.0.0.0", :port => 5432} 10 | 11 | opts = OptionParser.new do |opts| 12 | opts.banner = <<-EOF 13 | Usage: 14 | proxymachine -c [-h ] [-p ] 15 | 16 | Options: 17 | EOF 18 | 19 | opts.on("-cCONFIG", "--config CONFIG", "Configuration file") do |x| 20 | options[:config] = x 21 | end 22 | 23 | opts.on("-hHOST", "--host HOST", "Hostname to bind. Default 0.0.0.0") do |x| 24 | options[:host] = x 25 | end 26 | 27 | opts.on("-pPORT", "--port PORT", "Port to listen on. Default 5432") do |x| 28 | options[:port] = x 29 | end 30 | end 31 | 32 | opts.parse! 33 | 34 | load(options[:config]) 35 | name = options[:config].split('/').last.chomp('.rb') 36 | ProxyMachine.run(name, options[:host], options[:port]) 37 | rescue Exception => e 38 | if e.instance_of?(SystemExit) 39 | raise 40 | else 41 | LOGGER.info 'Uncaught exception' 42 | LOGGER.info e.message 43 | LOGGER.info e.backtrace.join("\n") 44 | end 45 | end -------------------------------------------------------------------------------- /examples/git.rb: -------------------------------------------------------------------------------- 1 | # This is a config file for ProxyMachine. It pulls the username out of 2 | # the Git stream and can proxy to different locations based on that value 3 | # Run with `proxymachine -c examples/git.rb` 4 | 5 | class GitRouter 6 | # Look at the routing table and return the correct address for +name+ 7 | # Returns ":" e.g. "ae8f31c.example.com:9418" 8 | def self.lookup(name) 9 | LOGGER.info "Proxying for user #{name}" 10 | "localhost:9418" 11 | end 12 | end 13 | 14 | # Perform content-aware routing based on the stream data. Here, the 15 | # header information from the Git protocol is parsed to find the 16 | # username and a lookup routine is run on the name to find the correct 17 | # backend server. If no match can be made yet, do nothing with the 18 | # connection yet. 19 | proxy do |data| 20 | if data =~ %r{^....git-upload-pack /([\w\.\-]+)/[\w\.\-]+\000host=(.+)\000} 21 | name, host = $1, $2 22 | { :remote => GitRouter.lookup(name) } 23 | else 24 | { :noop => true } 25 | end 26 | end -------------------------------------------------------------------------------- /examples/long.rb: -------------------------------------------------------------------------------- 1 | # To try out the graceful exit via SIGQUIT, start up a proxymachine with this 2 | # configuration, and run the following curl command a few times: 3 | # curl http://localhost:5432/ubuntu-releases/9.10/ubuntu-9.10-beta-alternate-amd64.iso \ 4 | # -H "Host: mirrors.cat.pdx.edu" > /dev/null 5 | # Then send a SIGQUIT to the process and stop the long downloads one by one. 6 | 7 | proxy do |data| 8 | { :remote => "mirrors.cat.pdx.edu:80" } 9 | end -------------------------------------------------------------------------------- /examples/transparent.rb: -------------------------------------------------------------------------------- 1 | proxy do |data| 2 | # p data 3 | { :remote => "google.com:80" } 4 | end -------------------------------------------------------------------------------- /lib/proxymachine.rb: -------------------------------------------------------------------------------- 1 | require 'yaml' 2 | require 'eventmachine' 3 | require 'logger' 4 | require 'socket' 5 | 6 | require 'proxymachine/client_connection' 7 | require 'proxymachine/server_connection' 8 | 9 | LOGGER = Logger.new(STDOUT) 10 | 11 | class ProxyMachine 12 | VERSION = '1.2.4' 13 | 14 | MAX_FAST_SHUTDOWN_SECONDS = 10 15 | 16 | def self.update_procline 17 | $0 = "proxymachine #{VERSION} - #{@@name} #{@@listen} - #{self.stats} cur/max/tot conns" 18 | end 19 | 20 | def self.stats 21 | "#{@@counter}/#{@@maxcounter}/#{@@totalcounter}" 22 | end 23 | 24 | def self.count 25 | @@counter 26 | end 27 | 28 | def self.incr 29 | @@totalcounter += 1 30 | @@counter += 1 31 | @@maxcounter = @@counter if @@counter > @@maxcounter 32 | self.update_procline 33 | @@counter 34 | end 35 | 36 | def self.decr 37 | @@counter -= 1 38 | if $server.nil? 39 | LOGGER.info "Waiting for #{@@counter} connections to finish." 40 | end 41 | self.update_procline 42 | EM.stop if $server.nil? and @@counter == 0 43 | @@counter 44 | end 45 | 46 | def self.set_router(block) 47 | @@router = block 48 | end 49 | 50 | def self.router 51 | @@router 52 | end 53 | 54 | def self.graceful_shutdown(signal) 55 | EM.stop_server($server) if $server 56 | LOGGER.info "Received #{signal} signal. No longer accepting new connections." 57 | LOGGER.info "Waiting for #{ProxyMachine.count} connections to finish." 58 | $server = nil 59 | EM.stop if ProxyMachine.count == 0 60 | end 61 | 62 | def self.fast_shutdown(signal) 63 | EM.stop_server($server) if $server 64 | LOGGER.info "Received #{signal} signal. No longer accepting new connections." 65 | LOGGER.info "Maximum time to wait for connections is #{MAX_FAST_SHUTDOWN_SECONDS} seconds." 66 | LOGGER.info "Waiting for #{ProxyMachine.count} connections to finish." 67 | $server = nil 68 | EM.stop if ProxyMachine.count == 0 69 | Thread.new do 70 | sleep MAX_FAST_SHUTDOWN_SECONDS 71 | exit! 72 | end 73 | end 74 | 75 | def self.set_connect_error_callback(&block) 76 | @@connect_error_callback = block 77 | end 78 | 79 | def self.connect_error_callback 80 | @@connect_error_callback 81 | end 82 | 83 | def self.set_inactivity_error_callback(&block) 84 | @@inactivity_error_callback = block 85 | end 86 | 87 | def self.inactivity_error_callback 88 | @@inactivity_error_callback 89 | end 90 | 91 | def self.run(name, host, port) 92 | @@totalcounter = 0 93 | @@maxcounter = 0 94 | @@counter = 0 95 | @@name = name 96 | @@listen = "#{host}:#{port}" 97 | @@connect_error_callback ||= proc { |remote| } 98 | @@inactivity_error_callback ||= proc { |remote| } 99 | self.update_procline 100 | EM.epoll 101 | 102 | EM.run do 103 | ProxyMachine::ClientConnection.start(host, port) 104 | trap('QUIT') do 105 | self.graceful_shutdown('QUIT') 106 | end 107 | trap('TERM') do 108 | self.fast_shutdown('TERM') 109 | end 110 | trap('INT') do 111 | self.fast_shutdown('INT') 112 | end 113 | end 114 | end 115 | end 116 | 117 | module Kernel 118 | def proxy(&block) 119 | ProxyMachine.set_router(block) 120 | end 121 | 122 | def proxy_connect_error(&block) 123 | ProxyMachine.set_connect_error_callback(&block) 124 | end 125 | 126 | def proxy_inactivity_error(&block) 127 | ProxyMachine.set_inactivity_error_callback(&block) 128 | end 129 | end 130 | -------------------------------------------------------------------------------- /lib/proxymachine/client_connection.rb: -------------------------------------------------------------------------------- 1 | class ProxyMachine 2 | class ClientConnection < EventMachine::Connection 3 | def self.start(host, port) 4 | $server = EM.start_server(host, port, self) 5 | LOGGER.info "Listening on #{host}:#{port}" 6 | LOGGER.info "Send QUIT to quit after waiting for all connections to finish." 7 | LOGGER.info "Send TERM or INT to quit after waiting for up to 10 seconds for connections to finish." 8 | end 9 | 10 | def post_init 11 | LOGGER.info "Accepted #{peer}" 12 | @buffer = [] 13 | @remote = nil 14 | @tries = 0 15 | @connected = false 16 | @connect_timeout = nil 17 | @inactivity_timeout = nil 18 | ProxyMachine.incr 19 | end 20 | 21 | def peer 22 | @peer ||= 23 | begin 24 | port, ip = Socket.unpack_sockaddr_in(get_peername) 25 | "#{ip}:#{port}" 26 | end 27 | end 28 | 29 | def receive_data(data) 30 | if !@connected 31 | @buffer << data 32 | establish_remote_server if @remote.nil? 33 | end 34 | rescue => e 35 | close_connection 36 | LOGGER.error "#{e.class} - #{e.message}" 37 | end 38 | 39 | # Called when new data is available from the client but no remote 40 | # server has been established. If a remote can be established, an 41 | # attempt is made to connect and proxy to the remote server. 42 | def establish_remote_server 43 | fail "establish_remote_server called with remote established" if @remote 44 | commands = ProxyMachine.router.call(@buffer.join) 45 | LOGGER.info "#{peer} #{commands.inspect}" 46 | close_connection unless commands.instance_of?(Hash) 47 | if remote = commands[:remote] 48 | m, host, port = *remote.match(/^(.+):(.+)$/) 49 | @remote = [host, port] 50 | if data = commands[:data] 51 | @buffer = [data] 52 | end 53 | if reply = commands[:reply] 54 | send_data(reply) 55 | end 56 | @connect_timeout = commands[:connect_timeout] 57 | @inactivity_timeout = commands[:inactivity_timeout] 58 | connect_to_server 59 | elsif close = commands[:close] 60 | if close == true 61 | close_connection 62 | else 63 | send_data(close) 64 | close_connection_after_writing 65 | end 66 | elsif commands[:noop] 67 | # do nothing 68 | else 69 | close_connection 70 | end 71 | end 72 | 73 | # Connect to the remote server 74 | def connect_to_server 75 | fail "connect_server called without remote established" if @remote.nil? 76 | host, port = @remote 77 | LOGGER.info "Establishing new connection with #{host}:#{port}" 78 | @server_side = ServerConnection.request(host, port, self) 79 | @server_side.pending_connect_timeout = @connect_timeout 80 | @server_side.comm_inactivity_timeout = @inactivity_timeout 81 | end 82 | 83 | # Called by the server side immediately after the server connection was 84 | # successfully established. Send any buffer we've accumulated and start 85 | # raw proxying. 86 | def server_connection_success 87 | LOGGER.info "Successful connection to #{@remote.join(':')}" 88 | @connected = true 89 | @buffer.each { |data| @server_side.send_data(data) } 90 | @buffer = [] 91 | proxy_incoming_to(@server_side, 10240) 92 | end 93 | 94 | # Called by the server side when a connection could not be established, 95 | # either due to a hard connection failure or to a connection timeout. 96 | # Leave the client connection open and retry the server connection up to 97 | # 10 times. 98 | def server_connection_failed 99 | @server_side = nil 100 | if @connected 101 | LOGGER.error "Connection with #{@remote.join(':')} was terminated prematurely." 102 | close_connection 103 | ProxyMachine.connect_error_callback.call(@remote.join(':')) 104 | elsif @tries < 10 105 | @tries += 1 106 | LOGGER.warn "Retrying connection with #{@remote.join(':')} (##{@tries})" 107 | EM.add_timer(0.1) { connect_to_server } 108 | else 109 | LOGGER.error "Connect #{@remote.join(':')} failed after ten attempts." 110 | close_connection 111 | ProxyMachine.connect_error_callback.call(@remote.join(':')) 112 | end 113 | end 114 | 115 | # Called by the server when an inactivity timeout is detected. The timeout 116 | # argument is the configured inactivity timeout in seconds as a float; the 117 | # elapsed argument is the amount of time that actually elapsed since 118 | # connecting but not receiving any data. 119 | def server_inactivity_timeout(timeout, elapsed) 120 | LOGGER.error "Disconnecting #{@remote.join(':')} after #{elapsed}s of inactivity (> #{timeout.inspect})" 121 | @server_side = nil 122 | close_connection 123 | ProxyMachine.inactivity_error_callback.call(@remote.join(':')) 124 | end 125 | 126 | def unbind 127 | @server_side.close_connection_after_writing if @server_side 128 | ProxyMachine.decr 129 | end 130 | 131 | # Proxy connection has been lost 132 | def proxy_target_unbound 133 | @server_side = nil 134 | end 135 | end 136 | end 137 | -------------------------------------------------------------------------------- /lib/proxymachine/server_connection.rb: -------------------------------------------------------------------------------- 1 | class ProxyMachine 2 | class ServerConnection < EventMachine::Connection 3 | def self.request(host, port, client_side) 4 | EventMachine.connect(host, port, self, client_side) 5 | end 6 | 7 | def initialize(conn) 8 | @client_side = conn 9 | @connected = false 10 | @data_received = false 11 | @timeout = nil 12 | end 13 | 14 | def receive_data(data) 15 | fail "receive_data called after raw proxy enabled" if @data_received 16 | @data_received = true 17 | @client_side.send_data(data) 18 | proxy_incoming_to(@client_side, 10240) 19 | end 20 | 21 | def connection_completed 22 | @connected = Time.now 23 | @timeout = comm_inactivity_timeout || 0.0 24 | @client_side.server_connection_success 25 | end 26 | 27 | def unbind 28 | now = Time.now 29 | if @client_side.error? 30 | # the client side disconnected while we were in progress with 31 | # the server. do nothing. 32 | LOGGER.info "Client closed while server connection in progress. Dropping." 33 | elsif !@connected 34 | # a connection error or timeout occurred 35 | @client_side.server_connection_failed 36 | elsif !@data_received 37 | if @timeout > 0.0 && (elapsed = now - @connected) >= @timeout 38 | # EM aborted the connection due to an inactivity timeout 39 | @client_side.server_inactivity_timeout(@timeout, elapsed) 40 | else 41 | # server disconnected soon after connecting without sending data 42 | # treat this like a failed server connection 43 | @client_side.server_connection_failed 44 | end 45 | else 46 | @client_side.close_connection_after_writing 47 | end 48 | end 49 | end 50 | end 51 | -------------------------------------------------------------------------------- /proxymachine.gemspec: -------------------------------------------------------------------------------- 1 | ## This is the rakegem gemspec template. Make sure you read and understand 2 | ## all of the comments. Some sections require modification, and others can 3 | ## be deleted if you don't need them. Once you understand the contents of 4 | ## this file, feel free to delete any comments that begin with two hash marks. 5 | ## You can find comprehensive Gem::Specification documentation, at 6 | ## http://docs.rubygems.org/read/chapter/20 7 | Gem::Specification.new do |s| 8 | s.specification_version = 2 if s.respond_to? :specification_version= 9 | s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version= 10 | s.rubygems_version = '1.3.5' 11 | 12 | ## Leave these as is they will be modified for you by the rake gemspec task. 13 | ## If your rubyforge_project name is different, then edit it and comment out 14 | ## the sub! line in the Rakefile 15 | s.name = 'proxymachine' 16 | s.version = '1.2.4' 17 | s.date = '2011-02-01' 18 | s.rubyforge_project = 'proxymachine' 19 | 20 | ## Make sure your summary is short. The description may be as long 21 | ## as you like. 22 | s.summary = "ProxyMachine is a simple content aware (layer 7) TCP routing proxy." 23 | s.description = "ProxyMachine is a simple content aware (layer 7) TCP routing proxy written in Ruby with EventMachine." 24 | 25 | ## List the primary authors. If there are a bunch of authors, it's probably 26 | ## better to set the email to an email list or something. If you don't have 27 | ## a custom homepage, consider using your GitHub URL or the like. 28 | s.authors = ["Tom Preston-Werner"] 29 | s.email = 'tom@mojombo.com' 30 | s.homepage = 'http://github.com/mojombo/proxymachine' 31 | 32 | ## This gets added to the $LOAD_PATH so that 'lib/NAME.rb' can be required as 33 | ## require 'NAME.rb' or'/lib/NAME/file.rb' can be as require 'NAME/file.rb' 34 | s.require_paths = %w[lib] 35 | 36 | ## If your gem includes any executables, list them here. 37 | s.executables = ["proxymachine"] 38 | s.default_executable = 'proxymachine' 39 | 40 | ## Specify any RDoc options here. You'll want to add your README and 41 | ## LICENSE files to the extra_rdoc_files list. 42 | s.rdoc_options = ["--charset=UTF-8"] 43 | s.extra_rdoc_files = %w[README.md LICENSE] 44 | 45 | ## List your runtime dependencies here. Runtime dependencies are those 46 | ## that are needed for an end user to actually USE your code. 47 | s.add_runtime_dependency(%q, [">= 0.12.10"]) 48 | 49 | ## List your development dependencies here. Development dependencies are 50 | ## those that are only needed during development 51 | s.add_development_dependency(%q, ["~> 0.8.7"]) 52 | s.add_development_dependency(%q, ["~> 2.11.3"]) 53 | s.add_development_dependency(%q, ["~> 1.5.2"]) 54 | 55 | ## Leave this section as-is. It will be automatically generated from the 56 | ## contents of your Git repository via the gemspec task. DO NOT REMOVE 57 | ## THE MANIFEST COMMENTS, they are used as delimiters by the task. 58 | # = MANIFEST = 59 | s.files = %w[ 60 | Gemfile 61 | History.txt 62 | LICENSE 63 | README.md 64 | Rakefile 65 | bin/proxymachine 66 | examples/git.rb 67 | examples/long.rb 68 | examples/transparent.rb 69 | lib/proxymachine.rb 70 | lib/proxymachine/client_connection.rb 71 | lib/proxymachine/server_connection.rb 72 | proxymachine.gemspec 73 | test/configs/simple.rb 74 | test/proxymachine_test.rb 75 | test/test_helper.rb 76 | ] 77 | # = MANIFEST = 78 | 79 | ## Test files will be grabbed from the file list. Make sure the path glob 80 | ## matches what you actually use. 81 | s.test_files = s.files.select { |path| path =~ /^test\/.*_test\.rb$/ } 82 | end 83 | -------------------------------------------------------------------------------- /test/configs/simple.rb: -------------------------------------------------------------------------------- 1 | LOGGER = Logger.new(File.new('/dev/null', 'w')) 2 | 3 | proxy do |data| 4 | if data == 'a' 5 | { :remote => "localhost:9980" } 6 | elsif data == 'b' 7 | { :remote => "localhost:9981" } 8 | elsif data == 'c' 9 | { :remote => "localhost:9980", :data => 'ccc' } 10 | elsif data == 'd' 11 | { :close => 'ddd' } 12 | elsif data == 'e' * 2048 13 | { :noop => true } 14 | elsif data == 'e' * 2048 + 'f' 15 | { :remote => "localhost:9980" } 16 | elsif data == 'g' 17 | { :remote => "localhost:9980", :data => 'g2', :reply => 'g3-' } 18 | elsif data == 'connect reject' 19 | { :remote => "localhost:9989" } 20 | elsif data == 'inactivity' 21 | { :remote => "localhost:9980", :data => 'sleep 3', :inactivity_timeout => 1 } 22 | else 23 | { :close => true } 24 | end 25 | end 26 | 27 | ERROR_FILE = File.expand_path('../../proxy_error', __FILE__) 28 | 29 | proxy_connect_error do |remote| 30 | File.open(ERROR_FILE, 'wb') { |fd| fd.write("connect error: #{remote}") } 31 | end 32 | 33 | proxy_inactivity_error do |remote| 34 | File.open(ERROR_FILE, 'wb') { |fd| fd.write("activity error: #{remote}") } 35 | end 36 | -------------------------------------------------------------------------------- /test/proxymachine_test.rb: -------------------------------------------------------------------------------- 1 | require 'test_helper' 2 | 3 | def assert_proxy(host, port, send, recv) 4 | sock = TCPSocket.new(host, port) 5 | sock.write(send) 6 | assert_equal recv, sock.read 7 | sock.close 8 | end 9 | 10 | class ProxymachineTest < Test::Unit::TestCase 11 | def setup 12 | @proxy_error_file = "#{File.dirname(__FILE__)}/proxy_error" 13 | end 14 | 15 | def teardown 16 | File.unlink(@proxy_error_file) rescue nil 17 | end 18 | 19 | should "handle simple routing" do 20 | assert_proxy('localhost', 9990, 'a', '9980:a') 21 | assert_proxy('localhost', 9990, 'b', '9981:b') 22 | end 23 | 24 | should "handle connection closing" do 25 | sock = TCPSocket.new('localhost', 9990) 26 | sock.write('xxx') 27 | assert_equal nil, sock.read(1) 28 | sock.close 29 | end 30 | 31 | should "handle rewrite routing" do 32 | assert_proxy('localhost', 9990, 'c', '9980:ccc') 33 | end 34 | 35 | should "handle rewrite closing" do 36 | assert_proxy('localhost', 9990, 'd', 'ddd') 37 | end 38 | 39 | should "handle data plus reply" do 40 | assert_proxy('localhost', 9990, 'g', 'g3-9980:g2') 41 | end 42 | 43 | should "handle noop" do 44 | sock = TCPSocket.new('localhost', 9990) 45 | sock.write('e' * 2048) 46 | sock.flush 47 | sock.write('f') 48 | assert_equal '9980:' + 'e' * 2048 + 'f', sock.read 49 | sock.close 50 | end 51 | 52 | should "call proxy_connect_error when a connection is rejected" do 53 | sock = TCPSocket.new('localhost', 9990) 54 | sock.write('connect reject') 55 | sock.flush 56 | assert_equal "", sock.read 57 | sock.close 58 | assert_equal "connect error: localhost:9989", File.read(@proxy_error_file) 59 | end 60 | 61 | should "call proxy_inactivity_error when initial read times out" do 62 | sock = TCPSocket.new('localhost', 9990) 63 | sent = Time.now 64 | sock.write('inactivity') 65 | sock.flush 66 | assert_equal "", sock.read 67 | assert_operator Time.now - sent, :>=, 1.0 68 | assert_equal "activity error: localhost:9980", File.read(@proxy_error_file) 69 | sock.close 70 | end 71 | 72 | should "not consider client disconnect a server error" do 73 | sock = TCPSocket.new('localhost', 9990) 74 | sock.write('inactivity') 75 | sock.close 76 | sleep 3.1 77 | assert !File.exist?(@proxy_error_file) 78 | end 79 | end 80 | -------------------------------------------------------------------------------- /test/test_helper.rb: -------------------------------------------------------------------------------- 1 | require 'rubygems' 2 | require 'test/unit' 3 | require 'shoulda' 4 | require 'socket' 5 | 6 | $LOAD_PATH.unshift(File.join(File.dirname(__FILE__), '..', 'lib')) 7 | $LOAD_PATH.unshift(File.dirname(__FILE__)) 8 | require 'proxymachine' 9 | 10 | # A simple echo server to use in tests 11 | module EventMachine 12 | module Protocols 13 | class TestConnection < Connection 14 | def self.start(host, port) 15 | @@port = port 16 | EM.start_server(host, port, self) 17 | end 18 | 19 | def receive_data(data) 20 | sleep $1.to_f if data =~ /^sleep (.*)/ 21 | send_data("#{@@port}:#{data}") 22 | close_connection_after_writing 23 | end 24 | end 25 | end 26 | end 27 | 28 | def harikari(ppid) 29 | Thread.new do 30 | loop do 31 | begin 32 | Process.kill(0, ppid) 33 | rescue 34 | exit 35 | end 36 | sleep 1 37 | end 38 | end 39 | end 40 | 41 | ppid = Process.pid 42 | 43 | # Start the simple proxymachine 44 | fork do 45 | harikari(ppid) 46 | load(File.join(File.dirname(__FILE__), *%w[configs simple.rb])) 47 | ProxyMachine.run('simple', 'localhost', 9990) 48 | end 49 | 50 | # Start two test daemons 51 | [9980, 9981].each do |port| 52 | fork do 53 | harikari(ppid) 54 | EM.run do 55 | EventMachine::Protocols::TestConnection.start('localhost', port) 56 | end 57 | end 58 | end 59 | --------------------------------------------------------------------------------