Merge PR #9881: Protect some I/O routines from SIGALRM

Ack-by: SkySkimmer
Reviewed-by: ejgallego
Ack-by: maximedenes

5 files changed, 36 insertions, 5 deletions

diff --git a/clib/cThread.ml b/clib/cThread.ml
index 9e0319e8f8..5fa44b1eec 100644
--- a/clib/cThread.ml
+++ b/clib/cThread.ml
@@ -100,10 +100,10 @@ let thread_friendly_input_value ic =
 (* On the ocaml runtime used in some opam-for-windows version the
  * [Thread.sigmask] API raises Invalid_argument "not implemented",
  * hence we protect the call and turn the exception into a no-op *)
-let protect_sigalrm f x =
+let mask_sigalrm f x =
   begin try ignore(Thread.sigmask Unix.SIG_BLOCK [Sys.sigalrm])
   with Invalid_argument _ -> () end;
   f x
 
 let create f x =
-  Thread.create (protect_sigalrm f) x
+  Thread.create (mask_sigalrm f) x
diff --git a/ide/idetop.ml b/ide/idetop.ml
index f744ce2ee3..10b8a2cdc5 100644
--- a/ide/idetop.ml
+++ b/ide/idetop.ml
@@ -466,7 +466,7 @@ let print_xml =
   let m = Mutex.create () in
   fun oc xml ->
     Mutex.lock m;
-    try Xml_printer.print oc xml; Mutex.unlock m
+    try Control.protect_sigalrm (Xml_printer.print oc) xml; Mutex.unlock m
     with e -> let e = CErrors.push e in Mutex.unlock m; iraise e
 
 let slave_feeder fmt xml_oc msg =
diff --git a/lib/control.ml b/lib/control.ml
index ffb3584f1e..9054507e46 100644
--- a/lib/control.ml
+++ b/lib/control.ml
@@ -89,3 +89,21 @@ let timeout_fun_ref = ref timeout_fun
 let set_timeout f = timeout_fun_ref := f
 
 let timeout n f e = !timeout_fun_ref.timeout n f e
+
+let protect_sigalrm f x =
+  let timed_out = ref false in
+  let timeout_handler _ = timed_out := true in
+  try
+    let old_handler = Sys.signal Sys.sigalrm (Sys.Signal_handle timeout_handler) in
+    try
+      let res = f x in
+      Sys.set_signal Sys.sigalrm old_handler;
+      match !timed_out, old_handler with
+      | true, Sys.Signal_handle f -> f Sys.sigalrm; res
+      | _, _ -> res
+    with e ->
+      let e = Backtrace.add_backtrace e in
+      Sys.set_signal Sys.sigalrm old_handler;
+      Exninfo.iraise e
+  with Invalid_argument _ -> (* This happens on Windows, as handling SIGALRM does not seem supported *)
+    f x
diff --git a/lib/control.mli b/lib/control.mli
index 59e2a15158..640d41a4f7 100644
--- a/lib/control.mli
+++ b/lib/control.mli
@@ -29,3 +29,14 @@ val timeout : int -> ('a -> 'b) -> 'a -> exn -> 'b
    API and it is scheduled to go away. *)
 type timeout = { timeout : 'a 'b. int -> ('a -> 'b) -> 'a -> exn -> 'b }
 val set_timeout : timeout -> unit
+
+(** [protect_sigalrm f x] computes [f x], but if SIGALRM is received during that
+    computation, the signal handler is executed only once the computation is
+    terminated. Otherwise said, it makes the execution of [f] atomic w.r.t.
+    handling of SIGALRM.
+
+    This is useful for example to prevent the implementation of `Timeout` to
+    interrupt I/O routines, generating ill-formed output.
+
+*)
+val protect_sigalrm : ('a -> 'b) -> 'a -> 'b
diff --git a/stm/asyncTaskQueue.ml b/stm/asyncTaskQueue.ml
index 73b9ef7da0..d1bd3692ab 100644
--- a/stm/asyncTaskQueue.ml
+++ b/stm/asyncTaskQueue.ml
@@ -329,10 +329,12 @@ module Make(T : Task) () = struct
   let main_loop () =
     (* We pass feedback to master *)
     let slave_feeder oc fb =
-      Marshal.to_channel oc (RespFeedback (debug_with_pid fb)) []; flush oc in
+      Control.protect_sigalrm (fun () ->
+          Marshal.to_channel oc (RespFeedback (debug_with_pid fb)) []; flush oc) ()
+    in
     ignore (Feedback.add_feeder (fun x -> slave_feeder (Option.get !slave_oc) x));
     (* We ask master to allocate universe identifiers *)
-    UnivGen.set_remote_new_univ_id (bufferize (fun () ->
+    UnivGen.set_remote_new_univ_id (bufferize @@ Control.protect_sigalrm (fun () ->
       marshal_response (Option.get !slave_oc) RespGetCounterNewUnivLevel;
       match unmarshal_more_data (Option.get !slave_ic) with
       | MoreDataUnivLevel l -> l));