#!/usr/local/bin/ruby
# $Id: port-forward,v 1.3 2004/12/17 07:27:06 tommy Exp $
# 
# Copyright (C) 2004 TOMITA Masahiro
# tommy@tmtm.org
# 

require "socket"
require "syslog"
require "getopts"
require "timeout"

$version = "0.2"

def usage()
  STDERR.puts <<EOS
port-forward version #{$version}
Usage: port-forward [options] port server:port ...
options:
  --pid-file=filename
  --timeout=#
  --connect-timeout=#
  --syslog=facility
EOS
  exit 1
end

def log_error(msg)
  Syslog.err("%s", msg)
end

def log_warning(msg)
  Syslog.warning("%s", msg)
end

def log_notice(msg)
  Syslog.notice("%s", msg)
end

def log_info(msg)
  Syslog.info("%s", msg)
end

def forward(client, servers, i, timeout, ctimeout)
  begin
    ipaddr = client.peeraddr[3]
  rescue Errno::EINVAL, Errno::ENOTCONN
    exit
  end
  log_info "connect from #{ipaddr}"
  start = i
  begin
    if ctimeout then
      server = nil
      timeout(ctimeout) do
        server = TCPSocket.new(*servers[i].split(/:/))
      end
    else
      server = TCPSocket.new(*servers[i].split(/:/))
    end
    log_info("connect to #{servers[i]}")
  rescue StandardError, Timeout::Error
    log_warning("connection failed: #{servers[i]}: #{$!.to_s}")
    i = (i + 1) % servers.length
    if i == start then
      log_error("cannot connect to any server")
      log_error("disconnect to client")
      client.close
      exit
    end
    retry
  end

  loop do
    res, = IO.select([client, server], nil, nil, timeout)
    if res == nil then
      log_warning("timed out")
      break
    end
    if res[0] == client then
      r, w = client, server
    else
      r, w = server, client
    end
    begin
      data = r.sysread(1024)
    rescue EOFError
      log_info "disconnect from #{r == server ? "server" : "client"}"
      r.close
      break
    rescue
      log_error "#{r == server ? "server" : "client"}: #{$!.to_s}"
      break
    end
    begin
      w.syswrite(data)
    rescue
      log_error "#{w == server ? "server" : "client"}: #{$!.to_s}"
      break
    end
  end
  unless server.closed? then
    log_info "disconnect to server"
    server.close
  end
  unless client.closed? then
    log_info "disconnect to client"
    client.close
  end
end

unless getopts(nil, "pid-file:", "timeout:", "connect-timeout:", "syslog:local0") then
  usage
end

if ARGV.size < 2 then
  usage
end
ARGV[1..-1].each do |a|
  usage unless a.include? ":"
end

if fork then
  exit
end
Process.setsid

if $OPT_pid_file then
  begin
    File.open($OPT_pid_file, File::WRONLY|File::CREAT|File::EXCL) do |f|
      f.puts $$
    end
  rescue Errno::EEXIST
    STDERR.puts "#{$OPT_pid_file}: already exist"
    exit 1
  end
end

facility = eval("Syslog::LOG_#{$OPT_syslog.upcase}")
Syslog.open(File.basename($0), nil, facility)

Signal.trap(:TERM, "EXIT")
Signal.trap(:CHLD, "IGNORE")

log_notice "start"

myport, *servers = ARGV
timeout = $OPT_timeout ? $OPT_timeout.to_i : nil
ctimeout = $OPT_connect_timeout ? $OPT_connect_timeout.to_i : nil
begin
  sock = TCPServer.new(*myport.split(/:/))
  i = 0
  loop do
    begin
      s = sock.accept
    rescue Errno::EAGAIN, Errno::ECONNABORTED, Errno::EPROTO, Errno::ECONNREFUSED, Errno::ECONNRESET
      next
    end
    fork do
      sock.close
      forward s, servers, i, timeout, ctimeout
    end
    s.close
    i = (i + 1) % servers.length
  end
rescue
  log_error $!.to_s
ensure
  log_notice "stop"
  if $OPT_pid_file then
    File.unlink $OPT_pid_file rescue nil
  end
end
