From aebdb4923c6f998ecc042c3aaf3021752b608e28 Mon Sep 17 00:00:00 2001
From: zach <zach@dylibso.com>
Date: Mon, 13 Nov 2023 12:27:26 -0800
Subject: [PATCH] feat: add set_log_callback

---
 src/bindings.ml | 11 +++++++++++
 src/extism.ml   | 48 +++++++++++++++++++++++++++++++++++++-----------
 src/extism.mli  |  6 ++++++
 3 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/src/bindings.ml b/src/bindings.ml
index e43c9fc..57866b7 100644
--- a/src/bindings.ml
+++ b/src/bindings.ml
@@ -126,6 +126,17 @@ let extism_plugin_output_data =
 let extism_log_file =
   fn "extism_log_file" (string @-> string_opt @-> returning bool)
 
+let log_callback =
+  Foreign.dynamic_funptr ~runtime_lock:true
+    Ctypes.(ptr char @-> uintptr_t @-> returning void)
+
+let extism_log_callback name f =
+  let (module Callback) = log_callback in
+  let cb = Callback.of_fun f in
+  fn "extism_log_callback"
+    (string_opt @-> Callback.t @-> returning bool)
+    name cb
+
 let extism_version = fn "extism_version" (void @-> returning string)
 let extism_plugin_free = fn "extism_plugin_free" (plugin @-> returning void)
 
diff --git a/src/extism.ml b/src/extism.ml
index 8417479..b079f9c 100644
--- a/src/extism.ml
+++ b/src/extism.ml
@@ -30,17 +30,43 @@ let%test "with_plugin" =
 
 let%test _ = String.length (extism_version ()) > 0
 
+let parse_level =
+  Option.map (function
+    | `Error -> "error"
+    | `Warn -> "warn"
+    | `Info -> "info"
+    | `Debug -> "debug"
+    | `Trace -> "trace")
+
 let set_log_file ?level filename =
-  let level =
-    Option.map
-      (function
-        | `Error -> "error"
-        | `Warn -> "warn"
-        | `Info -> "info"
-        | `Debug -> "debug"
-        | `Trace -> "trace")
-      level
+  Bindings.extism_log_file filename (parse_level level)
+
+let callback = ref None
+
+let set_log_callback ?level f =
+  callback := Some f;
+  let fx s length =
+    f @@ Ctypes.string_from_ptr s ~length:(Ctypes.Uintptr.to_int length)
   in
-  Bindings.extism_log_file filename level
+  Bindings.extism_log_callback (parse_level level) fx
 
-let%test _ = set_log_file ~level:`Trace "stderr"
+let%test _ =
+  let log_file =
+    try
+      let _ = Unix.getenv "TEST_LOG_FILE" in
+      true
+    with Not_found -> false
+  in
+  if log_file then set_log_file ~level:`Trace "stderr"
+  else
+    let logs = ref [] in
+    let ok =
+      set_log_callback ~level:`Trace (fun line -> logs := line :: !logs)
+    in
+    let manifest = Manifest.(create [ Wasm.file "test/code.wasm" ]) in
+    let plugin = Plugin.of_manifest manifest |> Error.unwrap in
+    let _ =
+      Plugin.call Type.string Type.string plugin ~name:"count_vowels"
+        "this is a test"
+    in
+    ok && List.length !logs > 0
diff --git a/src/extism.mli b/src/extism.mli
index a9301d6..05e8b62 100644
--- a/src/extism.mli
+++ b/src/extism.mli
@@ -425,6 +425,12 @@ val set_log_file :
 (** Set the log file and level for all Extism plugins, the names [stdout] or
     [stderr] can be used to write to the terminal *)
 
+val set_log_callback :
+  ?level:[ `Error | `Warn | `Info | `Debug | `Trace ] ->
+  (string -> unit) ->
+  bool
+(** Set the log level and callback function *)
+
 (** Extism error type *)
 module Error : sig
   type t = [ `Msg of string ]