systemd-ml/lib/systemd.ml

let is_systemd_booted env =
  let fs = Eio.Stdenv.fs env in
  Eio.Path.( / ) fs "/run/systemd/system/" |> Eio.Path.is_directory 

let is_in_systemd () =
  begin match Option.bind (Sys.getenv_opt "SYSTEMD_EXEC_PID") int_of_string_opt with
  | None -> false
  | Some(pid) ->
      pid = Unix.getpid ()
  end

let invocation_id () =
  Sys.getenv "INVOCATION_ID"

let is_journald_attached () =
  let stderr = Unix.fstat Unix.stderr in
  let dev_ino = Format.sprintf "%d:%d" stderr.st_dev stderr.st_ino in
  Sys.getenv_opt "JOURNAL_STREAM" = Some dev_ino

module Dirs = struct
  let runtime_dir () =
    Sys.getenv "RUNTIME_DIRECTORY"

  let state_dir () =
    Sys.getenv "STATE_DIRECTORY"

  let cache_dir () =
    Sys.getenv "CACHE_DIRECTORY"

  let logs_dir () =
    Sys.getenv "LOGS_DIRECTORY"

  let configuration_dir () =
    Sys.getenv "CONFIGURATION_DIRECTORY"

  let credentials_dir () =
    Sys.getenv "CREDENTIALS_DIRECTORY"
end

module Notify = struct
  type t = {
    sock : Eio_unix.Fd.t;
    mono_clock : Eio.Time.Mono.ty Eio.Time.Mono.t;
  }

  module Private = struct
    external unix_code_of_unix_error : Unix.error -> int = "ocaml_systemd_unix_code_of_unix_error"
  end

  let make ~sw ~env =
  (* debatable whether we need a whole fd for this - reference code just opens and closes the socket
     each time, and this avoids keeping the fds open the whole time. however, this also means that
     if we run out of fds, notify will fail. currently, we just keep the fd open forever *)
    if (Eio.Stdenv.backend_id env <> "linux") then failwith "Backend must be Eio_linux";
    let orig_addr = Sys.getenv "NOTIFY_SOCKET" in
    let addr = Bytes.of_string orig_addr in
    begin if Bytes.starts_with ~prefix:("@"[@bytes]) addr then
      Bytes.set addr 0 '\x00'
    else if not (Bytes.starts_with ~prefix:("/"[@bytes]) addr) then
      failwith "Invalid address family"
    end;
    let sock_unix = Unix.socket ~cloexec:true Unix.PF_UNIX Unix.SOCK_DGRAM 0 in
    let sock = Eio_unix.Fd.of_unix ~sw ~seekable:false ~close_unix:true sock_unix in
    Eio_linux.Low_level.connect sock (Unix.ADDR_UNIX (String.of_bytes addr));
    { sock = sock; mono_clock = Eio.Stdenv.mono_clock env; }

  let notify t ?(fds=[]) msg =
    let buf = [(Cstruct.of_string msg)] in
    let sent = Eio_linux.Low_level.send_msg ~fds:fds t.sock buf in
    assert (sent = Cstruct.lenv buf)

  let ready ?(status=None) ?(mainpid=(-1)) t =
    let buf = Buffer.create 32 in
    Buffer.add_string buf "READY=1\n";
    begin match status with
    | Some status -> Printf.bprintf buf "STATUS=%s\n" status
    | None -> ()
    end;
    if mainpid > -1 then Printf.bprintf buf "MAINPID=%u\n" mainpid;
    notify t (Buffer.contents buf)

  let send_status t status =
    notify t ("STATUS=" ^ status ^ "\n")

  let send_status_errno t status errno =
    let errno = Private.unix_code_of_unix_error errno in
    let msg = Format.sprintf "STATUS=%s\nERRNO=%d\n" status errno in
    notify t msg

  let reloading t =
    (* this happens to use CLOCK_MONOTONIC on Eio_linux *)
    let now = Eio.Time.Mono.now t.mono_clock in
    let now_ns = Mtime.to_uint64_ns now in
    let now = Int64.div now_ns 1000L in
    let msg = Format.sprintf "RELOADING=1\nMONOTONIC_USEC=%Lu\n" now in
    notify t msg
    
  let stopping t =
    notify t "STOPPING=1\n"

  module NotifyAccess = struct
    type t =
    | None
    | Main
    | Exec
    | All

    let to_string = function
    | None -> "none"
    | Main -> "main"
    | Exec -> "exec"
    | All -> "all"
  end

  let update_access t acc =
    notify t ("NOTIFYACCESS=" ^ (NotifyAccess.to_string acc) ^ "\n")

  let extend_timeout t span =
    let nsec = Mtime.Span.to_uint64_ns span in
    let usec = Int64.div nsec 1000L in
    let msg = Format.sprintf "EXTEND_TIMEOUT_USEC=%Lu\n" usec in
    notify t msg

  let barrier t =
    Eio.Switch.run @@ fun sw ->
    let (p0, p1) = Unix.pipe ~cloexec:true () in
    let p0 = Eio_unix.Fd.of_unix ~sw ~seekable:false ~close_unix:true p0 in
    let p1 = Eio_unix.Fd.of_unix ~sw ~seekable:false ~close_unix:true p1 in
    notify t ~fds:[p1] "BARRIER=1\n";
    Eio_unix.Fd.close p1;
    (* the libsystemd source uses POLLHUP, which isn't exposed via Eio_linux at all, but POLLIN
       works to the same end here - it won't return until the pipe is closed (and no one should
       be writing anything to it) *)
    Eio_linux.Low_level.await_readable p0;
    Eio_unix.Fd.close p0
end

module Watchdog = struct
  let get_interval_opt () =
    match Option.bind (Sys.getenv_opt "WATCHDOG_USEC") Int64.of_string_opt with
    | None -> None
    | Some usec ->
        begin match Option.bind (Sys.getenv_opt "WATCHDOG_PID") int_of_string_opt with
        | None -> None
        | Some pid ->
          if pid = Unix.getpid () then
            let nsec = Int64.mul usec 1000L in
            Some (Mtime.Span.of_uint64_ns nsec)
          else
            None
        end

  let pet t =
    Notify.notify t "WATCHDOG=1\n"

  let trigger t =
    Notify.notify t "WATCHDOG=trigger\n"

  let set_watchdog_time t span =
    let nsec = Mtime.Span.to_uint64_ns span in
    let usec = Int64.div nsec 1000L in
    let msg = Format.sprintf "WATCHDOG_USEC=%Lu\n" usec in
    Notify.notify t msg
end

module Fdstore = struct
  module Private = struct
    let sd_listen_fds_start = 3
    external unsafe_make_fd : int -> Unix.file_descr = "%identity"
  end

  let is_valid_fdname s =
    let rec iter_chars i = begin
      if i = String.length s then
        true
      else
        let dec = String.get_utf_8_uchar s i in
        let chr = Uchar.utf_decode_uchar dec |> Uchar.to_int in
        let len = Uchar.utf_decode_length dec in
        if chr >= 32 && chr <= 126 && chr != 58 then
          iter_chars (i + len)
        else
          false
    end in
    iter_chars 0

  let get_fd_limit () =
    match Sys.getenv_opt "FDSTORE" with
    | None -> 0
    | Some limit -> int_of_string limit

  let store t name ?(poll=true) fd =
    if not (is_valid_fdname name) then failwith "Invalid fdstore name";
    let msg = Format.sprintf ("FDSTORE=1\nFDNAME=%s\n" ^^ (if poll then "" else "FDPOLL=0\n")) name in
    Notify.notify t ~fds:[fd] msg

  let remove t name =
    if not (is_valid_fdname name) then failwith "Invalid fdstore name";
    Notify.notify t (Format.sprintf "FDSTOREREMOVE=1\nFDNAME=%s\n" name)

  let listen_fds ~sw =
    match Option.bind (Sys.getenv_opt "LISTEN_PID") int_of_string_opt with
    | None -> []
    | Some(pid) ->
        if pid = Unix.getpid () then begin
          let fd_count = Sys.getenv "LISTEN_FDS" |> int_of_string in
          let fd_names = match Sys.getenv_opt "LISTEN_FDNAMES" with
          | None -> List.init fd_count (fun _ -> None)
          | Some names -> String.split_on_char ':' names |> List.map (fun a -> Some a) in
          assert (fd_count = List.length fd_names);
          Unix.putenv "LISTEN_FDS" "";
          Unix.putenv "LISTEN_FDNAMES" "";
          List.mapi (fun i name ->
            let fd_num = Private.sd_listen_fds_start + i in
            let fd = Private.unsafe_make_fd fd_num in
            Unix.set_close_on_exec fd;
            let fd = Eio_unix.Fd.of_unix ~sw ~seekable:false ~close_unix:true fd in
            (name, fd)
          ) fd_names
        end else []
end

(* we can't do memory pressure because of a lack of exposed pollpri :( *)

module MemPressure = struct
  module Private = struct
    external ext_memory_trim : unit -> unit = "ocaml_systemd_memory_trim"
  end

  let memory_trim () =
    Gc.compact ();
    Private.ext_memory_trim ()
end