Ruby code for pumping data between multiple pairs of IO streams

pump.rb

#!/usr/bin/env ruby

require 'pty'
require 'socket'

# pump([[reader, writer], ...])
#
# Description:
# pump() takes an array of arrays, with the first element of each being the read stream and the
# second being the write stream. pump() itself is synchronous but it uses select() to
# concurrently pump multiple streams. Each pair creates a unidirectional pipe; to get
# bidirectional pipe use two pairs.
# 
# Examples:
#
# echo STDIN to STDOUT:
#
#   pump([[STDIN, STDOUT]])
# 
# pipe a shell (or other command) to a TCP socket:
#
#   $reader, $writer = PTY.spawn(command)
#   $socket = TCPSocket.new(host, port)
#   pump([[$reader, $socket], [$socket, $writer]])

def pump(pairs)
  pipes = pairs.map {|pair| {
    :reader => pair[0],
    :writer => pair[1],
    :buffer => ""
  }}
  
  read_streams = pipes.map {|p| p[:reader] }
  all_streams = pipes.map {|p| [p[:reader], p[:writer]]}.flatten.uniq
  
  loop do
    ready = IO.select(
      read_streams,
      pipes.select {|p| !p[:buffer].empty? }.map {|p| p[:writer] },
      all_streams
    )
    # none ready
    next if ready.nil?
    # errors
    return ready[2] if ready[2].any?

    pipes.each do |p|
      if ready[0].include? p[:reader]
        p[:buffer] += p[:reader].read_nonblock(1024)
      end
      if ready[1].include? p[:writer]
        bytes = p[:writer].write_nonblock(p[:buffer])
        p[:buffer].slice!(0, bytes)
      end
    end
  end
end

def forward_command_to_host(command, host, port)
  begin
    $reader, $writer = PTY.spawn(command)
    $socket = TCPSocket.new(host, port)
    pump([[$reader, $socket], [$socket, $writer]])
  rescue
    puts "ERROR=#{$!}"
  end
end

if __FILE__ == $0
  HOST  =  ARGV[0] || 'localhost'
  PORT  = (ARGV[1] || 8080).to_i
  SHELL =  ARGV[2] ||'/bin/bash'

  puts "Forwarding #{SHELL} to #{HOST}:#{PORT}"
  forward_command_to_host(SHELL, HOST, PORT)
end