Source file ThunkIoDisk.ml

(** Disk I/O operations for thunks that uses the OCaml [unix] library and the
    [MlFront_ZipFile] library. *)

(** Remove directory recursively *)
let remove_directory_recursively ~return (dir : MlFront_Core.FilePath.t) =
  (* TODO: Retry when there are inevitable race conditions. *)
  try
    let rec aux path =
      if Sys.file_exists path then begin
        if Sys.is_directory path then begin
          Array.iter aux (Sys.readdir path |> Array.map (Filename.concat path));
          Unix.rmdir path
        end
        else Unix.unlink path
      end
    in
    aux (MlFront_Core.FilePath.to_string dir);
    return `Deleted
  with
  | Unix.Unix_error (e, fname, arg) ->
      return
        (`Error
           (Printf.sprintf "deleting directory `%s` had the error `%s` in `%s`"
              (MlFront_Core.FilePath.to_string dir)
              (Unix.error_message e)
              (if String.equal arg "" then fname else fname ^ " " ^ arg)))
  | Sys_error e ->
      return
        (`Error
           (Printf.sprintf "deleting directory `%s` had the error `%s`"
              (MlFront_Core.FilePath.to_string dir)
              e))

(** Make a directory recursively *)
let make_directory_recursively ~return (dir : MlFront_Core.FilePath.t) =
  let dir_s = MlFront_Core.FilePath.to_string dir in
  try
    if Sys.file_exists dir_s then
      (* optimization 1: short-circuit if the directory exists *)
      return `Created
    else
      let parent_dir_s = MlFront_Core.FilePath.(parent dir |> to_string) in
      if Sys.file_exists parent_dir_s then (
        (* optimization 2: parent directory exists *)
        Unix.mkdir dir_s 0o755;
        return `Created)
      else
        (* possibly expensive fallback: recreate from root *)
        let root_result =
          MlFront_Core.FilePath.(
            of_string
              (if is_absolute dir then root_noslash dir ^ slash dir else ""))
        in
        match root_result with
        | Error msg ->
            return
              (`Error (Printf.sprintf "invalid file path `%s`: %s" dir_s msg))
        | Ok root ->
            let rec aux fp = function
              | [] -> `Created
              | segment :: rest -> begin
                  match MlFront_Core.FilePath.append segment fp with
                  | Error msg ->
                      `Error
                        (Printf.sprintf
                           "`%s` can't be added to file path `%s`: %s" segment
                           (MlFront_Core.FilePath.to_string fp)
                           msg)
                  | Ok fp ->
                      let path_s = MlFront_Core.FilePath.to_string fp in
                      if not (Sys.file_exists path_s) then begin
                        Unix.mkdir path_s 0o755
                      end;
                      aux fp rest
                end
            in
            return (aux root (MlFront_Core.FilePath.rootless_segments dir))
  with
  | Unix.Unix_error (e, fname, arg) ->
      return
        (`Error
           (Printf.sprintf "creating directory `%s` had the error `%s` in `%s`"
              dir_s (Unix.error_message e)
              (if String.equal arg "" then fname else fname ^ " " ^ arg)))
  | Sys_error e ->
      return
        (`Error
           (Printf.sprintf "creating directory `%s` had the error `%s`" dir_s e))

let delete_local_file ~return local_file =
  try
    (* For Windows, can't write (or delete) without turning off
           read-only flag. *)
    (if Sys.win32 then try Unix.chmod local_file 0o644 with Sys_error _ -> ());
    Unix.unlink local_file;
    return `Deleted
  with
  | Unix.Unix_error (Unix.ENOENT, _, _) ->
      (* If the file does not exist, we do not raise an error. *)
      return `Deleted
  | Unix.Unix_error (e, fname, arg) ->
      return
        (`Error
           (Printf.sprintf "deleting file `%s` had the error `%s` in `%s`"
              local_file (Unix.error_message e)
              (if String.equal arg "" then fname else fname ^ " " ^ arg)))
  | Sys_error e ->
      if Sys.file_exists local_file then
        return
          (`Error
             (Printf.sprintf "deleting file `%s` had the error `%s`" local_file
                e))
      else return `Deleted

let checksum_local_file ?strip_carriage_returns ~algo ~return local_file =
  let m =
    match algo with
    | `Sha1 -> (module Digestif.SHA1 : Digestif.S)
    | `Sha256 -> (module Digestif.SHA256 : Digestif.S)
  in
  let module D = (val m) in
  let ctx = ref (D.init ()) in
  try
    In_channel.with_open_bin local_file (fun ic ->
        let buf = Bytes.create 32_768 in
        (* feed only the bytes that are not carriage returns *)
        let rec feed_no_cr_bytes i n =
          if i >= n then !ctx
          else if Bytes.get buf i = '\r' then feed_no_cr_bytes (i + 1) n
          else
            let j = ref (i + 1) in
            while !j < n && Bytes.get buf !j <> '\r' do
              incr j
            done;
            ctx := D.feed_bytes !ctx ~off:i ~len:(!j - i) buf;
            feed_no_cr_bytes !j n
        in
        let rec aux total =
          match In_channel.input ic buf 0 32_768 with
          | 0 ->
              let cksum = D.(get !ctx |> to_hex) in
              return (`Checksum (cksum, total))
          | n ->
              if strip_carriage_returns = Some () then
                ctx := feed_no_cr_bytes 0 n
              else ctx := D.feed_bytes !ctx ~off:0 ~len:n buf;
              aux (Int64.add total (Int64.of_int n))
        in
        aux 0L)
  with Sys_error e -> return (`Error e)

module Make (M : MlFront_Thunk.BuildConstraints.MONAD_PROMISE) = struct
  include MlFront_Thunk.ThunkIo.Make (M)

  (** [memory_limit].

      16 MiB is the maximum on 32-bit OCaml platforms for a single string.
      However, js_of_ocaml uses maximum JavaScript string size which is
      undocumented but at least 2^51.

      We'll use [2^31 - 1] as limit for now as that is the memory bound for
      32-bit signed C integers. *)
  let memory_limit = 2147483647L

  open struct
    let safe_is_directory s = try Sys.is_directory s with Sys_error _ -> false

    let sys_error_is_a_directory ~path s =
      if String.equal s "Is a directory" then
        (* macOS *)
        true
      else if
        String.ends_with ~suffix:"Permission denied" s
        && Sys.win32 && safe_is_directory path
      then
        (* Windows has errors like: "./src: Permission denied".
       We double-check with a racy check. *)
        true
      else false

    let default_open_bufsize = 16384

    let default_read_bufsize =
      match Sys.backend_type with
      | Sys.Other "js_of_ocaml" | Sys.Bytecode -> 16_384
      | Sys.Native | Sys.Other _ -> 1_048_576

    let disk_node_id = ref 0L
  end

  open struct
    let mk_env envmods =
      let envpairs =
        Unix.environment () |> Array.to_list
        |> List.map (Stringext.cut ~on:"=")
        |> List.filter_map Fun.id
      in
      Array.of_list
        (List.map
           (fun (k, v) -> k ^ "=" ^ v)
           (MlFront_Core.EnvMods.apply ~win32:Sys.win32 envmods envpairs))
  end

  let rec disk_dir origin =
    let origin_s = MlFront_Core.FilePath.to_string origin in
    let create_directory () =
      make_directory_recursively ~return:M.return origin
    in
    let delete_directory () =
      remove_directory_recursively ~return:M.return origin
    in
    let zip_directory ?intermediate ~staging_dir () =
      let pending_deletion_file = ref None in
      Fun.protect
        ~finally:(fun () ->
          if intermediate = None then
            try Option.iter Sys.remove !pending_deletion_file
            with Sys_error _ -> ())
        (fun () ->
          (* make a zip file *)
          let zipfile =
            Filename.temp_file
              ~temp_dir:(directory_origin staging_dir)
              "thunk" ".zip"
          in
          pending_deletion_file := Some zipfile;
          match MlFront_Core.FilePath.of_string zipfile with
          | Error msg ->
              M.pure (`Error (Printf.sprintf "invalid zipfile path: %s" msg))
          | Ok zipfile_fp -> begin
              match
                MlFront_ZipFile.ZipFile.zip_exn ~deterministic:()
                  ~srcdir:origin_s ~destzip:zipfile ()
              with
              | () ->
                  pending_deletion_file := None;
                  M.pure (`ZipFile (disk_file zipfile_fp))
              | exception MlFront_ZipFile.ZipFile.ZipError (_zipfile, message)
                ->
                  M.pure
                    (`Error (Printf.sprintf "zipping had error '%s'" message))
            end)
    in
    let spawn_in_directory ~command ~args ~envmods ~stdout ~stderr () =
      let env = mk_env envmods in
      let command_s = MlFront_Core.FilePath.to_string command in
      let args = Array.of_list (command_s :: args) in
      let in_channel, out_channel, err_channel =
        (ref None, ref None, ref None)
      in
      Fun.protect
        ~finally:(fun () ->
          (match !in_channel with Some ic -> Unix.close ic | None -> ());
          (match !out_channel with Some oc -> Unix.close oc | None -> ());
          match !err_channel with Some ec -> Unix.close ec | None -> ())
        (fun () ->
          (* /dev/null or NUL for stdin *)
          let in_fd =
            Unix.openfile
              (if Sys.win32 then "NUL" else "/dev/null")
              [ Unix.O_RDONLY ] 0o644
          in
          in_channel := Some in_fd;
          (* files for stdout and stderr *)
          let out_fd =
            Unix.openfile (file_origin stdout)
              [ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC; Unix.O_CLOEXEC ]
              0o644
          in
          out_channel := Some out_fd;
          let err_fd =
            Unix.openfile (file_origin stderr)
              [ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC; Unix.O_CLOEXEC ]
              0o644
          in
          err_channel := Some err_fd;
          (* create the process *)
          try
            (* for an instant before we create the process asynchronously,
               we have to change the directory. *)
            let pwd = Sys.getcwd () in
            Unix.chdir origin_s;
            let pid =
              Fun.protect
                ~finally:(fun () -> Unix.chdir pwd)
                (fun () ->
                  Unix.create_process_env command_s args env in_fd out_fd err_fd)
            in
            (* wait for process to finish *)
            let ret =
              match Unix.waitpid [] pid with
              | _childpid, Unix.WEXITED ec -> `Exited ec
              | _childpid, Unix.WSIGNALED s -> `Signaled s
              | _childpid, Unix.WSTOPPED s -> `Stopped s
            in
            M.pure ret
          with Unix.Unix_error (Unix.ENOENT, _, _) ->
            M.pure (`Error (Printf.sprintf "command not found: `%s`" command_s)))
    in
    let interactive_shell_in_directory ?promptname ~envmods () =
      let ( let* ) = M.bind in
      (* Common functions for using the optional prompt *)
      let template name =
        (* Avoid escaping of % in Printf. Instead use literal search/replace *)
        Stringext.replace_all ~pattern:"<PROMPTNAME>" ~with_:name
      in
      let add_if_prompt opts l =
        match promptname with
        | None -> l
        | Some pname -> l @ List.map (template pname) opts
      in
      (* Different logic based on Windows or Unix *)
      if Sys.win32 then
        let original_pwd = Sys.getcwd () in
        match MlFront_Core.FilePath.of_string original_pwd with
        | Error msg ->
            M.pure
              (`Error
                 (Printf.sprintf "invalid current working directory `%s`: %s"
                    original_pwd msg))
        | Ok original_pwd_fp -> (
            let absolute_origin =
              MlFront_Core.FilePath.to_string
                (MlFront_Core.FilePath.concat original_pwd_fp origin)
            in
            let env = mk_env envmods in
            (* NOTE: We can use [Sys.command], [Unix.create_process*] and [Unix.open_process*]
           but not [Unix.exec*] because the latter cooks the Windows Terminal.
           Confer https://github.com/ocaml/ocaml/pull/13879 *)
            let run_get_trimmed_stdout cmd =
              let stdout, stdin, stderr = Unix.open_process_full cmd env in
              Out_channel.close stdin;
              In_channel.close stderr;
              let output = In_channel.input_all stdout in
              In_channel.close stdout;
              String.trim output
            in
            let attempt :
                [ `Error of string
                | `Exited of int
                | `Signaled of int
                | `Stopped of int
                | `TryNext ] =
              `TryNext
            in
            let attempt_chdir_and_run_if_needed attempt command args =
              match attempt with
              | `Error e -> M.pure (`Error e)
              | `Exited ec -> M.pure (`Exited ec)
              | `Signaled ec -> M.pure (`Signaled ec)
              | `Stopped ec -> M.pure (`Stopped ec)
              | `TryNext ->
                  let full_command =
                    run_get_trimmed_stdout
                      (Printf.sprintf "where.exe %s" command)
                  in
                  if String.equal full_command "" then M.pure `TryNext
                  else begin
                    Unix.chdir origin_s;
                    let pid =
                      Fun.protect
                        ~finally:(fun () -> Unix.chdir original_pwd)
                        (fun () ->
                          Unix.create_process_env full_command
                            (Array.of_list (full_command :: args))
                            env Unix.stdin Unix.stdout Unix.stderr)
                    in
                    (* wait for process to finish *)
                    let ret =
                      match Unix.waitpid [] pid with
                      | _childpid, Unix.WEXITED ec -> `Exited ec
                      | _childpid, Unix.WSIGNALED s -> `Signaled s
                      | _childpid, Unix.WSTOPPED s -> `Stopped s
                    in
                    M.pure ret
                  end
            in

            (* PowerShell prompt ...
           https://learn.microsoft.com/en-us/powershell/module/microsoft.powershell.core/about/about_prompts?view=powershell-7.5#how-to-customize-the-prompt *)
            let powershell_prompt () =
              [
                Printf.sprintf
                  {|function prompt {
    $currentLocation = Get-Location
    $BaseDirectory = "%s"
    $PromptName = "<PROMPTNAME>"

    if ($currentLocation.Path.StartsWith($BaseDirectory)) {
        $relativePath = $currentLocation.Path.Substring($BaseDirectory.Length)
        if ($relativePath -eq "") {
            $PromptValue = "$PromptName"
        } else {
            $PromptValue = "$PromptName$relativePath"
        }
    } else {
        $PromptValue = $currentLocation.Path
    }

    Write-Host "PS $PromptValue>" -NoNewline -ForegroundColor 3
    return " "
} |}
                  absolute_origin;
              ]
              (* [@@warning "-unused-var"] *)
            in

            (* Powershell 7+. https://learn.microsoft.com/en-us/powershell/module/microsoft.powershell.core/about/about_pwsh?view=powershell-7.5 *)
            let* attempt =
              attempt_chdir_and_run_if_needed attempt "pwsh"
                ([ "-Interactive"; "-NoProfile"; "-NoExit"; "-Command" ]
                |> add_if_prompt (powershell_prompt ()))
            in
            (* Powershell 3+. https://learn.microsoft.com/en-us/powershell/module/microsoft.powershell.core/about/about_powershell_exe?view=powershell-5.1 *)
            let* attempt =
              attempt_chdir_and_run_if_needed attempt "powershell"
                ([ "-NoProfile"; "-NoExit"; "-Command" ]
                |> add_if_prompt (powershell_prompt ()))
            in
            (* Command Prompt. https://learn.microsoft.com/en-us/windows-server/administration/windows-commands/cmd *)
            let* attempt =
              attempt_chdir_and_run_if_needed attempt "cmd"
                [ "/d"; "/u"; "/e:on"; "/f:on" ]
            in
            (* Gather errors *)
            match attempt with
            | `Signaled ec -> M.pure (`Signaled ec)
            | `Stopped ec -> M.pure (`Stopped ec)
            | `Exited 0 -> M.pure (`Exited 0)
            | `Exited ec ->
                M.pure (`Error (Printf.sprintf "failed with exit code %d" ec))
            | `TryNext ->
                M.pure
                  (`Error
                     (Printf.sprintf
                        "Command Prompt `cmd.exe` interpreter was not found"))
            | `Error e -> M.pure (`Error e))
      else begin
        (* Unix shells do not have something like -WorkingDirectory. *)
        Unix.chdir origin_s;
        (* *)
        let shell, opts, envmods =
          let ew suffix = String.ends_with ~suffix in
          let set_if_prompt name value envmods =
            match promptname with
            | None -> envmods
            | Some pname ->
                (* Use [union] to not override the thunk envmods! *)
                MlFront_Core.EnvMods.(
                  union (add name (template pname value) empty) envmods)
          in
          match Sys.getenv_opt "SHELL" with
          | None | Some "" -> ("/bin/sh", [ "-i" ], envmods)
          | Some s when ew "bash" s ->
              (* GNU bash, version 5.2.37(1)-release (aarch64-apple-darwin23.4.0)
               requires [-i] after long options. *)
              ( s,
                [ "--norc"; "--noprofile"; "-i" ],
                set_if_prompt "PS1"
                  {|\[\033[32m\]<PROMPTNAME> \[\033[34m\]\W\[\033[0m\]\$ |}
                  envmods
                |> set_if_prompt "PROMPT_DIRTRIM" "2" )
          | Some s when ew "zsh" s ->
              ( s,
                [ "-i"; "--no-rcs"; "--no-globalrcs" ],
                set_if_prompt "PROMPT"
                  {|%F{green}<PROMPTNAME>%f %F{blue}%2~%f %# |} envmods )
          | Some s when ew "fish" s ->
              ( s,
                [ "-i"; "--no-config" ]
                |> add_if_prompt
                     [
                       "-C";
                       {|function fish_prompt; set_color green; echo -n '<PROMPTNAME> '; set_color blue; echo -n (prompt_pwd); set_color normal; echo -n '> '; end|};
                     ],
                envmods )
          | Some s when ew "tcsh" s ->
              ( s,
                [ "-i"; "-f" ]
                (* sigh ... no way in tcsh to set this without adding files ...
                
                   |> add_if_prompt
                     [
                       "-c";
                       {|set prompt = "<PROMPTNAME>:%c3/> "; exec tcsh -i -f|};
                     ] *),
                envmods )
          | Some s when ew "csh" s -> (s, [ "-i"; "-f" ], envmods)
          | Some s -> (s, [ "-i" ], envmods)
        in
        let env = mk_env envmods in
        Unix.execvpe shell (Array.of_list (shell :: opts)) env
      end
    in
    generic_dir ~origin:origin_s ~create_directory ~delete_directory
      ~zip_directory ~spawn_in_directory ~interactive_shell_in_directory ()

  (** [disk_file origin] creates a new file object that reads from the disk from
      the file [origin]. The file is {b read synchronously}, so use a different
      file object implementation for asynchronous reads. *)
  and disk_file origin : file_object =
    let origin_s = MlFront_Core.FilePath.to_string origin in
    let out_channels = Hashtbl.create 1 in
    let in_channels = Hashtbl.create 1 in
    let open_for_writing () =
      (* We flush _before_ the file is opened to ensure that
         the output is writable (it is not a directory, etc.).
         If the file is lazily opened, we can't distinguish
         input exceptions from output exceptions. *)
      let success = ref false in
      try
        let oc = Out_channel.open_bin origin_s in
        Fun.protect
          ~finally:(fun () -> if not !success then Out_channel.close oc)
          (fun () ->
            match Out_channel.flush oc with
            | () ->
                let idx = !disk_node_id in
                disk_node_id := Int64.succ !disk_node_id;
                Hashtbl.add out_channels idx oc;
                success := true;
                M.return (`Node idx)
            | exception Sys_error s
              when sys_error_is_a_directory ~path:origin_s s ->
                M.return (`IsDirectory (disk_dir origin)))
      with Sys_error s ->
        M.return
        @@ `Error
             (Format.asprintf "`%s` while opening `%s` for writing" s origin_s)
    in
    let open_for_reading () =
      (* We have to read at least one byte to make sure the file is not
         readable rather than lazily opened. Similar to flushing in
         open_for_writing to raise exceptions early. *)
      let success = ref false in
      try
        let ic = In_channel.open_bin origin_s in
        Fun.protect
          ~finally:(fun () -> if not !success then In_channel.close ic)
          (fun () ->
            let first_byte_queue = Queue.create () in
            (match In_channel.input_byte ic with
            | None -> ()
            | Some first_byte -> Queue.add first_byte first_byte_queue);
            let idx = !disk_node_id in
            disk_node_id := Int64.succ !disk_node_id;
            Hashtbl.add in_channels idx
              (ic, Bytes.create default_open_bufsize, first_byte_queue);
            success := true;
            M.return (`Node idx))
      with
      | Sys_error s when sys_error_is_a_directory ~path:origin_s s ->
          M.return (`IsDirectory (disk_dir origin))
      | Sys_error s ->
          M.return
          @@ `Error
               (Format.asprintf "`%s` while opening `%s` for reading" s origin_s)
    in
    let read_some idx =
      match Hashtbl.find_opt in_channels idx with
      | None ->
          M.return (`Error (Printf.sprintf "file `%s` is not open" origin_s))
      | Some (channel, bytes, first_byte_queue) -> begin
          match Queue.take_opt first_byte_queue with
          | Some first_byte ->
              Bytes.set_int8 bytes 0 first_byte;
              M.return (`Bytes (bytes, 0, 1))
          | None ->
              let bytes_read =
                In_channel.input channel bytes 0 (Bytes.length bytes)
              in
              if bytes_read = 0 then M.return `Eof
              else M.return (`Bytes (bytes, 0, bytes_read))
        end
    in
    let write_all idx str pos len =
      match Hashtbl.find_opt out_channels idx with
      | Some channel ->
          if pos = 0 && len = String.length str then
            Out_channel.output_string channel str
          else Out_channel.output channel (Bytes.unsafe_of_string str) pos len;
          M.return `WroteBytes
      | None ->
          M.return (`Error (Printf.sprintf "file `%s` is not open" origin_s))
    in
    let close idx =
      (match Hashtbl.find_opt out_channels idx with
      | Some channel ->
          Out_channel.close channel;
          (* All generated files should be executable per the file_object docs. *)
          Unix.chmod origin_s 0o755;
          Hashtbl.remove out_channels idx
      | None -> ());
      (match Hashtbl.find_opt in_channels idx with
      | Some (channel, _bytes, _first_byte_queue) ->
          In_channel.close channel;
          Hashtbl.remove in_channels idx
      | None -> ());
      M.return ()
    in
    let read_all () =
      let bytes_sz = default_read_bufsize in
      let buf = Buffer.create bytes_sz in
      let bytes = Bytes.create bytes_sz in
      M.pure
      @@
      try
        In_channel.with_open_bin origin_s (fun ic ->
            (* Loop and read up to 1 MB of data *)
            let rec read_loop total_read =
              if Int64.compare total_read memory_limit < 0 then
                let bytes_read = In_channel.input ic bytes 0 bytes_sz in
                if bytes_read = 0 then `Content (Buffer.contents buf)
                else (
                  Buffer.add_subbytes buf bytes 0 bytes_read;
                  read_loop (Int64.add total_read (Int64.of_int bytes_read)))
              else
                (* Exceeded memory. Too big! *)
                `ExceededSizeLimit memory_limit
            in
            read_loop 0L)
      with Sys_error s ->
        `Error
          (Printf.sprintf "`%s` while reading file `%s`" s
             (MlFront_Core.FilePath.show origin))
    in
    let delete_file () = delete_local_file ~return:M.return origin_s in
    let prepare_as_copy_destination () =
      (* For Windows, can't write without turning off read-only flag.
         In fact, you can still get Permission Denied even after turning
         off read-only flag, perhaps because Windows has a richer
         permissions model than POSIX. So we remove the file
         after turning off read-only *)
      let ( let* ) = M.bind in
      let* predelete_result =
        if Sys.win32 then begin
          delete_file ()
        end
        else M.return `Deleted
      in
      match predelete_result with
      | `Error e -> M.return (`Error e)
      | `Deleted -> M.return `Ready
    in
    let checksum_file ~strip_carriage_returns ~algo () =
      let strip_carriage_returns =
        if strip_carriage_returns then Some () else None
      in
      checksum_local_file ?strip_carriage_returns ~algo ~return:M.return
        origin_s
    in
    generic_file ~origin:origin_s ~is_local_file:true ~open_for_writing
      ~open_for_reading ~read_some ~write_all ~close ~read_all
      ~prepare_as_copy_destination ~delete_file ~checksum_file ()
end