use std::collections::HashMap;
use parking_lot::RwLock;
use crate::role::Role;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MethodClass {
Read,
Write,
Admin,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum RateBucket {
ReadLight,
ReadHeavy,
WriteLight,
WriteHeavy,
AdminOnly,
}
#[derive(Debug, Clone)]
pub struct MethodMeta {
pub name: &'static str,
pub class: MethodClass,
pub min_role: Role,
pub rate_bucket: RateBucket,
pub public_exposed: bool,
}
impl MethodMeta {
pub const fn read(name: &'static str, min_role: Role, bucket: RateBucket) -> Self {
Self {
name,
class: MethodClass::Read,
min_role,
rate_bucket: bucket,
public_exposed: matches!(min_role, Role::Explorer),
}
}
pub const fn write(name: &'static str, min_role: Role, bucket: RateBucket) -> Self {
Self {
name,
class: MethodClass::Write,
min_role,
rate_bucket: bucket,
public_exposed: false,
}
}
pub const fn admin(name: &'static str) -> Self {
Self {
name,
class: MethodClass::Admin,
min_role: Role::Admin,
rate_bucket: RateBucket::AdminOnly,
public_exposed: false,
}
}
}
#[derive(Debug, Default)]
pub struct MethodRegistry {
inner: RwLock<HashMap<&'static str, MethodMeta>>,
}
impl MethodRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&self, meta: MethodMeta) {
self.inner.write().insert(meta.name, meta);
}
pub fn get(&self, name: &str) -> Option<MethodMeta> {
self.inner.read().get(name).cloned()
}
pub fn register_all(&self, metas: impl IntoIterator<Item = MethodMeta>) {
let mut g = self.inner.write();
for m in metas {
g.insert(m.name, m);
}
}
pub fn len(&self) -> usize {
self.inner.read().len()
}
pub fn is_empty(&self) -> bool {
self.inner.read().is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn builders_set_public_exposed_correctly() {
let r = MethodMeta::read("healthz", Role::Explorer, RateBucket::ReadLight);
assert!(r.public_exposed);
assert_eq!(r.class, MethodClass::Read);
let r_admin = MethodMeta::read("get_slashing_db", Role::Admin, RateBucket::ReadLight);
assert!(!r_admin.public_exposed);
let w = MethodMeta::write("push_tx", Role::Explorer, RateBucket::WriteHeavy);
assert!(!w.public_exposed); assert_eq!(w.class, MethodClass::Write);
let a = MethodMeta::admin("stop_node");
assert!(!a.public_exposed);
assert_eq!(a.min_role, Role::Admin);
assert_eq!(a.rate_bucket, RateBucket::AdminOnly);
}
#[test]
fn registry_register_and_lookup() {
let r = MethodRegistry::new();
assert!(r.is_empty());
assert!(r.get("healthz").is_none());
r.register(MethodMeta::read(
"healthz",
Role::Explorer,
RateBucket::ReadLight,
));
assert_eq!(r.len(), 1);
let meta = r.get("healthz").unwrap();
assert_eq!(meta.name, "healthz");
assert_eq!(meta.class, MethodClass::Read);
}
#[test]
fn register_overwrites() {
let r = MethodRegistry::new();
r.register(MethodMeta::read("m", Role::Explorer, RateBucket::ReadLight));
r.register(MethodMeta::admin("m"));
assert_eq!(r.len(), 1);
assert_eq!(r.get("m").unwrap().class, MethodClass::Admin);
}
}