use std::{collections::BTreeMap, future::Future};
use serde::{Deserialize, Serialize};
use crate::{
error::{FraiseQLError, Result},
schema::{SessionVariableSource, SessionVariablesConfig},
};
pub const MAX_SESSION_HEADER_BYTES: usize = 1024;
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct SessionHeaders(BTreeMap<String, String>);
impl SessionHeaders {
pub fn resolve(
config: &SessionVariablesConfig,
values: &dyn Fn(&str) -> Vec<Option<String>>,
) -> Result<Self> {
let mut resolved = BTreeMap::new();
for mapping in &config.variables {
let SessionVariableSource::Header { header } = &mapping.source else {
continue;
};
let name = header.to_ascii_lowercase();
if resolved.contains_key(&name) {
continue;
}
let refuse = |why: &str| FraiseQLError::Validation {
message: format!(
"request header '{name}' (session variable '{}') {why}",
mapping.name
),
path: None,
};
let mut sent = values(&name).into_iter();
let Some(first) = sent.next() else {
continue;
};
if sent.next().is_some() {
return Err(refuse("was sent more than once"));
}
let Some(value) = first else {
return Err(refuse("is not valid UTF-8"));
};
if value.len() > MAX_SESSION_HEADER_BYTES {
return Err(refuse(&format!("is longer than {MAX_SESSION_HEADER_BYTES} bytes")));
}
resolved.insert(name, value);
}
Ok(Self(resolved))
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
#[must_use]
pub fn get(&self, header: &str) -> Option<&str> {
self.0.get(&header.to_ascii_lowercase()).map(String::as_str)
}
}
tokio::task_local! {
static SESSION_HEADERS: SessionHeaders;
}
pub async fn with_session_headers<F: Future>(headers: SessionHeaders, future: F) -> F::Output {
SESSION_HEADERS.scope(headers, future).await
}
pub fn with_session_headers_sync<R>(headers: SessionHeaders, f: impl FnOnce() -> R) -> R {
SESSION_HEADERS.sync_scope(headers, f)
}
#[must_use]
pub fn scoped_session_headers() -> Option<SessionHeaders> {
SESSION_HEADERS.try_with(Clone::clone).ok()
}
#[must_use]
pub fn scoped_session_header(header: &str) -> Option<String> {
SESSION_HEADERS.try_with(|h| h.get(header).map(str::to_string)).ok().flatten()
}
#[cfg(test)]
mod tests;