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.register(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
        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 = Hash.new
            @id_to_acceptor = Hash.new
            @mutex = Mutex.new
            @next_id = 1
            @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.  Once an object is registered,
        # it cannot be unregistered.  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 we want to be sure that the id will fit into
        # a 16-bit short.  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 register(obj, name, acceptor=nil)
            @mutex.synchronize do
                if @next_id >= Session::MAX_ID then
                    raise ArgumentError, "Object limit exceeded"
                end
                @name_to_id[name] = @next_id
                @id_to_acceptor[@next_id] = acceptor
                @id_to_object[@next_id] = obj
                @next_id = @next_id.succ()
            end
        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()
                obj = @id_to_object[object_id]
                acceptor = @id_to_acceptor[object_id]
                begin
                    if acceptor then
                        if !acceptor.call(session.socket) then
                            if type != Session::ONEWAY then
                                raise ArgumentError, "No authorization"
                            end
                        end
                    end
                    case type
                        when Session::REQUEST
                            retval = obj.__send__(*message)
                        when Session::REQUEST_BLOCK
                            retval = obj.__send__(*message) do |*i|
                                session.send_message(Session::YIELD, 0, i)
                            end
                        when Session::ONEWAY
                            begin; obj.__send__(*message); rescue Exception; end
                            next
                        when Session::RESOLVE
                            retval = @name_to_id[message]
                        when Session::SYNC
                            session.reply_sync(object_id)
                            next
                        else
                            raise ArgumentError, "Bad session request"
                    end
                    session.send_message(Session::RETVAL, 0, retval)
                rescue Exception
                    session.send_message(Session::EXCEPTION, 0, $!)
                end
            end
        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
                        return Object.new(@session, @mutex, message)
                    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

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

        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 message
                        when Session::YIELD      ; yield message
                        when Session::EXCEPTION  ; raise message
                        when Session::SYNC       ; @session.reply_sync(obj)
                        else                     ; raise RuntimeError
                    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)
                return nil
            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

        GOOD_FUNCTIONS = [
            :inspect, :class_variables, :instance_eval, :instance_variables,
            :to_a, :to_s
        ]

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

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

        RESPOND_FUNCTIONS = [
            [ :method,      "raise NameError" ],
            [ :respond_to?, "return false" ]
        ]

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

        # And make sure others never get called.
        BAD_FUNCTIONS.each do |method|
            undef_method method
        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.
        METHOD_FUNCTIONS.each do |method|
            eval %{
                def #{method}(*args)
                    retval = method_missing(:#{method}, *args)
                    retval.each do |item|
                        BAD_FUNCTIONS.each do |bad|
                            retval.delete(bad.to_s)
                        end
                    end
                    return retval
                end
            }
        end

        # Same here, except don't let the call go through in the first place.
        RESPOND_FUNCTIONS.each do |method, action|
            eval %{
                def #{method}(arg, *args)
                    BAD_FUNCTIONS.each do |bad|
                        if arg === bad.to_s then
                            eval("#{action}")
                        end
                    end
                    return method_missing(:#{method}, arg, *args)
                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)
            return [ type, obj, message ]
        end

        def finished()
            return @io.eof
        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 RuntimeError, "DRuby synchonization failed"
            end
        end

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

    end

end