#!/usr/local/bin/ruby -w
#
# vim:sw=2 ts=8:et sta:fdm=marker
#
#
# Copyright (c) 2005 Ariff Abdullah 
#        (skywizard@MyBSD.org.my) All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
#        $MyBSD$
#
# Date: Wed May 11 21:36:43 MYT 2005
#   OS: FreeBSD kasumi.MyBSD.org.my 5.4-STABLE i386
#

require 'socket'
require 'resolv'
require 'timeout'
require 'thread'
require 'fcntl'

Socket.do_not_reverse_lookup = true
Thread.abort_on_exception = true
STDOUT.sync = true
STDERR.sync = true
Giant = Mutex.new()

ipv6_support = (defined?(Socket::AF_INET6) && (Socket::AF_INET != Socket::AF_INET6))
if ipv6_support
  begin
    s = Socket.new(Socket::AF_INET6, Socket::SOCK_STREAM, Socket::IPPROTO_TCP)
  rescue Exception
    ipv6_support = false
  ensure
    s.close() if s && !s.closed?
  end
end
IPV6_SUPPORT = ipv6_support

class Resolv
  class DNS
    alias old_each_address each_address
    def each_address(name)
      if IPV6_SUPPORT
        each_resource(name, Resource::IN::AAAA) {|resource| yield resource.address}
      end
      each_resource(name, Resource::IN::A) {|resource| yield resource.address}
    end
  end
end

is_win32 = false

if /mswin32|cygwin|mingw|bccwin/ =~ RUBY_PLATFORM
  is_win32 = true
  require 'Win32API'
  getv = Win32API.new('kernel32.dll', 'GetVersionExA', 'P', 'L')
  info = [ 148, 0, 0, 0, 0 ].pack('V5') + "\0" * 128
  getv.call(info)
  if info.unpack('V5')[4] == 2
    module Win32
      module Resolv
        class << self
          private
          alias old_get_info get_info
          def get_info
            search = nil
            nameserver = []
            Registry::HKEY_LOCAL_MACHINE.open(TCPIP_NT) do |reg|
              begin
                slist = reg.read_s('SearchList')
                search = slist.split(/\s*,\s*|\s+/) unless slist.empty?
              rescue Registry::Error
              end

              if add_search = search.nil?
                search = []
                begin
                  nvdom = reg.read_s('NV Domain')
                  unless nvdom.empty?
                    @search = [ nvdom ]
                    if reg.read_i('UseDomainNameDevolution') != 0
                      if /^[\w\d]+\./ =~ nvdom
                        devo = $'
                      end
                    end
                  end
                rescue Registry::Error
                end
              end

              begin
                oreg = reg
                reg.open('Interfaces') do |reg|
                  reg.each_key do |iface,|
                    reg.open(iface) do |regif|
                      begin
                        [ 'NameServer', 'DhcpNameServer' ].each do |key|
                          ns = regif.read_s(key)
                          unless ns.empty?
                            nameserver.concat(ns.split(/\s*,\s*|\s+/))
                            break
                          end
                        end
                      rescue Registry::Error
                      end

                      if add_search
                        begin
                          [ 'Domain', 'DhcpDomain' ].each do |key|
                            dom = regif.read_s(key)
                            unless dom.empty?
                              search.concat(dom.split(/\s*,\s*|\s+/))
                              break
                            end
                          end
                        rescue Registry::Error
                        end
                      end
                    end
                  end
                end
              rescue Registry::Error
                begin
                  [ 'NameServer', 'DhcpNameServer' ].each do |key|
                    ns = oreg.read_s(key)
                    unless ns.empty?
                      nameserver.concat(ns.split(/\s*,\s*|\s+/))
                      break
                    end
                  end
                rescue Registry::Error
                end

                if add_search
                  begin
                    [ 'Domain', 'DhcpDomain' ].each do |key|
                      dom = oreg.read_s(key)
                      unless dom.empty?
                        search.concat(dom.split(/\s*,\s*|\s+/))
                        break
                      end
                    end
                  rescue Registry::Error
                  end
                end
              end
              search << devo if add_search and devo
            end
            [ search.uniq, nameserver.uniq ]
          end
        end
      end
    end
  else
    module Win32
      module Resolv
        class << self
          alias old_get_dhcpinfo get_dhcpinfo
          def get_dhcpinfo
            begin
              old_get_dhcpinfo()
            rescue Win32::Registry::Error
              [[], []]
            end
          end
        end
      end
    end
  end
  module Win32
    module Resolv
      class << self
        alias old_get_hosts_path get_hosts_path
        def get_hosts_path
          path = get_hosts_dir
          ['hosts', 'hosts.sam'].each do |f|
            hpath = File.join(path.gsub(/\\/, File::SEPARATOR), f)
            return hpath if File.exist?(hpath)
          end
        end
      end
    end
  end
  class Resolv
    class Hosts
      alias old_initialize initialize
      def initialize(*argv)
        argv[0] ||= Win32::Resolv.get_hosts_path()
        old_initialize(*argv)
      end
    end
    class DNS
      alias old_lazy_initialize lazy_initialize
      def lazy_initialize
        @mutex.synchronize {
          unless @initialized
            @config.lazy_initialize
            @requester = Requester::UnconnectedUDP.new
            @initialized = true
          end
        }
      end
    end
  end
end

class IO
  REQUEST_TIMEOUT = 120
  CONNECT_TIMEOUT = 30
  BLOCKSIZE       = 1 << 10
  if defined?(Fcntl::F_SETFL) &&
      defined?(Fcntl::F_GETFL) &&
      defined?(Fcntl::O_NONBLOCK)
    def nonblocking=(flag)
      fcntl(
        Fcntl::F_SETFL,
        (fcntl(Fcntl::F_GETFL, 0) & ~Fcntl::O_NONBLOCK)|(flag ? Fcntl::O_NONBLOCK : 0)
      )
    end
    def nonblocking?
      (fcntl(Fcntl::F_GETFL, 0) & Fcntl::O_NONBLOCK) != 0
    end
  else
    def nonblocking=(flag)
      false
    end
    def nonblocking?
      false
    end
  end
  def blocking=(flag)
    self.nonblocking = !flag
  end
  def blocking?
    !self.nonblocking?
  end
  attr_accessor :is_server, :readable, :writable, :buf
  attr_accessor :slave, :thread, :data
  attr_accessor :proxy, :forever
  attr_accessor :main, :connected, :retryconnect, :time
  attr_accessor :inet_family
end

class RetryConnect < RuntimeError
end

def log(arg, urgent = true)
  if urgent
    Giant.synchronize do
      p arg
    end
  end
end

def construct_header(buf)
  buf1, buf2 = buf.split(/\r?\n\r?\n/, 2)
  if /^(GET|HEAD|POST|CONNECT)\s+((http):\/\/)?([^:\/]+)(:(\d+))?(.+)?\s+HTTP\/(\d+\.\d+)(.*)?$/mi =~ buf1
    method   = $1
    protocol = $3
    host     = $4
    port     = ($6 || 80).to_i
    request  = $7 || '/'
    httpver  = $8
    header   = $9
    rethdr   = {}
    if header
      header.strip!
      header.split(/\r?\n/).each do |hdr|
        k, v = hdr.split(/:\s+/, 2)
        if k && v
          k.downcase!
          k.capitalize!
          rethdr[k] = v
        end
      end
    end
    {
      :method => method,
      :protocol => protocol,
      :host => host,
      :port => port,
      :request => request,
      :httpver => httpver,
      :header => rethdr,
      :leftover => buf2.empty? ? nil : buf2,
      :mode => :http
    }
  else
    nil
  end
end

class CacheResolver < Resolv
  CACHE_AGE = 5 * 60
  def initialize(age = CACHE_AGE)
    super()
    @cache = {}
    #@inprogress = {}
    @age = age
    @mtx = Mutex.new()
    @lastvisit = Time.now()
  end
  def getaddresses(host)
    ckey = host.downcase()
    tnow = Time.now()
    ret = nil
    @mtx.synchronize do
      @lastvisit = tnow
      ret = @cache[ckey]
      if ret
        ret = ret.first.dup()
      end
    end
    unless ret
      ret = super(host)
      @mtx.synchronize do
        @cache[ckey] = [ret.dup(), tnow]
      end
    end
=begin
    inprogth = nil
    cth = Thread.current
    @mtx.synchronize do
      @lastvisit = tnow
      inprogth = @inprogress[ckey]
      ret = @cache[ckey]
      if ret
        ret[1] = tnow
        ret = ret.first.dup()
      else
        @inprogress[ckey] = cth unless inprogth
      end
    end
    if !ret && inprogth
      catch :finish do
        cnt = 0
        while true
          @mtx.synchronize do
            ret = @cache[ckey]
            if ret
              ret[1] = tnow
              ret = ret.first.dup()
              throw :finish
            end
            inprogth = @inprogress[ckey]
            throw :finish unless inprogth && inprogth.alive?
          end
          cnt += 1
          throw :finish if cnt == 1000
          sleep(0.25)
        end
      end
    end
    unless ret
      begin
        ret = super(host)
      ensure
        @mtx.synchronize do
          @inprogress.delete(ckey)
          @cache[ckey] = [ret.dup(), tnow] if ret
        end
      end
    end
=end
    ret
  end
  def updatehost(host, addrs)
    tnow = Time.now()
    host = host.downcase()
    addrs = addrs.dup()
    @mtx.synchronize do
      @lastvisit = tnow
      @cache[host] = [addrs, tnow]
    end
    true
  end
  def flush(age = nil)
    age ||= @age
    tnow = Time.now()
    cnt = 0
    sz = 0
    @mtx.synchronize do
      sz = @cache.size()
      @cache.delete_if do |k, v|
        if tnow - v[1] > age
          #@inprogress.delete(k)
          cnt += 1
          true
        else
          false
        end
      end
    end
    [sz, cnt]
  end
  def stat
    @mtx.synchronize do
      @cache.size()
    end
  end
  def elapsed(age)
    age ||= @age
    tnow = Time.now()
    @mtx.synchronize do
      tnow - @lastvisit > age ? true : false
    end
  end
end

def close_sock(s)
  if s
    s.close() unless s.closed?
    s.thread.kill() if s.thread && s.thread.alive?
    $total_clients -= 1 if s.main
    s.main = false
    os = s.slave
    s.slave = nil
    s = os
    if s
      s.close() unless s.closed?
      s.thread.kill() if s.thread && s.thread.alive?
      $total_clients -= 1 if s.main
      s.main = false
      s.slave = nil
    end
  end
end

def resolve_client(s)
  s.readable = s.writable = false
  host = s.data[:host]
  rto = s.data[:rto]
  th = Thread.new do
    Thread.current[:__name__] = "Resolving: #{host}"
    try = 2
    begin
      addr = Timeout.timeout(rto) do
        $resolver.getaddresses(host)
      end
      raise Resolv::ResolvError, 'resolver error, retrying...' \
          if !addr || addr.empty?
      Thread.current[:addr] = addr
    rescue Timeout::Error, Resolv::ResolvError, IOError, \
        SystemCallError, SocketError  => err
      if err.class == Resolv::ResolvError && try > 0
        try -= 1
        log([host, "Resolver retry left: #{try}"], false)
        sleep(0.25)
        retry
      end
      log([host, 'Resolver Timeout']) if err.class == Timeout::Error
      Thread.current[:error] = err
    end
  end
  th[:connection] = :none
  th
end

def construct_proxy(s)
  if s.retryconnect
    s.close() unless s.closed?
    s = s.slave
    $resolver.updatehost(s.data[:host], s.data[:addr])
  end
  if s.thread
    s.thread.join() if s.thread.alive?
    raise s.thread[:error] if s.thread[:error]
    s.data[:addr] = s.thread[:addr]
    s.thread = nil
  end
  outbind = nil
  addr = s.data[:addr].shift()
  if /:/ =~ addr && IPV6_SUPPORT
    family = Socket::AF_INET6
    outbind = s.data[:ipv6]
    inet_family = :IPv6
  else
    family = Socket::AF_INET
    outbind = s.data[:ipv4]
    inet_family = :IPv4
  end
  begin
    sockaddr_in = Socket.pack_sockaddr_in(s.data[:port], addr)
    bindaddr_in = outbind ? Socket.pack_sockaddr_in(0, outbind) : nil
  rescue Exception => err
    raise SocketError.new(err.message)
  end
  prx = Socket.new(family, Socket::SOCK_STREAM, Socket::IPPROTO_TCP)
  if outbind
    begin
      prx.bind(bindaddr_in)
    rescue Exception => err
      log([s.time, s.data[:host], err.class, err.message, outbind])
    end
  end
  prx.nonblocking = true
  prx.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
  prx.binmode()
  prx.readable = prx.writable = prx.proxy = true
  prx.time = s.time
  prx.is_server = false
  prx.retryconnect = false
  prx.connected = false
  prx.main = false
  prx.thread = nil
  prx.buf = ''
  prx.data = s.data
  prx.slave = s
  prx.forever = false
  prx.inet_family = inet_family
  s.slave = prx
  cto = prx.data[:cto]
  prx.thread = Thread.new do
    Thread.current[:__name__] = "Connecting: #{prx.data[:host]}"
    begin
      Timeout.timeout(cto) do
        prx.connect(sockaddr_in)
      end
    rescue Timeout::Error, IOError, SystemCallError, SocketError => err
      unless prx.data[:addr].empty?
        err = RetryConnect.new()
      end
      Thread.current[:error] = err
      if err.class == Timeout::Error
        log([prx.time, prx.data[:host], 'Socket connect Timeout'])
      elsif err.class != RetryConnect
        log([prx.time, prx.data[:host], err.class, err.message])
      end
      prx.close() unless prx.closed?
      prx.slave.close() \
          if prx.slave && !prx.slave.closed? &&
              err.class != RetryConnect
    end
  end
  prx.thread[:connection] = :progress
  prx
end

def resync_proxy(prx)
  if prx.thread
    prx.thread.join() if prx.thread.alive?
    raise prx.thread[:error] if prx.thread[:error]
    prx.thread = nil
  end
  prx.connected = true
  prx.proxy = false
  s = prx.slave
  prx.time = s.time
  data = s.data
  s.data = prx.data = nil
  sxinfo = Socket.getnameinfo(s.to_io.getpeername,
                Socket::NI_NUMERICHOST|Socket::NI_NUMERICSERV)
  s.readable = s.writable = prx.readable = prx.writable = false
  if data[:mode] == :raw
    s.readable = prx.readable = s.forever = prx.forever = true
    log([s.time, sxinfo, 'Raw Proxy Connection', data[:host], data[:port]])
  else
    header = data[:header]
    header['Connection'] = 'close'
    if data[:method] != 'CONNECT'
      header['Proxy-connection'] = 'close'
    end
    if data[:method] == 'CONNECT'
      s.buf.replace('')
      s.buf << 'HTTP/1.1 200 OK' << "\r\n" <<
        'Connection: close' << "\r\n" <<
        'Content-Length: 0' << "\r\n\r\n"
      s.writable = true
      s.forever = prx.forever = true
      log([s.time, sxinfo, data[:method], data[:host],
            data[:port], prx.inet_family])
    else
      prx.buf.replace('')
      prx.buf << "#{data[:method]} #{data[:request]} HTTP/#{data[:httpver]}\r\n"
      header.each do |k, v|
        prx.buf << "#{k}: #{v}\r\n"
      end
      prx.buf << "\r\n"
      prx.buf << data[:leftover] if data[:leftover]
      prx.writable = true
      log([s.time, sxinfo, data[:method], data[:host], data[:request],
            data[:port], prx.inet_family])
    end
  end
end

def usage(errmsg = nil)
  STDERR.puts errmsg if errmsg
  STDERR.puts "Usage: #{File.basename($0)} [-l host/port] " <<
      "[-r host/port] [-i seconds] [-4|-6 host] [-cto seconds] [-rto seconds] " <<
      "[-b bytes] [-dc seconds] [-d] [-k]"
  STDERR.puts "\t-l   Bind/listen proxy server to this host/port"
  STDERR.puts "\t-r   Raw/tunnel proxy mode to this host/port"
  STDERR.puts "\t-i   Force client disconnection if idle more than <idle seconds>"
  STDERR.puts "\t-6   Bind outgoing IPv6 connection to this address/host"
  STDERR.puts "\t-4   Bind outgoing IPv4 connection to this address/host"
  STDERR.puts "\t-cto Timeout (seconds) for socket on " <<
      "connection attempt (Default: #{IO::CONNECT_TIMEOUT})"
  STDERR.puts "\t-rto Timeout (seconds) for socket request " <<
      "processing (Default: #{IO::REQUEST_TIMEOUT})"
  STDERR.puts "\t-b   Size of IO buffer"
  STDERR.puts "\t-dc  DNS cache lifetime (Default: " <<
      "#{CacheResolver::CACHE_AGE} seconds, - means forever)"
  STDERR.puts "\t-d   Run as daemon"
  STDERR.puts "\t     Pidfile: /var/run/pundek.pid"
  STDERR.puts "\t     Logfile: /var/log/pundek.log"
  STDERR.puts "\t-k   Kill pundek daemon. All pundek will be killed!"
  STDERR.puts "\t     Of course you can have multiple pundek daemons running,"
  STDERR.puts "\t     but in terms of killing it, ALL OF IT WILL BE KILLED!"
  STDERR.puts "\t     BwahahahAHhahHAhha!!!"
  exit(1)
end

rawproxy = ''
srvarg = ''
idlemax = ''
v6bind = ''
v4bind = ''
rtoarg = ''
ctoarg = ''
nextarg = nil
bufsz = ''
dcage = ''
daemon = false
pidfile = '/var/run/pundek.pid'
logfile = '/var/log/pundek.log'

ARGV.each do |arg|
  if arg == '-l'
    nextarg = srvarg
  elsif arg == '-r'
    nextarg = rawproxy
  elsif arg == '-i'
    nextarg = idlemax
  elsif arg == '-h'
    usage()
  elsif arg == '-6'
    if IPV6_SUPPORT
      nextarg = v6bind
    else
      STDERR.puts "WARNING: IPv6 not supported. Ignoring..."
      nextarg = nil
    end
  elsif arg == '-4'
    nextarg = v4bind
  elsif arg == '-cto'
    nextarg = ctoarg
  elsif arg == '-rto'
    nextarg = rtoarg
  elsif arg == '-b'
    nextarg = bufsz
  elsif arg == '-dc'
    nextarg = dcage
  elsif arg == '-d'
    usage('Background / Daemon mode not supported in this platform') \
        if is_win32
    daemon = true
  elsif arg == '-k'
    usage('Killing running process not supported in this platform') \
        if is_win32
    begin
      killed = false
      File.foreach(pidfile) do |l|
        l.chomp!
        if /^\d+$/ =~ l
          l = l.to_i
          begin
            Process.kill('SIGTERM', l)
            killed = true
          rescue Exception => xerr
            STDERR.puts "#{xerr.class}: #{xerr.message}"
          end
        end
      end
      if killed
        STDERR.puts 'Pundek killed. BE GONE, Pundek.'
      end
    rescue Exception => err
      STDERR.puts "#{err.class}: #{err.message}"
      STDERR.puts 'Pundek killing failed. Pundek betul...'
      exit(1)
    end
    exit(0)
  elsif nextarg
    nextarg.replace(arg.strip())
    nextarg = nil
  end
end

srvarg = nil if srvarg && srvarg.empty?
rawproxy = nil if rawproxy && rawproxy.empty?
idlemax = nil if idlemax && idlemax.empty?
v6bind = nil if v6bind && v6bind.empty?
v4bind = nil if v4bind && v4bind.empty?
ctoarg = (ctoarg.empty? || /^\-?\d+$/ !~ ctoarg) ? IO::CONNECT_TIMEOUT : ctoarg.to_i
if ctoarg < 1
  usage("Nonsensical connection timeout '#{ctoarg}'")
end
rtoarg = (rtoarg.empty? || /^\-?\d+$/ !~ rtoarg) ? IO::REQUEST_TIMEOUT : rtoarg.to_i
if rtoarg < 1
  usage("Nonsensical request timeout '#{rtoarg}'")
end
if bufsz.empty?
  bufsz = IO::BLOCKSIZE
else
  if /^\d+$/ =~ bufsz
    bufsz = bufsz.to_i
    if bufsz < 1
      usage("IO buffer size too small: '#{bufsz}'")
    elsif bufsz > (1 << 16)
      usage("IO buffer size too big: '#{bufsz}'")
    end
  else
    usage("Nonsensical io buffer size '#{bufsz}'")
  end
end

if dcage.empty?
  dcage = CacheResolver::CACHE_AGE
elsif /^\d+$/ =~ dcage
  dcage = dcage.to_i
  if dcage < 2
    usage("DNS cache lifetime too short: '#{dcage}'")
  end
elsif dcage == '-'
  dcage = nil
else
  usage("Nonsensical DNS cache lifetime '#{dcage}'")
end


srvhost, srvport = (srvarg || 'localhost/1234').split('/', 2)
srvhost ||= 'localhost'
srvport ||= '1234'
srvhost.strip!
srvport.strip!
srvhost = 'localhost' if srvhost.empty?
srvport = '1234' if srvport.empty? || /^\d+$/ !~ srvport
srvport = srvport.to_i

if rawproxy
  if /^([^\/]+)\/(\d+)$/ =~ rawproxy
    rawproxy = { :host => $1, :port => $2.to_i }
  else
    usage('Invalid raw proxy target')
  end
end

if idlemax
  if /^\d+/ =~ idlemax
    idlemax = idlemax.to_i
    if idlemax < 1
      usage('Idle seconds too short')
    end
  else
    usage('Invalid idle argument')
  end
end

# DAEMON
if daemon
  exit!(0) if fork()
  Process.setsid()
  exit!(0) if fork()
  Dir.chdir('/')
  STDIN.reopen('/dev/null', 'r+')
  %w[SIGTERM SIGKILL SIGQUIT].each do |sig|
    trap(sig) do exit(0) end
  end
  File.umask(077)
  logfd = File.open(logfile, 'a+')
  logfd.sync = true
  STDOUT.reopen(logfd)
  STDERR.reopen(logfd)
  at_exit do
    begin
      logfd.puts 'Pundek Terminated. Pundek betul...'
      logfd.flush()
      logfd.close()
    rescue Exception
    end
    begin
      allpids = []
      File.foreach(pidfile) do |l|
        l.chomp!
        if /^\d+$/ =~ l
          l = l.to_i
          allpids << l unless $$ == l
        end
      end
      if allpids.empty?
        File.unlink(pidfile)
      else
        File.open(pidfile, 'w') do |fd|
          allpids.each do |pid| fd.puts pid end
        end
      end
    rescue Exception
    end
  end
end

if (srvhost == '::' || srvhost == '::0.0.0.0') && !IPV6_SUPPORT
  srvhost = '0.0.0.0'
end

listeners_v6 = []
listeners_v4 = []
afspec = []
allsocks = []
rfd = []
wfd = []

if srvhost == '::'
  afspec << [Socket::AF_INET6, '::']
elsif srvhost == '::0.0.0.0'
  afspec << [Socket::AF_INET6, '::'] << [Socket::AF_INET, '0.0.0.0']
elsif srvhost =~ /:/ && IPV6_SUPPORT
  afspec << [Socket::AF_INET6, srvhost]
else
  afspec << [Socket::AF_UNSPEC, srvhost]
end

afspec.each do |spec, host|
  begin
    Socket.getaddrinfo(host, srvport, spec,
          Socket::SOCK_STREAM, nil, Socket::AI_PASSIVE).each do
              |strfamily, port, xhost, ip, family, bleh, blah|
      if IPV6_SUPPORT && family == Socket::AF_INET6
        listeners_v6 << [Socket::AF_INET6, ip, port, host]
      elsif family == Socket::AF_INET
        listeners_v4 << [Socket::AF_INET, ip, port, host]
      else
        raise RuntimeError, "Unknown family '#{strfamily}' (#{family})"
      end
    end
  rescue Exception => err
    STDERR.puts "#{err.class}: #{err.message} (#{host})"
  end
end

$total_clients = 0
$resolver = CacheResolver.new()

ipbind = { :ipv6 => nil, :ipv4 => nil }
if v6bind
  $resolver.getaddresses(v6bind).each do |addr|
    if /:/ =~ addr
      ipbind[:ipv6] = addr.downcase()
      break
    end
  end
end
if v4bind
  $resolver.getaddresses(v4bind).each do |addr|
    if /:/ !~ addr
      ipbind[:ipv4] = addr.downcase()
      break
    end
  end
end

listeners_v6.concat(listeners_v4).each do |family, ip, port, host|
  srv = nil
  begin
    srv = Socket.new(family, Socket::SOCK_STREAM, Socket::IPPROTO_TCP)
    srv.nonblocking = true
    srv.setsockopt(Socket::SOL_SOCKET, Socket::SO_REUSEADDR, 1)
    srv.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
    srv.binmode()
    srv.bind(Socket.pack_sockaddr_in(port, ip))
    srv.listen(10)
    srv.is_server = true
    srv.retryconnect = false
    srv.slave = nil
    srv.time = Time.now()
    srv.inet_family = (IPV6_SUPPORT && family == Socket::AF_INET6) ? :IPv6 : :IPv4
    allsocks << srv
  rescue Exception => err
    STDERR.puts "#{err.class}: #{err.message} [Host: #{host} / IP: #{ip}]"
    srv.close() if srv && !srv.closed?
  end
end

origsize = allsocks.size()

if allsocks.empty?
  STDERR.puts "Pundek startup failed. No server socket found. Pundek betul..."
  exit(1)
end

lasttrap = Time.now()

STDOUT.puts "[ Pundek Server Started... (Pid: #{$$}) ]"
if rawproxy
  STDOUT.puts "[ Raw Proxy Mode: #{rawproxy[:host]}/#{rawproxy[:port]} ]"
else
  STDOUT.puts "[ HTTP Proxy Mode ]"
end
STDOUT.puts "[ Connect Timeout: #{ctoarg} / Request Timeout: #{rtoarg} ]"
if idlemax
  STDOUT.puts "[ Maximum idle time: #{idlemax} ]"
end
STDOUT.puts "[ IO Buffer Size: #{bufsz} ]"
if dcage
  STDOUT.puts "[ DNS cache lifetime: #{dcage} second(s) ]"
else
  STDOUT.puts "[ DNS cache lifetime: forever ]"
end

if ipbind[:ipv6]
  STDOUT.puts "\tOutgoing IPv6 : #{ipbind[:ipv6]}"
end
if ipbind[:ipv4]
  STDOUT.puts "\tOutgoing IPv4 : #{ipbind[:ipv4]}"
end

allsocks.each do |sx|
  ip, port = Socket.getnameinfo(sx.to_io.getsockname,
                Socket::NI_NUMERICHOST|Socket::NI_NUMERICSERV)
  STDOUT.puts "Listening On #{ip} Port #{port}"
end

Thread.current[:__name__] = 'Main Thread'

trap('INT') do
  trapnow = Time.now()
  total_cl = $total_clients
  total_ch = $resolver.stat()
  puts ''
  Giant.synchronize do
    p [trapnow, "Total Clients: #{total_cl} / DNS cache: #{total_ch}"]
    p [trapnow, "Remaining sockets: #{allsocks.size() - origsize}"]
    p [trapnow, "Running thread: #{Thread.list.size()}"]
    Thread.list.each do |t|
      p [trapnow, t, t.alive?, t[:__name__]]
    end
    p allsocks
  end
  if lasttrap && trapnow - lasttrap < 0.2
    puts "Exit.."
    exit(0)
  end
  lasttrap = trapnow
end

def dns_flush(dcage)
  Thread.new do
    Thread.current[:__name__] = 'DNS Flush'
    dcsleep = dcage / 2.0
    log([Time.now(), "DNS Flush Thread started: age: #{dcage} / sleep: #{dcsleep}"])
    while true do
      if $resolver.elapsed(dcage)
        total, purged = $resolver.flush(dcage)
        if purged > 0
          log([Time.now(), "Flushing DNS cache: #{total} total / #{purged} purged"])
        end
      end
      if $resolver.stat() == 0
        log([Time.now(), 'DNS Flush Thread exit...'])
        break
      end
      sleep(dcsleep)
    end
  end
end

if idlemax
  basepolltime = idlemax / 2.0
  maxidle = idlemax
else
  basepolltime = nil
  maxidle = rtoarg
end
polltime = nil
tnow = Time.now()
gcdone = true
allsize = allsocks.size()

if daemon
  begin
    File.open(pidfile, 'a') do |fd| fd.puts $$ end
  rescue Exception => err
    STDERR.puts "#{err.class}: #{err.message}"
  end
end

GC.start()
dcth = nil

#ocnt = 0
#pcnt = 0

while true
  rfd.clear()
  wfd.clear()
  waitinglist = []
  allsocks.delete_if do |s|
    if tnow - s.time > maxidle &&
          !s.is_server && !s.retryconnect && (idlemax || !s.forever)
      close_sock(s)
      allsize -= 1
      true
    elsif s.retryconnect
      s.close() unless s.closed?
      prx = construct_proxy(s)
      rfd << prx
      wfd << prx
      waitinglist << prx
      true
    elsif s.closed?
      unless s.is_server
        s.thread.join() if s.thread && s.thread.alive?
        err = s.thread ? s.thread[:error] : nil
        s.thread = nil
        if err && err.class == RetryConnect
          s.close() unless s.closed?
          s.retryconnect = true
          prx = construct_proxy(s)
          rfd << prx
          wfd << prx
          waitinglist << prx
        else
          close_sock(s)
          allsize -= 1
        end
      end
      true
    elsif s.is_server
      rfd << s
      false
    elsif s.thread
      if s.thread.alive?
        case s.thread[:connection]
          when :none
            polltime = 0.1
          when :progress
            rfd << s
            wfd << s
        end
        false
      elsif s.thread[:error]
        err = s.thread[:error]
        log([s.time, 1, err.class, err.message])
        close_sock(s)
        allsize -= 1
        true
      elsif s.main
        begin
          prx = construct_proxy(s)
          rfd << prx
          wfd << prx
          waitinglist << prx
          allsize += 1
          false
        rescue IOError, EOFError, SystemCallError, SocketError
          close_sock(s)
          allsize -= 1
          true
        end
      else
        rfd << s if s.readable
        wfd << s if s.writable
        false
      end
    else
      rfd << s if s.readable
      wfd << s if s.writable
      false
    end
  end
  allsocks.concat(waitinglist) unless waitinglist.empty?
  begin
    r, w = IO.select(rfd, wfd, nil, polltime)
  rescue IOError
  end
  if dcth
    dcth.kill() if dcth.alive?
    dcth = nil
  end
  polltime = (allsize == origsize) ? nil : \
      (basepolltime && basepolltime < 5 ? basepolltime : 5)
  tnow = Time.now()
  gcdone = false if r || w
  (r || []).each do |s|
    if s.is_server
      s.time = tnow
      cl = nil
      begin
        cl = s.accept.first()
        cl.nonblocking = true
        cl.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
        cl.binmode()
        cl.readable = true
        cl.writable = false
        cl.slave = nil
        cl.buf = ''
        cl.thread = nil
        cl.proxy = false
        cl.is_server = false
        cl.main = true
        cl.connected = true
        cl.retryconnect = false
        cl.time = tnow
        cl.forever = false
        cl.inet_family = s.inet_family
        $total_clients += 1
        if rawproxy
          cl.readable = false
          cl.data = {
            :host => rawproxy[:host],
            :port => rawproxy[:port],
            :mode => :raw,
            :ipv6 => ipbind[:ipv6],
            :ipv4 => ipbind[:ipv4],
            :cto => ctoarg,
            :rto => rtoarg
          }
          cl.thread = resolve_client(cl)
          polltime = 0.1
        end
        allsocks << cl
        allsize += 1
      rescue IOError, EOFError, SystemCallError, SocketError
        close_sock(cl)
      end
    elsif s.proxy
      s.time = tnow
      begin
        resync_proxy(s)
      rescue RetryConnect
        s.retryconnect = true
      rescue IOError, EOFError, SystemCallError, SocketError
        close_sock(s)
      end
    elsif !s.slave && !rawproxy
      s.time = tnow
      begin
        s.buf << s.sysread(bufsz)
        if s.buf.size() > 7 && /^(GET|HEAD|POST|CONNECT)\s/i !~ s.buf
          log([s.time, 2, 'Invalid request'])
          close_sock(s)
        elsif /\r?\n\r?\n/ =~ s.buf
          rbuf = construct_header(s.buf)
          if rbuf
            rbuf[:ipv6] = ipbind[:ipv6]
            rbuf[:ipv4] = ipbind[:ipv4]
            rbuf[:cto] = ctoarg
            rbuf[:rto] = rtoarg
            s.data = rbuf
            s.thread = resolve_client(s)
            s.buf.replace('')
            polltime = 0.1
          else
            log([s.time, 3, 'Invalid Request'])
            close_sock(s)
          end
        elsif s.buf.size() > 8192
          log([s.time, 4, 'Request Too Long'])
          close_sock(s)
        end
      rescue IOError, EOFError, SystemCallError, SocketError
        close_sock(s)
      end
    elsif s.connected && s.readable && s.slave
      s.time = tnow
      s.slave.readable = false
      begin
        s.slave.buf << s.sysread(bufsz)
        s.readable = s.writable = s.slave.readable = false
        s.slave.writable = true
      rescue IOError, EOFError, SystemCallError, SocketError
        close_sock(s)
      end
    end
  end
  (w || []).each do |s|
    if s.proxy
      s.time = tnow
      begin
        resync_proxy(s)
      rescue RetryConnect
        s.retryconnect = true
      rescue IOError, EOFError, SystemCallError, SocketError
        close_sock(s)
      end
    elsif s.connected && s.writable && s.slave
      s.time = tnow
      s.slave.writable = false
      begin
        s.buf.slice!(0, s.syswrite(s.buf.slice(0, bufsz)))
        if s.buf.empty?
          s.writable = s.slave.writable = false
          s.readable = s.slave.readable = true
        end
      rescue IOError, EOFError, SystemCallError, SocketError
        close_sock(s)
      end
    end
  end
  if !gcdone && $total_clients == 0 && allsize == origsize
    gcdone = true
    polltime = nil
    GC.start()
    log([tnow, "Garbage Collect / Total object(s): #{ObjectSpace.each_object do true end}"])
    dcth = dns_flush(dcage) if dcage && $resolver.stat() > 0
  end
  #raise RuntimeError, 'allsocks.size() != allsize' if allsocks.size() != allsize
end