use osproxy_core::{EndpointKind, IndexName, Instant, PartitionId, PrincipalId, RequestId};
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Default)]
pub enum DiagLevel {
#[default]
Off,
Shape,
ShapeTiming,
ShapeRewriteDiff,
}
impl DiagLevel {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Off => "Off",
Self::Shape => "Shape",
Self::ShapeTiming => "ShapeTiming",
Self::ShapeRewriteDiff => "ShapeRewriteDiff",
}
}
#[must_use]
pub fn from_name(name: &str) -> Option<Self> {
match name {
"Off" => Some(Self::Off),
"Shape" => Some(Self::Shape),
"ShapeTiming" => Some(Self::ShapeTiming),
"ShapeRewriteDiff" => Some(Self::ShapeRewriteDiff),
_ => None,
}
}
}
#[derive(Clone, PartialEq, Eq, Debug, Default)]
pub struct DirectiveMatch {
pub tenant: Option<PartitionId>,
pub index: Option<IndexName>,
pub principal: Option<PrincipalId>,
pub endpoint: Option<EndpointKind>,
}
impl DirectiveMatch {
#[must_use]
pub fn all() -> Self {
Self::default()
}
#[must_use]
pub fn for_tenant(mut self, tenant: PartitionId) -> Self {
self.tenant = Some(tenant);
self
}
#[must_use]
pub fn for_index(mut self, index: IndexName) -> Self {
self.index = Some(index);
self
}
#[must_use]
pub fn for_principal(mut self, principal: PrincipalId) -> Self {
self.principal = Some(principal);
self
}
#[must_use]
pub fn for_endpoint(mut self, endpoint: EndpointKind) -> Self {
self.endpoint = Some(endpoint);
self
}
#[must_use]
pub fn matches(&self, attrs: &RequestAttrs<'_>) -> bool {
self.tenant.as_ref().is_none_or(|t| attrs.tenant == Some(t))
&& self
.index
.as_ref()
.is_none_or(|i| i.as_str() == attrs.index)
&& self.principal.as_ref().is_none_or(|p| p == attrs.principal)
&& self.endpoint.is_none_or(|e| e == attrs.endpoint)
}
}
#[derive(Clone, Copy, Debug)]
pub struct RequestAttrs<'a> {
pub tenant: Option<&'a PartitionId>,
pub index: &'a str,
pub principal: &'a PrincipalId,
pub endpoint: EndpointKind,
}
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct DiagnosticsDirective {
pub id: String,
pub match_: DirectiveMatch,
pub level: DiagLevel,
pub sample_per_mille: u16,
pub expires_at: Instant,
pub ring_buffer: bool,
pub capture: bool,
}
impl DiagnosticsDirective {
#[must_use]
pub fn level_if_applies(
&self,
attrs: &RequestAttrs<'_>,
now: Instant,
request: &RequestId,
) -> Option<DiagLevel> {
self.applies(attrs, now, request).then_some(self.level)
}
#[must_use]
fn applies(&self, attrs: &RequestAttrs<'_>, now: Instant, request: &RequestId) -> bool {
now < self.expires_at && self.match_.matches(attrs) && self.is_sampled(request)
}
#[must_use]
fn is_sampled(&self, request: &RequestId) -> bool {
if self.sample_per_mille >= 1000 {
return true;
}
if self.sample_per_mille == 0 {
return false;
}
let bucket = fnv1a(request.as_str().as_bytes()) % 1000;
u16::try_from(bucket).unwrap_or(u16::MAX) < self.sample_per_mille
}
}
#[derive(Clone, Debug, Default)]
pub struct DirectiveSet {
directives: Vec<DiagnosticsDirective>,
}
impl DirectiveSet {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn from_directives(directives: Vec<DiagnosticsDirective>) -> Self {
Self { directives }
}
#[must_use]
pub fn len(&self) -> usize {
self.directives.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.directives.is_empty()
}
#[must_use]
pub fn evaluate(
&self,
attrs: &RequestAttrs<'_>,
now: Instant,
request: &RequestId,
) -> DiagLevel {
self.directives
.iter()
.filter(|d| d.applies(attrs, now, request))
.map(|d| d.level)
.max()
.unwrap_or(DiagLevel::Off)
}
#[must_use]
pub fn wants_ring_buffer(
&self,
attrs: &RequestAttrs<'_>,
now: Instant,
request: &RequestId,
) -> bool {
self.directives
.iter()
.any(|d| d.ring_buffer && d.applies(attrs, now, request))
}
#[must_use]
pub fn wants_capture(
&self,
attrs: &RequestAttrs<'_>,
now: Instant,
request: &RequestId,
) -> bool {
self.directives
.iter()
.any(|d| d.capture && d.applies(attrs, now, request))
}
#[must_use]
pub fn introspect(&self, now: Instant) -> serde_json::Value {
let directives: Vec<serde_json::Value> = self
.directives
.iter()
.map(|d| {
let mut obj = serde_json::Map::new();
obj.insert("id".into(), d.id.clone().into());
obj.insert("level".into(), d.level.as_str().into());
if let Some(t) = &d.match_.tenant {
obj.insert("tenant".into(), t.as_str().into());
}
if let Some(i) = &d.match_.index {
obj.insert("index".into(), i.as_str().into());
}
if let Some(p) = &d.match_.principal {
obj.insert("principal".into(), p.as_str().into());
}
if let Some(e) = d.match_.endpoint {
obj.insert("endpoint".into(), e.as_str().into());
}
obj.insert("sample_per_mille".into(), d.sample_per_mille.into());
obj.insert("ring_buffer".into(), d.ring_buffer.into());
obj.insert("capture".into(), d.capture.into());
obj.insert("expired".into(), (now >= d.expires_at).into());
serde_json::Value::Object(obj)
})
.collect();
serde_json::json!({ "directives": directives })
}
}
pub trait DirectiveVerifier: Send + Sync {
fn verify(&self, header_value: &str) -> Option<DiagnosticsDirective>;
}
#[derive(Clone, Copy, Debug, Default)]
pub struct NoVerifier;
impl DirectiveVerifier for NoVerifier {
fn verify(&self, _header_value: &str) -> Option<DiagnosticsDirective> {
None
}
}
fn fnv1a(bytes: &[u8]) -> u64 {
let mut h = 0xcbf2_9ce4_8422_2325;
for &b in bytes {
h ^= u64::from(b);
h = u64::wrapping_mul(h, 0x0000_0100_0000_01b3);
}
h
}
#[cfg(test)]
#[path = "directive_tests.rs"]
mod tests;