=begin
= distributed Ruby --- dRuby 1.3.8
 	Copyright (c) 1999-2002 Masatoshi SEKI 
        You can redistribute it and/or modify it under the same terms as Ruby.
 DRb --- dRuby module.
 DRbUndumped --- Mixin class. 
 DRbProtocol --- Mixin class.
 DRbObject --- dRuby remote object.
 DRbConn --- 
 DRbIdConv --- id <--> obj converter.
 DRbServer --- dRuby message handler.
 DRbUnknown --- unknown class object.
=end

require 'socket'
require 'thread'
require 'fcntl'

module DRb
  class DRbIdConv
    def to_obj(ref)
      ObjectSpace._id2ref(ref)
    end
    
    def to_id(obj)
      obj.nil? ? nil : obj.__id__
    end
  end

  module DRbUndumped 
    def _dump(dummy)
      raise TypeError, 'can\'t dump'
    end
  end

  class DRbServerNotFound < RuntimeError; end

  class DRbUnknownError < RuntimeError
    def initialize(unknown)
      @unknown = unknown
      super(unknown.name)
    end
    attr_reader :unknown

    def self._load(s)
      Marshal::load(s)
    end
    
    def _dump(lv)
      Marshal::dump(@unknown)
    end
  end

  class DRbUnknown
    def initialize(err, buf)
      case err
      when /uninitialized constant (\S+)/
	@name = $1
      when /undefined class\/module (\S+)/
	@name = $1
      else
	@name = nil
      end
      @buf = buf
    end
    attr_reader :name, :buf

    def self._load(s)
      begin
	Marshal::load(s)
      rescue NameError, ArgumentError
	DRbUnknown.new($!, s)
      end
    end

    def _dump(lv)
      @buf
    end

    def reload
      self.class._load(@buf)
    end

    def exception
      DRbUnknownError.new(self)
    end
  end

  module DRbProtocol

    if defined? Socket::SOL_TCP
      SOL_TCP = Socket::SOL_TCP
    elsif defined? Socket::IPPROTO_TCP
      SOL_TCP = Socket::IPPROTO_TCP
    else
      SOL_TCP = 6 #FIXME
    end

    def parse_uri(uri)
      if uri =~ /^druby:\/\/(.*?):(\d+)(.*)$/
	host = $1
	port = $2.to_i
	option = $3
	[host, port, option]
      else
	raise RuntimeError, 'can\'t parse uri:' + uri
      end
    end

    def dump(obj, ary)
      obj = DRbObject.new(obj) if obj.kind_of? DRbUndumped
      begin
	str = Marshal::dump(obj)
      rescue
	ro = DRbObject.new(obj)
	str = Marshal::dump(ro)
      end
      ary.push([str.size].pack('N') + str) if ary
      return str
    end

    def load(soc)
      sz = soc.read(4)	# sizeof (N)
      raise TypeError, 'invalid header' if sz.nil? || sz.size < 4
      sz = sz.unpack('N')[0]
      raise TypeError, 'too large packet' if load_limit && load_limit < sz
      str = soc.read(sz)
      raise TypeError, 'premature marshal format(can\'t read)' if str.nil? || str.size < sz
      begin
	Marshal::load(str)
      rescue NameError, ArgumentError
	DRbUnknown.new($!, str)
      end
    end

    def send_request(soc, ref, msg_id, *arg, &b)
      ary = []
      dump(ref.__drbref, ary)
      dump(msg_id.id2name, ary)
      dump(arg.length, ary)
      arg.each do |e|
	dump(e, ary)
      end
      dump(b, ary)
      soc.write(ary.join)
    end
    
    def recv_request(soc)
      ref = load(soc)
      ro = ref.nil? ? DRb.front : DRb.to_obj(ref)
      msg = load(soc)
      argc = load(soc)
      if argc_limit && argc_limit < argc
	raise ArgumentError, 'too many arguments' 
      end
      argv = Array.new(argc)
      argc.times do |n|
	argv[n] = load(soc)
      end
      block = load(soc)
      return ro, msg, argv, block
    end

    def send_reply(soc, succ, result)
      ary = []
      dump(succ, ary)
      dump(result, ary)
      soc.write(ary.join)
    end

    def recv_reply(soc)
      succ = load(soc)
      result = load(soc)
      [succ, result]
    end

    def ro_to_obj(ro)
      (ro && ro.__drbref) ? DRb.to_obj(ro.__drbref) : DRb.front
    end

    def setup_sockopt(soc)
      soc.setsockopt(SOL_TCP, Socket::TCP_NODELAY, 1)
      soc.fcntl(Fcntl::F_SETFL, Fcntl::O_NONBLOCK) if defined? Fcntl::O_NONBLOCK
      soc.fcntl(Fcntl::F_SETFL, Fcntl::FD_CLOEXEC) if defined? Fcntl::FD_CLOEXEC
    end
  end

  class DRbObject
    def self._load(s)
      uri, ref = Marshal.load(s)
      if DRb.uri == uri
	return ref ? DRb.to_obj(ref) : DRb.front
      end

      it = self.new(nil)
      it.reinit(uri, ref)
      it
    end

    def _dump(lv)
      Marshal.dump([@uri, @ref])
    end

    def initialize(obj, uri=nil)
      @uri = uri || DRb.uri
      @ref = obj ? DRb.to_id(obj) : nil
    end

    def reinit(uri, ref)
      @uri = uri
      @ref = ref
    end

    def __drburi
      @uri
    end

    def __drbref
      @ref
    end

    undef :to_s
    undef :to_a
    undef :respond_to?

    def method_missing(msg_id, *a, &b)
      if @uri == DRb.uri
	obj = @ref ? DRb.to_obj(@ref) : DRb.front
	DRb.current_server.check_insecure_method(obj, msg_id)
	return obj.__send__(msg_id, *a, &b) 
      end

      succ, result = DRbConn.open(@uri) do |conn|
	conn.send_message(self, msg_id, *a, &b)
      end
      return result if succ
      unless DRbUnknown === result
	prefix = "(#{@uri}) "
	bt = []
	result.backtrace.each do |x|
	  break if /`__send__'$/ =~ x 
	  if /^\(druby:\/\// =~ x
	    bt.push(x)
	  else
	    bt.push(prefix + x)
	  end
	end
	raise result, result.message, bt + caller
      else
	raise result
      end
    end
  end

  class DRbConn
    include DRbProtocol

    POOL_SIZE = 16
    @mutex = Mutex.new
    @pool = []

    def self.open(remote_uri)
      begin
	conn = nil

	@mutex.synchronize do
	  #FIXME
	  new_pool = []
	  @pool.each do |c|
	    if conn.nil? and c.uri == remote_uri
	      conn = c if c.alive?
	    else
	      new_pool.push c
	    end
	  end
	  @pool = new_pool
	end

	conn = self.new(remote_uri) unless conn
	succ, result = yield(conn)
	return succ, result

      ensure
	@mutex.synchronize do
	  if @pool.size > POOL_SIZE or ! succ
	    conn.close if conn
	  else
	    @pool.unshift(conn)
	  end
	end
      end
    end

    def initialize(remote_uri)
      @uri = remote_uri
      @load_limit = DRb.current_server.load_limit
      @host, @port, @opt = parse_uri(remote_uri)
      untaint_remote
      @soc = TCPSocket.open(@host, @port)
      setup_sockopt(@soc)
    end
    attr_reader :load_limit
    attr_reader :uri

    def send_message(ref, msg_id, *arg, &block)
      send_request(@soc, ref, msg_id, *arg, &block)
      recv_reply(@soc)
    end

    def close
      @soc.close
      @soc = nil
    end

    def alive?
      # return false unless @soc
      if IO.select([@soc], nil, nil, 0)
	close
	return false
      end
      true
    end

    private
    def untaint_remote
      @host.untaint
      @port.untaint
      @opt.untaint
    end
  end

  class DRbServer
    include DRbProtocol

    @@acl = nil
    @@idconv = DRbIdConv.new
    @@secondary_server = nil
    @@argc_limit = nil
    @@load_limit = nil
    @@verbose = false

    def self.default_argc_limit(argc)
      @@argc_limit = argc
    end

    def self.default_load_limit(sz)
      @@load_limit = sz
    end

    def self.default_acl(acl)
      @@acl = acl
    end

    def self.default_id_conv(idconv)
      @@idconv = idconv
    end

    def self.verbose=(on)
      @@verbose = on
    end
    
    def self.verbose
      @@verbose
    end

    def initialize(uri=nil, front=nil, acl=nil, idconv=nil)
      uri = 'druby://:0' unless uri
      host, port, opt = parse_uri(uri)
      if host.size == 0
	@soc = TCPServer.open(port)
	host = Socket.gethostname
      else
	@soc = TCPServer.open(host, port)
      end
      setup_sockopt(@soc)
      port = @soc.addr[1] if @soc
      @uri = "druby://#{host}:#{port}"
      @front = front
      @acl = acl || @@acl
      @idconv = idconv || @@idconv
      @argc_limit = @@argc_limit
      @load_limit = @@load_limit
      @verbose = @@verbose

      @grp = ThreadGroup.new
      @thread = run

      Thread.exclusive do
	DRb.primary_server = self unless DRb.primary_server
      end
    end
    attr_reader :uri, :thread, :front
    attr_accessor :argc_limit, :load_limit
    attr_accessor :verbose

    def alive?
      @thread.alive?
    end

    def stop_service
      @thread.kill
    end

    def to_obj(ref)
      @idconv.to_obj(ref)
    end

    def to_id(obj)
      return nil if obj.__id__ == front.__id__
      @idconv.to_id(obj)
    end

    private
    def kill_sub_thread
      list = @grp.list
      while list.size > 0
	list.each do |th|
	  th.kill
	end
	list @grp.list
      end
    end

    def run
      Thread.start do
	begin
	  while true
	    proc
	  end
	ensure
	  @soc.close if @soc
	  kill_sub_thread
	end
      end
    end

    def allow?(s)
      @acl ? @acl.allow_socket?(s) : true
    end
    
    INSECURE_METHOD = [
      :__send__
    ]
    def insecure_method?(msg_id)
      INSECURE_METHOD.include?(msg_id)
    end

    def any_to_s(obj)
      obj.to_s rescue sprintf("#<%s:0x%lx>", obj.class, obj.__id__)      
    end

    def check_insecure_method(obj, msg_id)
      return true if Proc === obj && msg_id == :yield_call
      raise(ArgumentError, "#{any_to_s(msg_id)} is not a symbol") unless Symbol == msg_id.type
      raise(SecurityError, "insecure method `#{msg_id}'") if insecure_method?(msg_id)
      unless obj.respond_to?(msg_id)
	desc = any_to_s(obj)
	if desc.nil? || desc[0] == '#'
	  desc << ":#{obj.class}"
	end
	
	if obj.private_methods.include?(msg_id.to_s)
	  raise NameError, "private method `#{msg_id}' called for #{desc}"
	else
	  raise NameError, "undefined method `#{msg_id}' called for #{desc}"
	end
      end
      true
    end
    public :check_insecure_method

    def obj_send(obj, msg_id, *argv)
      if Proc === obj && msg_id == :yield_call
	if argv.size == 1
	  ary = argv
	else
	  ary = [argv]
	end
	# ary.collect(&obj)[0]
	it = ary.collect(&obj)
	raise(LocalJumpError, 'break from proc-closure') if it.nil?
	it[0]
      else
	obj.__send__(msg_id, *argv)
      end
    end

    def proc
      Thread.start(@soc.accept) do |s|
	loop do
	  begin
	    @grp.add Thread.current
	    Thread.current['DRb'] = { 'socket' => s , 'server' => self }
	    succ = nil
	    begin 
	      raise 'Forbidden' unless allow?(s)
	      obj, msg, argv, block = recv_request(s)
	      msg_id = msg.intern
	      check_insecure_method(obj, msg_id)
	      if block
		catch_retry = nil
		result = obj.__send__(msg_id, *argv) do |*x|
		  catch_retry = false
		  begin
		    block_value = block.yield_call(*x)
		  rescue LocalJumpError
		    if /^retry/ =~ $!.message	# retry from proc-closure
		      catch_retry = true
		    elsif /^break/ =~ $!.message	# break from proc-closure
		      break
		    else
		      p $!
		      raise $!
		    end
		  end
		  retry if catch_retry
		  block_value
		end
	      else
		result = obj_send(obj, msg_id, *argv)
	      end
	      succ = true
	    rescue StandardError, ScriptError
	      result = $!
	      if @verbose
		p $!
		$@.each do |x| puts x end
	      end
	      succ = false
	    end
	    send_reply(s, succ, result) rescue nil
	  ensure
	    s.flush
	    unless succ
	      s.close
	      return
	    end
	  end
	end
      end
    end
  end

  @primary_server = nil

  def start_service(uri=nil, front=nil, acl=nil)
    @primary_server = DRbServer.new(uri, front, acl)
  end
  module_function :start_service

  attr_accessor :primary_server
  module_function :primary_server=, :primary_server

  def current_server
    drb = Thread.current['DRb'] 
    server = (drb && drb['server']) ? drb['server'] : @primary_server 
    raise DRbServerNotFound unless server
    return server
  end
  module_function :current_server

  def stop_service
    @primary_server.stop_service if @primary_server
    @primary_server = nil
  end
  module_function :stop_service

  def uri
    current_server.uri
  end
  module_function :uri
  
  def front
    current_server.front
  end
  module_function :front

  def to_obj(ref)
    current_server.to_obj(ref)
  end
  def to_id(obj)
    current_server.to_id(obj)
  end
  module_function :to_id
  module_function :to_obj

  def thread
    @primary_server ? @primary_server.thread : nil
  end
  module_function :thread

  def install_id_conv(idconv)
    DRbServer.default_id_conv(idconv)
  end
  module_function :install_id_conv

  def install_acl(acl)
    DRbServer.default_acl(acl)
  end
  module_function :install_acl
end

DRbObject = DRb::DRbObject
DRbUndumped = DRb::DRbUndumped
DRbIdConv = DRb::DRbIdConv


