ryan.freumh.org
forked from
patrick.sirref.org/merry
Shells in OCaml
1(* Much of this code is from Eio_posix.
2
3 Copyright (C) 2021 Anil Madhavapeddy Copyright (C) 2022 Thomas Leonard
4
5 Permission to use, copy, modify, and distribute this software for any purpose
6 with or without fee is hereby granted, provided that the above copyright notice
7 and this permission notice appear in all copies.
8
9 THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
10 REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
11 FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
12 INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
13 LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
14 OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
15 PERFORMANCE OF THIS SOFTWARE. *)
16
17open Eio.Std
18
19module Process = struct
20 type t = {
21 pid : int;
22 exit_status : Unix.process_status Promise.t;
23 lock : Mutex.t;
24 }
25 (* When [lock] is unlocked, [exit_status] is resolved iff the process has been reaped. *)
26
27 let exit_status t = t.exit_status
28 let pid t = t.pid
29
30 module Fork_action = Eio_unix.Private.Fork_action
31
32 (* Read a (typically short) error message from a child process. *)
33 let rec read_response fd =
34 let buf = Bytes.create 256 in
35 match Eio_posix.Low_level.read fd buf 0 (Bytes.length buf) with
36 | 0 -> ""
37 | n -> Bytes.sub_string buf 0 n ^ read_response fd
38
39 let with_pipe fn =
40 Switch.run @@ fun sw ->
41 let r, w = Eio_posix.Low_level.pipe ~sw in
42 fn r w
43
44 let signal t signal =
45 (* We need the lock here so that one domain can't signal the process exactly as another is reaping it. *)
46 Mutex.lock t.lock;
47 Fun.protect ~finally:(fun () -> Mutex.unlock t.lock) @@ fun () ->
48 if not (Promise.is_resolved t.exit_status) then Unix.kill t.pid signal
49 (* else process has been reaped and t.pid is invalid *)
50
51 external eio_spawn :
52 Unix.file_descr -> Eio_unix.Private.Fork_action.c_action list -> int
53 = "caml_eio_posix_spawn"
54
55 (* Wait for [pid] to exit and then resolve [exit_status] to its status. *)
56 let reap t exit_status =
57 Eio.Condition.loop_no_mutex Eio_unix.Process.sigchld (fun () ->
58 Mutex.lock t.lock;
59 match Unix.waitpid [ WNOHANG ] t.pid with
60 | 0, _ ->
61 Mutex.unlock t.lock;
62 None (* Not ready; wait for next SIGCHLD *)
63 | p, status ->
64 assert (p = t.pid);
65 Promise.resolve exit_status status;
66 Mutex.unlock t.lock;
67 Some ())
68
69 let iter_switch ~f = function
70 | Merry.Types.Async -> ()
71 | Merry.Types.Switched sw -> f sw
72
73 let spawn ?delay_reap ~mode actions =
74 with_pipe @@ fun errors_r errors_w ->
75 Eio_unix.Private.Fork_action.with_actions actions @@ fun c_actions ->
76 iter_switch ~f:Switch.check mode;
77 let exit_status, set_exit_status = Promise.create () in
78 let t =
79 let pid =
80 Eio_unix.Fd.use_exn "errors-w" errors_w @@ fun errors_w ->
81 Eio.Private.Trace.with_span "spawn" @@ fun () ->
82 eio_spawn errors_w c_actions
83 in
84 Eio_unix.Fd.close errors_w;
85 { pid; exit_status; lock = Mutex.create () }
86 in
87 let () =
88 iter_switch
89 ~f:(fun sw ->
90 let hook =
91 Switch.on_release_cancellable sw (fun () ->
92 (* Kill process (if still running) *)
93 signal t Sys.sigkill;
94 (* The switch is being released, so either the daemon fiber got
95 cancelled or it hasn't started yet (and never will start). *)
96 if not (Promise.is_resolved t.exit_status) then
97 (* Do a (non-cancellable) waitpid here to reap the child. *)
98 reap t set_exit_status)
99 in
100 Fiber.fork_daemon ~sw (fun () ->
101 Option.iter Eio.Promise.await delay_reap;
102 reap t set_exit_status;
103 Switch.remove_hook hook;
104 `Stop_daemon))
105 mode
106 in
107 (* Check for errors starting the process. *)
108 match read_response errors_r with
109 | "" -> t (* Success! Execing the child closed [errors_w] and we got EOF. *)
110 | err -> failwith err
111end
112
113module Process_impl = struct
114 type t = Process.t
115 type tag = [ `Generic | `Unix ]
116
117 let pid = Process.pid
118
119 let await t =
120 match Eio.Promise.await @@ Process.exit_status t with
121 | Unix.WEXITED i -> `Exited i
122 | Unix.WSIGNALED i -> `Signaled i
123 | Unix.WSTOPPED _ -> assert false
124
125 let signal = Process.signal
126end
127
128let process =
129 let handler = Eio.Process.Pi.process (module Process_impl) in
130 fun proc -> Eio.Resource.T (proc, handler)
131
132let read_of_fd ~mode ~default ~to_close v =
133 match (mode, v) with
134 | Merry.Types.Async, _ | _, None -> default
135 | Merry.Types.Switched sw, Some f -> (
136 match Eio_unix.Resource.fd_opt f with
137 | Some fd -> fd
138 | None ->
139 let r, w = Eio_unix.pipe sw in
140 Fiber.fork ~sw (fun () ->
141 Eio.Flow.copy f w;
142 Eio.Flow.close w);
143 let r = Eio_unix.Resource.fd r in
144 to_close := r :: !to_close;
145 r)
146
147let write_of_fd ~mode ~default ~to_close v =
148 match (mode, v) with
149 | Merry.Types.Async, _ | _, None -> default
150 | Merry.Types.Switched sw, Some f -> (
151 match Eio_unix.Resource.fd_opt f with
152 | Some fd -> fd
153 | None ->
154 let r, w = Eio_unix.pipe sw in
155 Fiber.fork ~sw (fun () ->
156 Eio.Flow.copy r f;
157 Eio.Flow.close r);
158 let w = Eio_unix.Resource.fd w in
159 to_close := w :: !to_close;
160 w)
161
162let with_close_list fn =
163 let to_close = ref [] in
164 let close () = List.iter Eio_unix.Fd.close !to_close in
165 match fn to_close with
166 | x ->
167 close ();
168 x
169 | exception ex ->
170 let bt = Printexc.get_raw_backtrace () in
171 close ();
172 Printexc.raise_with_backtrace ex bt
173
174let get_executable ~args = function
175 | Some exe -> exe
176 | None -> (
177 match args with
178 | [] -> invalid_arg "Arguments list is empty and no executable given!"
179 | x :: _ -> x)
180
181let get_env = function Some e -> e | None -> Unix.environment ()
182
183external action_dups : unit -> Eio_unix.Private.Fork_action.fork_fn
184 = "eio_unix_fork_dups"
185
186let action_dups = action_dups ()
187
188let rec with_fds mapping k =
189 match mapping with
190 | [] -> k []
191 | (dst, src, _) :: xs ->
192 Eio_unix.Fd.use_exn "inherit_fds" src @@ fun src ->
193 with_fds xs @@ fun xs -> k ((dst, (Obj.magic src : int)) :: xs)
194
195let inherit_fds m =
196 let blocking =
197 m
198 |> List.filter_map (fun (dst, _, flags) ->
199 match flags with
200 | `Blocking -> Some (dst, true)
201 | `Nonblocking -> Some (dst, false)
202 | `Preserve_blocking -> None)
203 in
204 with_fds m @@ fun m ->
205 (* TODO: investigate -- the plan from Eio seems to also invert the list of redirections.
206 This is problematic for redirections, so we have copied the entire action here. *)
207 let plan = Eio_unix__.Inherit_fds.plan m |> List.rev in
208 Eio_unix.Private.Fork_action.
209 { run = (fun k -> k (Obj.repr (action_dups, plan, blocking))) }
210
211let spawn_unix () ?delay_reap ~mode ~fork_actions ?pgid ?uid ?gid ~env ~fds
212 ~executable ~cwd args =
213 let open Eio_posix in
214 let actions =
215 [
216 inherit_fds fds;
217 Low_level.Process.Fork_action.execve executable ~argv:(Array.of_list args)
218 ~env;
219 ]
220 in
221 let actions =
222 match pgid with
223 | None -> actions
224 | Some pgid -> Low_level.Process.Fork_action.setpgid pgid :: actions
225 in
226 let actions =
227 match uid with
228 | None -> actions
229 | Some uid -> Eio_unix.Private.Fork_action.setuid uid :: actions
230 in
231 let actions =
232 match gid with
233 | None -> actions
234 | Some gid -> Eio_unix.Private.Fork_action.setgid gid :: actions
235 in
236 let actions = actions @ fork_actions in
237 let with_actions cwd fn =
238 let ((dir, path) : Eio.Fs.dir_ty Eio.Path.t) = cwd in
239 match Eio_posix__.Fs.as_posix_dir dir with
240 | None -> Fmt.invalid_arg "cwd is not an OS directory!"
241 | Some dirfd ->
242 Switch.run ~name:"spawn_unix" @@ fun launch_sw ->
243 let cwd =
244 Eio_posix__.Err.run
245 (fun () ->
246 let flags = Low_level.Open_flags.(rdonly + directory) in
247 Low_level.openat ~sw:launch_sw ~mode:0 dirfd path flags)
248 ()
249 in
250 fn (Low_level.Process.Fork_action.fchdir cwd :: actions)
251 in
252 with_actions cwd @@ fun actions ->
253 process (Process.spawn ?delay_reap ~mode actions)
254
255let fd_equal_int fd i =
256 Eio_unix.Fd.use_exn "fd_equal_int" fd @@ fun ufd ->
257 let ufd_int = (Obj.magic ufd : int) in
258 Int.equal i ufd_int
259
260let pp_redirections ppf (i, fd, _) = Fmt.pf ppf "(%i,%a)" i Eio_unix.Fd.pp fd
261
262let run ~mode ?delay_reap _ ?stdin ?stdout ?stderr ?(fds = [])
263 ?(fork_actions = []) ~pgid ~cwd ?env ?executable args =
264 with_close_list @@ fun to_close ->
265 let check_fd n = function
266 | Merry.Types.Redirect (m, _, _) -> Int.equal n m
267 | Merry.Types.Close fd -> fd_equal_int fd n
268 in
269 let fd_exists n = List.exists (check_fd n) fds in
270 let std_fds =
271 (if fd_exists 0 then []
272 else
273 [
274 ( 0,
275 read_of_fd ~mode stdin ~default:Eio_unix.Fd.stdin ~to_close,
276 `Blocking );
277 ])
278 @ (if fd_exists 1 then []
279 else begin
280 [
281 ( 1,
282 write_of_fd ~mode stdout ~default:Eio_unix.Fd.stdout ~to_close,
283 `Blocking );
284 ]
285 end)
286 @
287 if fd_exists 2 then []
288 else
289 [
290 ( 2,
291 write_of_fd ~mode stderr ~default:Eio_unix.Fd.stderr ~to_close,
292 `Blocking );
293 ]
294 in
295 let need_close, fds =
296 List.fold_left
297 (fun (cs, fs) -> function
298 | Merry.Types.Redirect (a, b, c) -> (cs, (a, b, c) :: fs)
299 | Close fd -> (fd :: cs, fs))
300 ([], []) fds
301 |> fun (cs, fs) -> (List.rev cs, List.rev fs)
302 in
303 List.iter Eio_unix.Fd.close need_close;
304 let fds = std_fds @ fds in
305 let executable = get_executable executable ~args in
306 let env = get_env env in
307 spawn_unix ?delay_reap ~mode ~fork_actions ~cwd ~pgid ~fds ~env ~executable ()
308 args