fix(io): handle interrupted writes

jscomp/melstd/io.ml

@@ -39,6 +39,11 @@ let with_file_in_fd fn ~f =
   protectx (Unix.openfile fn [ O_RDONLY; O_CLOEXEC ] 0) ~f ~finally:Unix.close
 
 let read_file_exn =
+  (* We use 65536 because that is the size of OCaml's IO buffers. *)
+  let chunk_size = 65536 in
+  let file_too_big () =
+    failwith "read_file: file is larger than Sys.max_string_length"
+  in
   let read_all_unless_large =
     let rec eagerly_input_acc ic s ~pos ~len acc =
       if len <= 0 then acc
@@ -53,8 +58,6 @@ let read_file_exn =
       if Int.equal r len then Bytes.unsafe_to_string buf
       else Bytes.sub_string buf ~pos:0 ~len:r
     in
-    (* We use 65536 because that is the size of OCaml's IO buffers. *)
-    let chunk_size = 65536 in
     (* Generic function for channels such that seeking is unsupported or
        broken *)
     let read_all_generic t buffer =
@@ -97,46 +100,67 @@ let read_file_exn =
   let read_file_chan ?binary fn =
     match with_file_in fn ~f:read_all_unless_large ?binary with
     | Ok x -> x
-    | Error () ->
-        failwith "read_file: file is larger than Sys.max_string_length"
+    | Error () -> file_too_big ()
   in
   let read_all_fd =
+    let rec unix_read fd buf pos len =
+      match Unix.read fd buf pos len with
+      | bytes_read -> bytes_read
+      | exception Unix.Unix_error (EINTR, _, _) -> unix_read fd buf pos len
+    in
     let rec read fd buf pos left =
       match left with
-      | 0 -> `Ok
+      | 0 -> pos
       | left -> (
-          match Unix.read fd buf pos left with
-          | 0 -> `Eof
+          match unix_read fd buf pos left with
+          | 0 -> pos
           | n -> read fd buf (pos + n) (left - n))
     in
-    fun fd ->
-      match Unix.fstat fd with
-      | exception Unix.Unix_error (e, x, y) -> Error (`Unix (e, x, y))
-      | { Unix.st_size; _ } -> (
-          if Int.equal st_size 0 then Ok ""
-          else if st_size > Sys.max_string_length then Error `Too_big
+    let read_to_eof fd initial =
+      let probe = Bytes.create 1 in
+      match unix_read fd probe 0 1 with
+      | 0 -> initial
+      | _ ->
+          let initial_length = String.length initial in
+          if initial_length >= Sys.max_string_length then file_too_big ()
           else
-            let b = Bytes.create st_size in
-            match read fd b 0 st_size with
-            | exception Unix.Unix_error (e, x, y) -> Error (`Unix (e, x, y))
-            | `Eof -> Error `Retry
-            | `Ok -> Ok (Bytes.unsafe_to_string b))
+            let capacity =
+              if initial_length > Sys.max_string_length - chunk_size - 1 then
+                Sys.max_string_length
+              else initial_length + chunk_size + 1
+            in
+            let buffer = Buffer.create capacity in
+            Buffer.add_string buffer initial;
+            Buffer.add_char buffer (Bytes.get probe 0);
+            let chunk = Bytes.create chunk_size in
+            let rec loop () =
+              match unix_read fd chunk 0 chunk_size with
+              | 0 -> Buffer.contents buffer
+              | n ->
+                  if n > Sys.max_string_length - Buffer.length buffer then
+                    file_too_big ()
+                  else (
+                    Buffer.add_subbytes buffer chunk 0 n;
+                    loop ())
+            in
+            loop ()
+    in
+    fun fd ->
+      let { Unix.st_size; _ } = Unix.fstat fd in
+      if st_size > Sys.max_string_length then file_too_big ()
+      else
+        let b = Bytes.create st_size in
+        let bytes_read = read fd b 0 st_size in
+        if bytes_read < st_size then Bytes.sub_string b ~pos:0 ~len:bytes_read
+        else read_to_eof fd (Bytes.unsafe_to_string b)
   in
   match Sys.backend_type with
   | Other _ ->
       (* use slow path for JSOO *)
       fun ?(binary = true) fn -> read_file_chan ~binary fn
   | Native | Bytecode ->
       fun ?(binary = true) fn ->
-        if binary then
-          with_file_in_fd fn ~f:(fun fd ->
-              match read_all_fd fd with
-              | Ok s -> s
-              | Error `Retry -> read_file_chan ~binary fn
-              | Error `Too_big ->
-                  failwith
-                    "read_file: file is larger than Sys.max_string_length"
-              | Error (`Unix (e, c, s)) -> raise (Unix.Unix_error (e, c, s)))
+        if binary then with_file_in_fd fn ~f:read_all_fd
         else read_file_chan ~binary fn
 
 let read_file ?binary fn =
@@ -167,8 +191,10 @@ let with_file_out_fd ?(perm = default_out_perm) fn ~f =
 
 let rec write fd str ~off ~len =
   if len > 0 then
-    let written = Unix.single_write_substring fd str off len in
-    write fd str ~off:(off + written) ~len:(len - written)
+    match Unix.single_write_substring fd str off len with
+    | exception Unix.Unix_error (EINTR, _, _) -> write fd str ~off ~len
+    | 0 -> raise (Unix.Unix_error (EIO, "single_write", ""))
+    | written -> write fd str ~off:(off + written) ~len:(len - written)
 
 let write_file_exn =
   let write_file_fast ?(perm = default_out_perm) fn data =

test/unit-tests/test_io.ml

@@ -13,10 +13,40 @@ let test_write_filev_error () =
   | Error _ -> ()
   | Ok () -> Alcotest.fail "expected write_filev to return Error"
 
+let test_read_file_with_zero_reported_size () =
+  let path = Filename.temp_file "melange-io" ".fifo" in
+  Sys.remove path;
+  Unix.mkfifo path 0o600;
+  Fun.protect
+    ~finally:(fun () -> Sys.remove path)
+    (fun () ->
+      match Unix.fork () with
+      | 0 -> (
+          try
+            let oc = open_out_bin path in
+            output_string oc "contents";
+            close_out oc;
+            Unix._exit 0
+          with _ -> Unix._exit 1)
+      | pid -> (
+          let result = Io.read_file path in
+          let _, status = Unix.waitpid [] pid in
+          Alcotest.(check int)
+            "writer exited successfully" 0
+            (match status with
+            | WEXITED code -> code
+            | WSIGNALED _ | WSTOPPED _ -> 1);
+          match result with
+          | Ok contents ->
+              Alcotest.(check string) "contents" "contents" contents
+          | Error exn -> raise exn))
+
 let suite =
   [
     Alcotest.test_case "read_file returns errors" `Quick test_read_file_error;
     Alcotest.test_case "write_file returns errors" `Quick test_write_file_error;
     Alcotest.test_case "write_filev returns errors" `Quick
       test_write_filev_error;
+    Alcotest.test_case "read zero-size streams" `Quick
+      test_read_file_with_zero_reported_size;
   ]