use miette::{IntoDiagnostic, Result, WrapErr};
use smol_str::{SmolStr, ToSmolStr};
use std::collections::BTreeSet;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::str::FromStr;
use url::Url;
use super::Hydrant;
use crate::db::keys;
use crate::pds_meta::HostStatus;
/// unix sockets are owner-only unless told otherwise, since the api has no auth of its own.
pub const DEFAULT_UNIX_SOCKET_MODE: u32 = 0o600;
/// one address the api listens on.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ApiBind {
Tcp(SocketAddr),
Unix(PathBuf),
}
impl ApiBind {
pub fn tcp(&self) -> Option<SocketAddr> {
match self {
Self::Tcp(addr) => Some(*addr),
Self::Unix(_) => None,
}
}
}
impl From<SocketAddr> for ApiBind {
fn from(addr: SocketAddr) -> Self {
Self::Tcp(addr)
}
}
impl FromStr for ApiBind {
type Err = miette::Report;
/// parses `host:port`, or `unix:` followed by a socket path.
fn from_str(s: &str) -> Result<Self> {
match s.strip_prefix("unix:") {
Some("") => miette::bail!("`unix:` bind is missing a socket path"),
Some(path) => Ok(Self::Unix(path.into())),
None => s
.parse()
.map(Self::Tcp)
.map_err(|e| miette::miette!("invalid bind address `{s}`: {e}")),
}
}
}
impl std::fmt::Display for ApiBind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Tcp(addr) => write!(f, "{addr}"),
Self::Unix(path) => write!(f, "unix:{}", path.display()),
}
}
}
#[derive(Debug, Clone)]
pub struct ApiBinds {
pub(crate) head: ApiBind,
pub(crate) tail: Vec<ApiBind>,
pub(crate) unix_mode: u32,
pub(crate) tcp_management: bool,
}
impl ApiBinds {
pub fn new(bind: impl Into<ApiBind>) -> Self {
Self {
head: bind.into(),
tail: Vec::new(),
unix_mode: DEFAULT_UNIX_SOCKET_MODE,
tcp_management: false,
}
}
pub fn try_from_iter<I>(iter: I) -> Option<Self>
where
I: IntoIterator,
I::Item: Into<ApiBind>,
{
let mut iter = iter.into_iter().map(Into::into);
let head = iter.next()?;
Some(Self {
head,
tail: iter.collect(),
unix_mode: DEFAULT_UNIX_SOCKET_MODE,
tcp_management: false,
})
}
/// sets the permissions unix sockets are created with. only the `0o777` bits are used.
pub fn with_unix_mode(mut self, mode: u32) -> Self {
self.unix_mode = mode & 0o777;
self
}
/// serves the management endpoints on tcp binds too. they're unix-only by default, since
/// they have no auth and a unix socket's permissions are the only thing guarding them.
pub fn with_tcp_management(mut self, enabled: bool) -> Self {
self.tcp_management = enabled;
self
}
/// whether any bind serves the management endpoints.
pub fn serves_management(&self) -> bool {
self.tcp_management || self.iter().any(|bind| matches!(bind, ApiBind::Unix(_)))
}
pub fn iter(&self) -> impl Iterator<Item = &ApiBind> + '_ {
std::iter::once(&self.head).chain(self.tail.iter())
}
/// reads `HYDRANT_API_BIND`, `HYDRANT_API_SOCKET_MODE` and `HYDRANT_API_TCP_MANAGEMENT`
/// through `lookup`. an unset bind gives `default`, and `none` gives `None`, meaning no api
/// at all.
pub fn from_lookup(
lookup: impl Fn(&str) -> Option<String>,
default: Option<ApiBinds>,
) -> Result<Option<Self>> {
let binds = match lookup("HYDRANT_API_BIND") {
Some(raw) if raw.trim().eq_ignore_ascii_case("none") => None,
Some(raw) => Some(
raw.parse::<Self>()
.map_err(|e| miette::miette!("invalid HYDRANT_API_BIND: {e}"))?,
),
None => default,
};
let Some(mut binds) = binds else {
return Ok(None);
};
if let Some(mode) = lookup("HYDRANT_API_SOCKET_MODE") {
binds = binds.with_unix_mode(parse_socket_mode(&mode)?);
}
let tcp_management =
crate::config::env::read(&lookup, "HYDRANT_API_TCP_MANAGEMENT", binds.tcp_management)?;
Ok(Some(binds.with_tcp_management(tcp_management)))
}
}
/// octal permission bits like `0660`, with or without the leading `0` or `0o`.
fn parse_socket_mode(s: &str) -> Result<u32> {
let digits = s.trim();
let digits = digits.strip_prefix("0o").unwrap_or(digits);
u32::from_str_radix(digits, 8)
.ok()
.filter(|mode| !digits.starts_with('+') && *mode <= 0o777)
.ok_or_else(|| {
miette::miette!(
"invalid HYDRANT_API_SOCKET_MODE `{s}`: want octal permissions like 0660"
)
})
}
impl FromStr for ApiBinds {
type Err = miette::Report;
/// parses a comma-separated list of [`ApiBind`]s, the format of `HYDRANT_API_BIND`.
fn from_str(s: &str) -> Result<Self> {
let binds = s
.split(',')
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::parse::<ApiBind>)
.collect::<Result<Vec<_>>>()?;
Self::try_from_iter(binds).ok_or_else(|| miette::miette!("no addresses to bind to"))
}
}
#[derive(Debug, Clone)]
/// information about a host hydrant is consuming from.
pub struct Host {
/// hostname of the host.
pub name: SmolStr,
/// latest seq hydrant has processed from this host.
pub seq: i64,
/// the amount of accounts hydrant has seen from this host.
pub account_count: u64,
/// the status of this host in hydrant.
pub status: crate::pds_meta::HostStatus,
}
impl Hydrant {
/// get the status of a (firehose) host we are consuming from.
///
/// returns the seq we are on for this host.
pub async fn get_host_status(&self, hostname: &str) -> Result<Option<Host>> {
let state = self.state.clone();
let hostname = hostname.to_smolstr();
let state_closure = state.clone();
state
.db
.run(move |db| {
let key = keys::firehose_cursor_key(&hostname);
let mut seq = 0;
if let Some(cursor_bytes) = db.cursors.get(&key).into_diagnostic()? {
seq = i64::from_be_bytes(cursor_bytes.as_ref().try_into().into_diagnostic()?);
} else {
// if it has no cursor, check if it's explicitly tracked in hosts map
// or firehose tasks (recently added via API but no messages yet)
let meta = state_closure.pds_meta.load();
if !meta.hosts.contains_key(hostname.as_str()) {
// we should also allow it if it's an active firehose ingestor
let mut found_in_cursors = false;
state_closure.firehose_cursors.iter_sync(|u, _| {
if u.host_str() == Some(hostname.as_str()) {
found_in_cursors = true;
}
!found_in_cursors // continue if not found
});
if !found_in_cursors {
return Ok(None);
}
}
}
let account_count = db.get_count_sync(&keys::pds_account_count_key(&hostname));
let status = state_closure.pds_meta.load().status(&hostname);
Ok(Some(Host {
name: hostname,
seq,
account_count,
status,
}))
})
.await
}
/// enumerates all hosts hydrant is consuming from.
///
/// returns hosts enumerated in this pagination and the cursor to paginate from.
pub async fn list_hosts(
&self,
cursor: Option<&str>,
limit: usize,
) -> Result<(Vec<Host>, Option<SmolStr>)> {
let state = self.state.clone();
let cursor = cursor.map(str::to_string);
let state_closure = state.clone();
state
.db
.run(move |db| {
let start_bound = match &cursor {
Some(after) => std::ops::Bound::Included(keys::firehose_cursor_key(after)),
None => std::ops::Bound::Included(keys::FIREHOSE_CURSOR_PREFIX.to_vec()),
};
let prefix_end = {
let mut end = keys::FIREHOSE_CURSOR_PREFIX.to_vec();
*end.last_mut().unwrap() += 1;
end
};
let end_bound = std::ops::Bound::Excluded(prefix_end);
let mut db_hosts = Vec::new();
for item in db.cursors.range((start_bound, end_bound)) {
let (k, _) = item.into_inner().into_diagnostic()?;
let hostname = std::str::from_utf8(&k[keys::FIREHOSE_CURSOR_PREFIX.len()..])
.into_diagnostic()
.wrap_err("firehose cursor key contains non-utf8 hostname")?;
if cursor
.as_ref()
.is_some_and(|after| hostname <= after.as_str())
{
continue;
}
db_hosts.push(SmolStr::new(hostname));
if db_hosts.len() > limit {
break;
}
}
let mut meta_hosts = Vec::new();
{
let meta = state_closure.pds_meta.load();
for hostname in meta.hosts.keys() {
if cursor
.as_ref()
.is_some_and(|after| hostname.as_str() <= after.as_str())
{
continue;
}
meta_hosts.push(SmolStr::new(hostname));
}
}
meta_hosts.sort();
meta_hosts.truncate(limit + 1);
let mut merged = BTreeSet::new();
for host in db_hosts {
merged.insert(host);
}
for host in meta_hosts {
merged.insert(host);
}
let selected: Vec<SmolStr> = merged.into_iter().take(limit + 1).collect();
let mut hosts: Vec<Host> = Vec::with_capacity(selected.len().min(limit));
for hostname in selected.iter().take(limit) {
let seq = db
.cursors
.get(keys::firehose_cursor_key(hostname))
.into_diagnostic()?
.map(|v| {
v.as_ref()
.try_into()
.into_diagnostic()
.wrap_err("cursor value is not 8 bytes")
.map(i64::from_be_bytes)
})
.transpose()?
.unwrap_or(0);
let account_count = db.get_count_sync(&keys::pds_account_count_key(hostname));
let status = state_closure.pds_meta.load().status(hostname);
hosts.push(Host {
name: hostname.clone(),
seq,
account_count,
status,
});
}
let next_cursor = if selected.len() > limit {
hosts.last().map(|h| h.name.clone())
} else {
None
};
Ok((hosts, next_cursor))
})
.await
}
}
#[cfg(test)]
mod api_bind_tests {
use super::*;
#[test]
fn parses_tcp_and_unix_binds() {
assert_eq!(
"127.0.0.1:3000".parse::<ApiBind>().unwrap(),
ApiBind::Tcp("127.0.0.1:3000".parse().unwrap())
);
assert_eq!(
"[::1]:3000".parse::<ApiBind>().unwrap(),
ApiBind::Tcp("[::1]:3000".parse().unwrap())
);
assert_eq!(
"unix:/run/hydrant/api.sock".parse::<ApiBind>().unwrap(),
ApiBind::Unix("/run/hydrant/api.sock".into())
);
assert!("unix:".parse::<ApiBind>().is_err());
assert!("localhost:3000".parse::<ApiBind>().is_err());
assert!("/run/hydrant/api.sock".parse::<ApiBind>().is_err());
}
#[test]
fn display_round_trips() {
for s in ["127.0.0.1:3000", "[::1]:3000", "unix:/run/hydrant/api.sock"] {
assert_eq!(s.parse::<ApiBind>().unwrap().to_string(), s);
}
}
#[test]
fn parses_bind_lists() {
let binds: ApiBinds = " 127.0.0.1:3000, unix:/tmp/api.sock ,".parse().unwrap();
assert_eq!(
binds.iter().cloned().collect::<Vec<_>>(),
[
ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()),
ApiBind::Unix("/tmp/api.sock".into()),
]
);
assert!("".parse::<ApiBinds>().is_err());
assert!(" , ".parse::<ApiBinds>().is_err());
assert!("127.0.0.1:3000,nope".parse::<ApiBinds>().is_err());
}
fn lookup<'a>(pairs: &'a [(&str, &str)]) -> impl Fn(&str) -> Option<String> + 'a {
|key| {
pairs
.iter()
.find(|(k, _)| *k == key)
.map(|(_, v)| v.to_string())
}
}
#[test]
fn from_lookup_reads_bind_and_socket_mode() {
let default = ApiBinds::new(ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()));
let unset = ApiBinds::from_lookup(lookup(&[]), Some(default.clone())).unwrap();
assert_eq!(unset.unwrap().iter().count(), 1);
assert!(ApiBinds::from_lookup(lookup(&[]), None).unwrap().is_none());
for none in ["none", " NONE "] {
let binds =
ApiBinds::from_lookup(lookup(&[("HYDRANT_API_BIND", none)]), Some(default.clone()));
assert!(binds.unwrap().is_none(), "{none:?} should turn the api off");
}
let unix =
ApiBinds::from_lookup(lookup(&[("HYDRANT_API_BIND", "unix:/tmp/api.sock")]), None);
assert_eq!(unix.unwrap().unwrap().unix_mode, 0o600);
for (mode, want) in [
("0660", 0o660),
("660", 0o660),
("0o640", 0o640),
("0777", 0o777),
] {
let pairs = [
("HYDRANT_API_BIND", "unix:/tmp/api.sock"),
("HYDRANT_API_SOCKET_MODE", mode),
];
let binds = ApiBinds::from_lookup(lookup(&pairs), None)
.unwrap()
.unwrap();
assert_eq!(binds.unix_mode, want, "mode {mode}");
}
}
#[test]
fn management_is_unix_only_unless_enabled_for_tcp() {
let tcp = ApiBinds::new(ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()));
assert!(!tcp.tcp_management);
assert!(!tcp.serves_management());
assert!(tcp.clone().with_tcp_management(true).serves_management());
let both: ApiBinds = "127.0.0.1:3000,unix:/tmp/api.sock".parse().unwrap();
assert!(both.serves_management());
// the flag applies to the default binds too, not just ones set in HYDRANT_API_BIND
let enabled = ApiBinds::from_lookup(
lookup(&[("HYDRANT_API_TCP_MANAGEMENT", "true")]),
Some(tcp.clone()),
);
assert!(enabled.unwrap().unwrap().tcp_management);
let disabled = ApiBinds::from_lookup(
lookup(&[("HYDRANT_API_TCP_MANAGEMENT", "false")]),
Some(tcp),
);
assert!(!disabled.unwrap().unwrap().tcp_management);
}
#[test]
fn from_lookup_refuses_bad_values() {
let sock = ("HYDRANT_API_BIND", "unix:/tmp/api.sock");
for bad in [
&[("HYDRANT_API_BIND", "")][..],
&[("HYDRANT_API_BIND", "none,unix:/tmp/api.sock")],
&[sock, ("HYDRANT_API_SOCKET_MODE", "0800")],
&[sock, ("HYDRANT_API_SOCKET_MODE", "1777")],
&[sock, ("HYDRANT_API_SOCKET_MODE", "+660")],
&[sock, ("HYDRANT_API_SOCKET_MODE", "rw")],
&[sock, ("HYDRANT_API_TCP_MANAGEMENT", "yes")],
] {
assert!(
ApiBinds::from_lookup(lookup(bad), None).is_err(),
"{bad:?} should be refused"
);
}
}
#[test]
fn with_unix_mode_keeps_only_permission_bits() {
let binds = ApiBinds::new(ApiBind::Unix("/tmp/api.sock".into())).with_unix_mode(0o4660);
assert_eq!(binds.unix_mode, 0o660);
}
#[test]
fn binds_can_be_unix_only() {
let binds = ApiBinds::new(ApiBind::Unix("/tmp/api.sock".into()));
assert_eq!(binds.iter().filter_map(ApiBind::tcp).count(), 0);
assert_eq!(binds.iter().count(), 1);
}
}
#[cfg(test)]
mod host_listing_tests {
use super::*;
use crate::config::Config;
use crate::db::set_ks_count;
use crate::pds_meta::HostStatus;
fn test_config(path: &std::path::Path) -> Config {
Config {
database_path: path.to_path_buf(),
..Default::default()
}
}
#[tokio::test]
async fn list_hosts_includes_seeded_hosts_without_cursors() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let hydrant = Hydrant::new(test_config(tmp.path())).await?;
{
let state = hydrant.state.clone();
state
.db
.run(move |db| -> Result<()> {
let mut batch = db.inner.batch();
crate::db::pds_meta::set_status(
&mut batch,
&db.filter,
"offline.example",
HostStatus::Offline,
)?;
crate::db::pds_meta::set_status(
&mut batch,
&db.filter,
"active.example",
HostStatus::Active,
)?;
set_ks_count(
&mut batch,
db,
&keys::pds_account_count_key("offline.example"),
7,
);
set_ks_count(
&mut batch,
db,
&keys::pds_account_count_key("active.example"),
42,
);
db.cursors
.insert(
keys::firehose_cursor_key("active.example"),
123_i64.to_be_bytes(),
)
.into_diagnostic()?;
batch.commit().into_diagnostic()?;
db.persist()
})
.await?;
crate::pds_meta::PdsMeta::update_host(
&hydrant.state.pds_meta,
"offline.example",
|h| {
h.status = HostStatus::Offline;
},
);
crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "active.example", |h| {
h.status = HostStatus::Active;
});
hydrant.state.db.update_count("p|offline.example", 7);
hydrant.state.db.update_count("p|active.example", 42);
}
let (hosts, next) = hydrant.list_hosts(None, 10).await?;
assert!(next.is_none());
let offline = hosts
.iter()
.find(|h| h.name == "offline.example")
.expect("seeded host without cursor should be listed");
let active = hosts
.iter()
.find(|h| h.name == "active.example")
.expect("host with cursor should be listed");
assert_eq!(offline.seq, 0);
assert_eq!(offline.account_count, 7);
assert_eq!(offline.status, HostStatus::Offline);
assert_eq!(active.seq, 123);
assert_eq!(active.account_count, 42);
assert_eq!(active.status, HostStatus::Active);
let status = hydrant.get_host_status("active.example").await?;
assert_eq!(status.map(|host| host.seq), Some(123));
Ok(())
}
#[tokio::test]
async fn list_hosts_pagination_correctness() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let hydrant = Hydrant::new(test_config(tmp.path())).await?;
// Seed some in DB: host2, host4, host6
{
let state = hydrant.state.clone();
state
.db
.run(move |db| -> Result<()> {
db.cursors
.insert(keys::firehose_cursor_key("host2"), 2_i64.to_be_bytes())
.into_diagnostic()?;
db.cursors
.insert(keys::firehose_cursor_key("host4"), 4_i64.to_be_bytes())
.into_diagnostic()?;
db.cursors
.insert(keys::firehose_cursor_key("host6"), 6_i64.to_be_bytes())
.into_diagnostic()?;
db.persist()
})
.await?;
}
// Seed some in memory: host1, host3, host5
crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host1", |h| {
h.status = HostStatus::Active;
});
crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host3", |h| {
h.status = HostStatus::Active;
});
crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host5", |h| {
h.status = HostStatus::Active;
});
// Lexicographically sorted: host1, host2, host3, host4, host5, host6
// Page 1: limit = 2
let (hosts1, next1) = hydrant.list_hosts(None, 2).await?;
assert_eq!(hosts1.len(), 2);
assert_eq!(hosts1[0].name.as_str(), "host1");
assert_eq!(hosts1[1].name.as_str(), "host2");
assert_eq!(next1.as_deref(), Some("host2"));
// Page 2: limit = 2, cursor = host2
let (hosts2, next2) = hydrant.list_hosts(next1.as_deref(), 2).await?;
assert_eq!(hosts2.len(), 2);
assert_eq!(hosts2[0].name.as_str(), "host3");
assert_eq!(hosts2[1].name.as_str(), "host4");
assert_eq!(next2.as_deref(), Some("host4"));
// Page 3: limit = 2, cursor = host4
let (hosts3, next3) = hydrant.list_hosts(next2.as_deref(), 2).await?;
assert_eq!(hosts3.len(), 2);
assert_eq!(hosts3[0].name.as_str(), "host5");
assert_eq!(hosts3[1].name.as_str(), "host6");
assert!(next3.is_none());
Ok(())
}
}