Shells in OCaml
at main 308 lines 10 kB view raw
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