Source file backend.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
open! Core
open! Async

let name = "Unauthenticated Async RPC"

module Settings = struct
  type t =
    | Only_local
    | Allow_remote
  [@@deriving bin_io, sexp]

  let only_allow_local_workers_will_fail_to_connect_if_remote = Only_local
  let unsafe_allow_unauthenticated_remote_workers = Allow_remote
end

let serve
  ?max_message_size
  ?buffer_age_limit
  ?handshake_timeout
  ?heartbeat_config
  ~implementations
  ~initial_connection_state
  ~where_to_listen
  (settings : Settings.t)
  =
  let where_to_listen =
    match settings with
    | Allow_remote -> where_to_listen
    | Only_local ->
      let port =
        Tcp.Where_to_listen.address where_to_listen |> Socket.Address.Inet.port
      in
      Tcp.Where_to_listen.bind_to Localhost (On_port port)
  in
  let make_transport fd ~max_message_size =
    Rpc.Transport.of_fd ?buffer_age_limit fd ~max_message_size
  in
  Rpc.Connection.serve
    ?max_message_size
    ~make_transport
    ?handshake_timeout
    ?heartbeat_config
    ~implementations
    ~initial_connection_state
    ~where_to_listen
    ()
;;

(* This is necessary binding to [localhost] requires also specifying [localhost] when
   connecting. Specifying the current box's hostname or ip doesn't work *)
let handle_where_to_connect where_to_connect (settings : Settings.t) =
  match settings with
  | Allow_remote -> return where_to_connect
  | Only_local ->
    let%map remote_address = Tcp.Where_to_connect.remote_address where_to_connect in
    let remote_port = Socket.Address.Inet.port remote_address in
    Tcp.Where_to_connect.of_inet_address
      (Socket.Address.Inet.create Unix.Inet_addr.localhost ~port:remote_port)
;;

let client
  ?implementations
  ?max_message_size
  ?buffer_age_limit
  ?handshake_timeout
  ?heartbeat_config
  ?description
  settings
  where_to_connect
  =
  let%bind where_to_connect = handle_where_to_connect where_to_connect settings in
  let make_transport fd ~max_message_size =
    Rpc.Transport.of_fd ?buffer_age_limit fd ~max_message_size
  in
  Rpc.Connection.client
    ?implementations
    ?max_message_size
    ~make_transport
    ?handshake_timeout
    ?heartbeat_config
    ?description
    where_to_connect
  |> Deferred.Or_error.of_exn_result
;;

let with_client
  ?implementations
  ?max_message_size
  ?buffer_age_limit
  ?handshake_timeout
  ?heartbeat_config
  settings
  where_to_connect
  f
  =
  let%bind where_to_connect = handle_where_to_connect where_to_connect settings in
  let make_transport fd ~max_message_size =
    Rpc.Transport.of_fd ?buffer_age_limit fd ~max_message_size
  in
  Rpc.Connection.with_client
    ?implementations
    ?max_message_size
    ~make_transport
    ?handshake_timeout
    ?heartbeat_config
    where_to_connect
    f
  |> Deferred.Or_error.of_exn_result
;;