Add on_connect_hook to serve_repl for customized sessions (#33)

For now, this allows the server to set a custom module for the client to evaluate
commands in.  This can be useful if you want
- The client to always start evaluating code in the server's module, for debugging
- Some minimal level of isolation between clients (a new anonymous module
  could be created for each client)

Author: c42f

Project.toml

@@ -1,7 +1,7 @@
 name = "RemoteREPL"
 uuid = "1bd9f7bb-701c-4338-bec7-ac987af7c555"
 authors = ["Chris Foster <[email protected]> and contributors"]
-version = "0.2.11"
+version = "0.2.12"
 
 [deps]
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"

src/server.jl

@@ -4,10 +4,14 @@ using REPL
 using Logging
 
 mutable struct ServerSideSession
+    socket
     display_properties::Dict
     in_module::Module
 end
 
+Base.isopen(session::ServerSideSession) = isopen(session.socket)
+Base.close(session::ServerSideSession)  = close(session.socket)
+
 function send_header(io, ser_version=Serialization.ser_version)
     write(io, PROTOCOL_MAGIC, PROTOCOL_VERSION)
     write(io, UInt32(ser_version))
@@ -93,9 +97,7 @@ function eval_message(session, messageid, messagebody)
     end
 end
 
-function evaluate_requests(request_chan, response_chan)
-    session = ServerSideSession(Dict(), Main)
-
+function evaluate_requests(session, request_chan, response_chan)
     while true
         try
             request = take!(request_chan)
@@ -174,15 +176,16 @@ function serialize_responses(socket, response_chan)
     end
 end
 
-# Serve a remote REPL session to a single client over `socket`.
-function serve_repl_session(socket)
+# Serve a remote REPL session to a single client
+function serve_repl_session(session)
+    socket = session.socket
     send_header(socket)
     @sync begin
         request_chan = Channel(1)
         response_chan = Channel(1)
 
         repl_backend = @async try
-            evaluate_requests(request_chan, response_chan)
+            evaluate_requests(session, request_chan, response_chan)
         catch exc
             @error "RemoteREPL backend crashed" exception=exc,catch_backtrace()
         finally
@@ -210,14 +213,18 @@ function serve_repl_session(socket)
 end
 
 """
-    serve_repl([address=Sockets.localhost,] port=$DEFAULT_PORT)
+    serve_repl([address=Sockets.localhost,] port=$DEFAULT_PORT; [on_client_connect=nothing])
     serve_repl(server)
 
 Start a REPL server listening on interface `address` and `port`. In normal
 operation `serve_repl()` serves REPL clients indefinitely (ie., it does not
 return), so you will generally want to launch it using `@async serve_repl()` to
 do other useful work at the same time.
 
+The hook `on_client_connect` may be supplied to modify the `ServerSideSession`
+for a client after each client connects. This can be used to define the default
+module in which the client evaluates commands.
+
 If you want to be able to stop the server you can pass an already-listening
 `server` object (the result of `Sockets.listen()`). The server can then be
 cancelled from another task using `close(server)` as necessary to control the
@@ -230,34 +237,38 @@ be used on open networks or multi-user machines where other users aren't
 trusted. For open networks, use the default `address=Sockets.localhost` and the
 automatic ssh tunnel support provided by the client-side `connect_repl()`.
 """
-function serve_repl(address=Sockets.localhost, port::Integer=DEFAULT_PORT)
+function serve_repl(address=Sockets.localhost, port::Integer=DEFAULT_PORT; kws...)
     server = listen(address, port)
     try
-        serve_repl(server)
+        serve_repl(server; kws...)
     finally
         close(server)
     end
 end
-serve_repl(port::Integer) = serve_repl(Sockets.localhost, port)
+serve_repl(port::Integer; kws...) = serve_repl(Sockets.localhost, port; kws...)
 
-function serve_repl(server::Base.IOServer)
-    open_sockets = Set()
+function serve_repl(server::Base.IOServer; on_client_connect=nothing)
+    open_sessions = Set{ServerSideSession}()
     @sync try
         while isopen(server)
             socket = accept(server)
-            push!(open_sockets, socket)
-            peer=getpeername(socket)
+            session = ServerSideSession(socket, Dict(), Main)
+            push!(open_sessions, session)
+            peer = getpeername(socket)
             @async try
-                serve_repl_session(socket)
+                if !isnothing(on_client_connect)
+                    on_client_connect(session)
+                end
+                serve_repl_session(session)
             catch exc
-                if !(exc isa EOFError && !isopen(socket))
+                if !(exc isa EOFError && !isopen(session))
                     @warn "Something went wrong evaluating client command" #=
                         =# exception=exc,catch_backtrace()
                 end
             finally
                 @info "REPL client exited" peer
-                close(socket)
-                pop!(open_sockets, socket)
+                close(session)
+                pop!(open_sessions, session)
             end
             @info "REPL client opened a connection" peer
         end
@@ -269,8 +280,8 @@ function serve_repl(server::Base.IOServer)
         @error "Unexpected server failure" isopen(server) exception=exc,catch_backtrace()
         rethrow()
     finally
-        for socket in open_sockets
-            close(socket)
+        for session in open_sessions
+            close(session)
         end
     end
 end

test/runtests.jl

@@ -47,6 +47,28 @@ end
     @test repl_prompt_text(fake_conn("ABC", DEFAULT_PORT, is_open=false)) == "julia@ABC [disconnected]> "
 end
 
+function wait_conn(host, port, use_ssh; max_tries=4)
+    for i=1:max_tries
+        try
+            return RemoteREPL.Connection(host=host, port=port,
+                                         tunnel=use_ssh ? :ssh : :none,
+                                         ssh_opts=`-o StrictHostKeyChecking=no`)
+        catch exc
+            if i == max_tries
+                rethrow()
+            end
+            # Server not yet started - continue waiting
+            sleep(2)
+        end
+    end
+end
+
+function runcommand_unwrap(conn, cmdstr)
+    result = RemoteREPL.run_remote_repl_command(conn, IOBuffer(), cmdstr)
+    # Unwrap Text for testing purposes
+    return result isa Text ? result.content : result
+end
+
 # Connect to a non-default loopback address to test SSH integration
 test_interface = ip"127.111.111.111"
 
@@ -56,7 +78,7 @@ use_ssh = if "use_ssh=true" in ARGS
 elseif "use_ssh=false" in ARGS
     false
 else
-    # Autodetct
+    # Autodetect
     try
         socket = Sockets.connect(test_interface, 22)
         # https://tools.ietf.org/html/rfc4253#section-4.2
@@ -78,34 +100,15 @@ server_proc = run(`$(Base.julia_cmd()) -e "using Sockets; using RemoteREPL; serv
 
 try
 
-@testset "RemoteREPL.jl" begin
-    local conn = nothing
-    max_tries = 4
-    for i=1:max_tries
-        try
-            conn = RemoteREPL.Connection(host=test_interface, port=test_port,
-                                         tunnel=use_ssh ? :ssh : :none,
-                                         ssh_opts=`-o StrictHostKeyChecking=no`)
-            break
-        catch exc
-            if i == max_tries
-                rethrown()
-            end
-            # Server not yet started - continue waiting
-            sleep(2)
-        end
-    end
+@testset "Client server tests" begin
+    conn = wait_conn(test_interface, test_port, use_ssh)
     @assert isopen(conn)
 
     # Some basic tests of the transport and server side and partial client side.
     #
     # More full testing of the client code would requires some tricky mocking
     # of the REPL environment.
-    function runcommand(cmdstr)
-        result = RemoteREPL.run_remote_repl_command(conn, IOBuffer(), cmdstr)
-        # Unwrap Text for testing purposes
-        return result isa Text ? result.content : result
-    end
+    runcommand(cmdstr) = runcommand_unwrap(conn, cmdstr)
 
     @test runcommand("asdf = 42") == "42"
     @test runcommand("Main.asdf") == "42"
@@ -262,3 +265,21 @@ end
 finally
     kill(server_proc)
 end
+
+
+test_port = RemoteREPL.find_free_port(Sockets.localhost)
+server_proc = run(```$(Base.julia_cmd()) -e "using Sockets; using RemoteREPL; module EvalInMod ; end;
+                  serve_repl($test_port, on_client_connect=sess->sess.in_module=EvalInMod)"```, wait=false)
+try
+
+@testset "on_client_connect" begin
+    conn = wait_conn(test_interface, test_port, use_ssh)
+
+    runcommand(cmdstr) = runcommand_unwrap(conn, cmdstr)
+
+    @test runcommand("@__MODULE__") == "Main.EvalInMod"
+end
+
+finally
+    kill(server_proc)
+end