From ce6d9a1c183d87844fc981e439aec1c1f77b986d Mon Sep 17 00:00:00 2001
From: eutro <git@eutro.dev>
Date: Fri, 19 Apr 2024 21:30:51 +0100
Subject: [PATCH] Allow olly to attach to an external process

- The command line changes are nontrivial in that they cannot be
expressed with usual `Cmdliner` combinators. Attaching is done with an
`--attach` option that is incompatible with the normal `EXECUTABLE`
arguments

- `EXECUTABLE` is now a sequence of arguments, rather than a single
space-separated one. This makes the above mutual exclusion slightly
easier to represent, and is more in-line with existing tools (`perf`,
`gdb --args`).

- `Unix.kill pid 0` is used to check if the process is alive, (since
`wait` can only be used for child processes), this is valid and
intended usage according to the POSIX specification.
---
 lib/olly_common/cli.ml    | 62 +++++++++++++++++++++++++++++++++++----
 lib/olly_common/launch.ml | 24 ++++++++++++---
 2 files changed, 76 insertions(+), 10 deletions(-)

diff --git a/lib/olly_common/cli.ml b/lib/olly_common/cli.ml
index 10e0e64..66ace74 100644
--- a/lib/olly_common/cli.ml
+++ b/lib/olly_common/cli.ml
@@ -34,13 +34,63 @@ let help man_format cmds topic =
           `Ok (Manpage.print man_format Format.std_formatter page))
 
 let exec_args p =
-  let doc =
-    "Executable (and its arguments) to trace. If the executable takes\n\
-    \              arguments, wrap quotes around the executable and its \
-     arguments.\n\
-    \              For example, olly '<exec> <arg_1> <arg_2> ... <arg_n>'."
+  let exec_and_args, ea_docv =
+    let doc = "Executable and arguments to trace." in
+    let docv = "EXECUTABLE" in
+    (Arg.(value & pos_right (p - 1) string [] & info [] ~docv ~doc), docv)
   in
-  Arg.(required & pos p (some string) None & info [] ~docv:"EXECUTABLE" ~doc)
+  let attach_opt, ao_docv =
+    let doc =
+      "Attach to the process with the given PID. The directory containing the \
+       PID.events file may be specified. This option cannot be combined with \
+       EXECUTABLE."
+    in
+    let docv = "[directory:]pid" in
+    let parser str =
+      let exception Fail of string in
+      try
+        let dir, pid_str =
+          match String.rindex_opt str ':' with
+          | None -> (".", str)
+          | Some idx ->
+              ( String.sub str 0 idx,
+                String.sub str (idx + 1) (String.length str - idx - 1) )
+        in
+        let pid =
+          try int_of_string pid_str
+          with _ ->
+            raise (Fail (Printf.sprintf "expected integer pid, got %s" pid_str))
+        in
+        if not @@ Sys.file_exists dir then
+          raise (Fail (Printf.sprintf "directory %s does not exist" dir));
+        if not @@ Sys.is_directory dir then
+          raise (Fail (Printf.sprintf "file %s is not a directory" dir));
+        Ok (dir, pid)
+      with Fail msg -> Error msg
+    in
+    let printer fmt (dir, pid) =
+      match dir with
+      | "." -> Format.fprintf fmt "%d" pid
+      | _ -> Format.fprintf fmt "%s:%d" dir pid
+    in
+    let dir_and_pid_conv = Arg.conv' ~docv (parser, printer) in
+    ( Arg.(
+        value
+        & opt (some dir_and_pid_conv) None
+        & info [ "a"; "attach" ] ~docv ~doc),
+      docv )
+  in
+  let cat_docvs sep = Printf.sprintf "%s %s --attach=%s" ea_docv sep ao_docv in
+  let combine dir_and_pid args =
+    match (args, dir_and_pid) with
+    | [], Some (dir, pid) -> Ok (Launch.Attach (dir, pid))
+    | _ :: _, None -> Ok (Launch.Execute args)
+    | [], None ->
+        Error (Printf.sprintf "required %s is missing" (cat_docvs "or"))
+    | _ ->
+        Error (Printf.sprintf "more than one of %s specified" (cat_docvs "and"))
+  in
+  Term.(term_result' ~usage:true (const combine $ attach_opt $ exec_and_args))
 
 let main name commands =
   let help_cmd =
diff --git a/lib/olly_common/launch.ml b/lib/olly_common/launch.ml
index 242dec0..d86c59a 100644
--- a/lib/olly_common/launch.ml
+++ b/lib/olly_common/launch.ml
@@ -7,8 +7,9 @@ type subprocess = {
   close : unit -> unit;
 }
 
-let exec_process exec_args =
-  let argsl = String.split_on_char ' ' exec_args in
+type exec_config = Attach of string * int | Execute of string list
+
+let exec_process argsl =
   let executable_filename = List.hd argsl in
 
   (* TODO Set the temp directory. We should make this configurable. *)
@@ -44,6 +45,21 @@ let exec_process exec_args =
   in
   { alive; cursor; close }
 
+let attach_process dir pid =
+  let cursor = Runtime_events.create_cursor (Some (dir, pid)) in
+  let alive () =
+    try
+      Unix.kill pid 0;
+      true
+    with Unix.Unix_error (Unix.ESRCH, _, _) -> false
+  and close () = Runtime_events.free_cursor cursor in
+  { alive; cursor; close }
+
+let launch_process exec_args =
+  match exec_args with
+  | Execute argsl -> exec_process argsl
+  | Attach (dir, pid) -> attach_process dir pid
+
 let collect_events child callbacks =
   (* Read from the child process *)
   while child.alive () do
@@ -76,10 +92,10 @@ let empty_config =
     cleanup = (fun () -> ());
   }
 
-let olly config exec_args =
+let olly config (exec_args : exec_config) =
   config.init ();
   Fun.protect ~finally:config.cleanup (fun () ->
-      let child = exec_process exec_args in
+      let child = launch_process exec_args in
       Fun.protect ~finally:child.close (fun () ->
           let callbacks =
             let {