use crate::suite::Suite;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
pub const PROTOCOL_MAJOR: u32 = 1;
pub const PROTOCOL_MINOR: u32 = 0;
pub const RANK_UNKNOWN: u8 = u8::MAX;
fn is_zero(v: &u64) -> bool {
*v == 0
}
fn rank_unknown() -> u8 {
RANK_UNKNOWN
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ServiceRecord {
pub protocol_major: u32,
pub protocol_minor: u32,
#[serde(default)]
pub service: String,
pub role: String,
#[serde(default = "rank_unknown")]
pub rank: u8,
pub pid: u32,
pub endpoint: String,
#[serde(default)]
pub version: String,
#[serde(default)]
pub started_ms: u64,
#[serde(default, skip_serializing_if = "is_zero")]
pub process_started_ms: u64,
#[serde(default)]
pub engine_version: String,
}
impl ServiceRecord {
pub fn host_version(&self) -> &str {
if self.version.is_empty() {
&self.engine_version
} else {
&self.version
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TakeoverRequest {
pub role: String,
#[serde(default = "rank_unknown")]
pub rank: u8,
pub pid: u32,
pub requested_at_ms: u64,
}
pub fn record_path(suite: &Suite, service: &str) -> PathBuf {
suite.service_dir(service).join("service.json")
}
pub fn takeover_path(suite: &Suite, service: &str) -> PathBuf {
suite.service_dir(service).join("takeover.json")
}
fn read_json<T: serde::de::DeserializeOwned>(path: &Path) -> std::io::Result<Option<T>> {
match std::fs::read(path) {
Ok(bytes) => Ok(serde_json::from_slice(&bytes).ok()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(e),
}
}
pub fn try_read_record(suite: &Suite, service: &str) -> std::io::Result<Option<ServiceRecord>> {
read_json(&record_path(suite, service))
}
pub fn try_read_takeover(suite: &Suite, service: &str) -> std::io::Result<Option<TakeoverRequest>> {
read_json(&takeover_path(suite, service))
}
pub fn read_record(suite: &Suite, service: &str) -> Option<ServiceRecord> {
try_read_record(suite, service).ok().flatten()
}
pub fn read_takeover(suite: &Suite, service: &str) -> Option<TakeoverRequest> {
try_read_takeover(suite, service).ok().flatten()
}
pub fn read_live_record(suite: &Suite, service: &str) -> Option<ServiceRecord> {
try_read_live_record(suite, service).ok().flatten()
}
pub fn try_read_live_record(
suite: &Suite,
service: &str,
) -> std::io::Result<Option<ServiceRecord>> {
Ok(try_read_record(suite, service)?.filter(crate::discovery::record_process_matches))
}
pub fn write_takeover(suite: &Suite, service: &str, req: &TakeoverRequest) -> std::io::Result<()> {
crate::suite::make_private_dir(&suite.service_dir(service)).map_err(std::io::Error::other)?;
let path = takeover_path(suite, service);
write_atomic(&path, &serde_json::to_vec_pretty(req).unwrap_or_default())?;
crate::takeover_watch::poke(&path);
Ok(())
}
pub fn clear_takeover_if_owner(suite: &Suite, service: &str, pid: u32) {
if read_takeover(suite, service).is_some_and(|t| t.pid == pid) {
let _ = std::fs::remove_file(takeover_path(suite, service));
}
}
pub fn write_atomic(path: &Path, body: &[u8]) -> std::io::Result<()> {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
match rightkit_fs::atomic_write_private(path, body) {
Ok(()) => return Ok(()),
Err(e)
if cfg!(windows)
&& e.kind() == Some(std::io::ErrorKind::PermissionDenied)
&& std::time::Instant::now() < deadline =>
{
std::thread::sleep(std::time::Duration::from_millis(10));
}
Err(rightkit_fs::FsError::Io { source, .. }) => return Err(source),
Err(e) => return Err(std::io::Error::other(e)),
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Allowlist {
apps: BTreeMap<String, BTreeSet<String>>,
}
impl Allowlist {
pub fn new() -> Self {
Self::default()
}
pub fn allow<I, S>(mut self, app: &str, methods: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.apps
.entry(app.to_string())
.or_default()
.extend(methods.into_iter().map(Into::into));
self
}
fn entry(&self, app: &str) -> Option<&BTreeSet<String>> {
self.apps.get(app).or_else(|| self.apps.get("*"))
}
pub fn app_known(&self, app: &str) -> bool {
self.entry(app).is_some()
}
pub fn permits(&self, app: &str, method: &str) -> bool {
method == "ping"
|| self
.entry(app)
.is_some_and(|m| m.contains(method) || m.contains("*"))
}
pub fn methods_for(&self, app: &str) -> Vec<String> {
self.entry(app)
.map(|m| m.iter().cloned().collect())
.unwrap_or_default()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn suite() -> (tempfile::TempDir, Suite) {
let d = tempfile::tempdir().unwrap();
let s = Suite::at(d.path().join("suite")).unwrap();
(d, s)
}
fn unreadable(path: &Path) {
std::fs::create_dir_all(path).unwrap();
}
#[test]
fn absent_files_are_none_not_errors() {
let (_d, s) = suite();
assert!(try_read_record(&s, "demo").unwrap().is_none());
assert!(try_read_takeover(&s, "demo").unwrap().is_none());
assert!(try_read_live_record(&s, "demo").unwrap().is_none());
assert!(read_record(&s, "demo").is_none() && read_takeover(&s, "demo").is_none());
}
#[test]
fn unparseable_files_stay_none() {
let (_d, s) = suite();
std::fs::create_dir_all(s.service_dir("demo")).unwrap();
std::fs::write(record_path(&s, "demo"), b"not json").unwrap();
std::fs::write(takeover_path(&s, "demo"), b"{").unwrap();
assert!(try_read_record(&s, "demo").unwrap().is_none());
assert!(try_read_takeover(&s, "demo").unwrap().is_none());
}
#[test]
fn unreadable_files_are_errors_in_the_try_readers() {
let (_d, s) = suite();
unreadable(&record_path(&s, "demo"));
unreadable(&takeover_path(&s, "demo"));
let e = try_read_record(&s, "demo").unwrap_err();
assert_ne!(e.kind(), std::io::ErrorKind::NotFound);
let e = try_read_takeover(&s, "demo").unwrap_err();
assert_ne!(e.kind(), std::io::ErrorKind::NotFound);
assert!(try_read_live_record(&s, "demo").is_err());
assert!(read_record(&s, "demo").is_none() && read_takeover(&s, "demo").is_none());
clear_takeover_if_owner(&s, "demo", std::process::id());
assert!(takeover_path(&s, "demo").is_dir());
}
}