#!/usr/bin/env ruby

require 'stringio'
require 'yaml'
require 'digest/sha1'
require 'logger'
require 'fileutils'
require 'socket'
require 'cgi'
require 'rubygems'
require 'cmdparse'
require 'monitor'

def log(msg)
    $stderr.print msg,"\n"
end

def error(msg, exc=nil)
    if exc
        $stderr.print "ERROR: #{msg}: #{exc}\n"
        $stderr.puts exc.backtrace
    else
        $stderr.print "ERROR: #{msg}\n"
    end
end


# Modifies CGI so that we can use it.
class SCGIFixed < ::CGI
    public :env_table
    
    def initialize(params, data, out, *args)
        @env_table = params
        @args = *args
        @input = StringIO.new(data)
        @out = out
        super(*args)
    end
    def args
        @args
    end
    def env_table
        @env_table
    end
    def stdinput
        @input
    end
    def stdoutput
        @out
    end
end


class SCGIProcessor < Monitor

    def initialize(settings)
        @env = settings[:env] || "development"
        @debug = settings[:debug] || false
        @host = settings[:host] || "127.0.0.1"
        @port = settings[:port] || "9999"
        @children = settings[:children] || 1
        @pid_file = settings[:pid_file] || "children.yaml"
        @status_dir = settings[:status_dir] || "/tmp"
        @log_file = settings[:logfile] || "log/scgi.log"
        @maxconns = settings[:maxconns]
        @busy_msg = settings[:busy_msg] || "BUSY"
        @settings = settings
        @started = Time.now
        @conns = 0
        @total_conns = 0
        @errors = 0
        
        if @maxconns
            @maxconns = @maxconns.to_i
        else
            @maxconns = 2**30-1
        end
        
        if settings[:conns_second]
            @throttle_sleep = 1.0/settings[:conns_second].to_i
        end

        super()
    end
    
    def run
        ENV['RAILS_ENV'] = @env

        begin
            require_gem 'rails'
            require "config/environment"
        rescue Object
            error("loading rails environment", $!)
        end

        server = TCPServer.new(@host, @port)
        
        if @debug
            log("Listening for connections on #@host:#@port")
            listen(server)
        else
            childpids = []
            @children.to_i.times do
                # fork each child listening to the same port.  very simple yet effective way to spread the load
                # to multiple processes without using threads and still using high performance libevent
                begin
                    pid = fork do
                        $stderr = open(@log_file,"w")
                        $stderr.sync = false
                        listen(server)
                    end
                    childpids << pid
                    Process.detach(pid)
                rescue Object
                    error("Could not fork child processes.  Your system might not support fork.  Use -D instead.", $!)
                end
            end
                
            # tell the user what the sha1 is so they can check for modification later
            log("#@pid_file will have SHA1 #{Digest::SHA1.hexdigest(YAML.dump(childpids))}")
            log("Record this somewhere so you know if it was modified later by someone else.")
            # all children forked and the pids are now ready to write to the pid file
            open(@pid_file,"w") { |f| f.write(YAML.dump(childpids)) }
        end
    end

    
    def listen(socket)
        thread = Thread.new do
            while true
                handle_client(socket.accept)
                sleep @throttle_sleep if @throttle_sleep
                
                @total_conns += 1
            end
        end
        
        begin
            thread.join
        rescue Interrupt
            log("Shutting down from SIGINT.")
        rescue Object
            error("while listening for connections on #@host:#@port", $!)
        end
    end
    
    
    def handle_client(socket)
        Thread.new do
            begin
                synchronize { @conns += 1}
                
                len = ""
                # we only read 10 bytes of the length.  any request longer than this is invalid
                while len.length <= 10
                    c = socket.read(1)
                    if c == ':'
                        # found the terminal, len now has a length in it so read the payload
                        break
                    else
                        len << c
                    end
                end
                
                # we should now either have a payload length to get
                payload = socket.read(len.to_i)
                if (c = socket.read(1)) != ','
                    error("Malformed request, does not end with ','")
                else
                    read_header(socket, payload, @conns)
                end
            rescue IOError
                error("received IOError #$! when handling client.  Your web server doesn't like me.")
            rescue Object
                @errors += 1
                error("after accepting client #@host:#@port -- #{$!.class}", $!)
            ensure
                synchronize { @conns -= 1}
                socket.close if not socket.closed?
            end
        end

    end

    
    def read_header(socket, payload, conns)
        return if socket.closed?
        request = split_body(payload)
        if request and request["CONTENT_LENGTH"]
            length = request["CONTENT_LENGTH"].to_i
            if length > 0
                body = socket.read(length)
            else
                body = ""
            end

            if @conns > @maxconns
                socket.write("Content-type: text/plain\r\n\r\n")
                socket.write(@busy_msg)
            else
                process_request(request, body, socket)
            end
        end
    end

    
    def process_request(request, body, socket)
        return if socket.closed?
        cgi = SCGIFixed.new(request, body, socket)
        begin
            synchronize do
                # unfortuneatly, the dependencies.rb file is not thread safe and will throw exceptions
                # claiming that Dispatcher is not defined, or that other classes are missing.  We have
                # to sync the dispatch call to get around this.
                Dispatcher.dispatch(cgi, ActionController::CgiRequest::DEFAULT_SESSION_OPTIONS, cgi.stdoutput)
            end
        rescue IOError
            error("received IOError #$! when handling client.  Your web server doesn't like me.")
        rescue Object => rails_error
            error("calling Dispatcher.dispatch", rails_error)
        end
    end
    
    
    def split_body(data)
        result = {}
        el = data.split("\0")
        i = 0
        len = el.length
        while i < len
            result[el[i]] = el[i+1]
            i += 1
        end
    
        return result
    end
            
    def status
        pid = Process.pid
        open("#@status_dir/scgi_rails-status.#{pid}","w") do |f|
            status = { 
               'time' => Time.now,  'pid' => pid, 'settings' => @settings,
               'env' => @env, 'status_dir' => @status_dir, 'started' => @started,
               'max_conns' => @maxconns, 'total_conns' => @total_conns,
               'conns' => @conns, 'errors' => @errors, 'systimes' => Process.times
             }
            f.write(YAML.dump(status))
        end
    end
end


def signal_children(pidfile,  signal)
    if not File.exists? pidfile
        log("No #{pidfile} as specified.  Probably nothing running or wrong path.")
        exit 1
    end
        
    childpids = YAML.load_file(pidfile)
    childpids.each do |pid|
        begin
            log("Signaling pid #{pid}")
            Process.kill(signal, pid)
        rescue Object
            log("Couldn't send #{signal} signal to #{pid} pid.")
        end
    end
end


def make_command(parent, name, desc, options)
    cmd = CmdParse::Command.new(name, false )
    cmd.short_desc = desc
    settings = {}
    cmd.options = CmdParse::OptionParserWrapper.new do |opt|
        options.each do |short, long, info, symbol|
            opt.on(short, long, info) {|val| settings[symbol] = val}
        end
    end
    cmd.set_execution_block do |args|
        yield(settings, args)
    end
    parent.add_command(cmd)
end


cmd = CmdParse::CommandParser.new( true )
cmd.program_name = "scgi_rails"
cmd.program_version = [0, 2, 1]
cmd.options = CmdParse::OptionParserWrapper.new do |opt|
    opt.separator "Global options:"
    opt.on("--verbose", "Be verbose when outputting info") {|t| $verbose = true }
end

cmd.add_command( CmdParse::HelpCommand.new )
cmd.add_command( CmdParse::VersionCommand.new )

make_command(cmd, 'start', "Start Rails Application", 
[['-e','--env STRING','Rails environment', :env],
['-D','--[no-]debug', 'Do not fork children, stay in foreground.', :debug],
['-h','--host STRING', 'IP address to bind as server', :host],
['-p','--port NUMBER', 'Port to bind to', :port],
['-c','--children NUMBER', 'Number of children to start (not win32)', :children],
['-f','--pid-file PATH', 'Where to read the list of running children', :pid_file],
['-l','--log-file PATH', 'Use a different log from from log/scgi.log', :logfile],
['-t','--throttle NUMBER', 'Max conn/second to allow.', :conns_second],
['-m','--max-conns NUMBER', 'Max simultaneous connections before the busy message', :maxconns],
['-b','--busy-msg', 'Busy message given to clients over the max connections ("busy")', :busy_msg],
['-s','--status-dir PATH', 'Where to put the status files', :status_dir]]) do |settings, args|
    scgi = SCGIProcessor.new(settings)
    begin
        trap("HUP") { scgi.status }
    rescue Object
        error("Could not setup a SIGHUP handler.  You won't be able to get status.")
    end
    
    scgi.run
end


make_command(cmd, 'status', "Get status from all running children",
[['-s','--status-dir PATH', 'Where to put the status files', :status_dir],
['-f','--pid-file PATH', 'Where to read the list of running children', :pid_file]]) do |settings, args|
    signal_children(settings[:pid_file] || 'children.yaml', "HUP")
    log("Status files for each child should show up in the configured status directory (/tmp by default).")
end

make_command(cmd, 'stop', "Stop all running children",
[['-s','--sig SIGNAL', 'Where to put the status files', :signal],
['-n','--[no-]delete', 'Keep the children.yaml file rather than delete', :nodelete],
['-f','--pid-file PATH', 'Where to read the list of running children', :pid_file]]) do |settings, args|
    pid_file = settings[:pid_file] || "children.yaml"
    signal_children(pid_file, settings[:signal] || "INT")
    if not settings[:nodelete] and File.exist?(pid_file)
        File.unlink(pid_file)
    end
end

cmd.parse