From 9cf16098fb20d9252704bd4b8f1e53624dab0735 Mon Sep 17 00:00:00 2001 From: Calascibetta Romain Date: Sat, 4 Mar 2023 13:22:14 +0100 Subject: [PATCH] Add two events Pty/Set_env/Start_shell into the server and export them into the mirage layer --- lib/server.ml | 15 +++++------ lwt/awa_lwt.ml | 1 + mirage/awa_mirage.ml | 59 +++++++++++++++++++++++++++-------------- mirage/awa_mirage.mli | 28 +++++++++---------- mirage/dune | 2 +- test/awa_test_server.ml | 5 ++-- 6 files changed, 64 insertions(+), 46 deletions(-) diff --git a/lib/server.ml b/lib/server.ml index b3e99bc..af88834 100644 --- a/lib/server.ml +++ b/lib/server.ml @@ -22,6 +22,9 @@ type event = | Channel_data of (int32 * Cstruct.t) | Channel_eof of int32 | Disconnected of string + | Pty of (string * int32 * int32 * int32 * int32 * string) + | Set_env of (string * string) + | Start_shell of int32 type t = { client_version : string option; (* Without crlf *) @@ -253,12 +256,6 @@ let input_channel_request t recp_channel want_reply data = else make_noreply t in - let success t = - if want_reply then - make_reply t (Msg_channel_success recp_channel) - else - make_noreply t - in let event t event = if want_reply then make_reply_with_event t (Msg_channel_success recp_channel) event @@ -266,10 +263,10 @@ let input_channel_request t recp_channel want_reply data = make_event t event in let handle t c = function - | Pty_req _ -> success t + | Pty_req v -> event t (Pty v) | X11_req _ -> fail t - | Env (_key, _value) -> success t (* TODO implement me *) - | Shell -> fail t + | Env v -> event t (Set_env v) + | Shell -> event t (Start_shell c) | Exec cmd -> event t (Channel_exec (c, cmd)) | Subsystem cmd -> event t (Channel_subsystem (c, cmd)) | Window_change _ -> fail t diff --git a/lwt/awa_lwt.ml b/lwt/awa_lwt.ml index 47f57ef..33035b6 100644 --- a/lwt/awa_lwt.ml +++ b/lwt/awa_lwt.ml @@ -157,6 +157,7 @@ let rec nexus t fd server input_buffer = let c = { cmd; id; sshin_mbox; exec_thread } in let t = { t with channels = c :: t.channels } in nexus t fd server input_buffer + | _ -> nexus t fd server input_buffer let spawn_server server msgs fd exec_callback = let t = { exec_callback; diff --git a/mirage/awa_mirage.ml b/mirage/awa_mirage.ml index 45f6581..dfd627c 100644 --- a/mirage/awa_mirage.ml +++ b/mirage/awa_mirage.ml @@ -161,24 +161,26 @@ module Make (F : Mirage_flow.S) (T : Mirage_time.S) (M : Mirage_clock.MCLOCK) = | Sshout of (int32 * Cstruct.t) | Ssherr of (int32 * Cstruct.t) - type sshin_msg = [ - | `Data of Cstruct.t - | `Eof - ] - type channel = { - cmd : string; + cmd : string option; id : int32; - sshin_mbox : sshin_msg Lwt_mvar.t; + sshin_mbox : Cstruct.t Mirage_flow.or_eof Lwt_mvar.t; exec_thread : unit Lwt.t; } - type exec_callback = - string -> (* cmd *) - (unit -> sshin_msg Lwt.t) -> (* sshin *) - (Cstruct.t -> unit Lwt.t) -> (* sshout *) - (Cstruct.t -> unit Lwt.t) -> (* ssherr *) - unit Lwt.t + type request = + | Pty_req of { width : int32; height : int32; max_width : int32; max_height : int32; term : string } + | Pty_set of { width : int32; height : int32; max_width : int32; max_height : int32 } + | Set_env of { key : string; value : string } + | Channel of { command : string + ; ic : unit -> Cstruct.t Mirage_flow.or_eof Lwt.t + ; oc : Cstruct.t -> unit Lwt.t + ; ec : Cstruct.t -> unit Lwt.t } + | Shell of { ic : unit -> Cstruct.t Mirage_flow.or_eof Lwt.t + ; oc : Cstruct.t -> unit Lwt.t + ; ec : Cstruct.t -> unit Lwt.t } + + type exec_callback = request -> unit Lwt.t type t = { exec_callback : exec_callback; (* callback to run on exec *) @@ -285,6 +287,12 @@ module Make (F : Mirage_flow.S) (T : Mirage_time.S) (M : Mirage_clock.MCLOCK) = >>= fun server -> match event with | None -> nexus t fd server input_buffer (List.append pending_promises [ Lwt_mvar.take t.nexus_mbox ]) + | Some Awa.Server.Pty (term, width, height, max_width, max_height, _modes) -> + t.exec_callback (Pty_req { width; height; max_width; max_height; term; }) >>= fun () -> + nexus t fd server input_buffer pending_promises + | Some Awa.Server.Set_env (key, value) -> + t.exec_callback (Set_env { key; value; }) >>= fun () -> + nexus t fd server input_buffer pending_promises | Some Awa.Server.Disconnected _ -> Lwt_list.iter_p sshin_eof t.channels >>= fun () -> Lwt.return t @@ -298,17 +306,28 @@ module Make (F : Mirage_flow.S) (T : Mirage_time.S) (M : Mirage_clock.MCLOCK) = | None -> Lwt.return_unit) >>= fun () -> nexus t fd server input_buffer (List.append pending_promises [ Lwt_mvar.take t.nexus_mbox ]) - | Some Awa.Server.Channel_subsystem (id, cmd) (* same as exec *) - | Some Awa.Server.Channel_exec (id, cmd) -> + | Some Awa.Server.Channel_subsystem (id, command) (* same as exec *) + | Some Awa.Server.Channel_exec (id, command) -> (* Create an input box *) let sshin_mbox = Lwt_mvar.create_empty () in (* Create a callback for each mbox *) - let sshin () = Lwt_mvar.take sshin_mbox in - let sshout id buf = Lwt_mvar.put t.nexus_mbox (Sshout (id, buf)) in - let ssherr id buf = Lwt_mvar.put t.nexus_mbox (Ssherr (id, buf)) in + let ic () = Lwt_mvar.take sshin_mbox in + let oc id buf = Lwt_mvar.put t.nexus_mbox (Sshout (id, buf)) in + let ec id buf = Lwt_mvar.put t.nexus_mbox (Ssherr (id, buf)) in + (* Create the execution thread *) + let exec_thread = t.exec_callback (Channel { command; ic; oc= oc id; ec= ec id; }) in + let c = { cmd= Some command; id; sshin_mbox; exec_thread } in + let t = { t with channels = c :: t.channels } in + nexus t fd server input_buffer (List.append pending_promises [ Lwt_mvar.take t.nexus_mbox ]) + | Some (Awa.Server.Start_shell id) -> + let sshin_mbox = Lwt_mvar.create_empty () in + (* Create a callback for each mbox *) + let ic () = Lwt_mvar.take sshin_mbox in + let oc id buf = Lwt_mvar.put t.nexus_mbox (Sshout (id, buf)) in + let ec id buf = Lwt_mvar.put t.nexus_mbox (Ssherr (id, buf)) in (* Create the execution thread *) - let exec_thread = t.exec_callback cmd sshin (sshout id) (ssherr id) in - let c = { cmd; id; sshin_mbox; exec_thread } in + let exec_thread = t.exec_callback (Shell { ic; oc= oc id; ec= ec id; }) in + let c = { cmd= None; id; sshin_mbox; exec_thread } in let t = { t with channels = c :: t.channels } in nexus t fd server input_buffer (List.append pending_promises [ Lwt_mvar.take t.nexus_mbox ]) diff --git a/mirage/awa_mirage.mli b/mirage/awa_mirage.mli index e2f1536..2c439a5 100644 --- a/mirage/awa_mirage.mli +++ b/mirage/awa_mirage.mli @@ -26,19 +26,21 @@ module Make (F : Mirage_flow.S) (T : Mirage_time.S) (M : Mirage_clock.MCLOCK) : Awa.Hostkey.priv -> Awa.Ssh.channel_request -> FLOW.flow -> (flow, error) result Lwt.t - type t + type t - type sshin_msg = [ - | `Data of Cstruct.t - | `Eof - ] + type request = + | Pty_req of { width : int32; height : int32; max_width : int32; max_height : int32; term : string } + | Pty_set of { width : int32; height : int32; max_width : int32; max_height : int32 } + | Set_env of { key : string; value : string } + | Channel of { command : string + ; ic : unit -> Cstruct.t Mirage_flow.or_eof Lwt.t + ; oc : Cstruct.t -> unit Lwt.t + ; ec : Cstruct.t -> unit Lwt.t } + | Shell of { ic : unit -> Cstruct.t Mirage_flow.or_eof Lwt.t + ; oc : Cstruct.t -> unit Lwt.t + ; ec : Cstruct.t -> unit Lwt.t } - type exec_callback = - string -> (* cmd *) - (unit -> sshin_msg Lwt.t) -> (* sshin *) - (Cstruct.t -> unit Lwt.t) -> (* sshout *) - (Cstruct.t -> unit Lwt.t) -> (* ssherr *) - unit Lwt.t + type exec_callback = request -> unit Lwt.t val spawn_server : ?stop:Lwt_switch.t -> Awa.Server.t -> Awa.Ssh.message list -> F.flow -> exec_callback -> t Lwt.t @@ -62,6 +64,4 @@ module Make (F : Mirage_flow.S) (T : Mirage_time.S) (M : Mirage_clock.MCLOCK) : {b NOTE}: Even if the [ssh_channel_handler] is fulfilled, [spawn_server] continues to handle SSH channels. Only [stop] can really stop the internal SSH channels handler. *) - -end - with module FLOW = F +end with module FLOW = F diff --git a/mirage/dune b/mirage/dune index b471bea..b04f5aa 100644 --- a/mirage/dune +++ b/mirage/dune @@ -2,4 +2,4 @@ (name awa_mirage) (public_name awa-mirage) (wrapped false) - (libraries awa mirage-flow mirage-clock mirage-time duration lwt mtime logs)) + (libraries awa mirage-flow mirage-clock mirage-time duration lwt mtime)) diff --git a/test/awa_test_server.ml b/test/awa_test_server.ml index b25e639..b05c833 100644 --- a/test/awa_test_server.ml +++ b/test/awa_test_server.ml @@ -90,7 +90,7 @@ let rec serve t cmd = | Channel_subsystem (id, exec) (* same as exec *) | Channel_exec (id, exec) -> printf "channel exec %s\n%!" exec; - match exec with + begin match exec with | "suicide" -> let* _ = Driver.disconnect t in Ok () @@ -104,7 +104,8 @@ let rec serve t cmd = let* t = Driver.send_channel_data t id (Cstruct.of_string m) in printf "%s\n%!" m; let* t = Driver.disconnect t in - serve t cmd + serve t cmd end + | _ -> failwith "Invalid SSH event" let user_db = (* User foo auths by passoword *)