#!/usr/bin/env ruby -w

require 'socket'

def usage
    puts "Usage: #{$0} [options] cmd [args]"
    puts "options:"
    puts "  --local-port <port>      tcp port to listen on"
    puts "  --exit-on-core           abort if child process segfault's"
    puts "  --hang-timeout <secs>    timeout to detect process hangs"
    exit 0
end

opt_exit_on_core = false
opt_hang_timeout = 10.0
opt_local_port   = 8080

while ARGV.size > 0 and ARGV[0] =~ /^--/
    opt = ARGV.shift
    if opt == '--exit-on-core'
        opt_exit_on_core = true
        next
    end

    if opt == '--local-port'
        opt_local_port = ARGV.shift.to_i
        next
    end

    if opt == '--hang-timeout'
        opt_hang_timeout = ARGV.shift.to_f
        raise "Minimum hang timeout is 1.0 seconds" if opt_hang_timeout < 1.0
        next
    end

    raise "Unknown option #{opt}"
end

usage if ARGV.empty?

EXEC  = ARGV.shift
ARGS  = ARGV

class Time
    def self.step limit, step
        last = now = Time.now
        while now - last <= limit
            now = Time.now
            break if yield
            sleep step
        end
    end
end

puts "running #{EXEC} @ tcp/#{opt_local_port}"
server = TCPServer.new opt_local_port
loop do
    # ---
    # wait for incoming connections and launch proggy
    # ---
    connect = server.accept
    if (child = fork).nil?
        $stdout.reopen connect
        $stdin.reopen connect
        $stderr.reopen connect

        ENV['rinetd.peeraddr'] = connect.peeraddr[3]
        exec EXEC, *ARGS
    end

    # ---
    # wait upto to the hang timeout for process to die
    # ---
    Time.step(opt_hang_timeout, 0.1) {
        Process.waitpid(child, Process::WNOHANG)
        $?.nil? != true
    }

    # ---
    # if we timed out, forcefully kill the process
    # ---
    if $?.nil?
        puts "pid #{child} hung"
        Process.kill 'TERM', child
        Process.wait
    else
        if $?.coredump?
            exit 1 if opt_exit_on_core
            puts "pid #{child} coredump'd"
        end
    end
    connect.close
end