require 'socket'
require 'thread'

=begin
DRuby (not to be confused with dRuby, written by Masatoshi SEKI)
(C) Copyright 2001 Paul Brannan (cout at rm-f.net)

DRuby is a set of classes for providing distributed object support to a
Ruby program.  You may distribute and/or modify it under the same terms as
Ruby (see http://www.ruby-lang.org/en/LICENSE.txt).  Example:

Client 
------
client = DRuby::Client.new('localhost', 4242)
obj = client.resolve("foo")
puts obj.foo("foo")
obj.oneway(:foo, "foo")

Server
------
class Foo
    def foo(x); return x; end
end
obj = Foo.new
server = DRuby::Server.new('localhost', 4242)
server.bind(obj, "foo")
server.thread.join

You can do all sorts of cool things with DRuby, including passing blocks to
the functions, throwing exceptions and propogating them from server to
client, and more.  Unlike CORBA, where you must create an interface
definition and strictly adhere to it, DRuby uses marshalling, so you can
use almost any object with it.  But, alas, it is not as powerful as CORBA.

On a fast machine, you should expect around 4000 messages per second with
normal method calls, and up to 10000 messages per second with oneway calls.
My dual-processor machine gets 3500/6500 messages per second, though, so
YMMV.

The DRuby message format is broken into 3 components:
    [ msg_type, obj_id, message ]
For each msg_type, there is a specific format to the message.  Additionally,
certain msg_types are only valid when being sent to the server, and others
are valid only when being sent back to the client.  Here is a summary:

msg_type        send to     meaning of obj_id       msg format
----------------------------------------------------------------------------
REQUEST         server      obj to talk to          [:method, *args]
REQUEST_BLOCK   server      obj to talk to          [:method, *args]
ONEWAY          server      obj to talk to          [:method, *args]
RETVAL          client      always 0                retval
EXCEPTION       client      always 0                $!
YIELD           client      always 0                [value, value, ...]
RESOLVE         server      always 0                object_name
SYNC            either      0=request, 1=response   nil

Known bugs:
1) Oneway calls are really slow if mixed with regular calls.  I'm not sure
why.  A workaround is to use a separate connection for each type of call,
but sometimes that doesn't work either.
2) The server does not refuse connections for unauthorized addresses; it will
instead accept a connection and then immediately close it.
3) If the server is flooded with oneway calls and drops packets, it will
listen for data until it encounters a valid message.  If this happens, the
client will not be notified.
=end

module DRuby

public

    # The DRuby server class.  Like its drb equivalent, this class spawns off
    # a new thread which processes requests, allowing the server to do other
    # things while it is doing processing for a distributed object.  This
    # means, though, that all objects used with DRuby must be thread-safe.
    class Server

    public
        attr_reader :host, :port, :obj, :thread

        # Start a server on the given host and port.  An acceptor may be
        # specified to accept/reject connections; it should be a proc that
        # takes a Socket as an argument and returns true or false.
        def initialize(host, port, acceptor=nil)
            @host = host
            @port = port
            @name_to_id = Hash.new
            @id_to_object = Array.new
            @id_to_acceptor = Array.new
            @mutex = Mutex.new
            @next_id = 1
            @unused_ids = Array.new
            @thread = Thread.new do
                server = TCPServer.new(@host, @port)
                while(socket = server.accept)
                    if acceptor then
                        if !acceptor.call(socket) then
                            socket.close
                            next
                        end
                    end
                    socket.setsockopt(Socket::SOL_TCP, Socket::TCP_NODELAY, 1)
                    Thread.new(socket) do |socket|
                        session_loop(Session.new(socket))
                    end
                end
            end
        end

        # Register an object with the server.  The object will be given an
        # id of @next_id, and @next_id will be incremented.  We could use the
        # object's real id, but this is insecure.  The acceptor argument is
        # a proc that takes a Socket as an arugment, and returns true or false
        # depending on whether the request should be allowed.  The supplied
        # object must be thread-safe.
        def create_reference(obj, acceptor=nil)
            @mutex.synchronize do
                ref = register_private(obj, acceptor)
                Object_Reference.new(ref) #return
            end
        end

        # Finds an object in linear time and unregister it.  If you need a
        # faster function, write one and send me a patch.
        def delete_reference(obj)
            @mutex.synchronize do
                unregister_private(obj)
            end
            nil #return
        end

        # Register an object with the server and bind it to name
        def bind(obj, name, acceptor=nil)
            id = register_private(obj, acceptor)
            @name_to_id[name] = id
            nil #return
        end

        # Main server loop.  Wait for a REQUEST message, process it, and send
        # back a YIELD, EXCEPTION, or RETVAL message.  Note that we don't do
        # any thread synchronization here, because all registered objects are
        # required to be thread-safe, and the @id_to_object lookup is atomic
        # (and if it succeeds, the object is guaranteed to be fully
        # registered).
        def session_loop(session)
            while not session.finished()
                type, object_id, message = session.get_message()
                begin
                    obj = @id_to_object[object_id]
                    if (acceptor = @id_to_acceptor[object_id]) then
                        if !acceptor.call(session.socket) then
                            if type != Session::ONEWAY then
                                raise "No authorization"
                            end
                        end
                    end
                    case type # In order of most to least common
                        when Session::REQUEST
                            retval = obj.__send__(*message)
                        when Session::ONEWAY
                            begin; obj.__send__(*message); rescue Exception; end
                            next
                        when Session::REQUEST_BLOCK
                            retval = obj.__send__(*message) do |*i|
                                session.send_message(Session::YIELD, 0, i)
                            end
                        when Session::RESOLVE
                            retval = @name_to_id[message]
                        when Session::SYNC
                            session.reply_sync(object_id)
                            next
                        else
                            raise "Bad session request"
                    end
                    session.send_message(Session::RETVAL, 0, retval)
                rescue Exception
                    session.send_message(Session::EXCEPTION, 0, $!)
                end
            end
        end

    private
        def register_private(obj, acceptor)
            if @next_id >= Session::MAX_ID then
                if @unused_ids.size == 0 then
                    raise "Object limit exceeded"
                else
                    id = @unused_ids.pop
                end
            end
            @id_to_acceptor[@next_id] = acceptor
            @id_to_object[@next_id] = obj
            old_id = @next_id
            @next_id = @next_id.succ()
            old_id #return
        end

        def unregister_private(obj)
            delete_obj_from_array_private(@id_to_acceptor, @next_id)
            delete_obj_from_array_private(@id_to_object, @next_id)
        end

        def delete_obj_from_array_private(array, obj)
            index = array.index(obj)
            array[index] = nil unless index == nil
        end
    end

    # The DRuby client class.  A DRuby server must be started on the given
    # host and port before instantiating a DRuby client.
    class Client
        attr_reader :host, :port
        
        # Connect a client to the server on the given host and port.  The
        # user can specify sync=false to turn off synchronization and get a
        # 20% speed boost.
        def initialize(host, port, sync=true)
            @host = host
            @port = port
            @server = TCPSocket.open(@host, @port)
            @server.setsockopt(Socket::SOL_TCP, Socket::TCP_NODELAY, 1)
            @session = Session.new(@server)
            @mutex = sync ? Mutex.new : Null_Mutex.new
        end

        # Given a string, return a proxy object that will forward requests
        # for an object on the server with that name.
        def resolve(object_name)
            @mutex.synchronize do
                @session.send_message(Session::RESOLVE, 0, object_name)
                type, obj, message = @session.get_message()
                case type
                    when Session::RETVAL
                        Proxy_Object.new(@session, @mutex, message) #return
                    when Session::EXCEPTION     ; raise message
                    else                        ; raise RuntimeError
                end
            end
        end
    end

private

    # In case the user does not want synchronization.
    class Null_Mutex
        def synchronize
            yield
        end
    end

    # All the special functions we have to keep track of
    class Functions
        GOOD = [
            :inspect, :class_variables, :instance_eval, :instance_variables,
            :to_a, :to_s
        ]

        BAD = [
            :clone, :dup, :display
        ]

        METHOD = [
            :methods, :private_methods, :protected_methods, :public_methods,
            :singleton_methods
        ]

        RESPOND = [
            [ :method,      "raise NameError" ],
            [ :respond_to?, "false" ]
        ]
    end

    # A DRuby::Object_Reference is created on the server side to represent an
    # object in the system.  It can be returned from a server object to a
    # client object, at which point it is converted into a DRuby::Proxy_Object.
    class Object_Reference
        attr_reader :object_id

        def initialize(object_id)
            @object_id = object_id
        end
    end

    # A DRuby::Object acts as a proxy; it forwards most methods to the server
    # for execution.
    class Proxy_Object

    public

        def initialize(session, mutex, object_id)
            @object_id = object_id
            @session = session
            @mutex = mutex
        end

        # Client request handler.  The idea here is to send out a REQUEST
        # message, and get back a YIELD, EXCEPTION, or RETVAL message.
        def method_missing(method, *args)
            message = [ method, *args ]
            @mutex.synchronize do
                @session.send_message(
                    (block_given?) ? Session::REQUEST_BLOCK : Session::REQUEST,
                    @object_id,
                    message)
                loop do
                    type, obj, message = @session.get_message()
                    case type
                        when Session::RETVAL     ; return msg_to_obj(message)
                        when Session::YIELD      ; yield msg_to_obj(message)
                        when Session::EXCEPTION  ; raise message
                        when Session::SYNC       ; @session.reply_sync(obj)
                        else                     ; raise "Invalid msg type"
                    end
                end
            end
        end

        # This is a special function that lets you ignore the return value
        # of the calling function to get some extra speed, in exchange for not
        # knowing when the function will complete.  A future version of the
        # server should probably check whether a given function is allowed to
        # be called as a oneway, in order to reduce DoS attacks.
        def oneway(*message)
            @mutex.synchronize do
                @session.send_message(Session::ONEWAY, @object_id, message)
                nil #return
            end
        end

        # sync() will synchonize the client with the server (useful for
        # determining when a oneway call completes)
        def sync()
            @mutex.synchronize do
                @session.send_sync()
                @session.wait_sync()
            end
        end

        # Make sure certain methods get passed down the wire.
        Functions::GOOD.each do |method|
            eval %{
                def #{method}(*args)
                      method_missing(:#{method}, *args) #return
                end
            }
        end

        # And make sure others never get called.
        Functions::BAD.each do |method|
            eval %{
                def #{method}(*args)
                    raise(NameError,
                        "undefined method `#{method}' for " +
                        "\#<#{self.class}:#{self.id}>")
                end
            }
        end

        # And remove these function names from any method lists that get
        # returned; there's nothing we can do about people who decide to
        # return them from other functions.
        Functions::METHOD.each do |method|
            eval %{
                def #{method}(*args)
                    retval = method_missing(:#{method}, *args)
                    retval.each do |item|
                        Functions::BAD.each do |bad|
                            retval.delete(bad.to_s)
                        end
                    end
                    retval #return
                end
            }
        end

        # Same here, except don't let the call go through in the first place.
        Functions::RESPOND.each do |method, action|
            eval %{
                def #{method}(arg, *args)
                    Functions::BAD.each do |bad|
                        if arg === bad.to_s then
                            return eval("#{action}")
                        end
                    end
                    method_missing(:#{method}, arg, *args) #return
                end
            }
        end

    private

        def msg_to_obj(obj)
            if Object_Reference === obj then
                Proxy_Object.new(@session, @mutex, obj.object_id) #return
            else
                obj #return
            end
        end

    end

    # A DRuby session sends messages back and forth between server and client.
    # The message format is as follows:
    # +-----+-----+-----+-----+-----+-----+-----+-----+-----+---  +  ---+-----+
    # | MSG_START | msg. size | msg. type |  obj. id  |  marshalled msg/args  |
    # +-----+-----+-----+-----+-----+-----+-----+-----+-----+---  +  ---+-----+
    class Session
        REQUEST         = 0x1001
        REQUEST_BLOCK   = 0x1002
        ONEWAY          = 0x1003
        RETVAL          = 0x2001
        EXCEPTION       = 0x2002
        YIELD           = 0x2003
        RESOLVE         = 0x3001
        SYNC            = 0x4001
        MSG_START       = 0x4242

        MAX_ID = 2**16

        attr_reader :io
        alias_method :socket, :io

        def initialize(io)
            @io = io
        end

        def send_message(type, obj, message)
            data = Marshal.dump(message)
            header = [MSG_START, data.length, type, obj]
            @io.write(header.pack("vvvv"))
            @io.write(data)
        end

        def get_message()
            header = @io.read(8)
            magic, size, type, obj = header.unpack("vvvv")
            while magic != MSG_START
                header = @io.read(8)
                magic, size, type, obj = header.unpack("vvvv")
            end
            data = @io.read(size)
            message = Marshal.load(data)
            [ type, obj, message ] #return
        end

        def finished()
            @io.eof #return
        end
        
        def send_sync()
            send_message(Session::SYNC, 0, nil)
        end
        
        def wait_sync()
            sleep 1
            type, obj, message = get_message()
            if type != Session::SYNC && obj != 0 && message != nil then
                raise "DRuby synchonization failed"
            end
        end

        def reply_sync(value)
            if value == 0 then
                send_message(Session::SYNC, 1, nil)
            end
        end

    end
end