banawa-chat/unikernel.ml

94 lines
3 KiB
OCaml

open Lwt.Syntax
type state =
{ env : (string, string) Hashtbl.t
; sigwinch : (int * int) Lwt_condition.t
; mutable size : int * int
}
module K = struct
open Cmdliner
let port =
let doc = Arg.info ~doc:"The TCP port for listening for SSH connections" ["port"] in
Arg.(value & opt int 22 doc)
let hostkey =
let doc = Arg.info ~doc:"SSH host key" ["hostkey"] in
Arg.(required & opt (some string) None doc)
end
module Main (_ : Mirage_crypto_rng_mirage.S) (T : Mirage_time.S) (M : Mirage_clock.MCLOCK) (Stack : Tcpip.Stack.V4V6) = struct
module Ssh = Banawa_mirage.Make(Stack.TCP)(T)(M)
module Nottui' = Nottui_mirage.Make(T)
let buffer = Rb.make 1024
let buffer_var = Lwd.var buffer
let callback flow stop t ~username r =
match r with
| Ssh.Pty_req { width; height; _ } ->
t.size <- (Int32.to_int width, Int32.to_int height);
Lwt.return_unit
| Ssh.Pty_set { width; height; _ } ->
Lwt_condition.broadcast t.sigwinch
(Int32.to_int width, Int32.to_int height);
Lwt.return_unit
| Ssh.Set_env _ -> Lwt.return_unit
| Ssh.Channel { cmd; ic=_; oc=_; ec } ->
let* () =
ec (Printf.ksprintf Cstruct.of_string
"Thanks for logging in! Currently, %S is unsupported\r\n\
Check back later." cmd)
in
let* () = Lwt_switch.turn_off stop in
Stack.TCP.close flow
| Ssh.Shell { ic; oc; ec=_ } ->
let ic () =
let+ r = ic () in
match r with
| `Data cs -> `Data (Cstruct.map (function '\r' -> '\n' | c -> c) cs)
| `Eof -> `Eof
in
let cursor = Lwd.var (0, 0) in
let message m =
let msg = Message.make ~nickname:username m in
Lwd.set buffer_var (Rb.push buffer msg; buffer);
in
let quit () =
let msg = Message.msgf "%s tried to quit, but it is not implemented" username in
Lwd.set buffer_var (Rb.push buffer msg; buffer);
in
Lwd.set buffer_var
(Rb.push buffer (Message.msgf "Welcome, %s!" username); buffer);
let ui =
let ( let* ) x f = Lwd.bind x ~f in
let* prompt = Prompt.make ~quit ~message cursor in
let* window = Window.make buffer_var in
Lwd.return (Nottui.Ui.vcat [window; prompt])
in
Lwt.join [
Nottui'.run ~cursor (t.size, t.sigwinch) ui ic oc;
]
let start _random _time _mtime stack port hostkey =
let hostkey =
match Awa.Keys.of_string hostkey with
| Ok k -> k
| Error `Msg e ->
Logs.err (fun m -> m "%s" e); exit 1
in
let server, msgs = Awa.Server.make hostkey in
Stack.TCP.listen (Stack.tcp stack) ~port
(fun flow ->
let stop = Lwt_switch.create () in
let state =
{ env = Hashtbl.create 0x10
; sigwinch = Lwt_condition.create ()
; size = (0, 0)
}
in
let _ssh = Ssh.spawn_server ~stop server msgs flow (callback flow stop state) in
Lwt.return_unit);
fst (Lwt.wait ())
end