Logs_lwt.info (fun m -> m "Starting wsproxy") >>= fun () ->

let fd_sock = Lwt_unix.stdin in

let () = Lwt_unix.listen fd_sock 5 in

let rec loop () =

let ensure_close = function

| [] ->

Lwt.return_unit

| fds ->

Logs_lwt.warn (fun m -> m "Closing %d excess fds" (List.length fds))

>>= fun () ->

List.iter (fun fd -> try Unix.close fd with _ -> ()) fds ;

Lwt.return_unit

in

Lwt.catch

(fun () ->

Lwt_unix.accept fd_sock >>= fun (fd_sock', _) ->

(* Background thread per connection *)

let (_ : unit Lwt.t) =

let buffer = Bytes.make 16384 '\000' in

with_fd fd_sock' ~callback:(fun fd ->

let io_vectors = Lwt_unix.IO_vectors.create () in

Lwt_unix.IO_vectors.append_bytes io_vectors buffer 0 16384 ;

Lwt_unix.Versioned.recv_msg_2 ~socket:fd ~io_vectors)

>>= fun (len, newfds) ->

match newfds with

| [] ->

Logs_lwt.warn (fun m ->

m "No fd to start a connection: not proxying")

| ufd :: ufds ->

ensure_close ufds >>= fun () ->

with_fd (Lwt_unix.of_unix_file_descr ufd) ~callback:(fun fd ->

Logs_lwt.debug (fun m -> m "About to start connection")

>>= fun () ->

Lwt_unix.setsockopt fd Lwt_unix.SO_KEEPALIVE true ;

let msg = Bytes.(to_string @@ sub buffer 0 len) in

handler fd msg)

in

loop ())

(fun e ->

Logs_lwt.err (fun m -> m "Caught exception: %s" (Printexc.to_string e))

>>= fun () -> Lwt.return_unit)

>>= fun () -> loop ()

in

let open LwtWsIteratee in

let open Lwt_support in

( match protocol with

| "hixie76" ->

Logs_lwt.debug (fun m -> m "Old-style (hixie76) protocol") >>= fun () ->

Lwt.return (wsframe_old, wsunframe_old)

| "hybi10" ->

Logs_lwt.debug (fun m -> m "New-style (hybi10) protocol") >>= fun () ->

Lwt.return (wsframe, wsunframe)

| _ ->

Logs_lwt.warn (fun m -> m "Unknown protocol, fallback to hybi10")

>>= fun () -> Lwt.return (wsframe, wsunframe)

)

>>= fun (frame, unframe) ->

with_open_connection_fd addr ~callback:(fun localfd ->

let session_id = Uuidm.v `V4 |> Uuidm.to_string in

Logs_lwt.debug (fun m -> m "Starting proxy session %s" session_id)

>>= fun () ->

let thread1 =

lwt_fd_enumerator localfd (frame (writer (really_write fd) "thread1"))

>>= fun _ -> Lwt.return_unit

in

let thread2 =

lwt_fd_enumerator fd (unframe (writer (really_write localfd) "thread2"))

>>= fun _ -> Lwt.return_unit

in

(* closing the connection in one of the threads above in general leaves the other pending forever,

Lwt.choose [thread1; thread2] >>= fun () ->

let socket = Re.Str.regexp "^/var/run/xen/vnc-[0-9]+$"

let port = Re.Str.regexp "^[0-9]+$"

Logs_lwt.debug (fun m -> m "Got msg: '%s'" msg) >>= fun () ->

match Re.Str.(split @@ regexp "[:]") msg with

| ([protocol; _; path] | [protocol; path])

when Re.Str.string_match RX.socket path 0 ->

let addr = Unix.ADDR_UNIX path in

proxy sock addr protocol

| ([protocol; _; sport] | [protocol; sport])

when Re.Str.string_match RX.port sport 0 ->

let localhost = Unix.inet_addr_loopback in

let addr = Unix.ADDR_INET (localhost, int_of_string sport) in

proxy sock addr protocol

| _ ->

Logs_lwt.warn (fun m ->

let buf_fmt ~like =

let b = Buffer.create 512 in

( Fmt.with_buffer ~like b

, fun () ->

let m = Buffer.contents b in

Buffer.reset b ; m )

in

let app, app_flush = buf_fmt ~like:Fmt.stdout in

let dst, dst_flush = buf_fmt ~like:Fmt.stderr in

let reporter = Logs_fmt.reporter ~app ~dst () in

let report src level ~over k msgf =

let k () =

let write () =

match level with

| Logs.App ->

Lwt_io.write Lwt_io.stdout (app_flush ())

| _ ->

Lwt_io.write Lwt_io.stderr (dst_flush ())

in

let unblock () = over () ; Lwt.return_unit in

Lwt.finalize write unblock |> Lwt.ignore_result ;

k ()

in

reporter.Logs.report src level ~over:(fun () -> ()) k msgf

in