|
| 1 | +require 'iruby/session_adapter' |
| 2 | +require 'iruby/session/mixin' |
| 3 | + |
| 4 | +require 'securerandom' |
| 5 | + |
| 6 | +module IRuby |
| 7 | + class Session |
| 8 | + include SessionSerialize |
| 9 | + |
| 10 | + def initialize(config, adapter_name=nil) |
| 11 | + @config = config |
| 12 | + @adapter = create_session_adapter(config, adapter_name) |
| 13 | + |
| 14 | + setup |
| 15 | + setup_sockets |
| 16 | + setup_heartbeat |
| 17 | + setup_security |
| 18 | + end |
| 19 | + |
| 20 | + attr_reader :adapter, :config |
| 21 | + |
| 22 | + def description |
| 23 | + "#{@adapter.name} session adapter" |
| 24 | + end |
| 25 | + |
| 26 | + def setup |
| 27 | + end |
| 28 | + |
| 29 | + def setup_sockets |
| 30 | + protocol, host = config.values_at('transport', 'ip') |
| 31 | + shell_port = config['shell_port'] |
| 32 | + iopub_port = config['iopub_port'] |
| 33 | + stdin_port = config['stdin_port'] |
| 34 | + |
| 35 | + @shell_socket, @shell_port = @adapter.make_router_socket(protocol, host, shell_port) |
| 36 | + @iopub_socket, @iopub_port = @adapter.make_pub_socket(protocol, host, iopub_port) |
| 37 | + @stdin_socket, @stdin_port = @adapter.make_router_socket(protocol, host, stdin_port) |
| 38 | + |
| 39 | + @sockets = { |
| 40 | + publish: @iopub_socket, |
| 41 | + reply: @shell_socket, |
| 42 | + stdin: @stdin_socket |
| 43 | + } |
| 44 | + end |
| 45 | + |
| 46 | + def setup_heartbeat |
| 47 | + protocol, host = config.values_at('transport', 'ip') |
| 48 | + hb_port = config['hb_port'] |
| 49 | + @hb_socket, @hb_port = @adapter.make_rep_socket(protocol, host, hb_port) |
| 50 | + @heartbeat_thread = Thread.start do |
| 51 | + begin |
| 52 | + # NOTE: this loop is copied from CZTop's old session code |
| 53 | + @adapter.heartbeat_loop(@hb_socket) |
| 54 | + rescue Exception => e |
| 55 | + IRuby.logger.fatal "Kernel heartbeat died: #{e.message}\n#{e.backtrace.join("\n")}" |
| 56 | + end |
| 57 | + end |
| 58 | + end |
| 59 | + |
| 60 | + def setup_security |
| 61 | + @session_id = SecureRandom.uuid |
| 62 | + unless config['key'].empty? || config['signature_scheme'].empty? |
| 63 | + unless config['signature_scheme'] =~ /\Ahmac-/ |
| 64 | + raise "Unknown signature_scheme: #{config['signature_scheme']}" |
| 65 | + end |
| 66 | + digest_algorithm = config['signature_scheme'][/\Ahmac-(.*)\Z/, 1] |
| 67 | + @hmac = OpenSSL::HMAC.new(config['key'], OpenSSL::Digest.new(digest_algorithm)) |
| 68 | + end |
| 69 | + end |
| 70 | + |
| 71 | + def send(socket_type, message_type, content) |
| 72 | + sock = check_socket_type(socket_type) |
| 73 | + idents = if socket_type == :reply && @last_recvd_msg |
| 74 | + @last_recvd_msg[:idents] |
| 75 | + else |
| 76 | + message_type == :stream ? "stream.#{content[:name]}" : message_type |
| 77 | + end |
| 78 | + header = { |
| 79 | + msg_type: message_type, |
| 80 | + msg_id: SecureRandom.uuid, |
| 81 | + username: 'kernel', |
| 82 | + session: @session_id, |
| 83 | + version: '5.0' |
| 84 | + } |
| 85 | + @adapter.send(sock, serialize(idents, header, content)) |
| 86 | + end |
| 87 | + |
| 88 | + def recv(socket_type) |
| 89 | + sock = check_socket_type(socket_type) |
| 90 | + data = @adapter.recv(sock) |
| 91 | + @last_recvd_msg = unserialize(data) |
| 92 | + end |
| 93 | + |
| 94 | + def recv_input |
| 95 | + sock = check_socket_type(:stdin) |
| 96 | + data = @adapter.recv(sock) |
| 97 | + unserialize(data)[:content]["value"] |
| 98 | + end |
| 99 | + |
| 100 | + private |
| 101 | + |
| 102 | + def check_socket_type(socket_type) |
| 103 | + case socket_type |
| 104 | + when :publish, :reply, :stdin |
| 105 | + @sockets[socket_type] |
| 106 | + else |
| 107 | + raise ArgumentError, "Invalid socket type #{socket_type}" |
| 108 | + end |
| 109 | + end |
| 110 | + |
| 111 | + def create_session_adapter(config, adapter_name) |
| 112 | + adapter_class = SessionAdapter.select_adapter_class(adapter_name) |
| 113 | + adapter_class.new(config) |
| 114 | + end |
| 115 | + end |
| 116 | +end |
0 commit comments