Source file CCUnix.ml

(* This file is free software, part of containers. See file "license" for more details. *)

(** {1 High-level Functions on top of Unix} *)

type 'a or_error = ('a, string) result
type 'a gen = unit -> 'a option

(** {2 Calling Commands} *)

let int_of_process_status = function
  | Unix.WEXITED i | Unix.WSIGNALED i | Unix.WSTOPPED i -> i

let str_exists s p =
  let rec f s p i =
    if i = String.length s then
      false
    else
      p s.[i] || f s p (i + 1)
  in
  f s p 0

let rec iter_gen f g =
  match g () with
  | None -> ()
  | Some x ->
    f x;
    iter_gen f g

let finally_ f x ~h =
  try
    let y = f x in
    ignore (h ());
    y
  with e ->
    ignore (h ());
    raise e

(* print a string, but escaped if required *)
let escape_str s =
  if
    str_exists s (function
      | ' ' | '"' | '\'' | '\n' | '\t' -> true
      | _ -> false)
  then (
    let buf = Buffer.create (String.length s) in
    Buffer.add_char buf '\'';
    String.iter
      (function
        | '\'' -> Buffer.add_string buf "'\\''"
        | c -> Buffer.add_char buf c)
      s;
    Buffer.add_char buf '\'';
    Buffer.contents buf
  ) else
    s

let read_all ?(size = 1024) ic =
  let buf = ref (Bytes.create size) in
  let len = ref 0 in
  try
    while true do
      (* resize *)
      if !len = Bytes.length !buf then buf := Bytes.extend !buf 0 !len;
      assert (Bytes.length !buf > !len);
      let n = input ic !buf !len (Bytes.length !buf - !len) in
      len := !len + n;
      if n = 0 then raise Exit
      (* exhausted *)
    done;
    assert false (* never reached*)
  with Exit -> Bytes.sub_string !buf 0 !len

type call_result =
  < stdout : string
  ; stderr : string
  ; status : Unix.process_status
  ; errcode : int  (** Extracted from status *) >

let kbprintf' buf fmt k = Printf.kbprintf k buf fmt

let call_full_inner ?(bufsize = 2048) ?(stdin = `Str "")
    ?(env = Unix.environment ()) ~f cmd =
  (* render the command *)
  let buf = Buffer.create 256 in
  kbprintf' buf cmd (fun buf ->
      let cmd = Buffer.contents buf in
      let oc, ic, errc = Unix.open_process_full cmd env in
      (* send stdin *)
      (match stdin with
      | `Str s -> output_string ic s
      | `Gen g -> iter_gen (output_string ic) g);
      close_out ic;
      (* read out and err *)
      let out = read_all ~size:bufsize oc in
      let err = read_all ~size:bufsize errc in
      let status = Unix.close_process_full (oc, ic, errc) in
      f (out, err, status))

let call_full ?bufsize ?stdin ?env cmd =
  call_full_inner ?bufsize ?stdin ?env cmd ~f:(fun (out, err, status) ->
      object
        method stdout = out
        method stderr = err
        method status = status
        method errcode = int_of_process_status status
      end)

let call ?bufsize ?stdin ?env cmd =
  call_full_inner ?bufsize ?stdin ?env cmd ~f:(fun (out, err, status) ->
      out, err, int_of_process_status status)

let call_stdout ?bufsize ?stdin ?env cmd =
  call_full_inner ?bufsize ?stdin ?env cmd ~f:(fun (out, _err, _status) -> out)

type line = string

type async_call_result =
  < stdout : line gen
  ; stderr : line gen
  ; stdin : line -> unit (* send a line *)
  ; close_in : unit (* close stdin *)
  ; close_err : unit
  ; close_out : unit
  ; close_all : unit (* close all 3 channels *)
  ; wait : Unix.process_status (* block until the process ends *)
  ; wait_errcode : int
  (* block until the process ends, then extract errcode *) >

let async_call ?(env = Unix.environment ()) cmd =
  (* render the command *)
  let buf = Buffer.create 256 in
  kbprintf' buf cmd (fun buf ->
      let cmd = Buffer.contents buf in
      let oc, ic, errc = Unix.open_process_full cmd env in
      object (self)
        method stdout () = try Some (input_line oc) with End_of_file -> None
        method stderr () = try Some (input_line errc) with End_of_file -> None

        method stdin l =
          output_string ic l;
          output_char ic '\n'

        method close_in = close_out ic
        method close_out = close_in oc
        method close_err = close_in errc

        method close_all =
          close_out ic;
          close_in oc;
          close_in errc;
          ()

        method wait = Unix.close_process_full (oc, ic, errc)
        method wait_errcode = int_of_process_status self#wait
      end)

let stdout x = x#stdout
let stderr x = x#stderr
let status x = x#status
let errcode x = x#errcode

let with_in ?(mode = 0o644) ?(flags = []) file ~f =
  let fd = Unix.openfile file (Unix.O_RDONLY :: flags) mode in
  let ic = Unix.in_channel_of_descr fd in
  finally_ f ic ~h:(fun () -> Unix.close fd)

let with_out ?(mode = 0o644) ?(flags = [ Unix.O_CREAT; Unix.O_TRUNC ]) file ~f =
  let fd = Unix.openfile file (Unix.O_WRONLY :: flags) mode in
  let oc = Unix.out_channel_of_descr fd in
  finally_ f oc ~h:(fun () ->
      flush oc;
      Unix.close fd)

let with_process_in cmd ~f =
  let ic = Unix.open_process_in cmd in
  finally_ f ic ~h:(fun () -> ignore (Unix.close_process_in ic))

let with_process_out cmd ~f =
  let oc = Unix.open_process_out cmd in
  finally_ f oc ~h:(fun () -> ignore (Unix.close_process_out oc))

type process_full =
  < stdin : out_channel
  ; stdout : in_channel
  ; stderr : in_channel
  ; close : Unix.process_status >

let with_process_full ?env cmd ~f =
  let env =
    match env with
    | None -> Unix.environment ()
    | Some e -> e
  in
  let oc, ic, err = Unix.open_process_full cmd env in
  let close = lazy (Unix.close_process_full (oc, ic, err)) in
  let p =
    object
      method stdin = ic
      method stdout = oc
      method stderr = err
      method close = Lazy.force close
    end
  in
  finally_ f p ~h:(fun () -> p#close)

let with_connection addr ~f =
  let ic, oc = Unix.open_connection addr in
  finally_ (fun () -> f ic oc) () ~h:(fun () -> Unix.shutdown_connection ic)

(* make sure that we are a session leader; that is, our children die if we die *)
let ensure_session_leader =
  let thunk =
    lazy (if (not Sys.win32) && not Sys.cygwin then ignore (Unix.setsid ()))
  in
  fun () -> Lazy.force thunk

exception ExitServer

(* version of {!Unix.establish_server} that doesn't fork *)
let establish_server sockaddr ~f =
  let sock =
    Unix.socket (Unix.domain_of_sockaddr sockaddr) Unix.SOCK_STREAM 0
  in
  Unix.setsockopt sock Unix.SO_REUSEADDR true;
  Unix.bind sock sockaddr;
  Unix.listen sock 5;
  let continue = ref true in
  while !continue do
    try
      let s, _ = Unix.accept sock in
      let ic = Unix.in_channel_of_descr s in
      let oc = Unix.out_channel_of_descr s in
      ignore (f ic oc)
    with ExitServer -> continue := false
  done

(** {2 Locking} *)

let with_file_lock ~kind filename f =
  let lock_file =
    Unix.openfile filename [ Unix.O_CREAT; Unix.O_WRONLY ] 0o644
  in
  let lock_action =
    match kind with
    | `Read -> Unix.F_RLOCK
    | `Write -> Unix.F_LOCK
  in
  Unix.lockf lock_file lock_action 0;
  try
    let x = f () in
    Unix.lockf lock_file Unix.F_ULOCK 0;
    Unix.close lock_file;
    x
  with e ->
    Unix.lockf lock_file Unix.F_ULOCK 0;
    Unix.close lock_file;
    raise e

module Infix = struct
  let ( ?| ) fmt = call_full fmt
  let ( ?|& ) fmt = async_call fmt
end

include Infix

(** {2 Temporary directory} *)

let rand_digits_ =
  let st = lazy (Random.State.make_self_init ()) in
  fun () ->
    let rand = Random.State.bits (Lazy.force st) land 0xFFFFFF in
    Printf.sprintf "%06x" rand

let rmdir_ dir = try ignore (Sys.command ("rm -r " ^ dir) : int) with _ -> ()

let with_temp_dir ?(mode = 0o700) ?dir pat (f : string -> 'a) : 'a =
  let dir =
    match dir with
    | Some d -> d
    | None -> Filename.get_temp_dir_name ()
  in
  let raise_err msg = raise (Sys_error msg) in
  let rec loop count =
    if count < 0 then
      raise_err "mk_temp_dir: too many failing attempts"
    else (
      let dir = Filename.concat dir (pat ^ rand_digits_ ()) in
      match Unix.mkdir dir mode with
      | () -> finally_ f dir ~h:(fun () -> rmdir_ dir)
      | exception Unix.Unix_error (Unix.EEXIST, _, _) -> loop (count - 1)
      | exception Unix.Unix_error (Unix.EINTR, _, _) -> loop count
      | exception Unix.Unix_error (e, _, _) ->
        raise_err ("mk_temp_dir: " ^ Unix.error_message e)
    )
  in
  loop 1000