Add Process.Env

lib_eio/process.ml

@@ -47,6 +47,100 @@ type 'tag mgr_ty = [ `Process_mgr | `Platform of 'tag ]
 type 'a mgr = 'a r
  constraint 'a = [> [> `Generic] mgr_ty]
 
+module Env = struct
+  let on_windows = (Sys.os_type = "Win32")
+
+  module Name = struct
+    type t = string
+
+    let normalise =
+      if on_windows then String.uppercase_ascii
+      else Fun.id
+
+    let compare x y =
+      String.compare (normalise x) (normalise y)
+
+    let starts_with ~prefix =
+      let prefix = normalise prefix in
+      fun x -> String.starts_with ~prefix (normalise x)
+
+    let validate t =
+      let bad_char = function
+        | '\000' | '=' -> true
+        | _ -> false
+      in
+      if t = "" || String.exists bad_char t then
+        Fmt.invalid_arg "Invalid environment variable name %S" t
+  end
+
+  module M = Map.Make(Name)
+
+  type t = string array
+
+  let of_array = Fun.id
+  let to_array = Fun.id
+  let empty = [| |]
+
+  let validate_value value =
+    if String.contains value '\000' then
+      Fmt.invalid_arg "Invalid environment variable value %S" value
+
+  let validate_binding (name, value) =
+    Name.validate name;
+    Option.iter validate_value value
+
+  let entry name value =
+    Printf.sprintf "%s=%s" name value
+
+  let get_opt name t =
+    Name.validate name;
+    let prefix = name ^ "=" in
+    Array.find_opt (Name.starts_with ~prefix) t
+    |> Option.map (fun e ->
+        let i = String.length prefix in
+        String.sub e i (String.length e - i)
+      )
+
+  let override bindings t =
+    List.iter validate_binding bindings;
+    let all_bindings = M.of_list bindings in
+    let bindings = ref all_bindings in
+    let updated =
+      Array.to_list t
+      |> List.filter_map (fun e ->
+          match String.index e '=' with
+          | exception Not_found -> Some e       (* Not a normal k=v entry *)
+          | i ->
+            let name = String.sub e 0 i in
+            match M.find_opt name all_bindings with
+            | None -> Some e                    (* We're not changing this *)
+            | Some x -> 
+              if M.mem name !bindings then (
+                bindings := M.remove name !bindings;
+                match x with
+                | None -> None                    (* Remove existing entry *)
+                | Some v -> Some (entry name v)   (* Update existing entry *)
+              ) else None                         (* Remove duplicate entry *)
+        )
+    in
+    let extra =
+      M.to_list !bindings
+      |> List.filter_map (function
+          | _, None -> None                     (* Remove entry that wasn't there anyway *)
+          | k, Some v ->
+            Some (entry k v)                    (* Add new entry *)
+        )
+    in
+    Array.of_list (updated @ extra)
+
+  let of_bindings xs =
+    override (List.map (fun (k, v) -> (k, Some v)) xs) empty
+
+  let pp f t =
+    Fmt.pf f "[@[<v>%a@]]"
+      (Fmt.array ~sep:Fmt.cut (Fmt.fmt "%S")) t
+end
+
 module Pi = struct
   module type PROCESS = sig
     type t

lib_eio/process.mli

@@ -52,6 +52,63 @@ type 'a mgr = 'a r
  constraint 'a = [> [> `Generic] mgr_ty]
 (** A process manager capable of spawning new processes. *)
 
+module Env : sig
+  (** A list of environment variable entries.
+
+      By convention, the entries are strings of the form "name=value",
+      each name is unique, order doesn't matter, and "" is not a valid name.
+
+      Also, due to the representation:
+      - Names cannot contain the '=' character.
+      - Neither names nor values can contain '\000'.
+
+      On Windows, name comparison is (ASCII) case-insensitive.
+      The convention of always using uppercase ASCII names for environment
+      variables will avoid different behaviour across platforms. *)
+
+  type t = string array
+  (** Note: this type is currently exposed for backwards compatibility and will
+      likely be made abstract in the future. The array is intended to be immutable. *)
+
+  val empty : t
+  (** An environment with no bindings. *)
+
+  val of_bindings : (string * string) list -> t
+  (** [of_bindings xs] is a new environment containing only the bindings in [xs].
+
+      This just adds [xs] to {!empty} using {!override}. *)
+
+  val get_opt : string -> t -> string option
+  (** [get_opt name t] is the value of [name] in [t], or [None] if there is no such binding.
+
+      @raise Invalid_argument if [name] is not a valid name. *)
+
+  val override : (string * string option) list -> t -> t
+  (** [override bindings t] is a new environment which is like [t]
+      except that the updates in [bindings] have been applied.
+
+      Each entry in [bindings] is a [(name, new_value)] pair.
+      [new_value] can be [None] to remove [name] (ignored if [name] is not present).
+
+      If there are several bindings for the same name in [bindings] then the last one is used.
+      If there are several bindings for an updated name in [t] then all are removed first.
+
+      @raise Invalid_argument if any binding is invalid. *)
+
+  val of_array : string array -> t
+  (** Create a [t] from e.g. the results of {!Unix.environment}.
+
+      The bindings are used as-is and need not conform to the conventions (e.g.
+      they may contain duplicate names).
+
+      Note: the array is assumed to be immutable. *)
+
+  val to_array : t -> string array
+  (** [to_array t] gets [t] as an array. The array should be treated as immutable. *)
+
+  val pp : t Fmt.t
+end
+
 (** {2 Processes} *)
 
 val pid : _ t -> int
@@ -82,7 +139,7 @@ val spawn :
   ?stdin:_ Flow.source ->
   ?stdout:_ Flow.sink ->
   ?stderr:_ Flow.sink ->
-  ?env:string array ->
+  ?env:Env.t ->
   ?executable:string ->
   string list -> 'tag ty r
 (** [spawn ~sw mgr args] creates a new child process that is connected to the switch [sw].
@@ -109,7 +166,7 @@ val run :
   ?stdout:_ Flow.sink ->
   ?stderr:_ Flow.sink ->
   ?is_success:(int -> bool) ->
-  ?env:string array ->
+  ?env:Env.t ->
   ?executable:string ->
   string list -> unit
 (** [run] does {!spawn} followed by {!await_exn}, with the advantage that if the process fails then
@@ -127,7 +184,7 @@ val parse_out :
   ?stdin:_ Flow.source ->
   ?stderr:_ Flow.sink ->
   ?is_success:(int -> bool) ->
-  ?env:string array ->
+  ?env:Env.t ->
   ?executable:string ->
   string list -> 'a
 (** [parse_out mgr parser args] runs [args] and parses the child's stdout with [parser].
@@ -183,7 +240,7 @@ module Pi : sig
       ?stdin:Flow.source_ty r ->
       ?stdout:Flow.sink_ty r ->
       ?stderr:Flow.sink_ty r ->
-      ?env:string array ->
+      ?env:Env.t ->
       ?executable:string ->
       string list ->
       tag ty r

lib_eio/unix/fork_action.ml

@@ -22,6 +22,7 @@ external make_string_array : int -> c_array = "eio_unix_make_string_array"
 external action_execve : unit -> fork_fn = "eio_unix_fork_execve"
 let action_execve = action_execve ()
 let execve path ~argv ~env =
+  let env = Eio.Process.Env.to_array env in
   let argv_c_array = make_string_array (Array.length argv) in
   let env_c_array = make_string_array (Array.length env) in
   { run = fun k -> k (Obj.repr (action_execve, path, argv_c_array, argv, env_c_array, env)) }

lib_eio/unix/fork_action.mli

@@ -32,7 +32,7 @@ val with_actions : t list -> (c_action list -> 'a) -> 'a
 
 (** {2 Actions} *)
 
-val execve : string -> argv:string array -> env:string array -> t
+val execve : string -> argv:string array -> env:Eio.Process.Env.t -> t
 (** See [execve(2)].
 
     This replaces the current executable,

lib_eio/unix/process.ml

@@ -67,7 +67,7 @@ let get_executable ~args = function
 
 let get_env = function
   | Some e -> e
-  | None -> Unix.environment ()
+  | None -> Unix.environment () |> Eio.Process.Env.of_array
 
 let translate_execve_error ~executable f =
   try f () with
@@ -98,7 +98,7 @@ module Pi = struct
       ?uid:int ->
       ?gid:int ->
       ?login_tty:Fd.t ->
-      env:string array ->
+      env:Eio.Process.Env.t ->
       fds:(int * Fd.t * Fork_action.blocking) list ->
       executable:string ->
       string list ->
@@ -126,7 +126,7 @@ module Make_mgr (X : sig
     ?uid:int ->
     ?gid:int ->
     ?login_tty:Fd.t ->
-    env:string array ->
+    env:Eio.Process.Env.t ->
     fds:(int * Fd.t * Fork_action.blocking) list ->
     executable:string ->
     string list ->

lib_eio/unix/process.mli

@@ -24,7 +24,7 @@ module Pi : sig
       ?uid:int ->
       ?gid:int ->
       ?login_tty:Fd.t ->
-      env:string array ->
+      env:Eio.Process.Env.t ->
       fds:(int * Fd.t * Fork_action.blocking) list ->
       executable:string ->
       string list ->
@@ -50,7 +50,7 @@ module Make_mgr (X : sig
     ?uid:int ->
     ?gid:int ->
     ?login_tty:Fd.t ->
-    env:string array ->
+    env:Eio.Process.Env.t ->
     fds:(int * Fd.t * Fork_action.blocking) list ->
     executable:string ->
     string list ->
@@ -66,7 +66,7 @@ val spawn_unix :
     ?gid:int ->
     ?login_tty:Fd.t ->
     fds:(int * Fd.t * Fork_action.blocking) list ->
-    ?env:string array ->
+    ?env:Eio.Process.Env.t ->
     ?executable:string ->
     string list ->
     ty r

lib_eio_linux/tests/spawn.md

@@ -5,7 +5,10 @@
 ```ocaml
 open Eio.Std
 
+module Env = Eio.Process.Env
 module Process = Eio_linux.Low_level.Process
+
+let default_env = Unix.environment () |> Env.of_array
 ```
 
 ## Spawning processes
@@ -18,7 +21,7 @@ Setting environment variables:
   let child = Process.spawn ~sw Process.Fork_action.[
     execve "/usr/bin/env"
       ~argv:[| "env" |]
-      ~env:[| "FOO=bar" |];
+      ~env:(Env.of_bindings ["FOO", "bar"]);
   ] in
   Promise.await (Process.exit_status child);;
 FOO=bar
@@ -34,7 +37,7 @@ Changing directory:
     chdir "/";
     execve "/usr/bin/env"
       ~argv:[| "env"; "pwd" |]
-      ~env:(Unix.environment ())
+      ~env:default_env
   ] in
   Promise.await (Process.exit_status child);;
 /
@@ -58,7 +61,7 @@ Changing directory using a file descriptor:
     fchdir root;
     execve "/usr/bin/env"
       ~argv:[| "env"; "pwd" |]
-      ~env:(Unix.environment ())
+      ~env:default_env
   ] in
   Promise.await (Process.exit_status child);;
 /
@@ -73,7 +76,7 @@ Exit status:
   let child = Process.spawn ~sw Process.Fork_action.[
     execve "/usr/bin/env"
       ~argv:[| "env"; "false" |]
-      ~env:(Unix.environment ())
+      ~env:default_env
   ] in
   Promise.await (Process.exit_status child);;
 - : Unix.process_status = Unix.WEXITED 1
@@ -88,7 +91,7 @@ Failure starting child:
     chdir "/idontexist";
     execve "/usr/bin/env"
       ~argv:[| "env"; "pwd" |]
-      ~env:(Unix.environment ())
+      ~env:default_env
   ]
 Exception: Unix.Unix_error(Unix.ENOENT, "chdir", "")
 ```
@@ -102,7 +105,7 @@ Signalling a running child:
     Process.spawn ~sw Process.Fork_action.[
       execve "/usr/bin/env"
         ~argv:[| "env"; "sleep"; "1000" |]
-        ~env:(Unix.environment ())
+        ~env:default_env
     ]
   in
   Process.signal child Sys.sigkill;
@@ -122,7 +125,7 @@ Signalling an exited child does nothing:
     Process.spawn ~sw Process.Fork_action.[
       execve "/usr/bin/env"
         ~argv:[| "env" |]
-        ~env:[| "FOO=bar" |];
+        ~env:(Env.of_bindings ["FOO", "bar"]);
     ]
   in
   ignore (Promise.await (Process.exit_status child) : Unix.process_status);