package containers

  1. Overview
  2. Docs
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Source file CCUnix.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348

(* 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

(*$T
  escape_str "foo" = "foo"
  escape_str "foo bar" = "'foo bar'"
  escape_str "fo'o b'ar" = "'fo'\\''o b'\\''ar'"
*)

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 *)
       begin match stdin with
         | `Str s -> output_string ic s
         | `Gen g -> iter_gen (output_string ic) g
       end;
       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)

(*$T
  call_full ~stdin:(`Str "abc") "cat" |> stdout = "abc"
  call_full "echo %s" (escape_str "a'b'c") |> stdout = "a'b'c\n"
  call_full "echo %s" "a'b'c" |> stdout = "abc\n"
*)

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)

(*$T
  call_stdout ~stdin:(`Str "abc") "cat" = "abc"
  call_stdout "echo %s" (escape_str "a'b'c") = "a'b'c\n"
  call_stdout "echo %s" "a'b'c" = "abc\n"
*)

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

(*$R
  let m = 200 in
  let n = 50 in
  let write_atom filename s =
    with_file_lock ~kind:`Write filename
      (fun () ->
        CCIO.with_out ~flags:[Open_append; Open_creat]
          filename (fun oc -> output_string oc s; flush oc))
  in
  let f filename =
    for j=1 to m do
      write_atom filename "foo\n"
    done
  in
  CCIO.File.with_temp ~prefix:"containers_" ~suffix:".txt"
    (fun filename ->
        let a = Array.init n (fun _ -> Thread.create f filename) in
        Array.iter Thread.join a;
        let lines = CCIO.with_in filename CCIO.read_lines_l in
        assert_equal ~printer:string_of_int (n * m) (List.length lines);
        assert_bool "all valid" (List.for_all ((=) "foo") lines))
*)

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

(*$R
  let filename = with_temp_dir "test_containers"
      (fun dir ->
         let name = Filename.concat dir "test" in
         CCIO.with_out name (fun oc -> output_string oc "content"; flush oc);
         assert_bool ("file exists:"^name) (Sys.file_exists name);
         name)
  in
  assert_bool ("file does not exist"^filename) (not (Sys.file_exists filename));
  ()
*)
OCaml

Innovation. Community. Security.