Source file extism.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
module Manifest = Extism_manifest
module Error = Error
module Plugin = Plugin
module Function = Function
module Host_function = Host_function
module Type = Type
module Val = Val
module Val_type = Val_type

let extism_version = Bindings.extism_version

let with_plugin f p =
  Fun.protect ~finally:(fun () -> Plugin.free p) (fun () -> f p)

let%test "with_plugin" =
  let manifest = Manifest.(create [ Wasm.file "test/code.wasm" ]) in
  let plugin = Plugin.of_manifest manifest |> Error.unwrap in
  let b =
    with_plugin
      (fun plugin ->
        Plugin.call Type.string Type.string plugin ~name:"count_vowels"
          "this is a test"
        |> Error.unwrap = "{\"count\":4,\"total\":4,\"vowels\":\"aeiouAEIOU\"}")
      plugin
  in
  Plugin.free plugin;
  Gc.minor ();
  Gc.full_major ();
  b

let%test _ = String.length (extism_version ()) > 0

let parse_level =
  Option.map (function
    | `Error -> "error"
    | `Warn -> "warn"
    | `Info -> "info"
    | `Debug -> "debug"
    | `Trace -> "trace"
    | `Filter f -> f)

let set_log_file ?level filename =
  Bindings.extism_log_file filename (parse_level level)

type drain_logs = unit -> unit

let mk_drain_logs f () =
  let fx s length =
    let s = Ctypes.string_from_ptr s ~length:(Ctypes.Uintptr.to_int length) in
    f s
  in
  Bindings.extism_log_drain fx

let set_log_custom ?level f =
  if Bindings.extism_log_custom (parse_level level) then
    let x = mk_drain_logs f in
    let () = Gc.finalise_last (mk_drain_logs f) x in
    x
  else Error.throw (`Msg "Unable to set custom logging")

let%test _ =
  let log_file =
    try String.length @@ Unix.getenv "TEST_LOG_FILE" > 0
    with Not_found -> false
  in
  if log_file then set_log_file ~level:`Trace "stderr"
  else
    let count = ref 0 in
    let _drain_logs =
      set_log_custom ~level:`Trace (fun s ->
          incr count;
          print_string s)
    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
    (* Make sure logs are drained when drain_logs is garbage collected *)
    Gc.minor ();
    Gc.full_major ();
    !count > 0