#!/usr/bin/env ruby

require 'rubygems'
require 'daemons'
require 'geoip'
require 'socket'
require 'fcntl'
require "optparse"

ENV["RAILS_ENV"] ||= "production"
require File.dirname(__FILE__)+'/../config/environment'

Rails.configuration.log_level = :info # Disable debug
ActiveRecord::Base.allow_concurrency = true

ENV["PATH"] = "/usr/local/bin/:/opt/local/bin:#{ENV["PATH"]}"

BASE_PATH = File.expand_path(GitoriousConfig['repository_base_path'])

TIMEOUT = 30
MAX_CHILDREN = 30
$children_reaped = 0
$children_active = 0

module Git
  class Daemon
    include Daemonize
    
    SERVICE_REGEXP = /^(git\-upload\-pack)\s(.+)\x00host=([\w\.\-]+)/.freeze
    
    def initialize(options)
      @options = options
    end
    
    def start
      if @options[:daemonize]
        daemonize(@options[:logfile])
        File.open(@options[:pidfile], "w") do |f| 
          f.write(Process.pid)
        end
      end
      @socket = TCPServer.new(@options[:host], @options[:port])
      @socket.setsockopt(Socket::SOL_SOCKET,Socket::SO_REUSEADDR, !!@options[:reuseaddr])
      @socket.fcntl(Fcntl::F_SETFD, Fcntl::FD_CLOEXEC)
      log(Process.pid, "Listening on #{@options[:host]}:#{@options[:port]}...")
      ActiveRecord::Base.verify_active_connections! if @options[:daemonize]
      run
    end
      
    def run
      while session = accept_socket
        connections = $children_active - $children_reaped
        if connections > MAX_CHILDREN
          log(Process.pid, "too many active children #{connections}/#{MAX_CHILDREN}")
          session.close
          next
        end
        
        run_service(session)
      end
    end
    
    def run_service(session)
      $children_active += 1
      ip_family, port, name, ip = session.peeraddr
      
      line = receive_data(session)
      
      if line =~ SERVICE_REGEXP
        start_time = Time.now
        service = $1
        base_path = $2
        host = $3
      
        path = File.expand_path("#{BASE_PATH}/#{base_path}")
        log(Process.pid, "Connection from #{ip} for #{base_path.inspect}")
        
        repository = nil        
        begin
          ActiveRecord::Base.verify_active_connections!
          repository = ::Repository.find_by_path(path)
        rescue => e
          log(Process.pid, "AR error: #{e.class.name} #{e.message}:\n #{e.backtrace.join("\n  ")}")
        end
        
        unless repository
          log(Process.pid, "Cannot find repository: #{path}")
          write_error_message(session, "Cannot find repository: #{base_path}")
          $children_active -= 1
          session.close
          return
        end
        
        real_path = File.expand_path(repository.full_repository_path)
        log(Process.pid, "#{ip} wants #{path.inspect} => #{real_path.inspect}")
        
        if real_path.index(BASE_PATH) != 0 || !File.directory?(real_path)
          log(Process.pid, "Invalid path: #{real_path}")
          write_error_message(session, "Cannot find repository: #{base_path}")
          session.close
          $children_active -= 1
          return
        end
      
        if !File.exist?(File.join(real_path, "git-daemon-export-ok"))
          session.close
          $children_active -= 1
          return
        end
        
        if ip_family == "AF_INET6"
          repository.cloned_from(ip)
        else
          geoip = GeoIP.new(File.join(RAILS_ROOT, "data", "GeoIP.dat"))
          localization = geoip.country(ip)
          repository.cloned_from(ip, localization[3], localization[5], 'git')
        end
      
        Dir.chdir(real_path) do
          cmd = "git-upload-pack --strict --timeout=#{TIMEOUT} ."
          
          child_pid = fork do
            log(Process.pid, "Deferred in #{'%0.5f' % (Time.now - start_time)}s")
            
            $stdout.reopen(session)
            $stdin.reopen(session)
            $stderr.reopen("/dev/null")
            
            exec(cmd)
            # FIXME; we don't ever get here since we exec(), so reaped count may be incorrect 
            $children_reaped += 1
            exit!
          end
        end rescue Errno::EAGAIN
      else
        # $stderr.puts "Invalid request from #{ip}: #{line.inspect}"
        $children_active -= 1
      end
      session.close
    end
  
    def handle_stop(signal)
      @socket.close
      log(Process.pid, "Received #{signal}, exiting..")
      exit 0
    end
  
    def handle_cld
      loop do
        pid = nil
        begin
          pid = Process.wait(-1, Process::WNOHANG)
        rescue Errno::ECHILD
          break
        end
        
        if pid && $?
          $children_reaped += 1
          log(pid, "Disconnected. (status=#{$?.exitstatus})") if pid > 0
          if $children_reaped == $children_active
            $children_reaped = 0
            $children_active = 0 
          end
          
          next
        end
        break
      end
    end
  
    def log(pid, msg)
      $stderr.puts "#{Time.now.strftime("%Y-%m-%d %H:%M:%S")} [#{pid}] #{msg}"
    end
    
    def write_error_message(session, msg)
      message = ["\n----------------------------------------------"]
      message << msg
      message << "----------------------------------------------\n"
      write_into_sideband(session, message.join("\n"), 2)
    end
    
    def write_into_sideband(session, message, channel)
      msg = "%s%s" % [channel.chr, message]
      session.write("%04x%s" % [msg.length+4, msg])
    end
    
    def accept_socket
      if RUBY_VERSION < '1.9'
        @socket.accept
      else
        begin
          @socket.accept_nonblock
        rescue Errno::EAGAIN, Errno::EWOULDBLOCK, Errno::ECONNABORTED, Errno::EPROTO, Errno::EINTR => e
          if IO.select([@socket])
            retry
          else
            raise e
          end
        end
      end
    end
    
    def receive_data(session)
      if RUBY_VERSION < '1.9'
        read_data(session)
      else
        read_data_nonblock(session)
      end
    end
    
    def read_data(session)
      size_string = session.recv(4)
      return "" if !size_string
      size = size_string.to_i(16)
      return "" unless size > 4
      session.recv(size - 4)
    rescue Errno::ECONNRESET
      return ""
    end
    
    def read_data_nonblock(session)
      begin
        size_string = session.recv_nonblock(4)
        return "" if !size_string
        size = size_string.to_i(16)
        return "" unless size > 4
        session.recv_nonblock(size - 4)
      rescue Errno::EAGAIN, Errno::EWOULDBLOCK, Errno::ECONNABORTED, Errno::EPROTO, Errno::EINTR
        if IO.select([@socket])
          retry
        else
          return ""
        end
      end
    end
  
  end
end

options = {
  :port => 9418,
  :host => "0.0.0.0",
  :logfile => File.join(RAILS_ROOT, "log", "git-daemon.log"),
  :pidfile => File.join(RAILS_ROOT, "log", "git-daemon.pid"),
  :daemonize => false,
  :reuseaddr => true,
}

OptionParser.new do |opts|
  opts.banner = "Usage: #{$0} [options]"

  opts.on("-p", "--port=[port]", Integer, "Port to listen on", "Default: #{options[:port]}") do |o|
    options[:port] = o
  end

  opts.on("-a", "--address=[host]", "Host to listen on", "Default: #{options[:host]}") do |o|
    options[:host] = o
  end
  
  opts.on("-l", "--logfile=[file]", "File to log to", "Default: #{options[:logfile]}") do |o|
    options[:logfile] = o
  end
  
  opts.on("-P", "--pidfile=[file]", "PID file to use (if daemonized)", "Default: #{options[:pidfile]}") do |o|
    options[:pidfile] = o
  end
  
  opts.on("-d", "--daemonize", "Daemonize or run in foreground", "Default: #{options[:daemonize]}") do |o|
    options[:daemonize] = o
  end
  
  opts.on("-r", "--reuseaddr", "Re-use addresses", "Default: #{options[:reuseaddr].inspect}") do |o|
    options[:reuseaddr] = o
  end
  
  opts.on_tail("-h", "--help", "Show this help message.") do
    puts opts
    exit
  end
  
  # opts.on("-e", "--environment", "RAILS_ENV to use") do |o|
  #   options[:port] = o
  # end
end.parse!

@git_daemon = Git::Daemon.new(options)

trap("SIGKILL")  { @git_daemon.handle_stop("SIGKILL") }
trap("TERM")     { @git_daemon.handle_stop("TERM")    }
trap("SIGINT")   { @git_daemon.handle_stop("SIGINT")  }
trap("CLD")      { @git_daemon.handle_cld  }

@git_daemon.start