iruby 0.2.7 → 0.5.0

data/lib/iruby/session/cztop.rb

@@ -0,0 +1,70 @@
+require 'cztop'
+
+module IRuby
+  class Session
+    include SessionSerialize
+
+    def initialize(config)
+      connection = "#{config['transport']}://#{config['ip']}:%d"
+
+      reply_socket = CZTop::Socket::ROUTER.new(connection % config['shell_port'])
+      pub_socket = CZTop::Socket::PUB.new(connection % config['iopub_port'])
+      stdin_socket = CZTop::Socket::ROUTER.new(connection % config['stdin_port'])
+
+      Thread.new do
+        begin
+          hb_socket = CZTop::Socket::REP.new(connection % config['hb_port'])
+          loop do
+            message = hb_socket.receive
+            hb_socket << message
+          end
+        rescue Exception => e
+          IRuby.logger.fatal "Kernel heartbeat died: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+
+      @sockets = {
+        publish: pub_socket,
+        reply: reply_socket,
+        stdin: stdin_socket,
+      }
+
+      @session = SecureRandom.uuid
+      unless config['key'].to_s.empty? || config['signature_scheme'].to_s.empty?
+        raise 'Unknown signature scheme' unless config['signature_scheme'] =~ /\Ahmac-(.*)\Z/
+        @hmac = OpenSSL::HMAC.new(config['key'], OpenSSL::Digest.new($1))
+      end
+    end
+
+    def description
+      'old-stle session using cztop'
+    end
+
+    # Build and send a message
+    def send(socket, type, content)
+      idents =
+        if socket == :reply && @last_recvd_msg
+          @last_recvd_msg[:idents]
+        else
+          type == :stream ? "stream.#{content[:name]}" : type
+        end
+      header = {
+        msg_type: type,
+        msg_id:   SecureRandom.uuid,
+        username: 'kernel',
+        session:  @session,
+        version:  '5.0'
+      }
+      @sockets[socket] << serialize(idents, header, content)
+    end
+
+    # Receive a message and decode it
+    def recv(socket)
+      @last_recvd_msg = unserialize(@sockets[socket].receive)
+    end
+
+    def recv_input
+      unserialize(@sockets[:stdin].receive)[:content]["value"]
+    end
+  end
+end

data/lib/iruby/session/ffi_rzmq.rb

@@ -0,0 +1,87 @@
+require 'ffi-rzmq'
+
+module IRuby
+  class Session
+    include SessionSerialize
+
+    def initialize(config)
+      c = ZMQ::Context.new
+
+      connection = "#{config['transport']}://#{config['ip']}:%d"
+      reply_socket = c.socket(ZMQ::XREP)
+      reply_socket.bind(connection % config['shell_port'])
+
+      pub_socket = c.socket(ZMQ::PUB)
+      pub_socket.bind(connection % config['iopub_port'])
+
+      stdin_socket = c.socket(ZMQ::XREP)
+      stdin_socket.bind(connection % config['stdin_port'])
+
+      Thread.new do
+        begin
+          hb_socket = c.socket(ZMQ::REP)
+          hb_socket.bind(connection % config['hb_port'])
+          ZMQ::Device.new(hb_socket, hb_socket)
+        rescue Exception => e
+          IRuby.logger.fatal "Kernel heartbeat died: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+
+      @sockets = { 
+        publish: pub_socket, reply: reply_socket, stdin: stdin_socket
+      }
+      
+      @session = SecureRandom.uuid
+      unless config['key'].to_s.empty? || config['signature_scheme'].to_s.empty?
+        raise 'Unknown signature scheme' unless config['signature_scheme'] =~ /\Ahmac-(.*)\Z/
+        @hmac = OpenSSL::HMAC.new(config['key'], OpenSSL::Digest.new($1))
+      end
+    end
+
+    # Build and send a message
+    def send(socket, type, content)
+      idents =
+        if socket == :reply && @last_recvd_msg
+          @last_recvd_msg[:idents]
+        else
+          type == :stream ? "stream.#{content[:name]}" : type
+        end
+      header = {
+        msg_type: type,
+        msg_id:   SecureRandom.uuid,
+        username: 'kernel',
+        session:  @session,
+        version:  '5.0'
+      }
+      socket = @sockets[socket]
+      list = serialize(idents, header, content)
+      list.each_with_index do |part, i|
+        socket.send_string(part, i == list.size - 1 ? 0 : ZMQ::SNDMORE)
+      end
+    end
+
+    # Receive a message and decode it
+    def recv(socket)
+      socket = @sockets[socket]
+      msg = []
+      while msg.empty? || socket.more_parts?
+        begin
+          frame = ''
+          rc = socket.recv_string(frame)
+          ZMQ::Util.error_check('zmq_msg_send', rc)
+          msg << frame
+        rescue
+        end
+      end
+
+      @last_recvd_msg = unserialize(msg)
+    end
+
+    def recv_input 
+      last_recvd_msg = @last_recvd_msg
+      input = recv(:stdin)[:content]["value"]
+      @last_recvd_msg = last_recvd_msg
+      input
+    end
+  end
+end

data/lib/iruby/session/mixin.rb

@@ -0,0 +1,47 @@
+module IRuby
+  module SessionSerialize
+    DELIM = '<IDS|MSG>'
+
+    private
+
+    def serialize(idents, header, content)
+      msg = [MultiJson.dump(header),
+             MultiJson.dump(@last_recvd_msg ? @last_recvd_msg[:header] : {}),
+             '{}',
+             MultiJson.dump(content || {})]
+      frames = ([*idents].compact.map(&:to_s) << DELIM << sign(msg)) + msg
+      IRuby.logger.debug "Sent #{frames.inspect}"
+      frames
+    end
+
+    def unserialize(msg)
+      raise 'no message received' unless msg
+      frames = msg.to_a.map(&:to_s)
+      IRuby.logger.debug "Received #{frames.inspect}"
+
+      i = frames.index(DELIM)
+      idents, msg_list = frames[0..i-1], frames[i+1..-1]
+
+      minlen = 5
+      raise 'malformed message, must have at least #{minlen} elements' unless msg_list.length >= minlen
+      s, header, parent_header, metadata, content, buffers = *msg_list
+      raise 'Invalid signature' unless s == sign(msg_list[1..-1])
+      {
+        idents:        idents,
+        header:        MultiJson.load(header),
+        parent_header: MultiJson.load(parent_header),
+        metadata:      MultiJson.load(metadata),
+        content:       MultiJson.load(content),
+        buffers:       buffers
+      }
+    end
+
+    # Sign using HMAC
+    def sign(list)
+      return '' unless @hmac
+      @hmac.reset
+      list.each {|m| @hmac.update(m) }
+      @hmac.hexdigest
+    end
+  end
+end

data/lib/iruby/session_adapter.rb

@@ -0,0 +1,66 @@
+module IRuby
+  class SessionAdapterNotFound < RuntimeError; end
+
+  module SessionAdapter
+    class BaseAdapter
+      def self.available?
+        load_requirements
+        true
+      rescue LoadError
+        false
+      end
+
+      def initialize(config)
+        @config = config
+      end
+
+      def name
+        self.class.name[/::(\w+)Adapter\Z/, 1].downcase
+      end
+
+      def make_router_socket(protocol, host, port)
+        socket, port = make_socket(:ROUTER, protocol, host, port)
+        [socket, port]
+      end
+
+      def make_pub_socket(protocol, host, port)
+        socket, port = make_socket(:PUB, protocol, host, port)
+        [socket, port]
+      end
+
+      def make_rep_socket(protocol, host, port)
+        socket, port = make_socket(:REP, protocol, host, port)
+        [socket, port]
+      end
+    end
+
+    require_relative 'session_adapter/ffirzmq_adapter'
+    require_relative 'session_adapter/cztop_adapter'
+    require_relative 'session_adapter/pyzmq_adapter'
+
+    def self.select_adapter_class(name=nil)
+      classes = {
+        'ffi-rzmq' => SessionAdapter::FfirzmqAdapter,
+        'cztop' => SessionAdapter::CztopAdapter,
+        # 'pyzmq' => SessionAdapter::PyzmqAdapter
+      }
+      if (name ||= ENV.fetch('IRUBY_SESSION_ADAPTER', nil))
+        cls = classes[name]
+        unless cls.available?
+          if ENV['IRUBY_SESSION_ADAPTER']
+            raise SessionAdapterNotFound,
+                  "Session adapter `#{name}` from IRUBY_SESSION_ADAPTER is unavailable"
+          else
+            raise SessionAdapterNotFound,
+                  "Session adapter `#{name}` is unavailable"
+          end
+        end
+        return cls
+      end
+      classes.each_value do |cls|
+        return cls if cls.available?
+      end
+      raise SessionAdapterNotFound, "No session adapter is available"
+    end
+  end
+end

data/lib/iruby/session_adapter/cztop_adapter.rb

@@ -0,0 +1,45 @@
+module IRuby
+  module SessionAdapter
+    class CztopAdapter < BaseAdapter
+      def self.load_requirements
+        require 'cztop'
+      end
+
+      def send(sock, data)
+        sock << data
+      end
+
+      def recv(sock)
+        sock.receive
+      end
+
+      def heartbeat_loop(sock)
+        loop do
+          message = sock.receive
+          sock << message
+        end
+      end
+
+      private
+
+      def socket_type_class(type_symbol)
+        case type_symbol
+        when :ROUTER, :PUB, :REP
+          CZTop::Socket.const_get(type_symbol)
+        else
+          if CZTop::Socket.const_defined?(type_symbol)
+            raise ArgumentError, "Unsupported ZMQ socket type: #{type_symbol}"
+          else
+            raise ArgumentError, "Invalid ZMQ socket type: #{type_symbol}"
+          end
+        end
+      end
+
+      def make_socket(type_symbol, protocol, host, port)
+        uri = "#{protocol}://#{host}:#{port}"
+        socket_class = socket_type_class(type_symbol)
+        socket_class.new(uri)
+      end
+    end
+  end
+end

data/lib/iruby/session_adapter/ffirzmq_adapter.rb

@@ -0,0 +1,55 @@
+module IRuby
+  module SessionAdapter
+    class FfirzmqAdapter < BaseAdapter
+      def self.load_requirements
+        require 'ffi-rzmq'
+      end
+
+      def send(sock, data)
+        data.each_with_index do |part, i|
+          sock.send_string(part, i == data.size - 1 ? 0 : ZMQ::SNDMORE)
+        end
+      end
+
+      def recv(sock)
+        msg = []
+        while msg.empty? || sock.more_parts?
+          begin
+            frame = ''
+            rc = sock.recv_string(frame)
+            ZMQ::Util.error_check('zmq_msg_recv', rc)
+            msg << frame
+          rescue
+          end
+        end
+        msg
+      end
+
+      def heartbeat_loop(sock)
+        @heartbeat_device = ZMQ::Device.new(sock, sock)
+      end
+
+      private
+
+      def make_socket(type, protocol, host, port)
+        case type
+        when :ROUTER, :PUB, :REP
+          type = ZMQ.const_get(type)
+        else
+          if ZMQ.const_defined?(type)
+            raise ArgumentError, "Unsupported ZMQ socket type: #{type_symbol}"
+          else
+            raise ArgumentError, "Invalid ZMQ socket type: #{type_symbol}"
+          end
+        end
+        zmq_context.socket(type).tap do |sock|
+          sock.bind("#{protocol}://#{host}:#{port}")
+        end
+      end
+
+      def zmq_context
+        @zmq_context ||= ZMQ::Context.new
+      end
+    end
+  end
+end

data/lib/iruby/session_adapter/pyzmq_adapter.rb

@@ -0,0 +1,77 @@
+module IRuby
+  module SessionAdapter
+    class PyzmqAdapter < BaseAdapter
+
+      class << self
+        def load_requirements
+          require 'pycall'
+          import_pyzmq
+        end
+
+        def import_pyzmq
+          @zmq = PyCall.import_module('zmq')
+        rescue PyCall::PyError => error
+          raise LoadError, error.message
+        end
+
+        attr_reader :zmq
+      end
+
+      def make_router_socket(protocol, host, port)
+        make_socket(:ROUTER, protocol, host, port)
+      end
+
+      def make_pub_socket(protocol, host, port)
+        make_socket(:PUB, protocol, host, port)
+      end
+
+      def heartbeat_loop(sock)
+        PyCall.sys.path.append(File.expand_path('../pyzmq', __FILE__))
+        heartbeat = PyCall.import_module('iruby.heartbeat')
+        @heartbeat_thread = heartbeat.Heartbeat.new(sock)
+        @heartbeat_thread.start
+      end
+
+      private
+
+      def socket_type(type_symbol)
+        case type_symbol
+        when :ROUTER, :PUB, :REP
+          zmq[type_symbol]
+        else
+          raise ArgumentError, "Unknown ZMQ socket type: #{type_symbol}"
+        end
+      end
+
+      def make_socket(type_symbol, protocol, host, port)
+        type = socket_type(type_symbol)
+        sock = zmq_context.socket(type)
+        bind_socket(sock, protocol, host, port)
+        sock
+      end
+
+      def bind_socket(sock, protocol, host, port)
+        iface = "#{protocol}://#{host}"
+        case protocol
+        when 'tcp'
+          if port <= 0
+            port = sock.bind_to_random_port(iface)
+          else
+            sock.bind("#{iface}:#{port}")
+          end
+        else
+          raise ArgumentError, "Unsupported protocol: #{protocol}"
+        end
+        [sock, port]
+      end
+
+      def zmq_context
+        zmq.Context.instance
+      end
+
+      def zmq
+        self.class.zmq
+      end
+    end
+  end
+end