Source file env.ml

open Import

module Working_dir_spec = struct
  type t =
    | Inherit
    | Path of string
    | Physical of { path : string
                  ; fd   : Unix.file_descr
                  }
end

module Working_dir = struct
  (* Physical working directories are ref-counted so that the file descriptor can be
     closed as soon as it is not needed anymore. This allow to optimize a [chdir] in tail
     position. *)
  module Physical = struct
    type t =
      | Win32 (* Physical path is not supported on Windows *)
      | Unix of { fd            : Unix.file_descr
                ; mutable users : int
                ; mutex         : Mutex.t
                }

    let create_unix fd =
      Unix { fd; users = 1; mutex = Mutex.create () }

    let addref = function
      | Win32 -> ()
      | Unix u ->
        Mutex.lock u.mutex;
        let users = u.users + 1 in
        u.users <- users;
        Mutex.unlock u.mutex;
        (* One must not try to reference a physical path that has been released *)
        assert (users >= 2)

    let deref = function
      | Win32 -> ()
      | Unix u ->
        Mutex.lock u.mutex;
        let users = u.users - 1 in
        u.users <- users;
        Mutex.unlock u.mutex;
        assert (users >= 0);
        if users = 0 then try Unix.close u.fd with _ -> ()

    let map_unix t ~f =
      match t with
      | Win32 -> Win32
      | Unix u -> create_unix (f u.fd)
  end

  type t =
    { logical       : string
    ; physical      : Physical.t
    }

  let create (spec : Working_dir_spec.t) =
    let logical, physical =
      match spec with
      | Inherit -> (Sys.getcwd (), None)
      | Path p -> (p, None)
      | Physical { path; fd } -> (path, Some fd)
    in
    let physical : Physical.t =
      match physical, Sys.win32 with
      | Some _, true ->
        invalid_arg "Shexp_exec.Std.Env.create: cannot use a physical \
                     working directory on Windows";
      | Some fd, false ->
        Physical.create_unix fd
      | None, true ->
        Win32
      | None, false ->
        Physical.create_unix (Unix.openfile logical [O_CLOEXEC; O_RDONLY] 0)
    in
    { logical; physical }

  let split_relative_dir dir =
    let rec loop dir acc =
      let acc = Filename.basename dir :: acc in
      let dir = Filename.dirname dir in
      if dir = Filename.current_dir_name then
        acc
      else
        loop dir acc
    in
    loop dir []

  let chdir_logical dir relative_dir =
    let rec loop dir components ~symlinks_followed =
      match components with
      | [] -> dir
      | component :: rest ->
        if component = Filename.current_dir_name then
          loop dir rest ~symlinks_followed
        else if component = Filename.parent_dir_name then
          match Unix.readlink dir with
          | exception (Unix.Unix_error (EINVAL, _, _)) ->
            loop (Filename.dirname dir) rest ~symlinks_followed
          | target ->
            let symlinks_followed = symlinks_followed + 1 in
            if symlinks_followed = 1000 then raise Exit;
            (* The current dir is a symlink, we need to resolve it before we can interpret
               the ".." *)
            if Filename.is_relative target then
              loop dir (split_relative_dir target @ components) ~symlinks_followed
            else
              loop target components ~symlinks_followed
        else
          loop (dir ^/ component) rest ~symlinks_followed
    in
    let components = split_relative_dir relative_dir in
    try
      loop dir components ~symlinks_followed:0
    with Exit ->
      raise (Unix.Unix_error (ELOOP, "chdir", dir ^/ relative_dir))

  let chdir t dir =
    if dir = Filename.current_dir_name then
      t
    else
      let logical =
        if Filename.is_relative dir then
          chdir_logical t.logical dir
        else
          dir
      in
      let physical =
        Physical.map_unix t.physical ~f:(fun fd ->
          Posixat.openat ~dir:fd ~flags:[O_RDONLY] ~path:dir ~perm:0)
      in
      { logical; physical }

  let addref t = Physical.addref t.physical
  let deref t = Physical.deref t.physical
end

module Uenv = struct
  type t =
    { entries : string SMap.t
    ; (* The [PATH] variable pre-splitted *)
      path    : string list
    }

  let get t var = SMap.lookup var t.entries

  let split_path str =
    let len = String.length str in
    let is_sep c = c = ':' in
    let rec loop i j =
      if j = len then
        [String.sub str ~pos:i ~len:(j - i)]
      else
      if is_sep str.[j] then
        String.sub str ~pos:i ~len:(j - i) :: loop (j + 1) (j + 1)
      else
        loop i (j + 1)
    in
    loop 0 0

  let set t var value =
    let entries = SMap.add t.entries ~key:var ~data:value in
    if var = "PATH" then
      { entries; path = split_path value }
    else
      { t with entries }

  let unset t var =
    let entries = SMap.remove var t.entries in
    if var = "PATH" then
      { entries; path = [] }
    else
      { t with entries }

  let create initial ~pwd =
    let entries =
      match initial with
      | Some l ->
        List.fold_left l ~init:SMap.empty ~f:(fun acc (key, data) ->
          SMap.add acc ~key ~data)
      | None ->
        Array.fold_left (Unix.environment ()) ~init:SMap.empty ~f:(fun acc s ->
          let i = String.index s '=' in
          let key = String.sub s ~pos:0 ~len:i in
          let data = String.sub s ~pos:(i + 1) ~len:(String.length s - i - 1) in
          SMap.add acc ~key ~data)
    in
    let entries = SMap.add entries ~key:"PWD" ~data:pwd in
    let path =
      match SMap.find "PATH" entries with
      | exception Not_found -> []
      | s -> split_path s
    in
    { entries; path }
end

type t =
  { stdin        : Unix.file_descr
  ; stdout       : Unix.file_descr
  ; stderr       : Unix.file_descr
  ; cwd          : Working_dir.t
  ; unix_env     : Uenv.t
  }

let create ?(stdin=Unix.stdin) ?(stdout=Unix.stdout) ?(stderr=Unix.stderr)
      ?(cwd=Working_dir_spec.Inherit) ?unix_env () =
  let cwd = Working_dir.create cwd in
  let unix_env = Uenv.create unix_env ~pwd:cwd.logical in
  { stdin
  ; stdout
  ; stderr
  ; cwd
  ; unix_env
  }

let add_cwd_ref t = Working_dir.addref t.cwd
let deref_cwd t = Working_dir.deref t.cwd

let get_env t var = Uenv.get t.unix_env var
let set_env t var value = { t with unix_env = Uenv.set t.unix_env var value }
let unset_env t var = { t with unix_env = Uenv.unset t.unix_env var }

let set_env_many t bindings =
  let unix_env =
    List.fold_left bindings ~init:t.unix_env ~f:(fun uenv (var, value) ->
      Uenv.set uenv var value)
  in
  { t with unix_env }

let unset_env_many t bindings =
  let unix_env =
    List.fold_left bindings ~init:t.unix_env ~f:Uenv.unset
  in
  { t with unix_env }

let stdin  t = t.stdin
let stdout t = t.stdout
let stderr t = t.stderr

let set_stdin  t fd = { t with stdin  = fd }
let set_stdout t fd = { t with stdout = fd }
let set_stderr t fd = { t with stderr = fd }

let set_outputs t fd = { t with stdout = fd; stderr = fd }

let set_stdios t ~stdin ~stdout ~stderr = { t with stdin; stdout; stderr }

let get_stdio t (which : Std_io.t) =
  match which with
  | Stdin  -> t.stdin
  | Stdout -> t.stdout
  | Stderr -> t.stderr

let set_stdio t (which : Std_io.t) x =
  match which with
  | Stdin  -> set_stdin  t x
  | Stdout -> set_stdout t x
  | Stderr -> set_stderr t x

let cwd_logical t = t.cwd.logical

let chdir t dir =
  let cwd = Working_dir.chdir t.cwd dir in
  let unix_env = Uenv.set t.unix_env "PWD" cwd.logical in
  { t with cwd; unix_env }

let find_executable t exe =
  let rec loop = function
    | [] -> None
    | path :: rest ->
      let fn = path ^/ exe in
      if Sys.file_exists fn then
        Some fn
      else
        loop rest
  in
  if not (Filename.is_relative exe) then
    Some exe
  else if Filename.basename exe <> exe then
    Some exe
  else
    loop t.unix_env.path

type run_error =
  | Command_not_found

let spawn t ~prog ~args =
  match find_executable t prog with
  | None -> Error Command_not_found
  | Some real_prog ->
    let env =
      SMap.fold t.unix_env.entries ~init:[] ~f:(fun ~key ~data acc ->
        sprintf "%s=%s" key data :: acc)
      |> Spawn.Env.of_list
    in
    let cwd : Spawn.Working_dir.t =
      match t.cwd.physical with
      | Win32 -> Path t.cwd.logical
      | Unix u -> Fd u.fd
    in
    let pid =
      Spawn.spawn ()
        ~env
        ~cwd
        ~prog:real_prog
        ~argv:(prog :: args)
        ~stdin:t.stdin
        ~stdout:t.stdout
        ~stderr:t.stderr
    in
    Ok pid

type full_path =
  | Path of string
  | In_dir of Unix.file_descr * string

let mk_logical_path t path =
  if Filename.is_relative path then
    t.cwd.logical ^/ path
  else
    path

let full_path t path =
  match t.cwd.physical with
  | Win32 ->
    Path (mk_logical_path t path)
  | Unix u ->
    In_dir (u.fd, path)

type full_path2 =
  | Path of string * string
  | In_dir of Unix.file_descr * string * string

let full_path2 t path1 path2 =
  match t.cwd.physical with
  | Win32 ->
    Path (mk_logical_path t path1, mk_logical_path t path2)
  | Unix u ->
    In_dir (u.fd, path1, path2)

let open_file t ?(perm=0) ~flags path =
  match full_path t path with
  | Path path ->
    Unix.openfile
      path
      (O_CLOEXEC :: List.map ~f:Posixat.Open_flag.to_unix_open_flag_exn flags)
      perm
  | In_dir (dir,  path) ->
    Posixat.openat ~dir ~path ~perm
      ~flags:(O_CLOEXEC :: flags)

let close_noerr t = try Unix.close t with _ -> ()

let with_file t ?perm ~flags path ~f =
  let fd = open_file t path ~flags ?perm in
  match f fd with
  | x           -> Unix.close fd; x
  | exception e -> close_noerr fd; raise e

let mkdir_one t ~perm path =
  match full_path t path with
  | Path path ->
    Unix.mkdir path perm
  | In_dir (dir,  path) ->
    Posixat.mkdirat ~dir ~path ~perm

let rec mkdir t ?(perm=0o777) ?(p=false) path =
  if not p then
    mkdir_one t path ~perm
  else
    try
      mkdir_one t path ~perm
    with
    | Unix.Unix_error((EEXIST|EISDIR) , _, _) -> ()
    | Unix.Unix_error(ENOENT, _, _) when Filename.basename path <> path ->
      mkdir t (Filename.dirname path) ~perm ~p;
      mkdir_one t path ~perm

let chmod t path ~perm =
  match full_path t path with
  | Path path ->
    Unix.chmod path perm
  | In_dir (dir, path) ->
    Posixat.fchmodat ~dir ~path ~perm ~flags:[]

let chown t path ~uid ~gid =
  match full_path t path with
  | Path path ->
    Unix.chown path uid gid
  | In_dir (dir, path) ->
    Posixat.fchownat ~dir ~path ~uid ~gid ~flags:[]

let rm t path =
  match full_path t path with
  | Path path ->
    Unix.unlink path
  | In_dir (dir,  path) ->
    Posixat.unlinkat ~dir ~path ~flags:[]

let rmdir t path =
  match full_path t path with
  | Path path ->
    Unix.rmdir path
  | In_dir (dir,  path) ->
    Posixat.unlinkat ~dir ~path ~flags:[AT_REMOVEDIR]

let mkfifo t ?(perm=0o666) path =
  match full_path t path with
  | Path path ->
    Unix.mkfifo path perm
  | In_dir (dir,  path) ->
    Posixat.mkfifoat ~dir ~path ~perm

let link t oldpath newpath =
  match full_path2 t oldpath newpath with
  | Path (oldpath, newpath) ->
    Unix.link oldpath newpath
  | In_dir (dir,  oldpath, newpath) ->
    Posixat.linkat ~olddir:dir ~newdir:dir ~oldpath ~newpath ~flags:[]

let symlink t oldpath newpath =
  match full_path2 t oldpath newpath with
  | Path (oldpath, newpath) ->
    Unix.symlink oldpath newpath
  | In_dir (dir,  oldpath, newpath) ->
    Posixat.symlinkat ~newdir:dir ~oldpath ~newpath

let rename t oldpath newpath =
  match full_path2 t oldpath newpath with
  | Path (oldpath, newpath) ->
    Unix.rename oldpath newpath
  | In_dir (dir,  oldpath, newpath) ->
    Posixat.renameat ~olddir:dir ~newdir:dir ~oldpath ~newpath

let stat t path =
  match full_path t path with
  | Path path ->
    Unix.stat path
  | In_dir (dir,  path) ->
    Posixat.fstatat ~dir ~path ~flags:[]

let lstat t path =
  match full_path t path with
  | Path path ->
    Unix.lstat path
  | In_dir (dir,  path) ->
    Posixat.fstatat ~dir ~path ~flags:[AT_SYMLINK_NOFOLLOW]

let readlink t path =
  match full_path t path with
  | Path path ->
    Unix.readlink path
  | In_dir (dir,  path) ->
    Posixat.readlinkat ~dir ~path

let access t path mode =
  match full_path t path with
  | Path path ->
    Unix.access path mode
  | In_dir (dir,  path) ->
    Posixat.faccessat ~dir ~path ~mode ~flags:[]

let readdir t path =
  match full_path t path with
  | Path path ->
    Sys.readdir path |> Array.to_list
  | In_dir (dir,  path) ->
    let fd = Posixat.openat ~dir ~path ~flags:[O_RDONLY] ~perm:0 in
    match Posixat.fdopendir fd with
    | exception e -> close_noerr fd; raise e
    | dh ->
      let rec loop dh acc =
        match Unix.readdir dh with
        | "." | ".." -> loop dh acc
        | fname -> loop dh (fname :: acc)
        | exception End_of_file -> Unix.closedir dh; List.rev acc
        | exception e ->
          Unix.closedir dh;
          raise e
      in
      loop dh []