feat: publish Dever core with project scaffolding and versioned AI skills
This commit is contained in:
Antigravity committed
1 parent
70fa73cdc2
commit
ac458ffa7f
960 files changed
+221319
-1163
No files matched your search
@@ -0,0 +1,49 @@
|
||||
[package]
|
||||
name = "dever-runtime"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[features]
|
||||
default = []
|
||||
auth-crypto = ["dep:sha2", "dep:hmac", "dep:subtle"]
|
||||
crypto = ["auth-crypto", "dep:argon2", "dep:getrandom"]
|
||||
wire = ["dep:serde", "dep:serde_json", "serde_json/raw_value", "dep:base64", "dep:getrandom"]
|
||||
external = ["wire", "dep:sha2", "dep:dever-sandbox"]
|
||||
api = ["wire", "auth-crypto"]
|
||||
database = ["wire", "dep:getrandom"]
|
||||
sqlite = ["database", "dep:deadpool-sqlite"]
|
||||
postgres = ["database", "dep:deadpool-postgres", "dep:tokio-postgres", "dep:tokio-postgres-rustls"]
|
||||
|
||||
[dependencies]
|
||||
dever-sandbox = { path = "../dever-sandbox", optional = true }
|
||||
dec.workspace = true
|
||||
indexmap.workspace = true
|
||||
tokio.workspace = true
|
||||
futures-util.workspace = true
|
||||
bytes.workspace = true
|
||||
base64 = { workspace = true, optional = true }
|
||||
hyper.workspace = true
|
||||
hyper-util.workspace = true
|
||||
http-body-util.workspace = true
|
||||
tokio-tungstenite.workspace = true
|
||||
tokio-rustls.workspace = true
|
||||
rustls.workspace = true
|
||||
webpki-roots.workspace = true
|
||||
serde = { workspace = true, optional = true }
|
||||
serde_json = { workspace = true, optional = true }
|
||||
deadpool-sqlite = { workspace = true, optional = true }
|
||||
deadpool-postgres = { workspace = true, optional = true }
|
||||
tokio-postgres = { workspace = true, optional = true }
|
||||
tokio-postgres-rustls = { workspace = true, optional = true }
|
||||
getrandom = { workspace = true, optional = true }
|
||||
argon2 = { workspace = true, optional = true }
|
||||
sha2 = { workspace = true, optional = true }
|
||||
hmac = { workspace = true, optional = true }
|
||||
subtle = { workspace = true, optional = true }
|
||||
zeroize.workspace = true
|
||||
chrono.workspace = true
|
||||
@@ -0,0 +1,49 @@
|
||||
//! Safe value operations shared by the Rust backend and the native runtime ABI.
|
||||
|
||||
use crate::bytes::Bytes;
|
||||
use crate::number::DecimalValue;
|
||||
use crate::render::Render;
|
||||
|
||||
/// The finite Decimal128 interchange representation is little-endian bits,
|
||||
/// independent of the private Rust DecimalValue layout.
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub struct DecimalBits {
|
||||
pub low: u64,
|
||||
pub high: u64,
|
||||
}
|
||||
|
||||
impl DecimalBits {
|
||||
pub fn decode(self) -> Result<DecimalValue, &'static str> {
|
||||
let mut bytes = [0; 16];
|
||||
bytes[..8].copy_from_slice(&self.low.to_le_bytes());
|
||||
bytes[8..].copy_from_slice(&self.high.to_le_bytes());
|
||||
DecimalValue::from_bytes(bytes)
|
||||
}
|
||||
|
||||
pub fn encode(value: DecimalValue) -> Self {
|
||||
let bytes = value.to_bytes();
|
||||
Self {
|
||||
low: u64::from_le_bytes(bytes[..8].try_into().expect("Decimal low word")),
|
||||
high: u64::from_le_bytes(bytes[8..].try_into().expect("Decimal high word")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn text_from_utf8(bytes: &[u8]) -> Result<Vec<u8>, String> {
|
||||
Bytes::new(bytes.to_vec())
|
||||
.into_text()
|
||||
.map(String::into_bytes)
|
||||
}
|
||||
|
||||
pub fn int_to_text(value: i64) -> String {
|
||||
value.render()
|
||||
}
|
||||
|
||||
pub fn float_to_text(value: f64) -> String {
|
||||
value.render()
|
||||
}
|
||||
|
||||
pub fn decimal_to_text(value: DecimalBits) -> Result<String, &'static str> {
|
||||
value.decode().map(|value| value.render())
|
||||
}
|
||||
@@ -0,0 +1,581 @@
|
||||
use base64::Engine as _;
|
||||
use std::cell::RefCell;
|
||||
use std::future::Future;
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::bytes::Bytes;
|
||||
use crate::http::{self, Header, Request, Response};
|
||||
|
||||
pub use serde_json::{Map as Object, Value};
|
||||
mod input;
|
||||
pub mod upload;
|
||||
pub use input::{InputError, Inputs, QueryValue, text_bounds};
|
||||
|
||||
tokio::task_local! { static REQUEST: RefCell<RequestState>; }
|
||||
|
||||
struct RequestState {
|
||||
id: u64,
|
||||
method: String,
|
||||
path: String,
|
||||
client_address: Option<String>,
|
||||
headers: Vec<Header>,
|
||||
cookies: BTreeMap<String, crate::secret::Secret>,
|
||||
pending_response_headers: Vec<Header>,
|
||||
committed_response_headers: Vec<Header>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum SameSite {
|
||||
Strict,
|
||||
Lax,
|
||||
None,
|
||||
}
|
||||
|
||||
pub struct CookieOptions {
|
||||
pub path: String,
|
||||
pub domain: Option<String>,
|
||||
pub same_site: SameSite,
|
||||
pub secure: bool,
|
||||
pub http_only: bool,
|
||||
pub max_age: Option<i64>,
|
||||
}
|
||||
|
||||
/// 路由由编译器生成；运行时只负责既有 HTTP 引擎的配置与启动。
|
||||
pub async fn serve<F, Fut>(route: F) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<Response, String>> + Send + 'static,
|
||||
{
|
||||
if !crate::lifecycle::api_enabled() {
|
||||
return Ok(());
|
||||
}
|
||||
let settings = crate::config::http_settings()?;
|
||||
let listener = crate::net::listen(&settings.host, settings.port).await?;
|
||||
let limits = settings.limits;
|
||||
let components = crate::component::current();
|
||||
#[cfg(feature = "api")]
|
||||
let session = crate::application::current();
|
||||
let server = http::serve_incoming(
|
||||
move |incoming| {
|
||||
let route = route.clone();
|
||||
let settings = settings.clone();
|
||||
let request = async move {
|
||||
upload::scope(
|
||||
incoming.body,
|
||||
settings,
|
||||
scoped_request(incoming.head, route),
|
||||
)
|
||||
.await
|
||||
};
|
||||
#[cfg(feature = "api")]
|
||||
let request = crate::application::scope_resources(session.clone(), request);
|
||||
crate::component::scope_resources(components.clone(), request)
|
||||
},
|
||||
listener.clone(),
|
||||
limits,
|
||||
);
|
||||
tokio::pin!(server);
|
||||
tokio::select! {
|
||||
result = &mut server => result,
|
||||
_ = crate::lifecycle::shutdown_requested() => {
|
||||
listener.close()?;
|
||||
server.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The scope owns response metadata; no source-level Context value exists.
|
||||
pub async fn scoped_request<F, Fut>(request: Request, route: F) -> Result<Response, String>
|
||||
where
|
||||
F: FnOnce(Request) -> Fut,
|
||||
Fut: Future<Output = Result<Response, String>>,
|
||||
{
|
||||
let cookies = match read_cookies(&request.headers) {
|
||||
Ok(cookies) => cookies,
|
||||
Err(error) => return Ok(invalid_input(error)),
|
||||
};
|
||||
let state = RequestState {
|
||||
id: http::request_id(),
|
||||
method: request.method.clone(),
|
||||
path: request.target.split('?').next().unwrap_or("").to_owned(),
|
||||
client_address: http::client_address(),
|
||||
headers: request.headers.clone(),
|
||||
cookies,
|
||||
pending_response_headers: Vec::new(),
|
||||
committed_response_headers: Vec::new(),
|
||||
};
|
||||
REQUEST
|
||||
.scope(RefCell::new(state), async move {
|
||||
let mut response = route(request).await?;
|
||||
REQUEST.with(|state| {
|
||||
response
|
||||
.headers
|
||||
.append(&mut state.borrow_mut().committed_response_headers)
|
||||
});
|
||||
Ok(response)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Publish staged metadata only after the generated handler has serialized its
|
||||
/// response and committed its database transaction.
|
||||
pub fn commit_response_metadata() -> Result<(), String> {
|
||||
with_request_mut(|request| {
|
||||
request
|
||||
.committed_response_headers
|
||||
.append(&mut request.pending_response_headers);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn with_request<T>(read: impl FnOnce(&RequestState) -> Result<T, String>) -> Result<T, String> {
|
||||
REQUEST
|
||||
.try_with(|state| read(&state.borrow()))
|
||||
.map_err(|_| "API request capability is available only within an HTTP request".to_owned())?
|
||||
}
|
||||
|
||||
fn with_request_mut<T>(
|
||||
write: impl FnOnce(&mut RequestState) -> Result<T, String>,
|
||||
) -> Result<T, String> {
|
||||
REQUEST
|
||||
.try_with(|state| write(&mut state.borrow_mut()))
|
||||
.map_err(|_| "API request capability is available only within an HTTP request".to_owned())?
|
||||
}
|
||||
|
||||
pub fn request_id() -> Result<String, String> {
|
||||
with_request(|request| Ok(request.id.to_string()))
|
||||
}
|
||||
pub fn method() -> Result<String, String> {
|
||||
with_request(|request| Ok(request.method.clone()))
|
||||
}
|
||||
pub fn path() -> Result<String, String> {
|
||||
with_request(|request| Ok(request.path.clone()))
|
||||
}
|
||||
pub fn client_address() -> Result<Option<String>, String> {
|
||||
with_request(|request| Ok(request.client_address.clone()))
|
||||
}
|
||||
|
||||
pub fn header(name: &str) -> Result<Option<String>, String> {
|
||||
validate_header_name(name)?;
|
||||
if [
|
||||
"cookie",
|
||||
"authorization",
|
||||
"proxy-authorization",
|
||||
"set-cookie",
|
||||
"forwarded",
|
||||
"x-forwarded-for",
|
||||
"x-forwarded-host",
|
||||
"x-forwarded-proto",
|
||||
"x-real-ip",
|
||||
]
|
||||
.iter()
|
||||
.any(|blocked| name.eq_ignore_ascii_case(blocked))
|
||||
{
|
||||
return Err("sensitive request headers require a typed Secret reader".into());
|
||||
}
|
||||
with_request(|request| {
|
||||
let mut values = request
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|header| header.name.eq_ignore_ascii_case(name));
|
||||
let Some(value) = values.next() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if values.next().is_some() {
|
||||
return Err("duplicate request header".into());
|
||||
}
|
||||
let value = std::str::from_utf8(value.value.values())
|
||||
.map_err(|_| "request header is not UTF-8".to_owned())?;
|
||||
Ok(Some(value.to_owned()))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cookie(name: &str) -> Result<Option<String>, String> {
|
||||
validate_cookie_name(name)?;
|
||||
with_request(|request| {
|
||||
let Some(raw) = request.cookies.get(name) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if raw.0.starts_with(b"v1.") {
|
||||
return Err("Secret Cookie values require the typed Secret reader".into());
|
||||
}
|
||||
let value =
|
||||
std::str::from_utf8(&raw.0).map_err(|_| "Cookie value is not UTF-8".to_owned())?;
|
||||
Ok(Some(value.to_owned()))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn secret_cookie(name: &str) -> Result<Option<crate::secret::Secret>, String> {
|
||||
validate_cookie_name(name)?;
|
||||
with_request(|request| {
|
||||
let Some(raw) = request.cookies.get(name) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(encoded) = raw.0.strip_prefix(b"v1.") else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(encoded)
|
||||
.ok()
|
||||
.map(crate::secret::Secret::from_input))
|
||||
})
|
||||
}
|
||||
|
||||
fn read_cookies(headers: &[Header]) -> Result<BTreeMap<String, crate::secret::Secret>, InputError> {
|
||||
let mut cookies = BTreeMap::new();
|
||||
for header in headers
|
||||
.iter()
|
||||
.filter(|header| header.name.eq_ignore_ascii_case("cookie"))
|
||||
{
|
||||
let raw = header.value.values();
|
||||
if !raw.is_ascii() {
|
||||
return Err(InputError("Cookie header must be ASCII".into()));
|
||||
}
|
||||
for pair in raw.split(|byte| *byte == b';') {
|
||||
let pair = pair.trim_ascii();
|
||||
if pair.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let Some(separator) = pair.iter().position(|byte| *byte == b'=') else {
|
||||
return Err(InputError("invalid Cookie header".into()));
|
||||
};
|
||||
let (key, encoded) = pair.split_at(separator);
|
||||
let value = &encoded[1..];
|
||||
let value = if value.len() >= 2
|
||||
&& value.first() == Some(&b'"')
|
||||
&& value.last() == Some(&b'"')
|
||||
{
|
||||
&value[1..value.len() - 1]
|
||||
} else {
|
||||
value
|
||||
};
|
||||
let name = std::str::from_utf8(key).expect("ASCII Cookie header");
|
||||
validate_cookie_name(name).map_err(InputError)?;
|
||||
validate_cookie_value(value).map_err(InputError)?;
|
||||
if cookies
|
||||
.insert(
|
||||
name.to_owned(),
|
||||
crate::secret::Secret::from_input(value.to_vec()),
|
||||
)
|
||||
.is_some()
|
||||
{
|
||||
return Err(InputError("duplicate Cookie name".into()));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(cookies)
|
||||
}
|
||||
|
||||
pub fn response_header(name: &str, value: &str) -> Result<(), String> {
|
||||
validate_header_name(name)?;
|
||||
if [
|
||||
"content-length",
|
||||
"content-type",
|
||||
"content-encoding",
|
||||
"transfer-encoding",
|
||||
"connection",
|
||||
"set-cookie",
|
||||
"host",
|
||||
"upgrade",
|
||||
"keep-alive",
|
||||
"trailer",
|
||||
"te",
|
||||
]
|
||||
.iter()
|
||||
.any(|blocked| name.eq_ignore_ascii_case(blocked))
|
||||
{
|
||||
return Err("response header is controlled by the HTTP runtime".into());
|
||||
}
|
||||
hyper::header::HeaderValue::from_str(value)
|
||||
.map_err(|_| "invalid response header value".to_owned())?;
|
||||
with_request_mut(|request| {
|
||||
request.pending_response_headers.push(Header {
|
||||
name: name.to_owned(),
|
||||
value: Bytes::from_text(value),
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn response_cookie(name: &str, value: &str, options: CookieOptions) -> Result<(), String> {
|
||||
write_cookie(name, value.as_bytes(), options, false)
|
||||
}
|
||||
|
||||
pub fn response_secret_cookie(
|
||||
name: &str,
|
||||
value: crate::secret::Secret,
|
||||
options: CookieOptions,
|
||||
) -> Result<(), String> {
|
||||
let length = base64::encoded_len(value.0.len(), false)
|
||||
.ok_or_else(|| "Secret Cookie value is too large".to_owned())?;
|
||||
let mut encoded = vec![0; length + 3];
|
||||
encoded[..3].copy_from_slice(b"v1.");
|
||||
let written = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.encode_slice(&value.0, &mut encoded[3..])
|
||||
.map_err(|_| "Secret Cookie encoding failed".to_owned())?;
|
||||
encoded.truncate(written + 3);
|
||||
write_cookie(name, &encoded, options, true)
|
||||
}
|
||||
|
||||
fn write_cookie(
|
||||
name: &str,
|
||||
value: &[u8],
|
||||
options: CookieOptions,
|
||||
secret: bool,
|
||||
) -> Result<(), String> {
|
||||
require_same_origin_for_cookie_write()?;
|
||||
validate_cookie_name(name)?;
|
||||
validate_cookie_value(value)?;
|
||||
if secret && (!options.secure || !options.http_only) {
|
||||
return Err("Secret cookies require Secure and HttpOnly".into());
|
||||
}
|
||||
if options.same_site == SameSite::None && !options.secure {
|
||||
return Err("SameSite=None requires Secure".into());
|
||||
}
|
||||
if !options.path.starts_with('/')
|
||||
|| !options
|
||||
.path
|
||||
.bytes()
|
||||
.all(|byte| (0x21..=0x7e).contains(&byte) && byte != b';')
|
||||
{
|
||||
return Err("invalid Cookie Path".into());
|
||||
}
|
||||
if let Some(domain) = &options.domain
|
||||
&& (domain.len() > 253
|
||||
|| domain.is_empty()
|
||||
|| domain.split('.').any(|label| {
|
||||
label.is_empty()
|
||||
|| label.len() > 63
|
||||
|| label.starts_with('-')
|
||||
|| label.ends_with('-')
|
||||
|| !label
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'-')
|
||||
}))
|
||||
{
|
||||
return Err("invalid Cookie Domain".into());
|
||||
}
|
||||
if options.max_age.is_some_and(|age| age < 0) {
|
||||
return Err("Cookie Max-Age must be nonnegative".into());
|
||||
}
|
||||
let mut encoded = Vec::with_capacity(name.len() + value.len() + 96);
|
||||
encoded.extend_from_slice(name.as_bytes());
|
||||
encoded.push(b'=');
|
||||
encoded.extend_from_slice(value);
|
||||
encoded.extend_from_slice(b"; Path=");
|
||||
encoded.extend_from_slice(options.path.as_bytes());
|
||||
if let Some(domain) = options.domain {
|
||||
encoded.extend_from_slice(b"; Domain=");
|
||||
encoded.extend_from_slice(domain.as_bytes());
|
||||
}
|
||||
encoded.extend_from_slice(match options.same_site {
|
||||
SameSite::Strict => b"; SameSite=Strict",
|
||||
SameSite::Lax => b"; SameSite=Lax",
|
||||
SameSite::None => b"; SameSite=None",
|
||||
});
|
||||
if options.secure {
|
||||
encoded.extend_from_slice(b"; Secure");
|
||||
}
|
||||
if options.http_only {
|
||||
encoded.extend_from_slice(b"; HttpOnly");
|
||||
}
|
||||
if let Some(age) = options.max_age {
|
||||
encoded.extend_from_slice(format!("; Max-Age={age}").as_bytes());
|
||||
}
|
||||
if encoded.len() > 4096 {
|
||||
return Err("Set-Cookie header exceeds 4096 bytes".into());
|
||||
}
|
||||
with_request_mut(|request| {
|
||||
request.pending_response_headers.push(Header {
|
||||
name: "set-cookie".into(),
|
||||
value: Bytes::new(encoded),
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn require_same_origin_for_cookie_write() -> Result<(), String> {
|
||||
with_request(|request| {
|
||||
if matches!(request.method.as_str(), "GET" | "HEAD" | "OPTIONS") {
|
||||
return Ok(());
|
||||
}
|
||||
#[cfg(feature = "api")]
|
||||
{
|
||||
let origin = crate::auth::cookie_origin()
|
||||
.map_err(|_| "site origin is required for unsafe Cookie writes".to_owned())?;
|
||||
validate_same_origin(&request.headers, &origin)
|
||||
}
|
||||
#[cfg(not(feature = "api"))]
|
||||
{
|
||||
Err("site authentication is required for unsafe Cookie writes".into())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "api")]
|
||||
pub(crate) fn validate_same_origin(
|
||||
headers: &[Header],
|
||||
expected: &crate::config::HttpOrigin,
|
||||
) -> Result<(), String> {
|
||||
let host = unique_header(headers, "host")?
|
||||
.ok_or_else(|| "unsafe Cookie request requires Host".to_owned())?;
|
||||
let origin = unique_header(headers, "origin")?
|
||||
.ok_or_else(|| "unsafe Cookie request requires Origin".to_owned())?;
|
||||
let origin = origin
|
||||
.parse::<hyper::Uri>()
|
||||
.map_err(|_| "invalid Origin".to_owned())?;
|
||||
if origin.path() != "/" || origin.query().is_some() {
|
||||
return Err("Origin must not contain a path or query".into());
|
||||
}
|
||||
let origin_authority = origin
|
||||
.authority()
|
||||
.ok_or_else(|| "invalid Origin".to_owned())?;
|
||||
let request_authority = host
|
||||
.parse::<hyper::http::uri::Authority>()
|
||||
.map_err(|_| "invalid Host".to_owned())?;
|
||||
let default_port = if expected.scheme() == "https" {
|
||||
443
|
||||
} else {
|
||||
80
|
||||
};
|
||||
if origin.scheme_str() != Some(expected.scheme())
|
||||
|| normalized_authority(origin_authority, default_port)
|
||||
!= (expected.host().to_owned(), expected.port())
|
||||
|| normalized_authority(&request_authority, default_port)
|
||||
!= (expected.host().to_owned(), expected.port())
|
||||
{
|
||||
return Err("Origin and Host must exactly match the configured site origin".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(feature = "api")]
|
||||
fn unique_header(headers: &[Header], name: &str) -> Result<Option<String>, String> {
|
||||
let mut values = headers
|
||||
.iter()
|
||||
.filter(|header| header.name.eq_ignore_ascii_case(name));
|
||||
let Some(value) = values.next() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if values.next().is_some() {
|
||||
return Err(format!("duplicate {name} header"));
|
||||
}
|
||||
let value = std::str::from_utf8(value.value.values())
|
||||
.map_err(|_| format!("{name} header is not UTF-8"))?;
|
||||
Ok(Some(value.to_owned()))
|
||||
}
|
||||
|
||||
#[cfg(feature = "api")]
|
||||
fn normalized_authority(
|
||||
authority: &hyper::http::uri::Authority,
|
||||
default_port: u16,
|
||||
) -> (String, u16) {
|
||||
(
|
||||
authority.host().to_ascii_lowercase(),
|
||||
authority.port_u16().unwrap_or(default_port),
|
||||
)
|
||||
}
|
||||
|
||||
fn validate_header_name(name: &str) -> Result<(), String> {
|
||||
hyper::header::HeaderName::from_bytes(name.as_bytes())
|
||||
.map(|_| ())
|
||||
.map_err(|_| "invalid HTTP header name".to_owned())
|
||||
}
|
||||
|
||||
fn validate_cookie_name(name: &str) -> Result<(), String> {
|
||||
if name.is_empty()
|
||||
|| !name.bytes().all(|byte| {
|
||||
byte.is_ascii_alphanumeric()
|
||||
|| matches!(
|
||||
byte,
|
||||
b'!' | b'#'
|
||||
..=b'\'' | b'*' | b'+' | b'-' | b'.' | b'^' | b'_' | b'`' | b'|' | b'~'
|
||||
)
|
||||
})
|
||||
{
|
||||
return Err("invalid Cookie name".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_cookie_value(value: &[u8]) -> Result<(), String> {
|
||||
if value
|
||||
.iter()
|
||||
.all(|byte| matches!(byte, 0x21 | 0x23..=0x2b | 0x2d..=0x3a | 0x3c..=0x5b | 0x5d..=0x7e))
|
||||
{
|
||||
Ok(())
|
||||
} else {
|
||||
Err("invalid Cookie value".into())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn success(data: crate::wire::Encoded) -> Response {
|
||||
envelope(200, 0, "ok", data)
|
||||
}
|
||||
|
||||
pub fn invalid_input(error: InputError) -> Response {
|
||||
envelope(400, 400, &error.0, crate::wire::Encoded::null())
|
||||
}
|
||||
|
||||
pub fn standard_error(status: i64) -> Response {
|
||||
let message = match status {
|
||||
400 => "invalid request",
|
||||
401 => "unauthorized",
|
||||
403 => "forbidden",
|
||||
404 => "not found",
|
||||
409 => "conflict",
|
||||
429 => "too many requests",
|
||||
_ => return internal_error(),
|
||||
};
|
||||
envelope(status, status, message, crate::wire::Encoded::null())
|
||||
}
|
||||
|
||||
#[cfg(feature = "api")]
|
||||
pub fn auth_error(error: crate::auth::Error) -> Response {
|
||||
match error {
|
||||
crate::auth::Error::Unauthorized => standard_error(401),
|
||||
crate::auth::Error::Forbidden => standard_error(403),
|
||||
crate::auth::Error::Internal => {
|
||||
crate::log::error("authentication runtime failed", Vec::new());
|
||||
internal_error()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn not_found() -> Response {
|
||||
envelope(404, 404, "not found", crate::wire::Encoded::null())
|
||||
}
|
||||
|
||||
pub fn method_not_allowed(methods: &[&str]) -> Response {
|
||||
let mut response = envelope(405, 405, "method not allowed", crate::wire::Encoded::null());
|
||||
response.headers.push(Header {
|
||||
name: "allow".into(),
|
||||
value: Bytes::from_string(methods.join(", ")),
|
||||
});
|
||||
response
|
||||
}
|
||||
|
||||
pub fn internal_error() -> Response {
|
||||
envelope(
|
||||
500,
|
||||
500,
|
||||
"internal server error",
|
||||
crate::wire::Encoded::null(),
|
||||
)
|
||||
}
|
||||
|
||||
fn envelope(status: i64, code: i64, message: &str, data: crate::wire::Encoded) -> Response {
|
||||
Response {
|
||||
status,
|
||||
headers: vec![Header {
|
||||
name: "content-type".into(),
|
||||
value: Bytes::from_text("application/json; charset=utf-8"),
|
||||
}],
|
||||
body: Bytes::from_string(format!(
|
||||
"{{\"code\":{code},\"message\":{},\"data\":{}}}",
|
||||
serde_json::to_string(message).expect("Text JSON encoding"),
|
||||
data.as_str()
|
||||
)),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt;
|
||||
|
||||
use crate::http::Request;
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct InputError(pub String);
|
||||
|
||||
impl fmt::Display for InputError {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for InputError {}
|
||||
|
||||
pub struct Inputs {
|
||||
values: BTreeMap<String, String>,
|
||||
query: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub enum QueryValue {
|
||||
Text,
|
||||
Int,
|
||||
Bool,
|
||||
Float,
|
||||
}
|
||||
|
||||
impl Inputs {
|
||||
pub(crate) fn from_fields(values: BTreeMap<String, String>) -> Self {
|
||||
Self {
|
||||
values,
|
||||
query: true,
|
||||
}
|
||||
}
|
||||
pub fn from_request(request: &Request) -> Result<Self, InputError> {
|
||||
let query = request
|
||||
.target
|
||||
.split_once('?')
|
||||
.map_or("", |(_, query)| query);
|
||||
match request.method.as_str() {
|
||||
"GET" | "DELETE" => {
|
||||
if !request.body.values().is_empty() {
|
||||
return Err(InputError(
|
||||
"GET and DELETE inputs must use the query string".into(),
|
||||
));
|
||||
}
|
||||
let mut values = BTreeMap::new();
|
||||
for pair in query.split('&').filter(|pair| !pair.is_empty()) {
|
||||
let (name, value) = pair.split_once('=').unwrap_or((pair, ""));
|
||||
let name = decode_query(name)?;
|
||||
let value = serde_json::to_string(&decode_query(value)?)
|
||||
.map_err(|error| InputError(error.to_string()))?;
|
||||
if values.insert(name.clone(), value).is_some() {
|
||||
return Err(InputError(format!("duplicate input '{name}'")));
|
||||
}
|
||||
}
|
||||
Ok(Self {
|
||||
values,
|
||||
query: true,
|
||||
})
|
||||
}
|
||||
"POST" | "PUT" => {
|
||||
if !query.is_empty() {
|
||||
return Err(InputError(
|
||||
"POST and PUT inputs must use a JSON object body".into(),
|
||||
));
|
||||
}
|
||||
let mut content_types = request
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|header| header.name.eq_ignore_ascii_case("content-type"));
|
||||
let content_type = content_types
|
||||
.next()
|
||||
.and_then(|header| std::str::from_utf8(header.value.values()).ok());
|
||||
if content_types.next().is_some()
|
||||
|| !content_type.is_some_and(|value| {
|
||||
value
|
||||
.split(';')
|
||||
.next()
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("application/json")
|
||||
})
|
||||
{
|
||||
return Err(InputError(
|
||||
"POST and PUT require Content-Type: application/json".into(),
|
||||
));
|
||||
}
|
||||
let body = std::str::from_utf8(request.body.values())
|
||||
.map_err(|_| InputError("JSON body must be UTF-8".into()))?;
|
||||
Self::from_json_object(body)
|
||||
}
|
||||
_ => Err(InputError("unsupported API request method".into())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_json_object(body: &str) -> Result<Self, InputError> {
|
||||
let node = crate::wire::parse(body)
|
||||
.map_err(|error| InputError(format!("invalid JSON object body: {error}")))?;
|
||||
let values = node
|
||||
.object()
|
||||
.map_err(|error| InputError(format!("invalid JSON object body: {error}")))?
|
||||
.iter()
|
||||
.map(|(name, value)| (name.clone(), value.raw().to_owned()))
|
||||
.collect();
|
||||
Ok(Self {
|
||||
values,
|
||||
query: false,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn raw_json(&mut self, name: &str) -> Result<String, InputError> {
|
||||
self.values.remove(name).ok_or_else(|| missing(name))
|
||||
}
|
||||
|
||||
pub fn optional_raw_json(&mut self, name: &str) -> Result<Option<String>, InputError> {
|
||||
Ok(self.values.remove(name))
|
||||
}
|
||||
|
||||
pub fn json(&mut self, name: &str, kind: QueryValue) -> Result<String, InputError> {
|
||||
let value = self.raw_json(name)?;
|
||||
self.convert_query(name, kind, value)
|
||||
}
|
||||
|
||||
pub fn optional_json(&mut self, name: &str, kind: QueryValue) -> Result<String, InputError> {
|
||||
let value = self
|
||||
.optional_raw_json(name)?
|
||||
.unwrap_or_else(|| "null".into());
|
||||
self.convert_query(name, kind, value)
|
||||
}
|
||||
|
||||
fn convert_query(
|
||||
&self,
|
||||
name: &str,
|
||||
kind: QueryValue,
|
||||
value: String,
|
||||
) -> Result<String, InputError> {
|
||||
if !self.query || value == "null" {
|
||||
return Ok(value);
|
||||
}
|
||||
let text: String =
|
||||
serde_json::from_str(&value).map_err(|error| InputError(error.to_string()))?;
|
||||
match kind {
|
||||
QueryValue::Text => Ok(value),
|
||||
QueryValue::Int => text
|
||||
.parse::<i64>()
|
||||
.map(|value| value.to_string())
|
||||
.map_err(|_| InputError(format!("input '{name}' must be Int"))),
|
||||
QueryValue::Bool => text
|
||||
.parse::<bool>()
|
||||
.map(|value| value.to_string())
|
||||
.map_err(|_| InputError(format!("input '{name}' must be Bool"))),
|
||||
QueryValue::Float => text
|
||||
.parse::<f64>()
|
||||
.ok()
|
||||
.filter(|value| value.is_finite())
|
||||
.map(|value| value.to_string())
|
||||
.ok_or_else(|| InputError(format!("input '{name}' must be Float"))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn text(&mut self, name: &str) -> Result<String, InputError> {
|
||||
self.optional_text(name)?.ok_or_else(|| missing(name))
|
||||
}
|
||||
|
||||
pub fn optional_text(&mut self, name: &str) -> Result<Option<String>, InputError> {
|
||||
self.take(name, "Text", |node, _| node.text().ok().map(str::to_owned))
|
||||
}
|
||||
|
||||
pub fn int(&mut self, name: &str) -> Result<i64, InputError> {
|
||||
self.optional_int(name)?.ok_or_else(|| missing(name))
|
||||
}
|
||||
|
||||
pub fn optional_int(&mut self, name: &str) -> Result<Option<i64>, InputError> {
|
||||
self.take(name, "Int", |node, query| {
|
||||
if query {
|
||||
node.text().ok()?.parse().ok()
|
||||
} else {
|
||||
node.int().ok()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub fn boolean(&mut self, name: &str) -> Result<bool, InputError> {
|
||||
self.optional_boolean(name)?.ok_or_else(|| missing(name))
|
||||
}
|
||||
|
||||
pub fn optional_boolean(&mut self, name: &str) -> Result<Option<bool>, InputError> {
|
||||
self.take(name, "Bool", |node, query| {
|
||||
if query {
|
||||
node.text().ok()?.parse().ok()
|
||||
} else {
|
||||
node.boolean().ok()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn take<T>(
|
||||
&mut self,
|
||||
name: &str,
|
||||
expected: &str,
|
||||
decode: impl FnOnce(&crate::wire::Node<'_>, bool) -> Option<T>,
|
||||
) -> Result<Option<T>, InputError> {
|
||||
match self.values.remove(name) {
|
||||
None => Ok(None),
|
||||
Some(value) if value == "null" => Ok(None),
|
||||
Some(value) => {
|
||||
let node = crate::wire::parse(&value).map_err(InputError)?;
|
||||
decode(&node, self.query)
|
||||
.map(Some)
|
||||
.ok_or_else(|| InputError(format!("input '{name}' must be {expected}")))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn finish(self) -> Result<(), InputError> {
|
||||
match self.values.keys().next() {
|
||||
None => Ok(()),
|
||||
Some(name) => Err(InputError(format!("unknown input '{name}'"))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn text_bounds(
|
||||
name: &str,
|
||||
value: Option<&str>,
|
||||
minimum: u32,
|
||||
maximum: Option<u32>,
|
||||
) -> Result<(), InputError> {
|
||||
let Some(value) = value else { return Ok(()) };
|
||||
let length = value.chars().count();
|
||||
if length < minimum as usize || maximum.is_some_and(|limit| length > limit as usize) {
|
||||
return Err(InputError(format!(
|
||||
"input '{name}' is outside Model Text length bounds"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn missing(name: &str) -> InputError {
|
||||
InputError(format!("input '{name}' is required"))
|
||||
}
|
||||
|
||||
fn decode_query(value: &str) -> Result<String, InputError> {
|
||||
let mut bytes = value.bytes();
|
||||
let mut decoded = Vec::with_capacity(value.len());
|
||||
while let Some(byte) = bytes.next() {
|
||||
decoded.push(match byte {
|
||||
b'+' => b' ',
|
||||
b'%' => {
|
||||
let high = bytes.next().and_then(|byte| (byte as char).to_digit(16));
|
||||
let low = bytes.next().and_then(|byte| (byte as char).to_digit(16));
|
||||
match (high, low) {
|
||||
(Some(high), Some(low)) => ((high << 4) | low) as u8,
|
||||
_ => return Err(InputError("query contains invalid percent encoding".into())),
|
||||
}
|
||||
}
|
||||
byte => byte,
|
||||
});
|
||||
}
|
||||
String::from_utf8(decoded).map_err(|_| InputError("query must contain UTF-8 text".into()))
|
||||
}
|
||||
@@ -0,0 +1,580 @@
|
||||
use std::cell::RefCell;
|
||||
use std::collections::BTreeMap;
|
||||
use std::future::Future;
|
||||
use std::io::{Seek, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use http_body_util::BodyExt;
|
||||
use hyper::body::{Body, Incoming};
|
||||
|
||||
use super::{InputError, Inputs, QueryValue};
|
||||
use crate::config::{HttpSettings, UploadLimits};
|
||||
use crate::http::Request;
|
||||
|
||||
tokio::task_local! { static BODY: RefCell<(Option<Incoming>, HttpSettings)>; }
|
||||
|
||||
pub(crate) async fn scope<T>(
|
||||
body: Incoming,
|
||||
settings: HttpSettings,
|
||||
future: impl Future<Output = T>,
|
||||
) -> T {
|
||||
BODY.scope(RefCell::new((Some(body), settings)), future)
|
||||
.await
|
||||
}
|
||||
|
||||
fn take_body() -> Result<Option<(Incoming, HttpSettings)>, InputError> {
|
||||
BODY.try_with(|state| {
|
||||
let mut state = state.borrow_mut();
|
||||
let body = state
|
||||
.0
|
||||
.take()
|
||||
.ok_or_else(|| invalid("request body was already consumed"))?;
|
||||
Ok((body, state.1.clone()))
|
||||
})
|
||||
.ok()
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Called only after route selection and authentication. Buffered unit-test requests
|
||||
/// have no Incoming owner and already carry their bounded body.
|
||||
pub async fn read_body(request: &mut Request) -> Result<(), InputError> {
|
||||
let Some((mut body, settings)) = take_body()? else {
|
||||
return Ok(());
|
||||
};
|
||||
let mut bytes = Vec::new();
|
||||
while let Some(frame) = body.frame().await {
|
||||
let frame = frame.map_err(|_| invalid("incomplete request body"))?;
|
||||
let chunk = frame
|
||||
.into_data()
|
||||
.map_err(|_| invalid("request trailers are not supported"))?;
|
||||
if chunk.len() > settings.limits.body_bytes as usize - bytes.len() {
|
||||
return Err(invalid("request body exceeds body_bytes"));
|
||||
}
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
request.body = crate::bytes::Bytes::new(bytes);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct Temporary {
|
||||
path: PathBuf,
|
||||
file: Mutex<Option<std::fs::File>>,
|
||||
}
|
||||
|
||||
impl Drop for Temporary {
|
||||
fn drop(&mut self) {
|
||||
// Close before unlinking, including on platforms that cannot unlink open files.
|
||||
self.file
|
||||
.get_mut()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.take();
|
||||
if let Err(error) = std::fs::remove_file(&self.path)
|
||||
&& error.kind() != std::io::ErrorKind::NotFound
|
||||
{
|
||||
crate::log::error("upload temporary cleanup failed", Vec::new());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Non-cloneable request resource. The path and file are never exposed to source code.
|
||||
pub struct Upload {
|
||||
temporary: Arc<Temporary>,
|
||||
filename: String,
|
||||
content_type: String,
|
||||
size: i64,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Upload {
|
||||
fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
output.write_str("Upload(<request resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl Upload {
|
||||
pub fn filename(&self) -> String {
|
||||
self.filename.clone()
|
||||
}
|
||||
pub fn content_type(&self) -> String {
|
||||
self.content_type.clone()
|
||||
}
|
||||
pub fn size(&self) -> i64 {
|
||||
self.size
|
||||
}
|
||||
pub fn close(self) {
|
||||
drop(self);
|
||||
}
|
||||
}
|
||||
|
||||
struct PendingStore {
|
||||
path: Option<PathBuf>,
|
||||
file: Option<std::fs::File>,
|
||||
key: crate::orm::Uuid,
|
||||
}
|
||||
|
||||
impl PendingStore {
|
||||
fn commit(mut self) -> crate::orm::Uuid {
|
||||
self.file.take();
|
||||
self.path.take();
|
||||
self.key
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PendingStore {
|
||||
fn drop(&mut self) {
|
||||
self.file.take();
|
||||
if let Some(path) = self.path.take()
|
||||
&& let Err(error) = std::fs::remove_file(path)
|
||||
&& error.kind() != std::io::ErrorKind::NotFound
|
||||
{
|
||||
crate::log::error("uncommitted upload storage cleanup failed", Vec::new());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Only generated storage Adapter calls reach this consuming sink. If cancellation
|
||||
/// drops the blocking result before delivery, PendingStore removes its target too.
|
||||
pub async fn store(upload: Upload) -> Result<crate::orm::Uuid, String> {
|
||||
let directory = BODY
|
||||
.try_with(|state| state.borrow().1.upload_directory.with_file_name("upload"))
|
||||
.map_err(|_| "upload storage requires a live API request".to_owned())?;
|
||||
let pending = tokio::task::spawn_blocking(move || {
|
||||
std::fs::create_dir_all(&directory)
|
||||
.map_err(|_| "cannot create upload storage directory")?;
|
||||
let mut random = [0u8; 16];
|
||||
getrandom::getrandom(&mut random).map_err(|_| "upload storage randomness unavailable")?;
|
||||
random[6] = (random[6] & 0x0f) | 0x40;
|
||||
random[8] = (random[8] & 0x3f) | 0x80;
|
||||
let key = crate::orm::Uuid::from_bytes(random);
|
||||
let path = directory.join(key.to_string());
|
||||
let mut options = std::fs::OpenOptions::new();
|
||||
options.write(true).create_new(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.mode(0o600);
|
||||
}
|
||||
let file = options
|
||||
.open(&path)
|
||||
.map_err(|_| "cannot create upload storage file")?;
|
||||
let mut pending = PendingStore {
|
||||
path: Some(path),
|
||||
file: Some(file),
|
||||
key,
|
||||
};
|
||||
let mut source = upload
|
||||
.temporary
|
||||
.file
|
||||
.lock()
|
||||
.map_err(|_| "upload source lock failed")?;
|
||||
let source = source.as_mut().ok_or("upload source is closed")?;
|
||||
source.rewind().map_err(|_| "cannot rewind upload source")?;
|
||||
let target = pending.file.as_mut().expect("pending target owns its file");
|
||||
std::io::copy(source, target).map_err(|_| "upload storage copy failed")?;
|
||||
target
|
||||
.sync_all()
|
||||
.map_err(|_| "upload storage flush failed")?;
|
||||
Ok::<_, String>(pending)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "upload storage worker failed".to_owned())??;
|
||||
Ok(pending.commit())
|
||||
}
|
||||
|
||||
pub struct Multipart {
|
||||
inputs: Inputs,
|
||||
uploads: BTreeMap<String, Upload>,
|
||||
}
|
||||
|
||||
impl Multipart {
|
||||
pub fn raw_json(&mut self, name: &str) -> Result<String, InputError> {
|
||||
self.inputs.raw_json(name)
|
||||
}
|
||||
pub fn optional_raw_json(&mut self, name: &str) -> Result<Option<String>, InputError> {
|
||||
self.inputs.optional_raw_json(name)
|
||||
}
|
||||
pub fn take(&mut self, name: &str) -> Result<Upload, InputError> {
|
||||
self.uploads
|
||||
.remove(name)
|
||||
.ok_or_else(|| invalid("missing upload field"))
|
||||
}
|
||||
pub fn json(&mut self, name: &str, kind: QueryValue) -> Result<String, InputError> {
|
||||
self.inputs.json(name, kind)
|
||||
}
|
||||
pub fn optional_json(&mut self, name: &str, kind: QueryValue) -> Result<String, InputError> {
|
||||
self.inputs.optional_json(name, kind)
|
||||
}
|
||||
pub fn finish(self) -> Result<(), InputError> {
|
||||
if !self.uploads.is_empty() {
|
||||
return Err(invalid("unknown upload field"));
|
||||
}
|
||||
self.inputs.finish()
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn read_multipart(request: &Request) -> Result<Multipart, InputError> {
|
||||
let Some((body, settings)) = take_body()? else {
|
||||
return Err(invalid("multipart requires an incoming HTTP body"));
|
||||
};
|
||||
if request.method != "POST" || request.target.contains('?') {
|
||||
return Err(invalid("multipart requires POST without query inputs"));
|
||||
}
|
||||
let mut types = request
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|header| header.name.eq_ignore_ascii_case("content-type"));
|
||||
let content_type = types
|
||||
.next()
|
||||
.and_then(|header| std::str::from_utf8(header.value.values()).ok())
|
||||
.ok_or_else(|| invalid("missing multipart Content-Type"))?;
|
||||
if types.next().is_some() {
|
||||
return Err(invalid("duplicate Content-Type"));
|
||||
}
|
||||
decode(
|
||||
body,
|
||||
content_type,
|
||||
&settings.upload_directory,
|
||||
settings.upload,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// The same streaming decoder is used by the HTTP owner and bounded body tests.
|
||||
pub async fn decode<B>(
|
||||
body: B,
|
||||
content_type: &str,
|
||||
directory: &Path,
|
||||
limits: UploadLimits,
|
||||
) -> Result<Multipart, InputError>
|
||||
where
|
||||
B: Body<Data = bytes::Bytes> + Unpin,
|
||||
{
|
||||
limits.validate().map_err(InputError)?;
|
||||
let (media, mut parameters) = parameters(content_type)?;
|
||||
let boundary = parameters
|
||||
.remove("boundary")
|
||||
.ok_or_else(|| invalid("missing multipart boundary"))?;
|
||||
if !media.eq_ignore_ascii_case("multipart/form-data")
|
||||
|| !parameters.is_empty()
|
||||
|| boundary.is_empty()
|
||||
|| boundary.len() > 70
|
||||
|| !boundary
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || b"'-_.,+/:=?()".contains(&byte))
|
||||
{
|
||||
return Err(invalid("invalid multipart Content-Type"));
|
||||
}
|
||||
let mut reader = Reader {
|
||||
body,
|
||||
pending: bytes::Bytes::new(),
|
||||
buffer: Vec::new(),
|
||||
total: 0,
|
||||
limit: limits.total_bytes,
|
||||
ended: false,
|
||||
};
|
||||
reader.exact(format!("--{boundary}\r\n").as_bytes()).await?;
|
||||
let delimiter = format!("\r\n--{boundary}").into_bytes();
|
||||
let mut uploads = BTreeMap::new();
|
||||
let mut fields = BTreeMap::new();
|
||||
let mut names = std::collections::BTreeSet::new();
|
||||
for part in 0..limits.parts {
|
||||
let headers = reader.until(b"\r\n\r\n", limits.header_bytes).await?;
|
||||
let (name, filename, content_type) = part_headers(&headers, limits.filename_bytes)?;
|
||||
if !names.insert(name.clone()) {
|
||||
return Err(invalid("duplicate multipart field"));
|
||||
}
|
||||
let mut upload = match filename {
|
||||
Some(filename) => Some(create(directory.to_path_buf(), filename, content_type).await?),
|
||||
None => {
|
||||
if fields.len() >= limits.fields {
|
||||
return Err(invalid("multipart fields exceed limit"));
|
||||
}
|
||||
None
|
||||
}
|
||||
};
|
||||
let mut field = Vec::new();
|
||||
let mut size = 0usize;
|
||||
loop {
|
||||
let (chunk, end) = reader.part_chunk(&delimiter).await?;
|
||||
let maximum = if upload.is_some() {
|
||||
limits.file_bytes
|
||||
} else {
|
||||
limits.field_bytes
|
||||
};
|
||||
if chunk.len() > maximum - size {
|
||||
return Err(invalid("multipart part exceeds byte limit"));
|
||||
}
|
||||
size += chunk.len();
|
||||
if let Some(upload) = &mut upload {
|
||||
let temporary = upload.temporary.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
temporary
|
||||
.file
|
||||
.lock()
|
||||
.map_err(|_| invalid("upload file lock failed"))?
|
||||
.as_mut()
|
||||
.ok_or_else(|| invalid("upload file closed"))?
|
||||
.write_all(&chunk)
|
||||
.map_err(|_| invalid("upload write failed"))
|
||||
})
|
||||
.await
|
||||
.map_err(|_| invalid("upload writer failed"))??;
|
||||
upload.size = size as i64;
|
||||
} else {
|
||||
field.extend_from_slice(&chunk);
|
||||
}
|
||||
if end {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if let Some(upload) = upload {
|
||||
uploads.insert(name, upload);
|
||||
} else {
|
||||
let value =
|
||||
String::from_utf8(field).map_err(|_| invalid("multipart scalar is not UTF-8"))?;
|
||||
fields.insert(
|
||||
name,
|
||||
serde_json::to_string(&value)
|
||||
.map_err(|_| invalid("multipart scalar encoding failed"))?,
|
||||
);
|
||||
}
|
||||
reader.fill(2).await?;
|
||||
if reader.buffer.starts_with(b"--") {
|
||||
reader.exact(b"--\r\n").await?;
|
||||
if reader.fill(1).await? {
|
||||
return Err(invalid("multipart epilogue is not supported"));
|
||||
}
|
||||
return Ok(Multipart {
|
||||
inputs: Inputs::from_fields(fields),
|
||||
uploads,
|
||||
});
|
||||
}
|
||||
reader.exact(b"\r\n").await?;
|
||||
if part + 1 == limits.parts {
|
||||
return Err(invalid("multipart parts exceed limit"));
|
||||
}
|
||||
}
|
||||
Err(invalid("incomplete multipart body"))
|
||||
}
|
||||
|
||||
async fn create(
|
||||
directory: PathBuf,
|
||||
filename: String,
|
||||
content_type: String,
|
||||
) -> Result<Upload, InputError> {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
std::fs::create_dir_all(&directory)
|
||||
.map_err(|_| invalid("cannot create upload directory"))?;
|
||||
let mut random = [0u8; 16];
|
||||
getrandom::getrandom(&mut random).map_err(|_| invalid("upload randomness unavailable"))?;
|
||||
let name = random
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect::<String>();
|
||||
let path = directory.join(format!("upload-{name}"));
|
||||
let mut options = std::fs::OpenOptions::new();
|
||||
options.read(true).write(true).create_new(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.mode(0o600);
|
||||
}
|
||||
let file = options
|
||||
.open(&path)
|
||||
.map_err(|_| invalid("cannot create upload file"))?;
|
||||
Ok(Upload {
|
||||
temporary: Arc::new(Temporary {
|
||||
path,
|
||||
file: Mutex::new(Some(file)),
|
||||
}),
|
||||
filename,
|
||||
content_type,
|
||||
size: 0,
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|_| invalid("upload owner failed"))?
|
||||
}
|
||||
|
||||
fn invalid(message: &str) -> InputError {
|
||||
InputError(message.to_owned())
|
||||
}
|
||||
|
||||
fn parameters(value: &str) -> Result<(&str, BTreeMap<String, String>), InputError> {
|
||||
let mut parts = value.split(';');
|
||||
let name = parts.next().unwrap_or("").trim();
|
||||
let mut parameters = BTreeMap::new();
|
||||
for part in parts {
|
||||
let (key, value) = part
|
||||
.trim()
|
||||
.split_once('=')
|
||||
.ok_or_else(|| invalid("invalid multipart parameter"))?;
|
||||
let value = value.trim();
|
||||
let value = value
|
||||
.strip_prefix('"')
|
||||
.and_then(|value| value.strip_suffix('"'))
|
||||
.unwrap_or(value);
|
||||
if value
|
||||
.bytes()
|
||||
.any(|byte| byte < 32 || byte == 127 || byte == b'"' || byte == b'\\')
|
||||
|| parameters
|
||||
.insert(key.trim().to_ascii_lowercase(), value.to_owned())
|
||||
.is_some()
|
||||
{
|
||||
return Err(invalid("invalid or duplicate multipart parameter"));
|
||||
}
|
||||
}
|
||||
Ok((name, parameters))
|
||||
}
|
||||
|
||||
fn part_headers(
|
||||
bytes: &[u8],
|
||||
filename_limit: usize,
|
||||
) -> Result<(String, Option<String>, String), InputError> {
|
||||
let headers = std::str::from_utf8(bytes).map_err(|_| invalid("invalid multipart headers"))?;
|
||||
let mut disposition = None;
|
||||
let mut content_type = None;
|
||||
for line in headers.split("\r\n") {
|
||||
let (name, value) = line
|
||||
.split_once(':')
|
||||
.ok_or_else(|| invalid("invalid multipart header"))?;
|
||||
if value.bytes().any(|byte| byte < 32 || byte == 127) {
|
||||
return Err(invalid("invalid multipart header value"));
|
||||
}
|
||||
match name.to_ascii_lowercase().as_str() {
|
||||
"content-disposition" if disposition.is_none() => disposition = Some(value.trim()),
|
||||
"content-type" if content_type.is_none() => {
|
||||
content_type = Some(value.trim().to_owned())
|
||||
}
|
||||
_ => return Err(invalid("unsupported or duplicate multipart header")),
|
||||
}
|
||||
}
|
||||
let (kind, mut parameters) =
|
||||
parameters(disposition.ok_or_else(|| invalid("missing Content-Disposition"))?)?;
|
||||
if !kind.eq_ignore_ascii_case("form-data") {
|
||||
return Err(invalid("expected form-data disposition"));
|
||||
}
|
||||
let name = parameters
|
||||
.remove("name")
|
||||
.ok_or_else(|| invalid("missing multipart field name"))?;
|
||||
if name.is_empty()
|
||||
|| name.len() > 128
|
||||
|| !name
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
|
||||
{
|
||||
return Err(invalid("invalid multipart field name"));
|
||||
}
|
||||
let filename = parameters.remove("filename");
|
||||
if !parameters.is_empty() {
|
||||
return Err(invalid("unsupported multipart disposition parameter"));
|
||||
}
|
||||
if filename.as_ref().is_some_and(|name| {
|
||||
name.is_empty()
|
||||
|| name.len() > filename_limit
|
||||
|| name == "."
|
||||
|| name == ".."
|
||||
|| name.contains(['/', '\\', ':'])
|
||||
}) {
|
||||
return Err(invalid("invalid upload filename"));
|
||||
}
|
||||
let content_type = content_type.unwrap_or_else(|| "application/octet-stream".into());
|
||||
Ok((name, filename, content_type))
|
||||
}
|
||||
|
||||
struct Reader<B> {
|
||||
body: B,
|
||||
pending: bytes::Bytes,
|
||||
buffer: Vec<u8>,
|
||||
total: usize,
|
||||
limit: usize,
|
||||
ended: bool,
|
||||
}
|
||||
|
||||
impl<B: Body<Data = bytes::Bytes> + Unpin> Reader<B> {
|
||||
async fn fill(&mut self, required: usize) -> Result<bool, InputError> {
|
||||
while self.buffer.len() < required && !self.ended {
|
||||
if self.pending.is_empty() {
|
||||
match self.body.frame().await {
|
||||
Some(Ok(frame)) => {
|
||||
self.pending = frame
|
||||
.into_data()
|
||||
.map_err(|_| invalid("multipart trailers are not supported"))?;
|
||||
if self.pending.len() > self.limit - self.total {
|
||||
return Err(invalid("multipart total exceeds byte limit"));
|
||||
}
|
||||
self.total += self.pending.len();
|
||||
}
|
||||
Some(Err(_)) => return Err(invalid("incomplete multipart body")),
|
||||
None => {
|
||||
self.ended = true;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
let count = self.pending.len().min(16_384);
|
||||
self.buffer.extend_from_slice(&self.pending.split_to(count));
|
||||
}
|
||||
Ok(self.buffer.len() >= required)
|
||||
}
|
||||
|
||||
async fn exact(&mut self, expected: &[u8]) -> Result<(), InputError> {
|
||||
if !self.fill(expected.len()).await? || !self.buffer.starts_with(expected) {
|
||||
return Err(invalid("invalid multipart delimiter"));
|
||||
}
|
||||
self.buffer.drain(..expected.len());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn until(&mut self, delimiter: &[u8], limit: usize) -> Result<Vec<u8>, InputError> {
|
||||
loop {
|
||||
if let Some(index) = self
|
||||
.buffer
|
||||
.windows(delimiter.len())
|
||||
.position(|part| part == delimiter)
|
||||
{
|
||||
if index > limit {
|
||||
return Err(invalid("multipart headers exceed limit"));
|
||||
}
|
||||
let bytes = self.buffer.drain(..index).collect();
|
||||
self.buffer.drain(..delimiter.len());
|
||||
return Ok(bytes);
|
||||
}
|
||||
if self.buffer.len() > limit + delimiter.len()
|
||||
|| !self.fill(self.buffer.len() + 1).await?
|
||||
{
|
||||
return Err(invalid("incomplete or oversized multipart headers"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn part_chunk(&mut self, delimiter: &[u8]) -> Result<(Vec<u8>, bool), InputError> {
|
||||
loop {
|
||||
if self.buffer.len() <= delimiter.len() + 2 && !self.ended {
|
||||
self.fill(16_384).await?;
|
||||
}
|
||||
let mut cursor = 0;
|
||||
while let Some(offset) = self.buffer[cursor..]
|
||||
.windows(delimiter.len())
|
||||
.position(|part| part == delimiter)
|
||||
{
|
||||
let index = cursor + offset;
|
||||
self.fill(index + delimiter.len() + 4).await?;
|
||||
let suffix = &self.buffer[index + delimiter.len()..];
|
||||
if suffix.starts_with(b"--\r\n") || suffix.starts_with(b"\r\n") {
|
||||
let chunk = self.buffer.drain(..index).collect();
|
||||
self.buffer.drain(..delimiter.len());
|
||||
return Ok((chunk, true));
|
||||
}
|
||||
cursor = index + 1;
|
||||
}
|
||||
if self.buffer.len() > delimiter.len() + 2 {
|
||||
let count = self.buffer.len() - delimiter.len() - 2;
|
||||
return Ok((self.buffer.drain(..count).collect(), false));
|
||||
}
|
||||
if !self.fill(self.buffer.len() + 1).await? {
|
||||
return Err(invalid("incomplete multipart file"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
//! Invocation-owned resources shared by generated application roots and requests.
|
||||
use std::cell::RefCell;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::config::{RuntimeProfile, Settings};
|
||||
|
||||
tokio::task_local! {
|
||||
static CURRENT: RefCell<Option<Arc<Session>>>;
|
||||
}
|
||||
|
||||
pub struct Session {
|
||||
settings: Arc<Settings>,
|
||||
lifecycle: std::sync::OnceLock<Arc<crate::lifecycle::Lifecycle>>,
|
||||
authorization: std::sync::OnceLock<bool>,
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
databases: Arc<crate::database::Session>,
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
tenant: std::sync::OnceLock<Arc<crate::tenant::Manager>>,
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
migrations: std::sync::Mutex<Vec<crate::database::Database>>,
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
jobs: std::sync::OnceLock<Arc<crate::job::resources::Resources>>,
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub fn load(
|
||||
profile: RuntimeProfile,
|
||||
bindings: &[(Option<&str>, &str, bool)],
|
||||
transactions: &[&[(Option<&str>, &str, bool)]],
|
||||
) -> Result<Arc<Self>, String> {
|
||||
let settings = Arc::new(Settings::load_for_executable()?);
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
let databases = Arc::new(crate::database::Session::from_settings(
|
||||
settings.clone(),
|
||||
profile,
|
||||
bindings,
|
||||
transactions,
|
||||
)?);
|
||||
#[cfg(not(any(feature = "sqlite", feature = "postgres")))]
|
||||
settings.prepare_databases(profile, bindings, transactions)?;
|
||||
Ok(Arc::new(Self {
|
||||
settings,
|
||||
lifecycle: std::sync::OnceLock::new(),
|
||||
authorization: std::sync::OnceLock::new(),
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
databases,
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
tenant: std::sync::OnceLock::new(),
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
migrations: std::sync::Mutex::new(Vec::new()),
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
jobs: std::sync::OnceLock::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn settings(&self) -> &Arc<Settings> {
|
||||
&self.settings
|
||||
}
|
||||
|
||||
pub fn initialize_authorization(&self, enabled: bool) -> Result<(), String> {
|
||||
self.authorization
|
||||
.set(enabled)
|
||||
.map_err(|_| "application authorization initialized twice".into())
|
||||
}
|
||||
|
||||
pub fn has_authorization(&self) -> bool {
|
||||
self.authorization.get().copied().unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(crate) fn lifecycle(&self) -> Option<Arc<crate::lifecycle::Lifecycle>> {
|
||||
self.lifecycle.get().cloned()
|
||||
}
|
||||
|
||||
pub(crate) fn set_lifecycle(
|
||||
&self,
|
||||
state: Arc<crate::lifecycle::Lifecycle>,
|
||||
) -> Result<(), String> {
|
||||
self.lifecycle
|
||||
.set(state)
|
||||
.map_err(|_| "application lifecycle initialized twice".into())
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub fn databases(&self) -> &Arc<crate::database::Session> {
|
||||
&self.databases
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub fn jobs(&self) -> Option<&Arc<crate::job::resources::Resources>> {
|
||||
self.jobs.get()
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub fn initialize_jobs(
|
||||
&self,
|
||||
jobs: Arc<crate::job::resources::Resources>,
|
||||
) -> Result<(), String> {
|
||||
self.jobs
|
||||
.set(jobs)
|
||||
.map_err(|_| "application Jobs initialized twice".into())
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub fn own_migration_database(&self, database: crate::database::Database) {
|
||||
self.migrations
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.push(database);
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub(crate) fn tenant(&self) -> Option<Arc<crate::tenant::Manager>> {
|
||||
self.tenant.get().cloned()
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub(crate) fn set_tenant(
|
||||
&self,
|
||||
manager: Arc<crate::tenant::Manager>,
|
||||
) -> Result<(), crate::orm::Error> {
|
||||
self.tenant.set(manager).map_err(|_| {
|
||||
crate::orm::Error::migration("tenant storage was initialized more than once")
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn scope<T>(self: Arc<Self>, future: impl Future<Output = T>) -> T {
|
||||
CURRENT.scope(RefCell::new(Some(self)), future).await
|
||||
}
|
||||
|
||||
pub async fn close(&self) -> Result<(), String> {
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
{
|
||||
self.databases.stop_accepting();
|
||||
let migrations = std::mem::take(
|
||||
&mut *self
|
||||
.migrations
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner()),
|
||||
);
|
||||
for database in &migrations {
|
||||
database.stop_accepting();
|
||||
}
|
||||
let mut migration_error = None;
|
||||
for database in migrations {
|
||||
if let Err(error) = database.close_pool().await {
|
||||
migration_error = Some(match migration_error {
|
||||
Some(previous) => crate::orm::Error::cause(previous, error),
|
||||
None => error,
|
||||
});
|
||||
}
|
||||
}
|
||||
let tenant = match self.tenant.get() {
|
||||
Some(manager) => manager.close().await,
|
||||
None => Ok(()),
|
||||
};
|
||||
let databases = self.databases.close().await;
|
||||
for result in [tenant, databases] {
|
||||
if let Err(error) = result {
|
||||
migration_error = Some(match migration_error {
|
||||
Some(previous) => previous.cause(error),
|
||||
None => error,
|
||||
});
|
||||
}
|
||||
}
|
||||
crate::log::flush();
|
||||
migration_error.map_or(Ok(()), |error| Err(error.to_string()))
|
||||
}
|
||||
#[cfg(not(any(feature = "sqlite", feature = "postgres")))]
|
||||
{
|
||||
crate::log::flush();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn current() -> Option<Arc<Session>> {
|
||||
CURRENT
|
||||
.try_with(|current| current.borrow().clone())
|
||||
.ok()
|
||||
.flatten()
|
||||
}
|
||||
|
||||
/// Installation occurs inside the generated root, before any application work.
|
||||
pub fn install(session: Arc<Session>) -> Result<(), String> {
|
||||
CURRENT
|
||||
.try_with(|current| {
|
||||
let mut current = current.borrow_mut();
|
||||
if current.is_some() {
|
||||
return Err("application session was initialized more than once".into());
|
||||
}
|
||||
*current = Some(session);
|
||||
Ok(())
|
||||
})
|
||||
.map_err(|_| "application session requires an application root".to_owned())?
|
||||
}
|
||||
|
||||
pub async fn scope_entry<T>(future: impl Future<Output = T>) -> T {
|
||||
CURRENT.scope(RefCell::new(None), future).await
|
||||
}
|
||||
|
||||
/// Carry application resources across a task boundary, never request identity.
|
||||
pub async fn scope_resources<T>(
|
||||
session: Option<Arc<Session>>,
|
||||
future: impl Future<Output = T>,
|
||||
) -> T {
|
||||
CURRENT.scope(RefCell::new(session), future).await
|
||||
}
|
||||
|
||||
pub fn with_resources<T>(session: Option<Arc<Session>>, apply: impl FnOnce() -> T) -> T {
|
||||
CURRENT.sync_scope(RefCell::new(session), apply)
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures_util::{Stream, StreamExt, stream};
|
||||
use tokio::sync::{Mutex, Semaphore};
|
||||
|
||||
use crate::resource::Pull;
|
||||
|
||||
type Producer<T> = Pin<Box<dyn Stream<Item = Pull<T>> + Send>>;
|
||||
|
||||
struct State<T> {
|
||||
producer: Mutex<Option<Producer<T>>>,
|
||||
closed: Semaphore,
|
||||
}
|
||||
|
||||
pub struct AsyncStream<T>(Arc<State<T>>);
|
||||
|
||||
impl<T> Clone for AsyncStream<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self(Arc::clone(&self.0))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> std::fmt::Debug for AsyncStream<T> {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("AsyncStream(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> AsyncStream<T> {
|
||||
pub fn new(producer: impl Stream<Item = Pull<T>> + Send + 'static) -> Self {
|
||||
Self(Arc::new(State {
|
||||
producer: Mutex::new(Some(Box::pin(producer))),
|
||||
closed: Semaphore::new(0),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn pull_next(&self) -> Pull<T> {
|
||||
let mut producer = self.0.producer.lock().await;
|
||||
let result = match producer.as_mut() {
|
||||
Some(producer) => tokio::select! {
|
||||
biased;
|
||||
_ = self.0.closed.acquire() => Pull::End,
|
||||
value = producer.next() => value.unwrap_or(Pull::End),
|
||||
},
|
||||
None => Pull::End,
|
||||
};
|
||||
if matches!(&result, Pull::Last(_) | Pull::End) {
|
||||
producer.take();
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn pull(&self) -> Option<T> {
|
||||
match self.pull_next().await {
|
||||
Pull::Item(value) | Pull::Last(value) => Some(value),
|
||||
Pull::End => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn close(&self) {
|
||||
self.0.closed.close();
|
||||
// 挂起的 pull 收到关闭信号后，会在其持有的锁内释放生产者。
|
||||
if let Ok(mut producer) = self.0.producer.try_lock() {
|
||||
producer.take();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn close_on_drop(&self) -> CloseOnDrop<T> {
|
||||
CloseOnDrop(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Send + 'static> AsyncStream<T> {
|
||||
pub fn from_values(values: impl IntoIterator<Item = T, IntoIter: Send + 'static>) -> Self {
|
||||
Self::new(stream::iter(values.into_iter().map(Pull::Item)))
|
||||
}
|
||||
|
||||
pub fn map<U: Send + 'static>(
|
||||
self,
|
||||
transform: impl FnMut(T) -> U + Send + 'static,
|
||||
) -> AsyncStream<U> {
|
||||
let owner = self.close_on_drop();
|
||||
AsyncStream::new(stream::unfold(
|
||||
(owner, transform),
|
||||
|(owner, mut transform)| async move {
|
||||
let result = match owner.0.pull_next().await {
|
||||
Pull::Item(value) => Pull::Item(transform(value)),
|
||||
Pull::Last(value) => Pull::Last(transform(value)),
|
||||
Pull::End => return None,
|
||||
};
|
||||
Some((result, (owner, transform)))
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CloseOnDrop<T>(AsyncStream<T>);
|
||||
|
||||
impl<T> Drop for CloseOnDrop<T> {
|
||||
fn drop(&mut self) {
|
||||
self.0.close();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,738 @@
|
||||
use std::future::Future;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use base64::Engine as _;
|
||||
use hmac::{Hmac, Mac};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::Sha256;
|
||||
|
||||
use crate::config::{AuthProvider, Settings};
|
||||
use crate::http::{Header, Request};
|
||||
use crate::secret::Secret;
|
||||
|
||||
const CLOCK_SKEW_SECONDS: i64 = 30;
|
||||
const MAX_TOKEN_BYTES: usize = 8 * 1024;
|
||||
const MAX_JWT_PART_BYTES: usize = 6 * 1024;
|
||||
const MAX_IDENTITY_TEXT_BYTES: usize = 256;
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub mod store;
|
||||
|
||||
tokio::task_local! { static INVOCATION: Invocation; }
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct Claims {
|
||||
pub subject: String,
|
||||
pub session: String,
|
||||
pub tenant: Option<String>,
|
||||
pub site: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct Identity {
|
||||
pub id: String,
|
||||
pub user_id: Option<i64>,
|
||||
pub tenant_id: Option<i64>,
|
||||
}
|
||||
|
||||
pub struct Prepared {
|
||||
site: String,
|
||||
provider_key: String,
|
||||
provider: AuthProvider,
|
||||
origin: Option<crate::config::HttpOrigin>,
|
||||
claims: Option<Claims>,
|
||||
}
|
||||
|
||||
impl Prepared {
|
||||
pub fn claims(&self) -> Option<&Claims> {
|
||||
self.claims.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum Error {
|
||||
Unauthorized,
|
||||
Forbidden,
|
||||
Internal,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for Error {
|
||||
fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
output.write_str(match self {
|
||||
Self::Unauthorized => "unauthorized",
|
||||
Self::Forbidden => "forbidden",
|
||||
Self::Internal => "authentication configuration error",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Invocation {
|
||||
site: String,
|
||||
provider_key: String,
|
||||
provider: AuthProvider,
|
||||
origin: Option<crate::config::HttpOrigin>,
|
||||
session: Option<String>,
|
||||
identity: Option<Identity>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct JwtHeader {
|
||||
alg: String,
|
||||
typ: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct JwtClaims {
|
||||
iss: String,
|
||||
sub: String,
|
||||
aud: String,
|
||||
sid: String,
|
||||
tenant: Option<String>,
|
||||
iat: i64,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
nbf: Option<i64>,
|
||||
exp: i64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct IssuedHeader {
|
||||
alg: &'static str,
|
||||
typ: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Eq, PartialEq)]
|
||||
enum CredentialKind {
|
||||
Authorization,
|
||||
Cookie,
|
||||
}
|
||||
|
||||
pub fn prepare(directory: &[&str], request: &Request, anonymous: bool) -> Result<Prepared, Error> {
|
||||
prepare_with_settings(
|
||||
&crate::config::current_settings(),
|
||||
directory,
|
||||
request,
|
||||
anonymous,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn prepare_with_settings(
|
||||
settings: &Settings,
|
||||
directory: &[&str],
|
||||
request: &Request,
|
||||
anonymous: bool,
|
||||
) -> Result<Prepared, Error> {
|
||||
let (site_key, site) = settings
|
||||
.site_for_directory(directory)
|
||||
.map_err(|_| Error::Internal)?
|
||||
.ok_or(Error::Internal)?;
|
||||
let provider = settings
|
||||
.auth()
|
||||
.provider(site.auth())
|
||||
.ok_or(Error::Internal)?;
|
||||
validate_request_host(request, site.hosts())?;
|
||||
|
||||
let authorization = bearer_token(request)?;
|
||||
// Public entries are anonymous for browser Cookie credentials. Expired or
|
||||
// revoked sessions therefore cannot prevent a fresh login. An explicit
|
||||
// Authorization credential remains strict.
|
||||
let cookie = if anonymous {
|
||||
None
|
||||
} else {
|
||||
cookie_token(request, provider.cookie())?
|
||||
};
|
||||
let credential = match (authorization, cookie) {
|
||||
(Some(_), Some(_)) => return Err(Error::Unauthorized),
|
||||
(Some(token), None) => Some((CredentialKind::Authorization, token)),
|
||||
(None, Some(token)) => Some((CredentialKind::Cookie, token)),
|
||||
(None, None) if anonymous => None,
|
||||
(None, None) => return Err(Error::Unauthorized),
|
||||
};
|
||||
if credential
|
||||
.as_ref()
|
||||
.is_some_and(|(kind, _)| *kind == CredentialKind::Cookie)
|
||||
{
|
||||
validate_cookie_origin(request, site.origin())?;
|
||||
}
|
||||
let claims = credential
|
||||
.map(|(_, token)| verify_jwt(&token, site_key, site.auth(), provider))
|
||||
.transpose()?;
|
||||
Ok(Prepared {
|
||||
site: site_key.to_owned(),
|
||||
provider_key: site.auth().to_owned(),
|
||||
provider: provider.clone(),
|
||||
origin: site.origin().cloned(),
|
||||
claims,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn scope<F>(
|
||||
prepared: Prepared,
|
||||
identity: Option<Identity>,
|
||||
future: F,
|
||||
) -> Result<F::Output, Error>
|
||||
where
|
||||
F: Future,
|
||||
{
|
||||
validate_identity(prepared.claims.as_ref(), identity.as_ref())?;
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
let execution = prepared
|
||||
.claims
|
||||
.as_ref()
|
||||
.zip(identity.as_ref())
|
||||
.map(|(claims, identity)| crate::job::ExecutionIdentity::User {
|
||||
tenant_id: identity.tenant_id,
|
||||
provider: prepared.provider_key.clone(),
|
||||
site: prepared.site.clone(),
|
||||
subject: claims.subject.clone(),
|
||||
session: claims.session.clone(),
|
||||
tenant_claim: claims.tenant.clone(),
|
||||
permission_key: String::new(),
|
||||
});
|
||||
let session = prepared
|
||||
.claims
|
||||
.as_ref()
|
||||
.map(|claims| claims.session.clone());
|
||||
let invocation = Invocation {
|
||||
site: prepared.site,
|
||||
provider_key: prepared.provider_key,
|
||||
provider: prepared.provider,
|
||||
origin: prepared.origin,
|
||||
session,
|
||||
identity,
|
||||
};
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
if let Some(tenant_id) = invocation
|
||||
.identity
|
||||
.as_ref()
|
||||
.and_then(|identity| identity.tenant_id)
|
||||
{
|
||||
let future = crate::tenant::scope(tenant_id, future);
|
||||
return Ok(match execution {
|
||||
Some(execution) => {
|
||||
INVOCATION
|
||||
.scope(invocation, crate::job::scope_execution(execution, future))
|
||||
.await
|
||||
}
|
||||
None => INVOCATION.scope(invocation, future).await,
|
||||
});
|
||||
}
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
if let Some(execution) = execution {
|
||||
return Ok(INVOCATION
|
||||
.scope(invocation, crate::job::scope_execution(execution, future))
|
||||
.await);
|
||||
}
|
||||
Ok(INVOCATION.scope(invocation, future).await)
|
||||
}
|
||||
|
||||
pub fn validate_job_source(provider_key: &str, site_key: &str) -> Result<(), Error> {
|
||||
let settings = crate::config::current_settings();
|
||||
let site = settings.sites().get(site_key).ok_or(Error::Internal)?;
|
||||
if site.auth() != provider_key || settings.auth().provider(provider_key).is_none() {
|
||||
return Err(Error::Internal);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn validate_job_identity(claims: &Claims, identity: &Identity) -> Result<(), Error> {
|
||||
validate_identity(Some(claims), Some(identity))
|
||||
}
|
||||
|
||||
pub fn site_key() -> Result<String, Error> {
|
||||
with_invocation(|invocation| Ok(invocation.site.clone()))
|
||||
}
|
||||
|
||||
pub fn id() -> Result<String, Error> {
|
||||
with_identity(|identity| Ok(identity.id.clone()))
|
||||
}
|
||||
|
||||
pub fn session() -> Result<String, Error> {
|
||||
with_invocation(|invocation| {
|
||||
if invocation.identity.is_none() {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
invocation.session.clone().ok_or(Error::Internal)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn user_id() -> Result<Option<i64>, Error> {
|
||||
with_identity(|identity| Ok(identity.user_id))
|
||||
}
|
||||
|
||||
pub fn tenant_id() -> Result<Option<i64>, Error> {
|
||||
with_identity(|identity| Ok(identity.tenant_id))
|
||||
}
|
||||
|
||||
pub fn owns_user(user_id: i64) -> Result<bool, Error> {
|
||||
with_identity(|identity| Ok(identity.user_id == Some(user_id)))
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub async fn require_components(
|
||||
components: &[&str],
|
||||
tenant_components: &[&str],
|
||||
) -> Result<(), Error> {
|
||||
if crate::config::current_settings().tenant().is_none() || components.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let tenant = tenant_id()?.ok_or(Error::Forbidden)?;
|
||||
crate::tenant::require_components(tenant, components, tenant_components)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
if error.kind() == crate::orm::ErrorKind::InvalidData {
|
||||
Error::Forbidden
|
||||
} else {
|
||||
crate::log::error(error.to_string(), Vec::new());
|
||||
Error::Internal
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub async fn authorize_permission(permission: &str, site: &str) -> Result<(), Error> {
|
||||
let user = user_id()?.ok_or(Error::Forbidden)?;
|
||||
let (_, database, actual_site) = authorization_context().await.map_err(|error| {
|
||||
crate::log::error(error, Vec::new());
|
||||
Error::Internal
|
||||
})?;
|
||||
if actual_site != site {
|
||||
return Err(Error::Forbidden);
|
||||
}
|
||||
match store::authorize(database, user, site, permission).await {
|
||||
Ok(true) => crate::job::bind_permission(permission).map_err(|error| {
|
||||
crate::log::error(error.to_string(), Vec::new());
|
||||
Error::Internal
|
||||
}),
|
||||
Ok(false) => Err(Error::Forbidden),
|
||||
Err(error) => {
|
||||
crate::log::error(error.to_string(), Vec::new());
|
||||
Err(Error::Internal)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub async fn permissions() -> Result<Vec<store::CatalogPermission>, String> {
|
||||
let (catalog, _, site) = authorization_context().await?;
|
||||
let mut permissions = store::list_permissions(catalog)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
permissions.retain(|permission| permission.site == site);
|
||||
Ok(permissions)
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub async fn save_role(
|
||||
id: &str,
|
||||
name: &str,
|
||||
all_permissions: bool,
|
||||
permission_keys: Vec<String>,
|
||||
) -> Result<(), String> {
|
||||
let (catalog, authorization, site) = authorization_context().await?;
|
||||
store::save_role(
|
||||
catalog,
|
||||
authorization,
|
||||
&store::Role {
|
||||
id: id.to_owned(),
|
||||
site,
|
||||
name: name.to_owned(),
|
||||
all_permissions,
|
||||
},
|
||||
&permission_keys,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub async fn grant_role(user_id: i64, role_id: &str) -> Result<(), String> {
|
||||
let (_, authorization, site) = authorization_context().await?;
|
||||
store::grant_role(authorization, user_id, &site, role_id)
|
||||
.await
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub async fn revoke_role(user_id: i64, role_id: &str) -> Result<(), String> {
|
||||
let (_, authorization, site) = authorization_context().await?;
|
||||
store::revoke_role(authorization, user_id, &site, role_id)
|
||||
.await
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub async fn disable_role(role_id: &str) -> Result<(), String> {
|
||||
let (_, authorization, site) = authorization_context().await?;
|
||||
store::disable_role(authorization, &site, role_id)
|
||||
.await
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
async fn authorization_context()
|
||||
-> Result<(crate::database::Database, crate::database::Database, String), String> {
|
||||
let settings = crate::config::current_settings();
|
||||
let connection = settings
|
||||
.tenant()
|
||||
.map(|tenant| tenant.database())
|
||||
.unwrap_or("default");
|
||||
let catalog =
|
||||
crate::database::database_for(Some(connection), "").map_err(|error| error.to_string())?;
|
||||
let authorization = match tenant_id().map_err(|error| error.to_string())? {
|
||||
Some(_) => crate::tenant::database(connection)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?,
|
||||
None => catalog.clone(),
|
||||
};
|
||||
let site = site_key().map_err(|error| error.to_string())?;
|
||||
Ok((catalog, authorization, site))
|
||||
}
|
||||
|
||||
pub fn issue(subject: &str, session: &str, tenant: Option<&str>) -> Result<Secret, String> {
|
||||
validate_identity_text("subject", subject).map_err(|error| error.to_string())?;
|
||||
validate_identity_text("session", session).map_err(|error| error.to_string())?;
|
||||
if let Some(tenant) = tenant {
|
||||
validate_identity_text("tenant", tenant).map_err(|error| error.to_string())?;
|
||||
}
|
||||
INVOCATION
|
||||
.try_with(|invocation| {
|
||||
let now = unix_seconds().map_err(|error| error.to_string())?;
|
||||
let ttl = i64::try_from(invocation.provider.ttl_seconds())
|
||||
.map_err(|_| Error::Internal.to_string())?;
|
||||
let claims = JwtClaims {
|
||||
iss: invocation.provider_key.clone(),
|
||||
sub: subject.to_owned(),
|
||||
aud: invocation.site.clone(),
|
||||
sid: session.to_owned(),
|
||||
tenant: tenant.map(str::to_owned),
|
||||
iat: now,
|
||||
nbf: None,
|
||||
exp: now
|
||||
.checked_add(ttl)
|
||||
.ok_or_else(|| Error::Internal.to_string())?,
|
||||
};
|
||||
encode_jwt(&claims, &invocation.provider)
|
||||
.map(Secret::from_input)
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.map_err(|_| {
|
||||
"authentication capability is available only within a site API entry".to_owned()
|
||||
})?
|
||||
}
|
||||
|
||||
pub fn issue_cookie(subject: &str, session: &str, tenant: Option<&str>) -> Result<(), String> {
|
||||
let token = issue(subject, session, tenant)?;
|
||||
INVOCATION
|
||||
.try_with(|invocation| {
|
||||
let max_age = i64::try_from(invocation.provider.ttl_seconds())
|
||||
.map_err(|_| "authentication configuration error".to_owned())?;
|
||||
crate::api::response_secret_cookie(
|
||||
invocation.provider.cookie(),
|
||||
token,
|
||||
session_cookie_options(Some(max_age)),
|
||||
)
|
||||
})
|
||||
.map_err(|_| {
|
||||
"authentication capability is available only within a site API entry".to_owned()
|
||||
})?
|
||||
}
|
||||
|
||||
pub fn clear_cookie() -> Result<(), String> {
|
||||
INVOCATION
|
||||
.try_with(|invocation| {
|
||||
crate::api::response_cookie(
|
||||
invocation.provider.cookie(),
|
||||
"",
|
||||
session_cookie_options(Some(0)),
|
||||
)
|
||||
})
|
||||
.map_err(|_| {
|
||||
"authentication capability is available only within a site API entry".to_owned()
|
||||
})?
|
||||
}
|
||||
|
||||
fn session_cookie_options(max_age: Option<i64>) -> crate::api::CookieOptions {
|
||||
crate::api::CookieOptions {
|
||||
path: "/".into(),
|
||||
domain: None,
|
||||
same_site: crate::api::SameSite::Lax,
|
||||
secure: true,
|
||||
http_only: true,
|
||||
max_age,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn cookie_origin() -> Result<crate::config::HttpOrigin, Error> {
|
||||
INVOCATION
|
||||
.try_with(|invocation| invocation.origin.clone().ok_or(Error::Internal))
|
||||
.map_err(|_| Error::Internal)?
|
||||
}
|
||||
|
||||
pub fn log_fields() -> Vec<crate::log::Field> {
|
||||
INVOCATION
|
||||
.try_with(|invocation| {
|
||||
let mut fields = vec![crate::log::Field {
|
||||
name: "site".into(),
|
||||
value: invocation.site.clone(),
|
||||
}];
|
||||
if let Some(identity) = &invocation.identity {
|
||||
fields.push(crate::log::Field {
|
||||
name: "auth_id".into(),
|
||||
value: identity.id.clone(),
|
||||
});
|
||||
if let Some(user_id) = identity.user_id {
|
||||
fields.push(crate::log::Field {
|
||||
name: "user_id".into(),
|
||||
value: user_id.to_string(),
|
||||
});
|
||||
}
|
||||
if let Some(tenant_id) = identity.tenant_id {
|
||||
fields.push(crate::log::Field {
|
||||
name: "tenant_id".into(),
|
||||
value: tenant_id.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
fields
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn validate_identity(claims: Option<&Claims>, identity: Option<&Identity>) -> Result<(), Error> {
|
||||
match (claims, identity) {
|
||||
(None, None) => return Ok(()),
|
||||
(Some(_), None) => return Err(Error::Unauthorized),
|
||||
(None, Some(_)) => return Err(Error::Internal),
|
||||
(Some(claims), Some(identity)) => {
|
||||
if identity.id != claims.subject
|
||||
|| claims.tenant.is_some() != identity.tenant_id.is_some()
|
||||
|| identity.user_id.is_some_and(|id| id <= 0)
|
||||
|| identity.tenant_id.is_some_and(|id| id <= 0)
|
||||
{
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
}
|
||||
}
|
||||
validate_identity_text("identity", &identity.expect("checked identity").id)
|
||||
}
|
||||
|
||||
fn validate_request_host(request: &Request, allowed: &[String]) -> Result<(), Error> {
|
||||
let host = unique_header(&request.headers, "host")?.ok_or(Error::Unauthorized)?;
|
||||
let host = std::str::from_utf8(host).map_err(|_| Error::Unauthorized)?;
|
||||
let host = crate::config::normalize_authority_host(host).map_err(|_| Error::Unauthorized)?;
|
||||
if !allowed.is_empty() && !allowed.iter().any(|allowed| allowed == &host) {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn bearer_token(request: &Request) -> Result<Option<Vec<u8>>, Error> {
|
||||
let Some(value) = unique_header(&request.headers, "authorization")? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !value.is_ascii() || value.len() > MAX_TOKEN_BYTES + 7 {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
let Some(separator) = value.iter().position(|byte| *byte == b' ') else {
|
||||
return Err(Error::Unauthorized);
|
||||
};
|
||||
let (scheme, token) = (&value[..separator], &value[separator + 1..]);
|
||||
if !scheme.eq_ignore_ascii_case(b"bearer")
|
||||
|| token.is_empty()
|
||||
|| token.len() > MAX_TOKEN_BYTES
|
||||
|| token.iter().any(|byte| byte.is_ascii_whitespace())
|
||||
{
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
Ok(Some(token.to_vec()))
|
||||
}
|
||||
|
||||
fn cookie_token(request: &Request, name: &str) -> Result<Option<Vec<u8>>, Error> {
|
||||
let mut selected = None;
|
||||
for header in request
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|header| header.name.eq_ignore_ascii_case("cookie"))
|
||||
{
|
||||
let raw = header.value.values();
|
||||
if !raw.is_ascii() {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
for pair in raw.split(|byte| *byte == b';') {
|
||||
let pair = pair.trim_ascii();
|
||||
let Some(separator) = pair.iter().position(|byte| *byte == b'=') else {
|
||||
return Err(Error::Unauthorized);
|
||||
};
|
||||
if &pair[..separator] != name.as_bytes() {
|
||||
continue;
|
||||
}
|
||||
if selected.is_some() {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
let encoded = &pair[separator + 1..];
|
||||
let encoded = if encoded.len() >= 2
|
||||
&& encoded.first() == Some(&b'"')
|
||||
&& encoded.last() == Some(&b'"')
|
||||
{
|
||||
&encoded[1..encoded.len() - 1]
|
||||
} else {
|
||||
encoded
|
||||
};
|
||||
let encoded = encoded.strip_prefix(b"v1.").ok_or(Error::Unauthorized)?;
|
||||
if encoded.len() > MAX_TOKEN_BYTES.saturating_mul(2) {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
let token = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(encoded)
|
||||
.map_err(|_| Error::Unauthorized)?;
|
||||
if token.is_empty() || token.len() > MAX_TOKEN_BYTES {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
selected = Some(token);
|
||||
}
|
||||
}
|
||||
Ok(selected)
|
||||
}
|
||||
|
||||
fn validate_cookie_origin(
|
||||
request: &Request,
|
||||
origin: Option<&crate::config::HttpOrigin>,
|
||||
) -> Result<(), Error> {
|
||||
if matches!(request.method.as_str(), "GET" | "HEAD" | "OPTIONS") {
|
||||
return Ok(());
|
||||
}
|
||||
let origin = origin.ok_or(Error::Internal)?;
|
||||
crate::api::validate_same_origin(&request.headers, origin).map_err(|_| Error::Forbidden)
|
||||
}
|
||||
|
||||
fn unique_header<'a>(headers: &'a [Header], name: &str) -> Result<Option<&'a [u8]>, Error> {
|
||||
let mut values = headers
|
||||
.iter()
|
||||
.filter(|header| header.name.eq_ignore_ascii_case(name));
|
||||
let value = values.next().map(|header| header.value.values());
|
||||
if values.next().is_some() {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn verify_jwt(
|
||||
token: &[u8],
|
||||
site: &str,
|
||||
provider_key: &str,
|
||||
provider: &AuthProvider,
|
||||
) -> Result<Claims, Error> {
|
||||
let text = std::str::from_utf8(token).map_err(|_| Error::Unauthorized)?;
|
||||
let mut parts = text.split('.');
|
||||
let header_part = parts.next().ok_or(Error::Unauthorized)?;
|
||||
let claims_part = parts.next().ok_or(Error::Unauthorized)?;
|
||||
let signature_part = parts.next().ok_or(Error::Unauthorized)?;
|
||||
if parts.next().is_some()
|
||||
|| [header_part, claims_part, signature_part]
|
||||
.iter()
|
||||
.any(|part| part.is_empty())
|
||||
{
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
let header = decode_part(header_part)?;
|
||||
let claims = decode_part(claims_part)?;
|
||||
let signature = decode_part(signature_part)?;
|
||||
let header_text = std::str::from_utf8(&header).map_err(|_| Error::Unauthorized)?;
|
||||
let claims_text = std::str::from_utf8(&claims).map_err(|_| Error::Unauthorized)?;
|
||||
crate::wire::parse(header_text).map_err(|_| Error::Unauthorized)?;
|
||||
crate::wire::parse(claims_text).map_err(|_| Error::Unauthorized)?;
|
||||
let header: JwtHeader = serde_json::from_slice(&header).map_err(|_| Error::Unauthorized)?;
|
||||
if header.alg != "HS256" || header.typ != "JWT" {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
let mut mac =
|
||||
Hmac::<Sha256>::new_from_slice(&provider.jwt_secret().0).map_err(|_| Error::Internal)?;
|
||||
mac.update(format!("{header_part}.{claims_part}").as_bytes());
|
||||
mac.verify_slice(&signature)
|
||||
.map_err(|_| Error::Unauthorized)?;
|
||||
|
||||
let claims: JwtClaims = serde_json::from_slice(&claims).map_err(|_| Error::Unauthorized)?;
|
||||
let now = unix_seconds()?;
|
||||
if claims.iss != provider_key
|
||||
|| claims.aud != site
|
||||
|| claims.iat <= 0
|
||||
|| claims.exp <= now - CLOCK_SKEW_SECONDS
|
||||
|| claims.iat > now + CLOCK_SKEW_SECONDS
|
||||
|| claims.nbf.is_some_and(|nbf| nbf > now + CLOCK_SKEW_SECONDS)
|
||||
|| claims.exp <= claims.iat
|
||||
|| claims.exp - claims.iat > provider.ttl_seconds() as i64 + CLOCK_SKEW_SECONDS
|
||||
{
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
validate_identity_text("subject", &claims.sub)?;
|
||||
validate_identity_text("session", &claims.sid)?;
|
||||
if let Some(tenant) = &claims.tenant {
|
||||
validate_identity_text("tenant", tenant)?;
|
||||
}
|
||||
Ok(Claims {
|
||||
subject: claims.sub,
|
||||
session: claims.sid,
|
||||
tenant: claims.tenant,
|
||||
site: claims.aud,
|
||||
})
|
||||
}
|
||||
|
||||
fn encode_jwt(claims: &JwtClaims, provider: &AuthProvider) -> Result<Vec<u8>, Error> {
|
||||
let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
let header = serde_json::to_vec(&IssuedHeader {
|
||||
alg: "HS256",
|
||||
typ: "JWT",
|
||||
})
|
||||
.map_err(|_| Error::Internal)?;
|
||||
let claims = serde_json::to_vec(claims).map_err(|_| Error::Internal)?;
|
||||
let signed = format!("{}.{}", engine.encode(header), engine.encode(claims));
|
||||
let mut mac =
|
||||
Hmac::<Sha256>::new_from_slice(&provider.jwt_secret().0).map_err(|_| Error::Internal)?;
|
||||
mac.update(signed.as_bytes());
|
||||
let signature = mac.finalize().into_bytes();
|
||||
Ok(format!("{signed}.{}", engine.encode(signature)).into_bytes())
|
||||
}
|
||||
|
||||
fn decode_part(part: &str) -> Result<Vec<u8>, Error> {
|
||||
if part.len() > MAX_JWT_PART_BYTES {
|
||||
return Err(Error::Unauthorized);
|
||||
}
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(part)
|
||||
.map_err(|_| Error::Unauthorized)
|
||||
}
|
||||
|
||||
fn unix_seconds() -> Result<i64, Error> {
|
||||
let seconds = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_err(|_| Error::Internal)?
|
||||
.as_secs();
|
||||
i64::try_from(seconds).map_err(|_| Error::Internal)
|
||||
}
|
||||
|
||||
fn validate_identity_text(_field: &str, value: &str) -> Result<(), Error> {
|
||||
if value.is_empty()
|
||||
|| value.len() > MAX_IDENTITY_TEXT_BYTES
|
||||
|| value.chars().any(char::is_control)
|
||||
{
|
||||
Err(Error::Unauthorized)
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn with_invocation<T>(read: impl FnOnce(&Invocation) -> Result<T, Error>) -> Result<T, Error> {
|
||||
INVOCATION.try_with(read).map_err(|_| Error::Unauthorized)?
|
||||
}
|
||||
|
||||
fn with_identity<T>(read: impl FnOnce(&Identity) -> Result<T, Error>) -> Result<T, Error> {
|
||||
with_invocation(|invocation| {
|
||||
invocation
|
||||
.identity
|
||||
.as_ref()
|
||||
.ok_or(Error::Unauthorized)
|
||||
.and_then(read)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,646 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::database::{Database, Executor, Sql, Transaction};
|
||||
use crate::orm::{self, Error, Row, Value};
|
||||
|
||||
const OWNER_ROLE_PREFIX: &str = "_dever_owner:";
|
||||
const ROLE_SCHEMA_VERSION: i64 = 2;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub struct Permission {
|
||||
pub key: &'static str,
|
||||
pub component: &'static str,
|
||||
pub domain: &'static str,
|
||||
pub site: &'static str,
|
||||
pub action: &'static str,
|
||||
pub method: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct CatalogPermission {
|
||||
pub key: String,
|
||||
pub component: String,
|
||||
pub domain: String,
|
||||
pub site: String,
|
||||
pub action: String,
|
||||
pub method: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct Role {
|
||||
pub id: String,
|
||||
pub site: String,
|
||||
pub name: String,
|
||||
pub all_permissions: bool,
|
||||
}
|
||||
|
||||
pub async fn initialize_catalog(database: Database) -> Result<(), Error> {
|
||||
require_control_database(&database)?;
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = initialize_catalog_in(&executor).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
pub async fn sync_catalog(database: Database, permissions: &[Permission]) -> Result<(), Error> {
|
||||
require_control_database(&database)?;
|
||||
validate_catalog(permissions)?;
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = sync_catalog_in(&executor, permissions).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
pub async fn list_permissions(database: Database) -> Result<Vec<CatalogPermission>, Error> {
|
||||
require_control_database(&database)?;
|
||||
let rows = database.query(Sql {
|
||||
sqlite: "SELECT key,component,domain,site,action,method FROM _dever_permission WHERE active=1 ORDER BY component,domain,site,action",
|
||||
postgres: "SELECT key,component,domain,site,action,method FROM _dever_permission WHERE active=1 ORDER BY component,domain,site,action",
|
||||
}, vec![]).await?;
|
||||
rows.into_iter().map(catalog_permission).collect()
|
||||
}
|
||||
|
||||
pub async fn initialize_roles(database: Database) -> Result<(), Error> {
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = initialize_roles_in(&executor).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
pub async fn authorize(
|
||||
database: Database,
|
||||
user_id: i64,
|
||||
site: &str,
|
||||
permission_key: &str,
|
||||
) -> Result<bool, Error> {
|
||||
validate_user_id(user_id)?;
|
||||
validate_name("site", site)?;
|
||||
validate_permission_key(permission_key, Some(site))?;
|
||||
let rows = database.query(Sql {
|
||||
sqlite: "SELECT 1 FROM _dever_auth_user_role AS ur JOIN _dever_auth_role AS r ON r.site=ur.site AND r.id=ur.role_id LEFT JOIN _dever_auth_role_permission AS rp ON rp.site=r.site AND rp.role_id=r.id AND rp.permission_key=?3 WHERE ur.user_id=?1 AND ur.site=?2 AND r.enabled=1 AND (r.all_permissions=1 OR rp.permission_key=?3) LIMIT 1",
|
||||
postgres: "SELECT 1 FROM _dever_auth_user_role AS ur JOIN _dever_auth_role AS r ON r.site=ur.site AND r.id=ur.role_id LEFT JOIN _dever_auth_role_permission AS rp ON rp.site=r.site AND rp.role_id=r.id AND rp.permission_key=$3 WHERE ur.user_id=$1 AND ur.site=$2 AND r.enabled=1 AND (r.all_permissions=1 OR rp.permission_key=$3) LIMIT 1",
|
||||
}, vec![number(user_id), text(site), text(permission_key)]).await?;
|
||||
Ok(!rows.is_empty())
|
||||
}
|
||||
|
||||
pub async fn save_role(
|
||||
catalog_database: Database,
|
||||
authorization_database: Database,
|
||||
role: &Role,
|
||||
permission_keys: &[String],
|
||||
) -> Result<(), Error> {
|
||||
require_control_database(&catalog_database)?;
|
||||
validate_role(role)?;
|
||||
let permissions = validate_role_permissions(&catalog_database, role, permission_keys).await?;
|
||||
let transaction = authorization_database.begin().await?;
|
||||
let executor = Executor::new(authorization_database, Some(&transaction))?;
|
||||
let result = save_role_in(&executor, role, &permissions).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
pub async fn grant_role(
|
||||
database: Database,
|
||||
user_id: i64,
|
||||
site: &str,
|
||||
role_id: &str,
|
||||
) -> Result<(), Error> {
|
||||
validate_user_id(user_id)?;
|
||||
validate_name("site", site)?;
|
||||
validate_mutable_role_id(role_id)?;
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = grant_role_in(&executor, user_id, site, role_id).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
pub async fn revoke_role(
|
||||
database: Database,
|
||||
user_id: i64,
|
||||
site: &str,
|
||||
role_id: &str,
|
||||
) -> Result<(), Error> {
|
||||
validate_user_id(user_id)?;
|
||||
validate_name("site", site)?;
|
||||
validate_mutable_role_id(role_id)?;
|
||||
database.execute(Sql {
|
||||
sqlite: "DELETE FROM _dever_auth_user_role WHERE user_id=?1 AND site=?2 AND role_id=?3",
|
||||
postgres: "DELETE FROM _dever_auth_user_role WHERE user_id=$1 AND site=$2 AND role_id=$3",
|
||||
}, vec![number(user_id), text(site), text(role_id)]).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn disable_role(database: Database, site: &str, role_id: &str) -> Result<(), Error> {
|
||||
validate_name("site", site)?;
|
||||
validate_mutable_role_id(role_id)?;
|
||||
let changed = database
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: "UPDATE _dever_auth_role SET enabled=0 WHERE id=?1 AND site=?2",
|
||||
postgres: "UPDATE _dever_auth_role SET enabled=0 WHERE id=$1 AND site=$2",
|
||||
},
|
||||
vec![text(role_id), text(site)],
|
||||
)
|
||||
.await?;
|
||||
if changed == 0 {
|
||||
return Err(Error::invalid_data("role does not exist"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn provision_owner(
|
||||
database: Database,
|
||||
user_id: i64,
|
||||
site: &str,
|
||||
) -> Result<String, Error> {
|
||||
validate_user_id(user_id)?;
|
||||
validate_name("site", site)?;
|
||||
let role_id = format!("{OWNER_ROLE_PREFIX}{site}");
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = async {
|
||||
executor.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_auth_role(id,site,name,all_permissions,enabled) VALUES (?1,?2,'Owner',1,1) ON CONFLICT(site,id) DO UPDATE SET name=excluded.name,all_permissions=1,enabled=1",
|
||||
postgres: "INSERT INTO _dever_auth_role(id,site,name,all_permissions,enabled) VALUES ($1,$2,'Owner',1,1) ON CONFLICT(site,id) DO UPDATE SET name=excluded.name,all_permissions=1,enabled=1",
|
||||
}, vec![text(&role_id), text(site)]).await?;
|
||||
grant_role_in(&executor, user_id, site, &role_id).await
|
||||
}.await;
|
||||
finish_transaction(transaction, result).await?;
|
||||
Ok(role_id)
|
||||
}
|
||||
|
||||
async fn initialize_catalog_in(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
lock_private_schema(executor).await?;
|
||||
executor.execute(Sql {
|
||||
sqlite: "CREATE TABLE IF NOT EXISTS _dever_permission (key TEXT PRIMARY KEY, component TEXT NOT NULL, domain TEXT NOT NULL, site TEXT NOT NULL, action TEXT NOT NULL, method TEXT NOT NULL, active BIGINT NOT NULL CHECK(active IN (0,1)), updated_at BIGINT NOT NULL)",
|
||||
postgres: "CREATE TABLE IF NOT EXISTS _dever_permission (key TEXT PRIMARY KEY, component TEXT NOT NULL, domain TEXT NOT NULL, site TEXT NOT NULL, action TEXT NOT NULL, method TEXT NOT NULL, active BIGINT NOT NULL CHECK(active IN (0,1)), updated_at BIGINT NOT NULL)",
|
||||
}, vec![]).await?;
|
||||
executor.execute(Sql {
|
||||
sqlite: "CREATE INDEX IF NOT EXISTS _dever_permission_catalog ON _dever_permission(active,component,domain,site,action)",
|
||||
postgres: "CREATE INDEX IF NOT EXISTS _dever_permission_catalog ON _dever_permission(active,component,domain,site,action)",
|
||||
}, vec![]).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn sync_catalog_in(executor: &Executor<'_>, permissions: &[Permission]) -> Result<(), Error> {
|
||||
initialize_catalog_in(executor).await?;
|
||||
let updated_at = crate::time::now().map_err(Error::migration)?;
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: "UPDATE _dever_permission SET active=0,updated_at=?1 WHERE active=1",
|
||||
postgres: "UPDATE _dever_permission SET active=0,updated_at=$1 WHERE active=1",
|
||||
},
|
||||
vec![number(updated_at)],
|
||||
)
|
||||
.await?;
|
||||
for permission in permissions {
|
||||
executor.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_permission(key,component,domain,site,action,method,active,updated_at) VALUES (?1,?2,?3,?4,?5,?6,1,?7) ON CONFLICT(key) DO UPDATE SET component=excluded.component,domain=excluded.domain,site=excluded.site,action=excluded.action,method=excluded.method,active=1,updated_at=excluded.updated_at",
|
||||
postgres: "INSERT INTO _dever_permission(key,component,domain,site,action,method,active,updated_at) VALUES ($1,$2,$3,$4,$5,$6,1,$7) ON CONFLICT(key) DO UPDATE SET component=excluded.component,domain=excluded.domain,site=excluded.site,action=excluded.action,method=excluded.method,active=1,updated_at=excluded.updated_at",
|
||||
}, vec![
|
||||
text(permission.key),
|
||||
text(permission.component),
|
||||
text(permission.domain),
|
||||
text(permission.site),
|
||||
text(permission.action),
|
||||
text(permission.method),
|
||||
number(updated_at),
|
||||
]).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn initialize_roles_in(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
lock_private_schema(executor).await?;
|
||||
executor.execute(Sql {
|
||||
sqlite: "CREATE TABLE IF NOT EXISTS _dever_auth_version (singleton INTEGER PRIMARY KEY CHECK(singleton=1), version BIGINT NOT NULL)",
|
||||
postgres: "CREATE TABLE IF NOT EXISTS _dever_auth_version (singleton BIGINT PRIMARY KEY CHECK(singleton=1), version BIGINT NOT NULL)",
|
||||
}, vec![]).await?;
|
||||
let versions = executor
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT version FROM _dever_auth_version",
|
||||
postgres: "SELECT version FROM _dever_auth_version",
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
if let Some(version) = versions.first() {
|
||||
if versions.len() != 1 || integer(version, 0)? != ROLE_SCHEMA_VERSION {
|
||||
return Err(Error::migration(
|
||||
"unsupported private authorization schema version",
|
||||
));
|
||||
}
|
||||
validate_role_schema_v2(executor).await?;
|
||||
create_role_indexes(executor).await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let role = private_table_exists(executor, "_dever_auth_role").await?;
|
||||
let permission = private_table_exists(executor, "_dever_auth_role_permission").await?;
|
||||
let user = private_table_exists(executor, "_dever_auth_user_role").await?;
|
||||
match (role, permission, user) {
|
||||
(false, false, false) => create_role_schema_v2(executor).await?,
|
||||
(true, true, true) => {
|
||||
let permission_site =
|
||||
private_column_exists(executor, "_dever_auth_role_permission", "site").await?;
|
||||
let user_site =
|
||||
private_column_exists(executor, "_dever_auth_user_role", "site").await?;
|
||||
match (permission_site, user_site) {
|
||||
(false, false) => migrate_role_schema_v1(executor).await?,
|
||||
(true, true) => {
|
||||
validate_role_schema_v2(executor).await?;
|
||||
create_role_indexes(executor).await?;
|
||||
}
|
||||
_ => {
|
||||
return Err(Error::migration(
|
||||
"inconsistent private authorization schema",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => return Err(Error::migration("incomplete private authorization schema")),
|
||||
}
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: "INSERT INTO _dever_auth_version(singleton,version) VALUES (1,?1)",
|
||||
postgres: "INSERT INTO _dever_auth_version(singleton,version) VALUES (1,$1)",
|
||||
},
|
||||
vec![number(ROLE_SCHEMA_VERSION)],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_role_schema_v2(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
for sql in [
|
||||
"CREATE TABLE IF NOT EXISTS _dever_auth_role (id TEXT NOT NULL, site TEXT NOT NULL, name TEXT NOT NULL, all_permissions BIGINT NOT NULL CHECK(all_permissions IN (0,1)), enabled BIGINT NOT NULL CHECK(enabled IN (0,1)), PRIMARY KEY(site,id))",
|
||||
"CREATE TABLE IF NOT EXISTS _dever_auth_role_permission (site TEXT NOT NULL, role_id TEXT NOT NULL, permission_key TEXT NOT NULL, PRIMARY KEY(site,role_id,permission_key), FOREIGN KEY(site,role_id) REFERENCES _dever_auth_role(site,id) ON DELETE CASCADE)",
|
||||
"CREATE TABLE IF NOT EXISTS _dever_auth_user_role (user_id BIGINT NOT NULL, site TEXT NOT NULL, role_id TEXT NOT NULL, PRIMARY KEY(user_id,site,role_id), FOREIGN KEY(site,role_id) REFERENCES _dever_auth_role(site,id) ON DELETE CASCADE)",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
create_role_indexes(executor).await
|
||||
}
|
||||
|
||||
async fn create_role_indexes(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
for sql in [
|
||||
"CREATE INDEX IF NOT EXISTS _dever_auth_role_site ON _dever_auth_role(site,enabled)",
|
||||
"CREATE INDEX IF NOT EXISTS _dever_auth_user_role_user ON _dever_auth_user_role(user_id,site,role_id)",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn migrate_role_schema_v1(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
for sql in [
|
||||
"CREATE TABLE _dever_auth_role_v2 (id TEXT NOT NULL, site TEXT NOT NULL, name TEXT NOT NULL, all_permissions BIGINT NOT NULL CHECK(all_permissions IN (0,1)), enabled BIGINT NOT NULL CHECK(enabled IN (0,1)), PRIMARY KEY(site,id))",
|
||||
"CREATE TABLE _dever_auth_role_permission_v2 (site TEXT NOT NULL, role_id TEXT NOT NULL, permission_key TEXT NOT NULL, PRIMARY KEY(site,role_id,permission_key), FOREIGN KEY(site,role_id) REFERENCES _dever_auth_role_v2(site,id) ON DELETE CASCADE)",
|
||||
"CREATE TABLE _dever_auth_user_role_v2 (user_id BIGINT NOT NULL, site TEXT NOT NULL, role_id TEXT NOT NULL, PRIMARY KEY(user_id,site,role_id), FOREIGN KEY(site,role_id) REFERENCES _dever_auth_role_v2(site,id) ON DELETE CASCADE)",
|
||||
"INSERT INTO _dever_auth_role_v2(id,site,name,all_permissions,enabled) SELECT id,site,name,all_permissions,enabled FROM _dever_auth_role",
|
||||
"INSERT INTO _dever_auth_role_permission_v2(site,role_id,permission_key) SELECT role.site,permission.role_id,permission.permission_key FROM _dever_auth_role_permission AS permission JOIN _dever_auth_role AS role ON role.id=permission.role_id",
|
||||
"INSERT INTO _dever_auth_user_role_v2(user_id,site,role_id) SELECT assignment.user_id,role.site,assignment.role_id FROM _dever_auth_user_role AS assignment JOIN _dever_auth_role AS role ON role.id=assignment.role_id",
|
||||
"DROP INDEX IF EXISTS _dever_auth_role_site",
|
||||
"DROP INDEX IF EXISTS _dever_auth_user_role_user",
|
||||
"DROP TABLE _dever_auth_role_permission",
|
||||
"DROP TABLE _dever_auth_user_role",
|
||||
"DROP TABLE _dever_auth_role",
|
||||
"ALTER TABLE _dever_auth_role_v2 RENAME TO _dever_auth_role",
|
||||
"ALTER TABLE _dever_auth_role_permission_v2 RENAME TO _dever_auth_role_permission",
|
||||
"ALTER TABLE _dever_auth_user_role_v2 RENAME TO _dever_auth_user_role",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
create_role_indexes(executor).await
|
||||
}
|
||||
|
||||
async fn validate_role_schema_v2(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
for table in [
|
||||
"_dever_auth_role",
|
||||
"_dever_auth_role_permission",
|
||||
"_dever_auth_user_role",
|
||||
] {
|
||||
if !private_table_exists(executor, table).await? {
|
||||
return Err(Error::migration("incomplete private authorization schema"));
|
||||
}
|
||||
}
|
||||
for table in ["_dever_auth_role_permission", "_dever_auth_user_role"] {
|
||||
if !private_column_exists(executor, table, "site").await? {
|
||||
return Err(Error::migration(
|
||||
"invalid private authorization schema version",
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn private_table_exists(executor: &Executor<'_>, table: &str) -> Result<bool, Error> {
|
||||
let rows = executor.query(Sql {
|
||||
sqlite: "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?1",
|
||||
postgres: "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema=current_schema() AND table_name=$1",
|
||||
}, vec![text(table)]).await?;
|
||||
Ok(integer(&orm::required_row(rows, "authorization table lookup")?, 0)? == 1)
|
||||
}
|
||||
|
||||
async fn private_column_exists(
|
||||
executor: &Executor<'_>,
|
||||
table: &str,
|
||||
column: &str,
|
||||
) -> Result<bool, Error> {
|
||||
let rows = executor.query(Sql {
|
||||
sqlite: "SELECT COUNT(*) FROM pragma_table_info(?1) WHERE name=?2",
|
||||
postgres: "SELECT COUNT(*) FROM information_schema.columns WHERE table_schema=current_schema() AND table_name=$1 AND column_name=$2",
|
||||
}, vec![text(table), text(column)]).await?;
|
||||
Ok(integer(&orm::required_row(rows, "authorization column lookup")?, 0)? == 1)
|
||||
}
|
||||
|
||||
async fn save_role_in(
|
||||
executor: &Executor<'_>,
|
||||
role: &Role,
|
||||
permission_keys: &BTreeSet<String>,
|
||||
) -> Result<(), Error> {
|
||||
executor.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_auth_role(id,site,name,all_permissions,enabled) VALUES (?1,?2,?3,?4,1) ON CONFLICT(site,id) DO UPDATE SET name=excluded.name,all_permissions=excluded.all_permissions",
|
||||
postgres: "INSERT INTO _dever_auth_role(id,site,name,all_permissions,enabled) VALUES ($1,$2,$3,$4,1) ON CONFLICT(site,id) DO UPDATE SET name=excluded.name,all_permissions=excluded.all_permissions",
|
||||
}, vec![
|
||||
text(&role.id),
|
||||
text(&role.site),
|
||||
text(&role.name),
|
||||
number(i64::from(role.all_permissions)),
|
||||
]).await?;
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: "DELETE FROM _dever_auth_role_permission WHERE site=?1 AND role_id=?2",
|
||||
postgres: "DELETE FROM _dever_auth_role_permission WHERE site=$1 AND role_id=$2",
|
||||
},
|
||||
vec![text(&role.site), text(&role.id)],
|
||||
)
|
||||
.await?;
|
||||
for permission_key in permission_keys {
|
||||
executor.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_auth_role_permission(site,role_id,permission_key) VALUES (?1,?2,?3)",
|
||||
postgres: "INSERT INTO _dever_auth_role_permission(site,role_id,permission_key) VALUES ($1,$2,$3)",
|
||||
}, vec![text(&role.site), text(&role.id), text(permission_key)]).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn grant_role_in(
|
||||
executor: &Executor<'_>,
|
||||
user_id: i64,
|
||||
site: &str,
|
||||
role_id: &str,
|
||||
) -> Result<(), Error> {
|
||||
let rows = executor
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT enabled FROM _dever_auth_role WHERE id=?1 AND site=?2",
|
||||
postgres: "SELECT enabled FROM _dever_auth_role WHERE id=$1 AND site=$2 FOR SHARE",
|
||||
},
|
||||
vec![text(role_id), text(site)],
|
||||
)
|
||||
.await?;
|
||||
let role = orm::required_row(rows, "grant role")?;
|
||||
if integer(&role, 0)? != 1 {
|
||||
return Err(Error::invalid_data("disabled role cannot be granted"));
|
||||
}
|
||||
executor.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_auth_user_role(user_id,site,role_id) VALUES (?1,?2,?3) ON CONFLICT DO NOTHING",
|
||||
postgres: "INSERT INTO _dever_auth_user_role(user_id,site,role_id) VALUES ($1,$2,$3) ON CONFLICT DO NOTHING",
|
||||
}, vec![number(user_id), text(site), text(role_id)]).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn validate_role_permissions(
|
||||
catalog_database: &Database,
|
||||
role: &Role,
|
||||
permission_keys: &[String],
|
||||
) -> Result<BTreeSet<String>, Error> {
|
||||
if role.all_permissions && !permission_keys.is_empty() {
|
||||
return Err(Error::invalid_data(
|
||||
"all-permissions role cannot store individual permissions",
|
||||
));
|
||||
}
|
||||
let requested = permission_keys.iter().cloned().collect::<BTreeSet<_>>();
|
||||
if requested.len() != permission_keys.len() {
|
||||
return Err(Error::invalid_data("role permissions contain duplicates"));
|
||||
}
|
||||
for permission_key in &requested {
|
||||
validate_permission_key(permission_key, Some(&role.site))?;
|
||||
}
|
||||
if requested.is_empty() {
|
||||
return Ok(requested);
|
||||
}
|
||||
let available = catalog_database
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT key FROM _dever_permission WHERE site=?1 AND active=1",
|
||||
postgres: "SELECT key FROM _dever_permission WHERE site=$1 AND active=1",
|
||||
},
|
||||
vec![text(&role.site)],
|
||||
)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| row_text(&row, 0))
|
||||
.collect::<Result<BTreeSet<_>, _>>()?;
|
||||
if let Some(unknown) = requested.iter().find(|key| !available.contains(*key)) {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"permission '{unknown}' is not active for site '{}'",
|
||||
role.site
|
||||
)));
|
||||
}
|
||||
Ok(requested)
|
||||
}
|
||||
|
||||
fn validate_catalog(permissions: &[Permission]) -> Result<(), Error> {
|
||||
let mut keys = BTreeSet::new();
|
||||
for permission in permissions {
|
||||
validate_name("component", permission.component)?;
|
||||
validate_name("domain", permission.domain)?;
|
||||
validate_name("site", permission.site)?;
|
||||
validate_name("action", permission.action)?;
|
||||
if !matches!(permission.method, "GET" | "POST" | "PUT" | "DELETE") {
|
||||
return Err(Error::invalid_data(
|
||||
"permission method must be GET, POST, PUT, or DELETE",
|
||||
));
|
||||
}
|
||||
let expected = format!(
|
||||
"{}.{}.{}.{}",
|
||||
permission.component, permission.domain, permission.site, permission.action
|
||||
);
|
||||
if permission.key != expected {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"permission key '{}' must be '{expected}'",
|
||||
permission.key
|
||||
)));
|
||||
}
|
||||
if !keys.insert(permission.key) {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"duplicate permission key '{}'",
|
||||
permission.key
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_role(role: &Role) -> Result<(), Error> {
|
||||
validate_mutable_role_id(&role.id)?;
|
||||
validate_name("site", &role.site)?;
|
||||
validate_text("role name", &role.name, 128)
|
||||
}
|
||||
|
||||
fn validate_mutable_role_id(role_id: &str) -> Result<(), Error> {
|
||||
validate_text("role id", role_id, 128)?;
|
||||
if role_id.starts_with(OWNER_ROLE_PREFIX) {
|
||||
return Err(Error::invalid_data("reserved role id"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_permission_key(key: &str, expected_site: Option<&str>) -> Result<(), Error> {
|
||||
if key.len() > 512 {
|
||||
return Err(Error::invalid_data("permission key exceeds 512 bytes"));
|
||||
}
|
||||
let parts = key.split('.').collect::<Vec<_>>();
|
||||
if parts.len() != 4 || parts.iter().any(|part| !valid_name(part)) {
|
||||
return Err(Error::invalid_data(
|
||||
"permission key must contain component.domain.site.action",
|
||||
));
|
||||
}
|
||||
if expected_site.is_some_and(|site| parts[2] != site) {
|
||||
return Err(Error::invalid_data(
|
||||
"permission key belongs to another site",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_name(label: &str, value: &str) -> Result<(), Error> {
|
||||
if !valid_name(value) {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"{label} must use lower_snake_case"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn valid_name(value: &str) -> bool {
|
||||
!value.is_empty()
|
||||
&& value.len() <= 128
|
||||
&& value.starts_with(|character: char| character.is_ascii_lowercase())
|
||||
&& value.split('_').all(|word| {
|
||||
!word.is_empty()
|
||||
&& word
|
||||
.chars()
|
||||
.all(|character| character.is_ascii_lowercase() || character.is_ascii_digit())
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_text(label: &str, value: &str, maximum: usize) -> Result<(), Error> {
|
||||
if value.is_empty() || value.len() > maximum || value.chars().any(char::is_control) {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"{label} must contain 1..={maximum} non-control bytes"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_user_id(user_id: i64) -> Result<(), Error> {
|
||||
if user_id <= 0 {
|
||||
return Err(Error::invalid_data("user id must be positive"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn require_control_database(database: &Database) -> Result<(), Error> {
|
||||
if database.key().tenant_id().is_some() {
|
||||
return Err(Error::invalid_data(
|
||||
"permission catalog requires the global control database",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn lock_private_schema(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
executor
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT 1",
|
||||
postgres: "SELECT pg_advisory_xact_lock(1684371045, 1635087476) IS NULL",
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn catalog_permission(row: Row) -> Result<CatalogPermission, Error> {
|
||||
Ok(CatalogPermission {
|
||||
key: row_text(&row, 0)?,
|
||||
component: row_text(&row, 1)?,
|
||||
domain: row_text(&row, 2)?,
|
||||
site: row_text(&row, 3)?,
|
||||
action: row_text(&row, 4)?,
|
||||
method: row_text(&row, 5)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn row_text(row: &Row, index: usize) -> Result<String, Error> {
|
||||
orm::text(row.get(index)?.clone())
|
||||
}
|
||||
|
||||
fn integer(row: &Row, index: usize) -> Result<i64, Error> {
|
||||
orm::int(row.get(index)?.clone())
|
||||
}
|
||||
|
||||
fn text(value: impl Into<String>) -> Value {
|
||||
Value::Text(value.into())
|
||||
}
|
||||
|
||||
fn number(value: i64) -> Value {
|
||||
Value::Int(value)
|
||||
}
|
||||
|
||||
async fn finish_transaction<T>(
|
||||
transaction: Transaction,
|
||||
result: Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
match result {
|
||||
Ok(value) => {
|
||||
transaction.commit().await?;
|
||||
Ok(value)
|
||||
}
|
||||
Err(error) => match transaction.rollback().await {
|
||||
Ok(()) => Err(error),
|
||||
Err(rollback) => Err(error.cause(rollback)),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Bytes(bytes::Bytes);
|
||||
|
||||
impl Bytes {
|
||||
pub fn new(bytes: Vec<u8>) -> Self {
|
||||
Self(bytes.into())
|
||||
}
|
||||
pub fn values(&self) -> &[u8] {
|
||||
&self.0
|
||||
}
|
||||
pub fn into_values(self) -> BytesIntoValues {
|
||||
self.0.into_iter()
|
||||
}
|
||||
pub(crate) fn from_http(bytes: bytes::Bytes) -> Self {
|
||||
Self(bytes)
|
||||
}
|
||||
pub(crate) fn into_http(self) -> bytes::Bytes {
|
||||
self.0
|
||||
}
|
||||
pub fn from_text(text: &str) -> Self {
|
||||
Self::new(text.as_bytes().to_vec())
|
||||
}
|
||||
pub fn from_string(text: String) -> Self {
|
||||
Self::new(text.into_bytes())
|
||||
}
|
||||
pub fn to_text(&self) -> Result<String, String> {
|
||||
decode_text(self.0.to_vec())
|
||||
}
|
||||
pub fn into_text(self) -> Result<String, String> {
|
||||
decode_text(self.0.into())
|
||||
}
|
||||
pub fn length(&self) -> i64 {
|
||||
self.0.len() as i64
|
||||
}
|
||||
pub fn at(&self, index: i64) -> Option<i64> {
|
||||
usize::try_from(index)
|
||||
.ok()
|
||||
.and_then(|index| self.0.get(index))
|
||||
.map(|value| i64::from(*value))
|
||||
}
|
||||
pub fn slice(&self, start: i64, end: i64) -> Result<Self, &'static str> {
|
||||
let start = usize::try_from(start).map_err(|_| "byte range is out of bounds")?;
|
||||
let end = usize::try_from(end).map_err(|_| "byte range is out of bounds")?;
|
||||
if start > end || end > self.0.len() {
|
||||
return Err("byte range is out of bounds");
|
||||
}
|
||||
Ok(Self(self.0.slice(start..end)))
|
||||
}
|
||||
pub fn concat(self, other: &Self) -> Self {
|
||||
if other.values().is_empty() {
|
||||
return self;
|
||||
}
|
||||
if self.values().is_empty() {
|
||||
return other.clone();
|
||||
}
|
||||
let mut bytes = match self.0.try_into_mut() {
|
||||
Ok(bytes) => bytes,
|
||||
Err(bytes) => {
|
||||
let mut result = bytes::BytesMut::with_capacity(bytes.len() + other.0.len());
|
||||
result.extend_from_slice(&bytes);
|
||||
result
|
||||
}
|
||||
};
|
||||
bytes.extend_from_slice(other.values());
|
||||
Self(bytes.freeze())
|
||||
}
|
||||
pub fn from_ints(values: &[i64]) -> Result<Self, &'static str> {
|
||||
values
|
||||
.iter()
|
||||
.map(|value| u8::try_from(*value).map_err(|_| "byte value must be between 0 and 255"))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Self::new)
|
||||
}
|
||||
}
|
||||
|
||||
pub type BytesIntoValues = bytes::buf::IntoIter<bytes::Bytes>;
|
||||
|
||||
fn decode_text(bytes: Vec<u8>) -> Result<String, String> {
|
||||
String::from_utf8(bytes).map_err(|error| error.to_string())
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
pub const MAX_CHANNEL_CAPACITY: usize = 65_536;
|
||||
|
||||
struct State<T> {
|
||||
queue: Mutex<Queue<T>>,
|
||||
available_items: Arc<Semaphore>,
|
||||
available_slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
struct Queue<T> {
|
||||
values: VecDeque<T>,
|
||||
closed: bool,
|
||||
}
|
||||
|
||||
pub struct Channel<T>(Arc<State<T>>);
|
||||
|
||||
impl<T> std::fmt::Debug for Channel<T> {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("Channel(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Clone for Channel<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self(Arc::clone(&self.0))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Send + 'static> Channel<T> {
|
||||
pub fn stream(self) -> crate::async_stream::AsyncStream<T> {
|
||||
crate::async_stream::AsyncStream::new(futures_util::stream::unfold(
|
||||
self,
|
||||
|channel| async move {
|
||||
let value = channel
|
||||
.receive()
|
||||
.await
|
||||
.expect("channel state lock poisoned")?;
|
||||
Some((crate::resource::Pull::Item(value), channel))
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Channel<T> {
|
||||
pub fn new(capacity: i64) -> Result<Self, &'static str> {
|
||||
let capacity = usize::try_from(capacity)
|
||||
.ok()
|
||||
.filter(|capacity| (1..=MAX_CHANNEL_CAPACITY).contains(capacity))
|
||||
.ok_or("channel capacity must be between 1 and 65536")?;
|
||||
Ok(Self(Arc::new(State {
|
||||
queue: Mutex::new(Queue {
|
||||
values: VecDeque::new(),
|
||||
closed: false,
|
||||
}),
|
||||
available_items: Arc::new(Semaphore::new(0)),
|
||||
available_slots: Arc::new(Semaphore::new(capacity)),
|
||||
})))
|
||||
}
|
||||
|
||||
pub async fn send(&self, value: T) -> Result<(), String> {
|
||||
let slot = Arc::clone(&self.0.available_slots)
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| "channel is closed".to_owned())?;
|
||||
let mut queue = self
|
||||
.0
|
||||
.queue
|
||||
.lock()
|
||||
.map_err(|_| "channel state lock failed".to_owned())?;
|
||||
if queue.closed {
|
||||
return Err("channel is closed".to_owned());
|
||||
}
|
||||
queue.values.push_back(value);
|
||||
slot.forget();
|
||||
drop(queue);
|
||||
self.0.available_items.add_permits(1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn receive(&self) -> Result<Option<T>, String> {
|
||||
let item = Arc::clone(&self.0.available_items).acquire_owned().await;
|
||||
let value = self
|
||||
.0
|
||||
.queue
|
||||
.lock()
|
||||
.map_err(|_| "channel state lock failed".to_owned())?
|
||||
.values
|
||||
.pop_front();
|
||||
if let Ok(item) = item {
|
||||
item.forget();
|
||||
if value.is_some() {
|
||||
self.0.available_slots.add_permits(1);
|
||||
}
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub async fn close(&self) -> Result<(), String> {
|
||||
let mut queue = self
|
||||
.0
|
||||
.queue
|
||||
.lock()
|
||||
.map_err(|_| "channel state lock failed".to_owned())?;
|
||||
if queue.closed {
|
||||
return Ok(());
|
||||
}
|
||||
queue.closed = true;
|
||||
drop(queue);
|
||||
self.0.available_slots.close();
|
||||
self.0.available_items.close();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
use std::hash::Hash;
|
||||
use std::sync::Arc;
|
||||
|
||||
use indexmap::IndexMap;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct List<T>(Arc<Vec<T>>);
|
||||
|
||||
impl<T> List<T> {
|
||||
pub fn new(values: Vec<T>) -> Self {
|
||||
Self(Arc::new(values))
|
||||
}
|
||||
pub fn values(&self) -> &[T] {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Clone> List<T> {
|
||||
pub fn into_first(self) -> Option<T> {
|
||||
match Arc::try_unwrap(self.0) {
|
||||
Ok(values) => values.into_iter().next(),
|
||||
Err(values) => values.first().cloned(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn into_values(self) -> IntoValues<T> {
|
||||
match Arc::try_unwrap(self.0) {
|
||||
Ok(values) => IntoValues::Owned(values.into_iter()),
|
||||
Err(values) => IntoValues::Shared { values, index: 0 },
|
||||
}
|
||||
}
|
||||
|
||||
pub fn append(mut self, value: T) -> Self {
|
||||
Arc::make_mut(&mut self.0).push(value);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// 独占列表转移元素；共享列表按需复制，提前停止不复制剩余元素。
|
||||
pub enum IntoValues<T> {
|
||||
Owned(std::vec::IntoIter<T>),
|
||||
Shared { values: Arc<Vec<T>>, index: usize },
|
||||
}
|
||||
|
||||
impl<T: Clone> Iterator for IntoValues<T> {
|
||||
type Item = T;
|
||||
|
||||
fn next(&mut self) -> Option<T> {
|
||||
match self {
|
||||
Self::Owned(values) => values.next(),
|
||||
Self::Shared { values, index } => {
|
||||
let value = values.get(*index)?.clone();
|
||||
*index += 1;
|
||||
Some(value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||
let remaining = match self {
|
||||
Self::Owned(values) => values.len(),
|
||||
Self::Shared { values, index } => values.len() - index,
|
||||
};
|
||||
(remaining, Some(remaining))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Clone> ExactSizeIterator for IntoValues<T> {}
|
||||
|
||||
impl<T: PartialEq> PartialEq for List<T> {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.values().len() == other.values().len()
|
||||
&& self
|
||||
.values()
|
||||
.iter()
|
||||
.zip(other.values())
|
||||
.all(|(left, right)| left == right)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct MapEntry<K, V> {
|
||||
pub key: K,
|
||||
pub value: V,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Map<K, V>(Arc<IndexMap<K, V>>);
|
||||
|
||||
impl<K: Clone + Eq + Hash, V: Clone> Map<K, V> {
|
||||
pub fn new(entries: Vec<(K, V)>) -> Result<Self, &'static str> {
|
||||
let mut map = IndexMap::with_capacity(entries.len());
|
||||
for (key, value) in entries {
|
||||
if map.insert(key, value).is_some() {
|
||||
return Err("duplicate Map key");
|
||||
}
|
||||
}
|
||||
Ok(Self(Arc::new(map)))
|
||||
}
|
||||
|
||||
pub fn from_unique(entries: Vec<(K, V)>) -> Self {
|
||||
let mut map = IndexMap::with_capacity(entries.len());
|
||||
for (key, value) in entries {
|
||||
debug_assert!(map.insert(key, value).is_none());
|
||||
}
|
||||
Self(Arc::new(map))
|
||||
}
|
||||
|
||||
pub fn get(&self, key: &K) -> Option<V> {
|
||||
self.0.get(key).cloned()
|
||||
}
|
||||
|
||||
pub fn into_get(self, key: &K) -> Option<V> {
|
||||
match Arc::try_unwrap(self.0) {
|
||||
Ok(mut entries) => entries.swap_remove(key),
|
||||
Err(entries) => entries.get(key).cloned(),
|
||||
}
|
||||
}
|
||||
|
||||
// 将小型 COW 包装内联，使调用方的循环能一起优化所有权检查与插入。
|
||||
#[inline]
|
||||
pub fn put(mut self, key: K, value: V) -> Self {
|
||||
Arc::make_mut(&mut self.0).insert(key, value);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn remove(mut self, key: &K) -> Self {
|
||||
Arc::make_mut(&mut self.0).shift_remove(key);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn entries(&self) -> List<MapEntry<K, V>> {
|
||||
cloned_entries(&self.0)
|
||||
}
|
||||
|
||||
pub fn into_entries(self) -> List<MapEntry<K, V>> {
|
||||
match Arc::try_unwrap(self.0) {
|
||||
Ok(entries) => List::new(
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| MapEntry { key, value })
|
||||
.collect(),
|
||||
),
|
||||
Err(entries) => cloned_entries(&entries),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn cloned_entries<K: Clone, V: Clone>(entries: &IndexMap<K, V>) -> List<MapEntry<K, V>> {
|
||||
List::new(
|
||||
entries
|
||||
.iter()
|
||||
.map(|(key, value)| MapEntry {
|
||||
key: key.clone(),
|
||||
value: value.clone(),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
impl<K, V> Map<K, V> {
|
||||
pub fn pairs(&self) -> impl Iterator<Item = (&K, &V)> {
|
||||
self.0.iter()
|
||||
}
|
||||
}
|
||||
|
||||
impl<K: PartialEq, V: PartialEq> PartialEq for Map<K, V> {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.0.len() == other.0.len()
|
||||
&& self
|
||||
.pairs()
|
||||
.zip(other.pairs())
|
||||
.all(|((ak, av), (bk, bv))| ak == bk && av == bv)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,831 @@
|
||||
//! Supervised process boundary for external Adapter implementations.
|
||||
use std::collections::{BTreeMap, BTreeSet, VecDeque};
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||
use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command};
|
||||
use tokio::sync::{Mutex, mpsc, oneshot};
|
||||
|
||||
mod command;
|
||||
|
||||
const VERSION: &str = "dever-component-1";
|
||||
const COMMAND_CAPACITY: usize = 64;
|
||||
const START_LIMIT: usize = 3;
|
||||
const START_WINDOW: Duration = Duration::from_secs(10);
|
||||
const SHUTDOWN_GRACE: Duration = Duration::from_secs(2);
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Definition {
|
||||
pub key: String,
|
||||
pub ecosystem: String,
|
||||
pub entry: String,
|
||||
pub port: String,
|
||||
pub adapter: String,
|
||||
pub schema: String,
|
||||
pub capabilities: Vec<String>,
|
||||
pub operations: Vec<String>,
|
||||
pub setting: Option<String>,
|
||||
pub timeout_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Reply {
|
||||
Result(String),
|
||||
Error { identity: String, payload: String },
|
||||
}
|
||||
|
||||
struct Managed {
|
||||
definition: Definition,
|
||||
state: Mutex<State>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct State {
|
||||
sender: Option<mpsc::Sender<Request>>,
|
||||
stop: Option<oneshot::Sender<()>>,
|
||||
task: Option<crate::task::Task<()>>,
|
||||
starts: VecDeque<Instant>,
|
||||
stopping: bool,
|
||||
}
|
||||
|
||||
struct Request {
|
||||
operation: String,
|
||||
payload: String,
|
||||
response: oneshot::Sender<Result<Reply, String>>,
|
||||
command_lease: Option<command::RequestLease>,
|
||||
}
|
||||
|
||||
tokio::task_local! {
|
||||
static CURRENT: Option<Arc<Session>>;
|
||||
}
|
||||
|
||||
/// Worker handles belong to one runtime invocation, including partial startup.
|
||||
pub(crate) struct Session {
|
||||
scope: crate::task::Scope,
|
||||
registry: Mutex<Registry>,
|
||||
shutdown: Mutex<()>,
|
||||
pub(crate) settings: std::sync::OnceLock<Result<Arc<crate::config::Settings>, String>>,
|
||||
#[cfg(feature = "external")]
|
||||
pub(crate) resources: std::sync::OnceLock<crate::external::PreparedBundle>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Registry {
|
||||
workers: BTreeMap<String, Arc<Managed>>,
|
||||
stopping: bool,
|
||||
closed: bool,
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub(crate) fn new(scope: crate::task::Scope) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
scope,
|
||||
registry: Mutex::new(Registry::default()),
|
||||
shutdown: Mutex::new(()),
|
||||
settings: std::sync::OnceLock::new(),
|
||||
#[cfg(feature = "external")]
|
||||
resources: std::sync::OnceLock::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn current() -> Option<Arc<Session>> {
|
||||
CURRENT.try_with(Clone::clone).ok().flatten()
|
||||
}
|
||||
|
||||
fn session() -> Result<Arc<Session>, String> {
|
||||
current().ok_or_else(|| "external Adapter requires a Dever invocation".into())
|
||||
}
|
||||
|
||||
/// One settings snapshot owns both Adapter selection and deployment grants.
|
||||
pub fn adapter_settings() -> Result<Arc<crate::config::Settings>, String> {
|
||||
#[cfg(feature = "api")]
|
||||
if let Some(application) = crate::application::current() {
|
||||
return Ok(application.settings().clone());
|
||||
}
|
||||
session()?
|
||||
.settings
|
||||
.get_or_init(|| crate::config::Settings::load_adapter_settings().map(Arc::new))
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn scope_resources<T>(
|
||||
session: Option<Arc<Session>>,
|
||||
future: impl Future<Output = T>,
|
||||
) -> impl Future<Output = T> {
|
||||
CURRENT.scope(session, future)
|
||||
}
|
||||
|
||||
pub(crate) fn with_resources<T>(session: Option<Arc<Session>>, apply: impl FnOnce() -> T) -> T {
|
||||
CURRENT.sync_scope(session, apply)
|
||||
}
|
||||
|
||||
pub async fn start(definition: Definition) -> Result<(), String> {
|
||||
validate_definition(&definition)?;
|
||||
let session = session()?;
|
||||
let mut registered = session.registry.lock().await;
|
||||
if registered.stopping || registered.closed {
|
||||
return Err("external Adapter invocation is stopped".into());
|
||||
}
|
||||
if registered.workers.contains_key(&definition.key) {
|
||||
return Err("external Adapter initialized twice".into());
|
||||
}
|
||||
let managed = Arc::new(Managed {
|
||||
definition,
|
||||
state: Mutex::new(State::default()),
|
||||
});
|
||||
// Register before awaiting startup so cancellation cannot orphan the owner.
|
||||
registered
|
||||
.workers
|
||||
.insert(managed.definition.key.clone(), managed.clone());
|
||||
drop(registered);
|
||||
managed.ensure_started().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn call(
|
||||
key: &str,
|
||||
operation: &str,
|
||||
payload: &crate::wire::Encoded,
|
||||
) -> Result<Reply, String> {
|
||||
let session = session()?;
|
||||
let registered = session.registry.lock().await;
|
||||
if registered.stopping || registered.closed {
|
||||
return Err("external Adapter invocation is stopped".into());
|
||||
}
|
||||
let managed = registered
|
||||
.workers
|
||||
.get(key)
|
||||
.cloned()
|
||||
.ok_or_else(|| "external Adapter is not initialized".to_owned())?;
|
||||
drop(registered);
|
||||
if !managed
|
||||
.definition
|
||||
.operations
|
||||
.iter()
|
||||
.any(|name| name == operation)
|
||||
{
|
||||
return Err("external Adapter operation is not declared by its Port".into());
|
||||
}
|
||||
for _ in 0..2 {
|
||||
let sender = managed.ensure_started().await?;
|
||||
let (response, result) = oneshot::channel();
|
||||
let lease = (managed.definition.ecosystem == "command")
|
||||
.then(command::call_lease)
|
||||
.transpose()?;
|
||||
let request = Request {
|
||||
operation: operation.to_owned(),
|
||||
payload: payload.as_str().to_owned(),
|
||||
response,
|
||||
command_lease: lease.as_ref().map(|lease| lease.request()),
|
||||
};
|
||||
if sender.send(request).await.is_ok() {
|
||||
return result
|
||||
.await
|
||||
.unwrap_or_else(|_| Err("external Adapter worker stopped".into()));
|
||||
}
|
||||
}
|
||||
Err("external Adapter worker is unavailable".into())
|
||||
}
|
||||
|
||||
pub async fn shutdown() -> Result<(), String> {
|
||||
shutdown_session(false).await
|
||||
}
|
||||
|
||||
pub(crate) async fn close() -> Result<(), String> {
|
||||
shutdown_session(true).await
|
||||
}
|
||||
|
||||
async fn shutdown_session(terminal: bool) -> Result<(), String> {
|
||||
let session = session()?;
|
||||
let _shutdown = session.shutdown.lock().await;
|
||||
let managed = {
|
||||
let mut registered = session.registry.lock().await;
|
||||
registered.stopping = true;
|
||||
registered.closed |= terminal;
|
||||
registered.workers.values().cloned().collect::<Vec<_>>()
|
||||
};
|
||||
let mut first = None;
|
||||
for component in managed {
|
||||
if let Err(error) = component.stop().await {
|
||||
first = Some(match first {
|
||||
Some(previous) => format!("{previous}; worker cleanup: {error}"),
|
||||
None => error,
|
||||
});
|
||||
}
|
||||
session
|
||||
.registry
|
||||
.lock()
|
||||
.await
|
||||
.workers
|
||||
.remove(&component.definition.key);
|
||||
}
|
||||
if !terminal {
|
||||
session.registry.lock().await.stopping = false;
|
||||
}
|
||||
first.map_or(Ok(()), Err)
|
||||
}
|
||||
|
||||
impl Managed {
|
||||
async fn ensure_started(&self) -> Result<mpsc::Sender<Request>, String> {
|
||||
let mut state = self.state.lock().await;
|
||||
if state.stopping {
|
||||
return Err("external Adapter worker is stopped".into());
|
||||
}
|
||||
if let Some(sender) = &state.sender
|
||||
&& !sender.is_closed()
|
||||
{
|
||||
return Ok(sender.clone());
|
||||
}
|
||||
if state.task.is_some() {
|
||||
let _ = crate::task::wait_resource(&mut state.task).await;
|
||||
}
|
||||
state.sender = None;
|
||||
state.stop = None;
|
||||
if self.definition.ecosystem != "command" {
|
||||
let now = Instant::now();
|
||||
while state
|
||||
.starts
|
||||
.front()
|
||||
.is_some_and(|started| now.duration_since(*started) >= START_WINDOW)
|
||||
{
|
||||
state.starts.pop_front();
|
||||
}
|
||||
if state.starts.len() == START_LIMIT {
|
||||
return Err("external Adapter worker restart limit exceeded".into());
|
||||
}
|
||||
state.starts.push_back(now);
|
||||
}
|
||||
let (sender, receiver) = mpsc::channel(COMMAND_CAPACITY);
|
||||
let (stop, stopping) = oneshot::channel();
|
||||
let (ready, started) = oneshot::channel();
|
||||
let definition = self.definition.clone();
|
||||
let session = session()?;
|
||||
let task = crate::task::run_in(&session.scope, async move {
|
||||
supervise(definition, receiver, stopping, ready).await
|
||||
})?;
|
||||
state.stop = Some(stop);
|
||||
state.task = Some(task);
|
||||
match started.await {
|
||||
Ok(Ok(())) => {
|
||||
state.sender = Some(sender.clone());
|
||||
Ok(sender)
|
||||
}
|
||||
Ok(Err(error)) => {
|
||||
let cleanup = crate::task::wait_resource(&mut state.task).await;
|
||||
state.stop = None;
|
||||
match cleanup {
|
||||
Err(cleanup) if cleanup != error => {
|
||||
Err(format!("{error}; worker cleanup: {cleanup}"))
|
||||
}
|
||||
_ => Err(error),
|
||||
}
|
||||
}
|
||||
Err(_) => crate::task::wait_resource(&mut state.task)
|
||||
.await
|
||||
.and_then(|_| Err("external Adapter worker stopped during startup".into())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn stop(&self) -> Result<(), String> {
|
||||
let mut state = self.state.lock().await;
|
||||
state.stopping = true;
|
||||
state.sender = None;
|
||||
if let Some(stop) = state.stop.take() {
|
||||
let _ = stop.send(());
|
||||
}
|
||||
// Keep the task in its owner if an explicit shutdown caller is cancelled.
|
||||
crate::task::wait_resource(&mut state.task).await
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_definition(definition: &Definition) -> Result<(), String> {
|
||||
if !matches!(
|
||||
definition.ecosystem.as_str(),
|
||||
"exec" | "pip" | "npm" | "go" | "command"
|
||||
) {
|
||||
return Err("invalid external Adapter ecosystem".into());
|
||||
}
|
||||
if definition.key.is_empty()
|
||||
|| definition.port.is_empty()
|
||||
|| definition.adapter.is_empty()
|
||||
|| definition.schema.is_empty()
|
||||
|| definition.operations.is_empty()
|
||||
|| definition.timeout_ms == 0
|
||||
{
|
||||
return Err("incomplete external Adapter definition".into());
|
||||
}
|
||||
unique(&definition.capabilities, "capability")?;
|
||||
if definition
|
||||
.capabilities
|
||||
.iter()
|
||||
.any(|capability| !matches!(capability.as_str(), "network" | "file" | "process" | "gpu"))
|
||||
{
|
||||
return Err("external Adapter declares an unknown capability".into());
|
||||
}
|
||||
unique(&definition.operations, "operation")?;
|
||||
if definition.ecosystem == "command"
|
||||
&& (definition.operations.len() != 1 || definition.setting.is_some())
|
||||
{
|
||||
return Err("command Adapter requires one operation and no setting".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn unique(values: &[String], label: &str) -> Result<(), String> {
|
||||
let mut found = BTreeSet::new();
|
||||
if values
|
||||
.iter()
|
||||
.any(|value| value.is_empty() || !found.insert(value))
|
||||
{
|
||||
return Err(format!(
|
||||
"external Adapter has an invalid or duplicate {label}"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn supervise(
|
||||
definition: Definition,
|
||||
receiver: mpsc::Receiver<Request>,
|
||||
stopping: oneshot::Receiver<()>,
|
||||
ready: oneshot::Sender<Result<(), String>>,
|
||||
) -> Result<(), String> {
|
||||
if definition.ecosystem == "command" {
|
||||
return command::supervise(definition, receiver, stopping, ready).await;
|
||||
}
|
||||
let mut command = worker_command(&definition)?;
|
||||
command
|
||||
.arg("--dever-component")
|
||||
.env_clear()
|
||||
.stdin(std::process::Stdio::piped())
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
let mut child = command
|
||||
.spawn()
|
||||
.map_err(|error| format!("cannot start external Adapter worker: {error}"))?;
|
||||
let Some(stdin) = child.stdin.take() else {
|
||||
terminate_child(&mut child).await?;
|
||||
return Err("external Adapter worker has no stdin".into());
|
||||
};
|
||||
let Some(stdout) = child.stdout.take() else {
|
||||
terminate_child(&mut child).await?;
|
||||
return Err("external Adapter worker has no stdout".into());
|
||||
};
|
||||
let Some(stderr) = child.stderr.take() else {
|
||||
terminate_child(&mut child).await?;
|
||||
return Err("external Adapter worker has no stderr".into());
|
||||
};
|
||||
let protocol = protocol_loop(
|
||||
definition, &mut child, stdin, stdout, receiver, stopping, ready,
|
||||
);
|
||||
let logs = drain_stderr(stderr);
|
||||
let (protocol, logs) = tokio::join!(protocol, logs);
|
||||
protocol.and(logs)
|
||||
}
|
||||
|
||||
fn worker_command(definition: &Definition) -> Result<Command, String> {
|
||||
#[cfg(feature = "external")]
|
||||
{
|
||||
crate::external::worker_command(definition).map(Command::from)
|
||||
}
|
||||
#[cfg(not(feature = "external"))]
|
||||
{
|
||||
let _ = definition;
|
||||
Err("external Worker sandbox runtime feature is not enabled".into())
|
||||
}
|
||||
}
|
||||
|
||||
async fn protocol_loop(
|
||||
definition: Definition,
|
||||
child: &mut Child,
|
||||
mut input: ChildStdin,
|
||||
mut output: ChildStdout,
|
||||
mut requests: mpsc::Receiver<Request>,
|
||||
mut stopping: oneshot::Receiver<()>,
|
||||
ready: oneshot::Sender<Result<(), String>>,
|
||||
) -> Result<(), String> {
|
||||
let protocol = protocol_session(
|
||||
&definition,
|
||||
&mut input,
|
||||
&mut output,
|
||||
&mut requests,
|
||||
&mut stopping,
|
||||
ready,
|
||||
)
|
||||
.await;
|
||||
drop(input);
|
||||
drop(output);
|
||||
match protocol {
|
||||
Ok(()) => wait_for_exit(child).await,
|
||||
Err(error) => match terminate_child(child).await {
|
||||
Ok(()) => Err(error),
|
||||
Err(cleanup) => Err(format!("{error}; worker cleanup failed: {cleanup}")),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn protocol_session(
|
||||
definition: &Definition,
|
||||
input: &mut ChildStdin,
|
||||
output: &mut ChildStdout,
|
||||
requests: &mut mpsc::Receiver<Request>,
|
||||
stopping: &mut oneshot::Receiver<()>,
|
||||
ready: oneshot::Sender<Result<(), String>>,
|
||||
) -> Result<(), String> {
|
||||
let startup = handshake(definition, input, output).await;
|
||||
if let Err(error) = startup {
|
||||
let _ = ready.send(Err(error.clone()));
|
||||
return Err(error);
|
||||
}
|
||||
let _ = ready.send(Ok(()));
|
||||
let mut next_id = 1_u64;
|
||||
loop {
|
||||
let request = tokio::select! {
|
||||
biased;
|
||||
_ = &mut *stopping => return shutdown_protocol(input, output).await,
|
||||
request = requests.recv() => request,
|
||||
};
|
||||
let Some(mut request) = request else {
|
||||
return shutdown_protocol(input, output).await;
|
||||
};
|
||||
let id = next_id;
|
||||
next_id = next_id
|
||||
.checked_add(1)
|
||||
.ok_or_else(|| "external Adapter request id exhausted".to_owned())?;
|
||||
if request.response.is_closed() {
|
||||
continue;
|
||||
}
|
||||
let message = call_message(id, &request.operation, &request.payload)?;
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = &mut *stopping => {
|
||||
let _ = request.response.send(Err("external Adapter worker stopped".into()));
|
||||
return Err("external Adapter worker stopped while writing a call".into());
|
||||
}
|
||||
result = write_frame_timeout(input, &message, Duration::from_millis(definition.timeout_ms)) => result?,
|
||||
}
|
||||
let timeout = tokio::time::sleep(Duration::from_millis(definition.timeout_ms));
|
||||
tokio::pin!(timeout);
|
||||
let result = tokio::select! {
|
||||
biased;
|
||||
_ = &mut *stopping => {
|
||||
let _ = request.response.send(Err("external Adapter worker stopped".into()));
|
||||
cancel_request(input, output, id).await?;
|
||||
return shutdown_protocol(input, output).await;
|
||||
}
|
||||
frame = read_frame(output) => frame.and_then(|frame| parse_reply(frame, id)),
|
||||
_ = request.response.closed() => {
|
||||
cancel_request(input, output, id).await?;
|
||||
continue;
|
||||
}
|
||||
_ = &mut timeout => {
|
||||
let _ = request.response.send(Err("external Adapter call timed out".into()));
|
||||
cancel_request(input, output, id).await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
match result {
|
||||
Ok(reply) => {
|
||||
let _ = request.response.send(Ok(reply));
|
||||
}
|
||||
Err(error) => {
|
||||
requests.close();
|
||||
let _ = request.response.send(Err(error.clone()));
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn cancel_request(
|
||||
input: &mut ChildStdin,
|
||||
output: &mut ChildStdout,
|
||||
id: u64,
|
||||
) -> Result<(), String> {
|
||||
write_frame_timeout(input, &id_message("cancel", id)?, SHUTDOWN_GRACE).await?;
|
||||
let _ = parse_reply(read_frame_timeout(output, SHUTDOWN_GRACE).await?, id)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn shutdown_protocol(input: &mut ChildStdin, output: &mut ChildStdout) -> Result<(), String> {
|
||||
write_frame_timeout(input, &kind_message("shutdown")?, SHUTDOWN_GRACE).await?;
|
||||
let shutdown = read_frame_timeout(output, SHUTDOWN_GRACE).await?;
|
||||
expect_kind(&shutdown, "shutdown")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn wait_for_exit(child: &mut Child) -> Result<(), String> {
|
||||
match tokio::time::timeout(SHUTDOWN_GRACE, child.wait()).await {
|
||||
Ok(Ok(status)) if status.success() => Ok(()),
|
||||
Ok(Ok(status)) => Err(format!("external Adapter worker exited with {status}")),
|
||||
Ok(Err(error)) => Err(format!("cannot wait for external Adapter worker: {error}")),
|
||||
Err(_) => {
|
||||
terminate_child(child).await?;
|
||||
Err("external Adapter worker exceeded shutdown deadline".into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn terminate_child(child: &mut Child) -> Result<(), String> {
|
||||
match child.try_wait() {
|
||||
Ok(Some(_)) => return Ok(()),
|
||||
Ok(None) => {}
|
||||
Err(error) => return Err(format!("cannot inspect external Adapter worker: {error}")),
|
||||
}
|
||||
if let Err(error) = child.start_kill()
|
||||
&& !matches!(child.try_wait(), Ok(Some(_)))
|
||||
{
|
||||
return Err(format!("cannot stop external Adapter worker: {error}"));
|
||||
}
|
||||
match tokio::time::timeout(SHUTDOWN_GRACE, child.wait()).await {
|
||||
Ok(Ok(_)) => Ok(()),
|
||||
Ok(Err(error)) => Err(format!("cannot wait for external Adapter worker: {error}")),
|
||||
Err(_) => Err("external Adapter worker could not be reaped before its deadline".into()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handshake(
|
||||
definition: &Definition,
|
||||
input: &mut ChildStdin,
|
||||
output: &mut ChildStdout,
|
||||
) -> Result<(), String> {
|
||||
let timeout = Duration::from_millis(definition.timeout_ms);
|
||||
write_frame_timeout(input, &hello_message(definition)?, timeout).await?;
|
||||
let ready = read_frame_timeout(output, timeout).await?;
|
||||
validate_ready(definition, &ready)?;
|
||||
write_frame_timeout(input, &id_message("health", 0)?, timeout).await?;
|
||||
let health = read_frame_timeout(output, timeout).await?;
|
||||
let node = crate::wire::parse(&health)?;
|
||||
let fields = node.fields(&["kind", "id"])?;
|
||||
if text_field(fields, "kind")? != "health" || int_field(fields, "id")? != 0 {
|
||||
return Err("external Adapter returned an invalid health response".into());
|
||||
}
|
||||
require_fields(fields, &["kind", "id"])
|
||||
}
|
||||
|
||||
async fn drain_stderr(mut stderr: ChildStderr) -> Result<(), String> {
|
||||
let mut buffer = [0_u8; 8192];
|
||||
loop {
|
||||
match stderr.read(&mut buffer).await {
|
||||
Ok(0) => return Ok(()),
|
||||
Ok(_) => {}
|
||||
Err(error) => return Err(format!("cannot read external Adapter stderr: {error}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_frame(output: &mut (impl AsyncWrite + Unpin), body: &str) -> Result<(), String> {
|
||||
if body.is_empty() || body.len() > crate::wire::MAX_BYTES {
|
||||
return Err("component protocol frame exceeds byte limit".into());
|
||||
}
|
||||
let length =
|
||||
u32::try_from(body.len()).map_err(|_| "component protocol frame exceeds byte limit")?;
|
||||
output
|
||||
.write_all(&length.to_be_bytes())
|
||||
.await
|
||||
.map_err(|error| format!("cannot write component protocol frame header: {error}"))?;
|
||||
output
|
||||
.write_all(body.as_bytes())
|
||||
.await
|
||||
.map_err(|error| format!("cannot write component protocol frame body: {error}"))?;
|
||||
output
|
||||
.flush()
|
||||
.await
|
||||
.map_err(|error| format!("cannot flush component protocol frame: {error}"))
|
||||
}
|
||||
|
||||
async fn write_frame_timeout(
|
||||
output: &mut (impl AsyncWrite + Unpin),
|
||||
body: &str,
|
||||
timeout: Duration,
|
||||
) -> Result<(), String> {
|
||||
tokio::time::timeout(timeout, write_frame(output, body))
|
||||
.await
|
||||
.map_err(|_| "component protocol write timed out".to_owned())?
|
||||
}
|
||||
|
||||
async fn read_frame(input: &mut (impl AsyncRead + Unpin)) -> Result<String, String> {
|
||||
let mut length = [0_u8; 4];
|
||||
input
|
||||
.read_exact(&mut length)
|
||||
.await
|
||||
.map_err(|error| format!("cannot read component protocol frame header: {error}"))?;
|
||||
let length = usize::try_from(u32::from_be_bytes(length)).expect("u32 fits usize");
|
||||
if length == 0 || length > crate::wire::MAX_BYTES {
|
||||
return Err("component protocol frame exceeds byte limit".into());
|
||||
}
|
||||
let mut body = vec![0_u8; length];
|
||||
input
|
||||
.read_exact(&mut body)
|
||||
.await
|
||||
.map_err(|error| format!("cannot read complete component protocol frame: {error}"))?;
|
||||
String::from_utf8(body).map_err(|_| "component protocol frame is not UTF-8".into())
|
||||
}
|
||||
|
||||
async fn read_frame_timeout(
|
||||
input: &mut (impl AsyncRead + Unpin),
|
||||
timeout: Duration,
|
||||
) -> Result<String, String> {
|
||||
tokio::time::timeout(timeout, read_frame(input))
|
||||
.await
|
||||
.map_err(|_| "component protocol response timed out".to_owned())?
|
||||
}
|
||||
|
||||
fn hello_message(definition: &Definition) -> Result<String, String> {
|
||||
let mut writer = crate::wire::Encoder::default();
|
||||
writer.begin_object()?;
|
||||
field_text(&mut writer, "kind", "hello")?;
|
||||
field_text(&mut writer, "version", VERSION)?;
|
||||
field_text(&mut writer, "port", &definition.port)?;
|
||||
field_text(&mut writer, "schema", &definition.schema)?;
|
||||
field_text(&mut writer, "adapter", &definition.adapter)?;
|
||||
field_text_list(&mut writer, "capabilities", &definition.capabilities)?;
|
||||
field_text_list(&mut writer, "operations", &definition.operations)?;
|
||||
writer.key("setting")?;
|
||||
match &definition.setting {
|
||||
Some(setting) => writer.json(setting)?,
|
||||
None => writer.null()?,
|
||||
}
|
||||
writer.end()?;
|
||||
Ok(writer.finish()?.as_str().to_owned())
|
||||
}
|
||||
|
||||
fn call_message(id: u64, operation: &str, payload: &str) -> Result<String, String> {
|
||||
let mut writer = crate::wire::Encoder::default();
|
||||
writer.begin_object()?;
|
||||
field_text(&mut writer, "kind", "call")?;
|
||||
field_int(&mut writer, "id", id)?;
|
||||
field_text(&mut writer, "operation", operation)?;
|
||||
writer.key("payload")?;
|
||||
writer.json(payload)?;
|
||||
writer.end()?;
|
||||
Ok(writer.finish()?.as_str().to_owned())
|
||||
}
|
||||
|
||||
fn id_message(kind: &str, id: u64) -> Result<String, String> {
|
||||
let mut writer = crate::wire::Encoder::default();
|
||||
writer.begin_object()?;
|
||||
field_text(&mut writer, "kind", kind)?;
|
||||
field_int(&mut writer, "id", id)?;
|
||||
writer.end()?;
|
||||
Ok(writer.finish()?.as_str().to_owned())
|
||||
}
|
||||
|
||||
fn kind_message(kind: &str) -> Result<String, String> {
|
||||
let mut writer = crate::wire::Encoder::default();
|
||||
writer.begin_object()?;
|
||||
field_text(&mut writer, "kind", kind)?;
|
||||
writer.end()?;
|
||||
Ok(writer.finish()?.as_str().to_owned())
|
||||
}
|
||||
|
||||
fn field_text(writer: &mut crate::wire::Encoder, name: &str, value: &str) -> Result<(), String> {
|
||||
writer.key(name)?;
|
||||
writer.text(value)
|
||||
}
|
||||
|
||||
fn field_int(writer: &mut crate::wire::Encoder, name: &str, value: u64) -> Result<(), String> {
|
||||
let value = i64::try_from(value).map_err(|_| "component protocol request id exhausted")?;
|
||||
writer.key(name)?;
|
||||
writer.int(value)
|
||||
}
|
||||
|
||||
fn field_text_list(
|
||||
writer: &mut crate::wire::Encoder,
|
||||
name: &str,
|
||||
values: &[String],
|
||||
) -> Result<(), String> {
|
||||
writer.key(name)?;
|
||||
writer.begin_array()?;
|
||||
for value in values {
|
||||
writer.text(value)?;
|
||||
}
|
||||
writer.end()
|
||||
}
|
||||
|
||||
fn validate_ready(definition: &Definition, body: &str) -> Result<(), String> {
|
||||
let node = crate::wire::parse(body)?;
|
||||
let names = [
|
||||
"kind",
|
||||
"version",
|
||||
"port",
|
||||
"schema",
|
||||
"adapter",
|
||||
"capabilities",
|
||||
"operations",
|
||||
];
|
||||
let fields = node.fields(&names)?;
|
||||
require_fields(fields, &names)?;
|
||||
if text_field(fields, "kind")? != "ready"
|
||||
|| text_field(fields, "version")? != VERSION
|
||||
|| text_field(fields, "port")? != definition.port
|
||||
|| text_field(fields, "schema")? != definition.schema
|
||||
|| text_field(fields, "adapter")? != definition.adapter
|
||||
|| text_list_field(fields, "capabilities")? != definition.capabilities
|
||||
|| text_list_field(fields, "operations")? != definition.operations
|
||||
{
|
||||
return Err("external Adapter handshake does not match its compiled Port schema".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_reply(body: String, expected_id: u64) -> Result<Reply, String> {
|
||||
let node = crate::wire::parse(&body)?;
|
||||
let fields = node.object()?;
|
||||
let kind = text_field(fields, "kind")?;
|
||||
let id = int_field(fields, "id")?;
|
||||
if id != expected_id {
|
||||
return Err("component protocol response has an unknown or out-of-order request id".into());
|
||||
}
|
||||
match kind {
|
||||
"result" => {
|
||||
require_fields(fields, &["kind", "id", "payload"])?;
|
||||
Ok(Reply::Result(raw_field(fields, "payload")?.to_owned()))
|
||||
}
|
||||
"error" => {
|
||||
require_fields(fields, &["kind", "id", "error", "payload"])?;
|
||||
let identity = text_field(fields, "error")?;
|
||||
if identity.is_empty() {
|
||||
return Err("component protocol error identity is empty".into());
|
||||
}
|
||||
Ok(Reply::Error {
|
||||
identity: identity.to_owned(),
|
||||
payload: raw_field(fields, "payload")?.to_owned(),
|
||||
})
|
||||
}
|
||||
_ => Err("component protocol returned an unknown message kind".into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn expect_kind(body: &str, expected: &str) -> Result<(), String> {
|
||||
let node = crate::wire::parse(body)?;
|
||||
let fields = node.fields(&["kind"])?;
|
||||
require_fields(fields, &["kind"])?;
|
||||
if text_field(fields, "kind")? == expected {
|
||||
Ok(())
|
||||
} else {
|
||||
Err("component protocol returned an unexpected message kind".into())
|
||||
}
|
||||
}
|
||||
|
||||
fn require_fields(
|
||||
fields: &BTreeMap<String, crate::wire::Node<'_>>,
|
||||
names: &[&str],
|
||||
) -> Result<(), String> {
|
||||
if fields.len() != names.len() || names.iter().any(|name| !fields.contains_key(*name)) {
|
||||
return Err("component protocol message has missing or extra fields".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn text_field<'a>(
|
||||
fields: &'a BTreeMap<String, crate::wire::Node<'a>>,
|
||||
name: &str,
|
||||
) -> Result<&'a str, String> {
|
||||
fields
|
||||
.get(name)
|
||||
.ok_or_else(|| "component protocol message is missing a field".to_owned())?
|
||||
.text()
|
||||
}
|
||||
|
||||
fn int_field(fields: &BTreeMap<String, crate::wire::Node<'_>>, name: &str) -> Result<u64, String> {
|
||||
let value = fields
|
||||
.get(name)
|
||||
.ok_or_else(|| "component protocol message is missing a field".to_owned())?
|
||||
.int()?;
|
||||
u64::try_from(value).map_err(|_| "component protocol request id must be non-negative".into())
|
||||
}
|
||||
|
||||
fn raw_field<'a>(
|
||||
fields: &'a BTreeMap<String, crate::wire::Node<'a>>,
|
||||
name: &str,
|
||||
) -> Result<&'a str, String> {
|
||||
fields
|
||||
.get(name)
|
||||
.map(crate::wire::Node::raw)
|
||||
.ok_or_else(|| "component protocol message is missing a field".to_owned())
|
||||
}
|
||||
|
||||
fn text_list_field<'a>(
|
||||
fields: &'a BTreeMap<String, crate::wire::Node<'a>>,
|
||||
name: &str,
|
||||
) -> Result<Vec<&'a str>, String> {
|
||||
let values = fields
|
||||
.get(name)
|
||||
.ok_or_else(|| "component protocol message is missing a field".to_owned())?
|
||||
.list()?;
|
||||
let mut result = Vec::with_capacity(values.len());
|
||||
let mut unique = BTreeSet::new();
|
||||
for value in values {
|
||||
let value = value.text()?;
|
||||
if !unique.insert(value) {
|
||||
return Err("component protocol list contains a duplicate value".into());
|
||||
}
|
||||
result.push(value);
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
@@ -0,0 +1,342 @@
|
||||
//! One ordinary child per request, owned by the existing Adapter supervisor.
|
||||
use super::{Definition, Reply, Request, adapter_settings, terminate_child, worker_command};
|
||||
use crate::bytes::Bytes;
|
||||
use crate::config::CommandLimits;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
mod observation;
|
||||
use observation::Observation;
|
||||
|
||||
#[cfg(feature = "external")]
|
||||
const ARGUMENT_BYTES: usize = 65_536;
|
||||
#[cfg(feature = "external")]
|
||||
const STDIN_BYTES: usize = 8 * 1024 * 1024;
|
||||
|
||||
type CompletionSlot = std::sync::Arc<std::sync::Mutex<Option<crate::task::ScopedCompletion>>>;
|
||||
|
||||
pub(super) struct CallLease(CompletionSlot);
|
||||
pub(super) struct RequestLease(CompletionSlot);
|
||||
|
||||
pub(super) fn call_lease() -> Result<CallLease, String> {
|
||||
Ok(CallLease(std::sync::Arc::new(std::sync::Mutex::new(Some(
|
||||
crate::task::scoped_completion()?,
|
||||
)))))
|
||||
}
|
||||
|
||||
impl CallLease {
|
||||
pub(super) fn request(&self) -> RequestLease {
|
||||
RequestLease(self.0.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CallLease {
|
||||
fn drop(&mut self) {
|
||||
// Queued calls have no process to drain. An executing supervisor has
|
||||
// already taken the guard and retains it until the child is reaped.
|
||||
self.0
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.take();
|
||||
}
|
||||
}
|
||||
|
||||
impl RequestLease {
|
||||
fn begin(self) -> Option<crate::task::ScopedCompletion> {
|
||||
self.0
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.take()
|
||||
}
|
||||
}
|
||||
|
||||
struct Input {
|
||||
#[cfg(feature = "external")]
|
||||
arguments: Vec<String>,
|
||||
stdin: Option<Bytes>,
|
||||
}
|
||||
|
||||
struct Execution {
|
||||
reply: Result<Reply, String>,
|
||||
stopped: bool,
|
||||
cleanup_error: Option<String>,
|
||||
}
|
||||
|
||||
pub(super) async fn supervise(
|
||||
definition: Definition,
|
||||
mut requests: mpsc::Receiver<Request>,
|
||||
mut stopping: oneshot::Receiver<()>,
|
||||
ready: oneshot::Sender<Result<(), String>>,
|
||||
) -> Result<(), String> {
|
||||
let prepared: Result<CommandLimits, String> = (|| {
|
||||
let limits = adapter_settings()?.command_limits(&definition.port)?;
|
||||
// Validate resources and grants without starting the user's executable.
|
||||
worker_command(&definition)?;
|
||||
Ok(limits)
|
||||
})();
|
||||
let limits = match prepared {
|
||||
Ok(limits) => limits,
|
||||
Err(error) => {
|
||||
let _ = ready.send(Err(error.clone()));
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let _ = ready.send(Ok(()));
|
||||
loop {
|
||||
let request = tokio::select! {
|
||||
biased;
|
||||
_ = &mut stopping => break,
|
||||
request = requests.recv() => request,
|
||||
};
|
||||
let Some(mut request) = request else { break };
|
||||
if request.response.is_closed() {
|
||||
continue;
|
||||
}
|
||||
let Some(completion) = request.command_lease.take().and_then(RequestLease::begin) else {
|
||||
continue;
|
||||
};
|
||||
let execution = execute(&definition, limits, &mut request, &mut stopping).await;
|
||||
// execute has dropped I/O futures and reaped the child before replying.
|
||||
drop(completion);
|
||||
let _ = request.response.send(execution.reply);
|
||||
if let Some(error) = execution.cleanup_error {
|
||||
return Err(error);
|
||||
}
|
||||
if execution.stopped {
|
||||
break;
|
||||
}
|
||||
}
|
||||
requests.close();
|
||||
while let Some(request) = requests.recv().await {
|
||||
let _ = request
|
||||
.response
|
||||
.send(Err("external command Adapter stopped".into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(feature = "external")]
|
||||
fn input(payload: &str) -> Result<Input, String> {
|
||||
let node = crate::wire::parse(payload)?;
|
||||
let fields = node.fields(&["args", "stdin"])?;
|
||||
let arguments = fields
|
||||
.get("args")
|
||||
.ok_or("command args are missing")?
|
||||
.list()?;
|
||||
let mut bytes = 0_usize;
|
||||
let arguments = arguments
|
||||
.iter()
|
||||
.map(|value| {
|
||||
let value = value.text()?;
|
||||
if value.contains('\0') {
|
||||
return Err("command argument contains NUL".into());
|
||||
}
|
||||
// Include terminators so empty arguments also consume the budget.
|
||||
bytes = bytes.saturating_add(value.len()).saturating_add(1);
|
||||
if bytes > ARGUMENT_BYTES {
|
||||
return Err("command arguments exceed the 65536 byte limit".into());
|
||||
}
|
||||
Ok(value.to_owned())
|
||||
})
|
||||
.collect::<Result<Vec<_>, String>>()?;
|
||||
let stdin = fields
|
||||
.get("stdin")
|
||||
.filter(|value| !value.is_null())
|
||||
.map(|value| value.bytes())
|
||||
.transpose()?;
|
||||
if stdin
|
||||
.as_ref()
|
||||
.is_some_and(|value| value.values().len() > STDIN_BYTES)
|
||||
{
|
||||
return Err("command stdin exceeds the 8388608 byte limit".into());
|
||||
}
|
||||
Ok(Input { arguments, stdin })
|
||||
}
|
||||
|
||||
#[cfg(feature = "external")]
|
||||
fn spawn(definition: &Definition, payload: &str) -> Result<(Input, Child, Observation), String> {
|
||||
let input = input(payload)?;
|
||||
let (command, observation) = crate::external::observed_command(definition)?;
|
||||
let observation = Observation::new(observation)?;
|
||||
let mut command = tokio::process::Command::from(command);
|
||||
// Resource arguments were mapped by the sandbox owner. Call arguments
|
||||
// are opaque values, including absolute sandbox paths and shell syntax.
|
||||
command
|
||||
.args(&input.arguments)
|
||||
.env_clear()
|
||||
.stdin(std::process::Stdio::piped())
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
let child = command
|
||||
.spawn()
|
||||
.map_err(|error| format!("cannot start external command: {error}"))?;
|
||||
// Dropping Command closes the parent's copies of mapped pipe descriptors.
|
||||
drop(command);
|
||||
Ok((input, child, observation))
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "external"))]
|
||||
fn spawn(_: &Definition, _: &str) -> Result<(Input, Child, Observation), String> {
|
||||
Err("external command sandbox runtime feature is not enabled".into())
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
definition: &Definition,
|
||||
limits: CommandLimits,
|
||||
request: &mut Request,
|
||||
stopping: &mut oneshot::Receiver<()>,
|
||||
) -> Execution {
|
||||
let (input, mut child, mut observation) = match spawn(definition, &request.payload) {
|
||||
Ok(prepared) => prepared,
|
||||
Err(error) => {
|
||||
return Execution {
|
||||
reply: Err(error),
|
||||
stopped: false,
|
||||
cleanup_error: None,
|
||||
};
|
||||
}
|
||||
};
|
||||
let wrapper = child
|
||||
.id()
|
||||
.expect("newly spawned command launcher has a PID");
|
||||
let pipes = (child.stdin.take(), child.stdout.take(), child.stderr.take());
|
||||
let (Some(stdin), Some(stdout), Some(stderr)) = pipes else {
|
||||
let cleanup_error = cleanup(&mut child, &mut observation, wrapper).await.err();
|
||||
return Execution {
|
||||
reply: Err(cleanup_error
|
||||
.clone()
|
||||
.unwrap_or_else(|| "external command has incomplete pipes".into())),
|
||||
stopped: false,
|
||||
cleanup_error,
|
||||
};
|
||||
};
|
||||
let mut stopped = false;
|
||||
let result = {
|
||||
let completion = async {
|
||||
let (_, _, (stdout, stderr), status) = tokio::try_join!(
|
||||
observation.started(wrapper),
|
||||
write_stdin(stdin, input.stdin),
|
||||
capture(stdout, stderr, limits.output_limit),
|
||||
async {
|
||||
child
|
||||
.wait()
|
||||
.await
|
||||
.map_err(|error| format!("cannot wait for external command: {error}"))
|
||||
},
|
||||
)?;
|
||||
let code = status
|
||||
.code()
|
||||
.ok_or("external command sandbox terminated without an exit code")?;
|
||||
output(code, stdout, stderr)
|
||||
};
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = stopping => {
|
||||
stopped = true;
|
||||
Err("external command Adapter stopped".into())
|
||||
}
|
||||
_ = request.response.closed() => Err("external command call cancelled".into()),
|
||||
result = tokio::time::timeout(std::time::Duration::from_millis(limits.timeout_ms), completion) => {
|
||||
result.unwrap_or_else(|_| Err("external command call timed out".into()))
|
||||
}
|
||||
}
|
||||
};
|
||||
let cleanup_error = cleanup(&mut child, &mut observation, wrapper).await.err();
|
||||
let reply = match (result, &cleanup_error) {
|
||||
(result, None) => result,
|
||||
(Ok(_), Some(cleanup)) => Err(cleanup.clone()),
|
||||
(Err(error), Some(cleanup)) => Err(format!("{error}; command cleanup failed: {cleanup}")),
|
||||
};
|
||||
Execution {
|
||||
reply,
|
||||
stopped,
|
||||
cleanup_error,
|
||||
}
|
||||
}
|
||||
|
||||
async fn cleanup(
|
||||
child: &mut Child,
|
||||
observation: &mut Observation,
|
||||
wrapper: u32,
|
||||
) -> Result<(), String> {
|
||||
let namespace = tokio::time::timeout(super::SHUTDOWN_GRACE, observation.terminate(wrapper))
|
||||
.await
|
||||
.unwrap_or_else(|_| Err("external command namespace exceeded its cleanup deadline".into()));
|
||||
// Reap the owned launcher even if namespace observation itself failed.
|
||||
// Both stages have independent finite deadlines; neither error is lost.
|
||||
let launcher = terminate_child(child).await;
|
||||
match (namespace, launcher) {
|
||||
(Err(namespace), Err(launcher)) => {
|
||||
Err(format!("{namespace}; launcher cleanup failed: {launcher}"))
|
||||
}
|
||||
(namespace, launcher) => namespace.and(launcher),
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_stdin(mut stdin: ChildStdin, input: Option<Bytes>) -> Result<(), String> {
|
||||
if let Some(input) = input {
|
||||
match stdin.write_all(input.values()).await {
|
||||
Ok(()) => {}
|
||||
// Ordinary programs may intentionally stop reading before EOF.
|
||||
Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => return Ok(()),
|
||||
Err(error) => return Err(format!("cannot write external command stdin: {error}")),
|
||||
}
|
||||
}
|
||||
drop(stdin);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn capture(
|
||||
mut stdout: ChildStdout,
|
||||
mut stderr: ChildStderr,
|
||||
limit: usize,
|
||||
) -> Result<(Vec<u8>, Vec<u8>), String> {
|
||||
let mut output = (Vec::new(), Vec::new());
|
||||
let mut stdout_buffer = [0; 8192];
|
||||
let mut stderr_buffer = [0; 8192];
|
||||
let mut stdout_open = true;
|
||||
let mut stderr_open = true;
|
||||
while stdout_open || stderr_open {
|
||||
let (read, is_stdout) = tokio::select! {
|
||||
read = stdout.read(&mut stdout_buffer), if stdout_open => (read, true),
|
||||
read = stderr.read(&mut stderr_buffer), if stderr_open => (read, false),
|
||||
};
|
||||
let count =
|
||||
read.map_err(|error| format!("cannot read external command output: {error}"))?;
|
||||
if count == 0 {
|
||||
if is_stdout {
|
||||
stdout_open = false;
|
||||
} else {
|
||||
stderr_open = false;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if count > limit.saturating_sub(output.0.len() + output.1.len()) {
|
||||
return Err("external command output exceeds output_limit".into());
|
||||
}
|
||||
if is_stdout {
|
||||
output.0.extend_from_slice(&stdout_buffer[..count]);
|
||||
} else {
|
||||
output.1.extend_from_slice(&stderr_buffer[..count]);
|
||||
}
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn output(code: i32, stdout: Vec<u8>, stderr: Vec<u8>) -> Result<Reply, String> {
|
||||
let mut encoder = crate::wire::Encoder::default();
|
||||
encoder.begin_object()?;
|
||||
encoder.key("output")?;
|
||||
encoder.begin_object()?;
|
||||
encoder.key("code")?;
|
||||
encoder.int(i64::from(code))?;
|
||||
encoder.key("stdout")?;
|
||||
encoder.bytes(&Bytes::new(stdout))?;
|
||||
encoder.key("stderr")?;
|
||||
encoder.bytes(&Bytes::new(stderr))?;
|
||||
encoder.end()?;
|
||||
encoder.end()?;
|
||||
Ok(Reply::Result(encoder.finish()?.as_str().to_owned()))
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
pub(super) struct Observation {
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
startup: tokio::net::unix::pipe::Receiver,
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
information: tokio::net::unix::pipe::Receiver,
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
information_bytes: Vec<u8>,
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
release: Option<std::fs::File>,
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
namespace: Option<tokio::io::unix::AsyncFd<dever_sandbox::CommandNamespace>>,
|
||||
}
|
||||
|
||||
impl Observation {
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
pub(super) fn new(observation: dever_sandbox::CommandObservation) -> Result<Self, String> {
|
||||
use tokio::net::unix::pipe::Receiver;
|
||||
Ok(Self {
|
||||
startup: Receiver::from_file(observation.startup).map_err(|error| error.to_string())?,
|
||||
information: Receiver::from_file(observation.namespace)
|
||||
.map_err(|error| error.to_string())?,
|
||||
information_bytes: Vec::new(),
|
||||
release: Some(observation.release),
|
||||
namespace: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "external", not(target_os = "linux")))]
|
||||
pub(super) fn new(_: dever_sandbox::CommandObservation) -> Result<Self, String> {
|
||||
Err("external commands require Linux".into())
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
async fn pin(&mut self, wrapper: u32) -> Result<(), String> {
|
||||
if self.namespace.is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
let mut buffer = [0; 1024];
|
||||
loop {
|
||||
let count = self
|
||||
.information
|
||||
.read(&mut buffer)
|
||||
.await
|
||||
.map_err(|error| format!("cannot read command namespace identity: {error}"))?;
|
||||
if count == 0 {
|
||||
break;
|
||||
}
|
||||
if self.information_bytes.len() + count > 4096 {
|
||||
return Err("command namespace identity exceeds its byte limit".into());
|
||||
}
|
||||
self.information_bytes.extend_from_slice(&buffer[..count]);
|
||||
}
|
||||
let text = std::str::from_utf8(&self.information_bytes)
|
||||
.map_err(|_| "command namespace identity is not UTF-8")?;
|
||||
let node = crate::wire::parse(text)?;
|
||||
let pid = node
|
||||
.object()?
|
||||
.get("child-pid")
|
||||
.ok_or("command namespace identity is missing")?
|
||||
.int()?;
|
||||
let namespace = dever_sandbox::CommandNamespace::pin(
|
||||
i32::try_from(pid).map_err(|_| "invalid command namespace PID")?,
|
||||
wrapper,
|
||||
)?;
|
||||
self.namespace = Some(
|
||||
tokio::io::unix::AsyncFd::new(namespace)
|
||||
.map_err(|error| format!("cannot observe command namespace exit: {error}"))?,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn started(&mut self, wrapper: u32) -> Result<(), String> {
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
{
|
||||
self.pin(wrapper).await?;
|
||||
// EOF releases bwrap only after its namespace is pinned. Neither
|
||||
// normal cancellation nor a very short process can race PID reuse.
|
||||
drop(self.release.take());
|
||||
let mut status = Vec::new();
|
||||
(&mut self.startup)
|
||||
.take(4096)
|
||||
.read_to_end(&mut status)
|
||||
.await
|
||||
.map_err(|error| format!("cannot read command startup status: {error}"))?;
|
||||
match status.as_slice() {
|
||||
b"R" => Ok(()),
|
||||
[] => Err("external command sandbox failed before starting its guard".into()),
|
||||
status => {
|
||||
let message = status
|
||||
.strip_prefix(b"RE")
|
||||
.or_else(|| status.strip_prefix(b"E"));
|
||||
Err(match message {
|
||||
Some(message) => format!(
|
||||
"external command startup failed: {}",
|
||||
String::from_utf8_lossy(message)
|
||||
),
|
||||
None => "invalid external command startup status".into(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(not(all(feature = "external", target_os = "linux")))]
|
||||
{
|
||||
let _ = wrapper;
|
||||
Err("external commands require Linux".into())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn terminate(&mut self, wrapper: u32) -> Result<(), String> {
|
||||
#[cfg(all(feature = "external", target_os = "linux"))]
|
||||
{
|
||||
self.pin(wrapper).await?;
|
||||
let namespace = self
|
||||
.namespace
|
||||
.as_ref()
|
||||
.expect("namespace pinned before cleanup");
|
||||
namespace.get_ref().terminate()?;
|
||||
// pidfd readiness follows kernel namespace teardown, unlike pipe
|
||||
// EOF, which occurs while the init is still closing its files.
|
||||
let _exited = namespace
|
||||
.readable()
|
||||
.await
|
||||
.map_err(|error| format!("cannot wait for command namespace exit: {error}"))?;
|
||||
Ok(())
|
||||
}
|
||||
#[cfg(not(all(feature = "external", target_os = "linux")))]
|
||||
{
|
||||
let _ = wrapper;
|
||||
Err("external commands require Linux".into())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Mutex, mpsc};
|
||||
|
||||
pub const MAX_WORKERS: usize = 256;
|
||||
const MAX_BATCH_ITEMS: usize = 256;
|
||||
const BATCHES_PER_WORKER: usize = 8;
|
||||
|
||||
pub fn worker_count(workers: i64) -> Result<usize, &'static str> {
|
||||
usize::try_from(workers)
|
||||
.ok()
|
||||
.filter(|count| (1..=MAX_WORKERS).contains(count))
|
||||
.ok_or("parallel worker count must be between 1 and 256")
|
||||
}
|
||||
|
||||
/// 单一生产者按需取值，有界队列防止 Stream 被提前收集；返回前回收全部工作线程。
|
||||
pub fn each<T: Send>(
|
||||
values: impl IntoIterator<Item = T>,
|
||||
workers: i64,
|
||||
apply: impl Fn(T) -> Result<(), String> + Sync,
|
||||
) -> Result<(), String> {
|
||||
each_typed(values, workers, apply)
|
||||
}
|
||||
|
||||
pub fn each_typed<T: Send, E: From<String> + Send>(
|
||||
values: impl IntoIterator<Item = T>,
|
||||
workers: i64,
|
||||
apply: impl Fn(T) -> Result<(), E> + Sync,
|
||||
) -> Result<(), E> {
|
||||
let workers = worker_count(workers).map_err(|error| E::from(error.to_owned()))?;
|
||||
run_typed(values, workers, |value, _| apply(value))
|
||||
}
|
||||
|
||||
/// 内存序列没有拉取副作用；限制每批元素数，同时给工作线程保留多批调度机会。
|
||||
pub fn each_slice<T: Sync>(
|
||||
values: &[T],
|
||||
workers: i64,
|
||||
apply: impl Fn(&T) -> Result<(), String> + Sync,
|
||||
) -> Result<(), String> {
|
||||
each_slice_typed(values, workers, apply)
|
||||
}
|
||||
|
||||
pub fn each_slice_typed<T: Sync, E: From<String> + Send>(
|
||||
values: &[T],
|
||||
workers: i64,
|
||||
apply: impl Fn(&T) -> Result<(), E> + Sync,
|
||||
) -> Result<(), E> {
|
||||
let workers = worker_count(workers).map_err(|error| E::from(error.to_owned()))?;
|
||||
let batch = batch_size(values.len(), workers);
|
||||
run_typed(values.chunks(batch), workers, |chunk, cancelled| {
|
||||
for value in chunk {
|
||||
if cancelled.load(Ordering::Acquire) {
|
||||
break;
|
||||
}
|
||||
apply(value)?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn batch_size(items: usize, workers: usize) -> usize {
|
||||
(items / (workers * BATCHES_PER_WORKER)).clamp(1, MAX_BATCH_ITEMS)
|
||||
}
|
||||
|
||||
fn run_typed<T: Send, E: From<String> + Send>(
|
||||
values: impl IntoIterator<Item = T>,
|
||||
workers: usize,
|
||||
apply: impl Fn(T, &AtomicBool) -> Result<(), E> + Sync,
|
||||
) -> Result<(), E> {
|
||||
let mut values = values.into_iter();
|
||||
let (minimum_items, maximum_items) = values.size_hint();
|
||||
let workers = match maximum_items {
|
||||
Some(maximum_items) if maximum_items == minimum_items => workers.min(maximum_items),
|
||||
_ => workers,
|
||||
};
|
||||
if workers == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let (sender, receiver) = mpsc::sync_channel(workers);
|
||||
let receiver = Mutex::new(receiver);
|
||||
let cancelled = AtomicBool::new(false);
|
||||
let failure = Mutex::new(None);
|
||||
let record_failure = |error: E| {
|
||||
cancelled.store(true, Ordering::Release);
|
||||
let mut failure = failure.lock().expect("worker failure lock");
|
||||
if failure.is_none() {
|
||||
*failure = Some(error);
|
||||
}
|
||||
};
|
||||
std::thread::scope(|scope| {
|
||||
let mut handles = Vec::new();
|
||||
for _ in 0..workers {
|
||||
let spawn = std::thread::Builder::new().spawn_scoped(scope, || {
|
||||
loop {
|
||||
let next = receiver.lock().expect("worker receiver lock").recv();
|
||||
let Ok(value) = next else { break };
|
||||
// 故障后继续清空队列直到发送端关闭，避免生产者卡在有界 send。
|
||||
if !cancelled.load(Ordering::Acquire) {
|
||||
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
||||
apply(value, &cancelled)
|
||||
})) {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(error)) => record_failure(error),
|
||||
Err(_) => record_failure(E::from("parallel worker panicked".into())),
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
match spawn {
|
||||
Ok(handle) => handles.push(handle),
|
||||
Err(error) => {
|
||||
record_failure(E::from(error.to_string()));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
while !cancelled.load(Ordering::Acquire) {
|
||||
let Some(value) = values.next() else { break };
|
||||
if sender.send(value).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
drop(sender);
|
||||
for handle in handles {
|
||||
if handle.join().is_err() {
|
||||
record_failure(E::from("parallel worker panicked".into()));
|
||||
}
|
||||
}
|
||||
});
|
||||
match failure.into_inner().expect("worker failure lock") {
|
||||
Some(error) => Err(error),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,50 @@
|
||||
use super::{ADAPTER_FIELDS, Settings};
|
||||
use serde::Deserialize;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct CommandLimits {
|
||||
pub timeout_ms: u64,
|
||||
pub output_limit: usize,
|
||||
}
|
||||
|
||||
impl Default for CommandLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
timeout_ms: 30_000,
|
||||
output_limit: 1024 * 1024,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
/// Deployment limits apply only to ordinary command Adapters.
|
||||
pub fn command_limits(&self, identity: &str) -> Result<CommandLimits, String> {
|
||||
let node = self
|
||||
.adapter
|
||||
.as_ref()
|
||||
.map(|raw| crate::wire::parse(raw.get()))
|
||||
.transpose()?;
|
||||
let binding = node
|
||||
.as_ref()
|
||||
.map(|node| node.object())
|
||||
.transpose()?
|
||||
.and_then(|fields| fields.get(identity));
|
||||
let command = binding
|
||||
.map(|node| node.fields(ADAPTER_FIELDS))
|
||||
.transpose()?
|
||||
.and_then(|fields| fields.get("command"));
|
||||
let limits = command
|
||||
.map(|node| serde_json::from_str::<CommandLimits>(node.raw()))
|
||||
.transpose()
|
||||
.map_err(|_| "Adapter command limits require timeout_ms and output_limit integers")?
|
||||
.unwrap_or_default();
|
||||
if !(1..=3_600_000).contains(&limits.timeout_ms) {
|
||||
return Err("Adapter command timeout_ms must be between 1 and 3600000".into());
|
||||
}
|
||||
if !(1..=8 * 1024 * 1024).contains(&limits.output_limit) {
|
||||
return Err("Adapter command output_limit must be between 1 and 8388608".into());
|
||||
}
|
||||
Ok(limits)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
//! The non-secret deployment contract consumed by the compiler.
|
||||
|
||||
use super::{Database, RuntimeProfile, Settings};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct CompilationBindings {
|
||||
pub databases: BTreeMap<String, DatabaseDriver>,
|
||||
pub providers: BTreeMap<String, String>,
|
||||
pub sites: BTreeMap<String, CompilationSite>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Serialize, Eq, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum DatabaseDriver {
|
||||
Sqlite,
|
||||
Postgres,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct CompilationSite {
|
||||
pub path: String,
|
||||
pub auth: String,
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
pub fn compilation_bindings(&self) -> CompilationBindings {
|
||||
CompilationBindings {
|
||||
databases: self
|
||||
.databases
|
||||
.iter()
|
||||
.map(|(name, database)| {
|
||||
(
|
||||
name.clone(),
|
||||
match database {
|
||||
Database::Sqlite { .. } => DatabaseDriver::Sqlite,
|
||||
Database::Postgres { .. } => DatabaseDriver::Postgres,
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
providers: self
|
||||
.auth
|
||||
.providers
|
||||
.iter()
|
||||
.map(|(key, provider)| (key.clone(), provider.verify.clone()))
|
||||
.collect(),
|
||||
sites: self
|
||||
.sites
|
||||
.iter()
|
||||
.map(|(key, site)| {
|
||||
(
|
||||
key.clone(),
|
||||
CompilationSite {
|
||||
path: site.path.clone(),
|
||||
auth: site.auth.clone(),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CompilationBindings {
|
||||
pub fn validate(&self) -> Result<(), String> {
|
||||
if !self.databases.is_empty() && !self.databases.contains_key("default") {
|
||||
return Err("database.default is required".into());
|
||||
}
|
||||
for name in self.databases.keys() {
|
||||
validate_database_name(name)?;
|
||||
}
|
||||
for (key, verify) in &self.providers {
|
||||
validate_provider(key, verify)?;
|
||||
}
|
||||
for (key, site) in &self.sites {
|
||||
validate_site(
|
||||
key,
|
||||
&site.path,
|
||||
&site.auth,
|
||||
self.providers.contains_key(&site.auth),
|
||||
)?;
|
||||
}
|
||||
validate_site_paths(
|
||||
self.sites
|
||||
.iter()
|
||||
.map(|(key, site)| (key.as_str(), site.path.as_str())),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn site_for_directory(
|
||||
&self,
|
||||
directory: &[&str],
|
||||
) -> Result<Option<(&str, &CompilationSite)>, String> {
|
||||
select_site(&self.sites, directory, |site| &site.path)
|
||||
}
|
||||
|
||||
pub fn resolve_database(
|
||||
&self,
|
||||
explicit: Option<&str>,
|
||||
package_root: &str,
|
||||
) -> Result<(&str, &DatabaseDriver), String> {
|
||||
resolve_database(&self.databases, explicit, package_root)
|
||||
}
|
||||
|
||||
pub fn profile(&self) -> RuntimeProfile {
|
||||
RuntimeProfile {
|
||||
sqlite: self
|
||||
.databases
|
||||
.values()
|
||||
.any(|driver| *driver == DatabaseDriver::Sqlite),
|
||||
postgres: self
|
||||
.databases
|
||||
.values()
|
||||
.any(|driver| *driver == DatabaseDriver::Postgres),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn resolve_database<'a, T>(
|
||||
databases: &'a BTreeMap<String, T>,
|
||||
explicit: Option<&str>,
|
||||
package_root: &str,
|
||||
) -> Result<(&'a str, &'a T), String> {
|
||||
let name = match explicit {
|
||||
Some(name) => name,
|
||||
None if databases.contains_key(package_root) => package_root,
|
||||
None => "default",
|
||||
};
|
||||
databases
|
||||
.get_key_value(name)
|
||||
.map(|(name, database)| (name.as_str(), database))
|
||||
.ok_or_else(|| format!("database connection '{name}' is not configured"))
|
||||
}
|
||||
|
||||
pub(super) fn select_site<'a, T>(
|
||||
sites: &'a BTreeMap<String, T>,
|
||||
directory: &[&str],
|
||||
path: impl Fn(&T) -> &str,
|
||||
) -> Result<Option<(&'a str, &'a T)>, String> {
|
||||
let mut matches = sites.iter().filter(|(_, site)| {
|
||||
let path = path(site);
|
||||
path.is_empty()
|
||||
|| path.split('/').count() <= directory.len()
|
||||
&& path
|
||||
.split('/')
|
||||
.zip(directory)
|
||||
.all(|(expected, actual)| expected == *actual)
|
||||
});
|
||||
let selected = matches.next();
|
||||
if matches.next().is_some() {
|
||||
return Err("API directory matches multiple sites".into());
|
||||
}
|
||||
Ok(selected.map(|(key, site)| (key.as_str(), site)))
|
||||
}
|
||||
|
||||
pub(super) fn validate_database_name(name: &str) -> Result<(), String> {
|
||||
if !super::valid_name(name) {
|
||||
return Err(format!(
|
||||
"database connection name '{name}' must use lower_snake_case"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn validate_provider(key: &str, verify: &str) -> Result<(), String> {
|
||||
if !super::valid_name(key) {
|
||||
return Err(format!("auth provider '{key}' must use lower_snake_case"));
|
||||
}
|
||||
let parts = verify.split('.').collect::<Vec<_>>();
|
||||
if parts.len() != 3 || parts.iter().any(|part| !super::valid_name(part)) {
|
||||
return Err(format!(
|
||||
"auth.providers.{key}.verify must be component.domain.function"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn validate_site(
|
||||
key: &str,
|
||||
path: &str,
|
||||
auth: &str,
|
||||
provider_exists: bool,
|
||||
) -> Result<(), String> {
|
||||
if !super::valid_name(key) {
|
||||
return Err(format!("site '{key}' must use lower_snake_case"));
|
||||
}
|
||||
if !path.is_empty() {
|
||||
if path.starts_with('/') || path.ends_with('/') {
|
||||
return Err(format!(
|
||||
"sites.{key}.path must be relative to the api directory"
|
||||
));
|
||||
}
|
||||
if path.split('/').any(|part| !super::valid_name(part)) {
|
||||
return Err(format!(
|
||||
"sites.{key}.path segments must use lower_snake_case"
|
||||
));
|
||||
}
|
||||
}
|
||||
if !provider_exists {
|
||||
return Err(format!(
|
||||
"sites.{key}.auth references unknown provider '{auth}'"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn validate_site_paths<'a>(
|
||||
sites: impl Iterator<Item = (&'a str, &'a str)>,
|
||||
) -> Result<(), String> {
|
||||
let mut entries = sites.collect::<Vec<_>>();
|
||||
entries.sort_by_key(|(_, path)| *path);
|
||||
for adjacent in entries.windows(2) {
|
||||
let [(left_key, left), (right_key, right)] = adjacent else {
|
||||
unreachable!()
|
||||
};
|
||||
if left.is_empty()
|
||||
|| left == right
|
||||
|| right
|
||||
.strip_prefix(*left)
|
||||
.is_some_and(|suffix| suffix.starts_with('/'))
|
||||
{
|
||||
return Err(format!(
|
||||
"sites.{left_key}.path and sites.{right_key}.path overlap; site ownership must be unambiguous"
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
//! Bounded, timezone-independent five-field UTC schedules shared with the compiler.
|
||||
use chrono::{Datelike, Timelike};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Cron {
|
||||
fields: [u64; 5],
|
||||
day_any: bool,
|
||||
weekday_any: bool,
|
||||
}
|
||||
|
||||
impl Cron {
|
||||
pub fn parse(text: &str) -> Result<Self, String> {
|
||||
if text.len() > 512 {
|
||||
return Err("cron exceeds 512 bytes".into());
|
||||
}
|
||||
let parts = text.split_whitespace().collect::<Vec<_>>();
|
||||
if parts.len() != 5 {
|
||||
return Err("cron requires five UTC fields".into());
|
||||
}
|
||||
let mut fields = [0; 5];
|
||||
for (index, (min, max)) in [(0, 59), (0, 23), (1, 31), (1, 12), (0, 6)]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
{
|
||||
fields[index] = field(parts[index], min, max)?;
|
||||
}
|
||||
let day_any = fields[2] == ((1u64 << 32) - 2);
|
||||
let weekday_any = fields[4] == 127;
|
||||
if weekday_any && !day_any {
|
||||
let possible = [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.any(|(month, days)| {
|
||||
fields[3] & (1 << (month + 1)) != 0
|
||||
&& fields[2] & ((1u64 << (days + 1)) - 2) != 0
|
||||
});
|
||||
if !possible {
|
||||
return Err("cron has no possible UTC calendar date".into());
|
||||
}
|
||||
}
|
||||
Ok(Self {
|
||||
fields,
|
||||
day_any,
|
||||
weekday_any,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn fingerprint(&self) -> String {
|
||||
format!(
|
||||
"cron-v1:{:?}:{}:{}",
|
||||
self.fields, self.day_any, self.weekday_any
|
||||
)
|
||||
}
|
||||
|
||||
pub fn matches_minute(&self, minute: i64) -> bool {
|
||||
let Some(seconds) = minute.checked_mul(60) else {
|
||||
return false;
|
||||
};
|
||||
let Some(date) = chrono::DateTime::from_timestamp(seconds, 0) else {
|
||||
return false;
|
||||
};
|
||||
let matches = |index: usize, value: u32| self.fields[index] & (1u64 << value) != 0;
|
||||
let day = matches(2, date.day());
|
||||
let weekday = matches(4, date.weekday().num_days_from_sunday());
|
||||
let calendar = match (self.day_any, self.weekday_any) {
|
||||
(false, false) => day || weekday,
|
||||
(true, false) => weekday,
|
||||
(false, true) => day,
|
||||
(true, true) => true,
|
||||
};
|
||||
matches(0, date.minute()) && matches(1, date.hour()) && matches(3, date.month()) && calendar
|
||||
}
|
||||
}
|
||||
|
||||
fn field(text: &str, minimum: u32, maximum: u32) -> Result<u64, String> {
|
||||
let mut mask = 0;
|
||||
for part in text.split(',') {
|
||||
let (range, step) = match part.split_once('/') {
|
||||
Some((range, step)) => (range, number(step, 1, maximum + 1)?),
|
||||
None => (part, 1),
|
||||
};
|
||||
let (first, last) = if range == "*" {
|
||||
(minimum, maximum)
|
||||
} else if let Some((first, last)) = range.split_once('-') {
|
||||
(
|
||||
number(first, minimum, maximum)?,
|
||||
number(last, minimum, maximum)?,
|
||||
)
|
||||
} else {
|
||||
let value = number(range, minimum, maximum)?;
|
||||
(value, if part.contains('/') { maximum } else { value })
|
||||
};
|
||||
if first > last {
|
||||
return Err("cron range is reversed".into());
|
||||
}
|
||||
for value in (first..=last).step_by(step as usize) {
|
||||
mask |= 1u64 << value;
|
||||
}
|
||||
}
|
||||
if mask == 0 {
|
||||
return Err("empty cron field".into());
|
||||
}
|
||||
Ok(mask)
|
||||
}
|
||||
|
||||
fn number(text: &str, minimum: u32, maximum: u32) -> Result<u32, String> {
|
||||
if text.is_empty() || !text.bytes().all(|byte| byte.is_ascii_digit()) {
|
||||
return Err("cron fields require unsigned decimal numbers".into());
|
||||
}
|
||||
text.parse::<u32>()
|
||||
.ok()
|
||||
.filter(|value| (minimum..=maximum).contains(value))
|
||||
.ok_or_else(|| "cron field is outside its allowed range".into())
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
use argon2::password_hash::{Error, SaltString};
|
||||
use argon2::{Algorithm, Argon2, Params, PasswordHash, PasswordHasher, PasswordVerifier, Version};
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::{Digest, Sha256};
|
||||
use subtle::ConstantTimeEq;
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
use crate::{bytes::Bytes, secret::Secret};
|
||||
|
||||
const MAX_TOKEN_BYTES: i64 = 1024;
|
||||
const INVALID_HASH: &str = "invalid or unsupported password hash";
|
||||
|
||||
pub fn token(length: i64) -> Result<Secret, String> {
|
||||
if !(16..=MAX_TOKEN_BYTES).contains(&length) {
|
||||
return Err("token length must be between 16 and 1024 bytes".into());
|
||||
}
|
||||
let mut bytes = Zeroizing::new(vec![0; length as usize]);
|
||||
getrandom::getrandom(&mut bytes).map_err(|_| "secure random source unavailable")?;
|
||||
Ok(Secret(bytes))
|
||||
}
|
||||
|
||||
pub fn password_hash(password: &Secret) -> Result<String, String> {
|
||||
let mut salt = [0u8; 16];
|
||||
getrandom::getrandom(&mut salt).map_err(|_| "secure random source unavailable")?;
|
||||
let salt = SaltString::encode_b64(&salt).map_err(|_| "password hashing failed")?;
|
||||
Argon2::new(Algorithm::Argon2id, Version::V0x13, Params::default())
|
||||
.hash_password(&password.0, &salt)
|
||||
.map(|hash| hash.to_string())
|
||||
.map_err(|_| "password hashing failed".into())
|
||||
}
|
||||
|
||||
pub fn password_verify(password: &Secret, encoded: &str) -> Result<bool, String> {
|
||||
// Bound parsing and resource costs before invoking the password verifier:
|
||||
// that verifier otherwise accepts PHC-controlled Argon2 parameters.
|
||||
if encoded.len() > 256 {
|
||||
return Err(INVALID_HASH.into());
|
||||
}
|
||||
let hash = PasswordHash::new(encoded).map_err(|_| INVALID_HASH)?;
|
||||
let params = Params::try_from(&hash).map_err(|_| INVALID_HASH)?;
|
||||
let limits = Params::default();
|
||||
if hash.algorithm.as_str() != "argon2id"
|
||||
|| hash.version != Some(19)
|
||||
|| hash.salt.is_none()
|
||||
|| hash.hash.is_none()
|
||||
|| params.m_cost() > limits.m_cost()
|
||||
|| params.t_cost() > limits.t_cost()
|
||||
|| params.p_cost() > limits.p_cost()
|
||||
|| params.output_len().unwrap_or(Params::DEFAULT_OUTPUT_LEN) > Params::DEFAULT_OUTPUT_LEN
|
||||
|| !params.keyid().is_empty()
|
||||
|| !params.data().is_empty()
|
||||
{
|
||||
return Err(INVALID_HASH.into());
|
||||
}
|
||||
match Argon2::default().verify_password(&password.0, &hash) {
|
||||
Ok(()) => Ok(true),
|
||||
Err(Error::Password) => Ok(false),
|
||||
Err(_) => Err(INVALID_HASH.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sha256(bytes: &Bytes) -> Bytes {
|
||||
Bytes::new(Sha256::digest(bytes.values()).to_vec())
|
||||
}
|
||||
|
||||
pub fn hmac_sha256(key: &Secret, bytes: &Bytes) -> Bytes {
|
||||
let mut mac = Hmac::<Sha256>::new_from_slice(&key.0).expect("HMAC accepts any key length");
|
||||
mac.update(bytes.values());
|
||||
Bytes::new(mac.finalize().into_bytes().to_vec())
|
||||
}
|
||||
|
||||
/// Content comparison is constant-time for equal lengths; lengths are public.
|
||||
pub fn constant_time_eq(left: &Bytes, right: &Bytes) -> bool {
|
||||
left.values().ct_eq(right.values()).into()
|
||||
}
|
||||
@@ -0,0 +1,969 @@
|
||||
use crate::config::Database as DatabaseConfig;
|
||||
use crate::orm::{Error, ErrorKind, Row, RowStream, Value};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Database {
|
||||
key: DatabaseKey,
|
||||
backend: Backend,
|
||||
}
|
||||
|
||||
/// One application entry's database owners. Native process roots retain their
|
||||
/// existing global initialization; repeated embedded roots own this session.
|
||||
pub struct Session {
|
||||
settings: Arc<crate::config::Settings>,
|
||||
databases: std::collections::BTreeMap<String, Database>,
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub fn load(
|
||||
bundled: crate::config::RuntimeProfile,
|
||||
bindings: &[(Option<&str>, &str, bool)],
|
||||
transactions: &[&[(Option<&str>, &str, bool)]],
|
||||
) -> Result<Self, String> {
|
||||
let settings = crate::config::Settings::load_for_executable()?;
|
||||
settings.validate_database_entry()?;
|
||||
Self::from_settings(Arc::new(settings), bundled, bindings, transactions)
|
||||
}
|
||||
|
||||
pub fn from_settings(
|
||||
settings: Arc<crate::config::Settings>,
|
||||
bundled: crate::config::RuntimeProfile,
|
||||
bindings: &[(Option<&str>, &str, bool)],
|
||||
transactions: &[&[(Option<&str>, &str, bool)]],
|
||||
) -> Result<Self, String> {
|
||||
settings.prepare_databases(bundled, bindings, transactions)?;
|
||||
let mut databases = std::collections::BTreeMap::new();
|
||||
#[cfg(feature = "sqlite")]
|
||||
for (name, database) in crate::sqlite::configured_databases(&settings)? {
|
||||
databases.insert(
|
||||
name.clone(),
|
||||
Database {
|
||||
key: DatabaseKey {
|
||||
connection: name,
|
||||
tenant_id: None,
|
||||
},
|
||||
backend: Backend::Sqlite(database),
|
||||
},
|
||||
);
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
for (name, database) in crate::postgres::configured_databases(&settings)? {
|
||||
databases.insert(
|
||||
name.clone(),
|
||||
Database {
|
||||
key: DatabaseKey {
|
||||
connection: name,
|
||||
tenant_id: None,
|
||||
},
|
||||
backend: Backend::Postgres(database),
|
||||
},
|
||||
);
|
||||
}
|
||||
Ok(Self {
|
||||
settings,
|
||||
databases,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn database(&self, explicit: Option<&str>, root: &str) -> Result<Database, Error> {
|
||||
let (name, _) = self
|
||||
.settings
|
||||
.resolve_database(explicit, root)
|
||||
.map_err(|message| Error::new(ErrorKind::Pool, message))?;
|
||||
self.databases.get(name).cloned().ok_or_else(|| {
|
||||
Error::new(
|
||||
ErrorKind::Pool,
|
||||
format!("database '{name}' is not initialized"),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn prepare(&self) -> Result<(), Error> {
|
||||
for database in self.databases.values() {
|
||||
database.prepare_pool().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Called after the entry's structured scope has drained, before its
|
||||
/// runtime is destroyed. Stop every pool before awaiting any driver.
|
||||
pub async fn close(&self) -> Result<(), Error> {
|
||||
self.stop_accepting();
|
||||
let mut failure: Option<Error> = None;
|
||||
for database in self.databases.values() {
|
||||
if let Err(error) = database.close_pool().await {
|
||||
failure = Some(match failure {
|
||||
Some(previous) => previous.cause(error),
|
||||
None => error,
|
||||
});
|
||||
}
|
||||
}
|
||||
failure.map_or(Ok(()), Err)
|
||||
}
|
||||
|
||||
pub(crate) fn stop_accepting(&self) {
|
||||
for database in self.databases.values() {
|
||||
database.stop_accepting();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum Backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Sqlite(Arc<crate::sqlite::Database>),
|
||||
#[cfg(feature = "postgres")]
|
||||
Postgres(Arc<crate::postgres::Database>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
|
||||
pub struct DatabaseKey {
|
||||
connection: String,
|
||||
tenant_id: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum ModelScope {
|
||||
Global,
|
||||
Tenant,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct StorageBinding {
|
||||
explicit: Option<&'static str>,
|
||||
package_root: &'static str,
|
||||
scope: ModelScope,
|
||||
}
|
||||
|
||||
impl StorageBinding {
|
||||
pub const fn new(
|
||||
explicit: Option<&'static str>,
|
||||
package_root: &'static str,
|
||||
scope: ModelScope,
|
||||
) -> Self {
|
||||
Self {
|
||||
explicit,
|
||||
package_root,
|
||||
scope,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn database(&self) -> Result<Database, Error> {
|
||||
let settings = crate::config::current_settings();
|
||||
if self.scope == ModelScope::Global || settings.tenant().is_none() {
|
||||
return database_for(self.explicit, self.package_root);
|
||||
}
|
||||
let (connection, _) = settings
|
||||
.resolve_database(self.explicit, self.package_root)
|
||||
.map_err(|error| Error::new(ErrorKind::Pool, error))?;
|
||||
crate::tenant::database(connection).await
|
||||
}
|
||||
|
||||
pub(crate) fn selector(&self) -> (Option<&str>, &str) {
|
||||
(self.explicit, self.package_root)
|
||||
}
|
||||
}
|
||||
|
||||
impl DatabaseKey {
|
||||
pub fn connection(&self) -> &str {
|
||||
&self.connection
|
||||
}
|
||||
pub fn tenant_id(&self) -> Option<i64> {
|
||||
self.tenant_id
|
||||
}
|
||||
}
|
||||
|
||||
pub enum Transaction {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Sqlite {
|
||||
key: DatabaseKey,
|
||||
lease: Arc<crate::sqlite::Database>,
|
||||
transaction: crate::sqlite::Transaction,
|
||||
},
|
||||
#[cfg(feature = "postgres")]
|
||||
Postgres {
|
||||
key: DatabaseKey,
|
||||
lease: Arc<crate::postgres::Database>,
|
||||
transaction: crate::postgres::Transaction,
|
||||
},
|
||||
}
|
||||
|
||||
/// An internal shared transaction owner for foreign coroutine frames. Sequential
|
||||
/// operations borrow the same transaction; finish takes it exactly once. No
|
||||
/// task-local ambient transaction is installed or inherited by spawned tasks.
|
||||
pub struct TransactionState(tokio::sync::Mutex<Option<Transaction>>);
|
||||
|
||||
pub struct TransactionLease<'a> {
|
||||
transaction: tokio::sync::MutexGuard<'a, Option<Transaction>>,
|
||||
completed: bool,
|
||||
}
|
||||
|
||||
impl std::ops::Deref for TransactionLease<'_> {
|
||||
type Target = Transaction;
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.transaction
|
||||
.as_ref()
|
||||
.expect("validated live transaction lease")
|
||||
}
|
||||
}
|
||||
|
||||
impl TransactionLease<'_> {
|
||||
/// The driver returned a known result. Cancellation before this point must
|
||||
/// discard the shared transaction even if another frame still owns an alias.
|
||||
pub fn complete(&mut self) {
|
||||
self.completed = true;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TransactionLease<'_> {
|
||||
fn drop(&mut self) {
|
||||
if !self.completed {
|
||||
self.transaction.take();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TransactionState {
|
||||
pub fn new(transaction: Transaction) -> Self {
|
||||
Self(tokio::sync::Mutex::new(Some(transaction)))
|
||||
}
|
||||
|
||||
pub async fn borrow(state: Option<&Self>) -> Result<Option<TransactionLease<'_>>, Error> {
|
||||
let Some(state) = state else {
|
||||
return Ok(None);
|
||||
};
|
||||
let guard = state.0.lock().await;
|
||||
if guard.is_none() {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Database,
|
||||
"transaction was already finished",
|
||||
));
|
||||
}
|
||||
Ok(Some(TransactionLease {
|
||||
transaction: guard,
|
||||
completed: false,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn finish(&self, commit: bool) -> Result<(), Error> {
|
||||
let transaction =
|
||||
self.0.lock().await.take().ok_or_else(|| {
|
||||
Error::new(ErrorKind::Database, "transaction was already finished")
|
||||
})?;
|
||||
if commit {
|
||||
transaction.commit().await
|
||||
} else {
|
||||
transaction.rollback().await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub enum Executor<'a> {
|
||||
Pool(Database),
|
||||
#[cfg(feature = "sqlite")]
|
||||
Sqlite(&'a crate::sqlite::Transaction),
|
||||
#[cfg(feature = "postgres")]
|
||||
Postgres(&'a crate::postgres::Transaction),
|
||||
}
|
||||
|
||||
pub struct Sql {
|
||||
pub sqlite: &'static str,
|
||||
pub postgres: &'static str,
|
||||
}
|
||||
|
||||
pub struct OwnedSql {
|
||||
pub sqlite: String,
|
||||
pub postgres: String,
|
||||
}
|
||||
|
||||
pub fn relation_sql(
|
||||
prefix: &'static str,
|
||||
before_limit: &'static str,
|
||||
suffix: &'static str,
|
||||
ids: usize,
|
||||
) -> OwnedSql {
|
||||
let sqlite = (1..=ids)
|
||||
.map(|index| format!("?{index}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let postgres = (1..=ids)
|
||||
.map(|index| format!("${index}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
OwnedSql {
|
||||
sqlite: format!("{prefix}{sqlite}{before_limit}?{}{suffix}", ids + 1),
|
||||
postgres: format!("{prefix}{postgres}{before_limit}${}{suffix}", ids + 1),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
||||
pub struct Field {
|
||||
pub name: String,
|
||||
pub ty: String,
|
||||
pub nullable: bool,
|
||||
pub generated: bool,
|
||||
pub default: Option<String>,
|
||||
pub rename_from: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
||||
pub struct Index {
|
||||
pub name: String,
|
||||
pub fields: Vec<String>,
|
||||
pub unique: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
||||
pub struct Migration {
|
||||
pub name: String,
|
||||
pub revision: String,
|
||||
pub drops: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
||||
pub struct Schema {
|
||||
pub model: String,
|
||||
pub table: String,
|
||||
pub revision: String,
|
||||
pub seed_revision: String,
|
||||
pub fields: Vec<Field>,
|
||||
pub indexes: Vec<Index>,
|
||||
pub migrations: Vec<Migration>,
|
||||
}
|
||||
|
||||
pub(crate) fn validate_migration_order(previous: &Schema, desired: &Schema) -> Result<(), Error> {
|
||||
let previous_names = previous
|
||||
.migrations
|
||||
.iter()
|
||||
.map(|migration| migration.name.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
let retained = desired
|
||||
.migrations
|
||||
.iter()
|
||||
.map(|migration| migration.name.as_str())
|
||||
.filter(|name| previous_names.contains(name))
|
||||
.collect::<Vec<_>>();
|
||||
if retained != previous_names {
|
||||
return Err(Error::migration(format!(
|
||||
"applied migrations for '{}' changed order",
|
||||
desired.model,
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub struct SchemaSql {
|
||||
pub sqlite: String,
|
||||
pub postgres: String,
|
||||
}
|
||||
|
||||
pub struct Seed {
|
||||
pub sql: Sql,
|
||||
pub parameters: Vec<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum MigrationPhase {
|
||||
Before,
|
||||
After,
|
||||
}
|
||||
|
||||
pub struct DataMigration {
|
||||
pub name: String,
|
||||
pub phase: MigrationPhase,
|
||||
pub sql: Sql,
|
||||
pub parameters: Vec<Value>,
|
||||
}
|
||||
|
||||
pub struct PostgresConstraint {
|
||||
pub name: String,
|
||||
pub definition: String,
|
||||
pub foreign_key: bool,
|
||||
}
|
||||
|
||||
pub struct PostgresColumn {
|
||||
pub name: String,
|
||||
pub definition: String,
|
||||
pub sql_type: String,
|
||||
pub default: Option<String>,
|
||||
}
|
||||
|
||||
pub struct Model {
|
||||
pub schema: Schema,
|
||||
pub create_table: SchemaSql,
|
||||
pub create_temporary_table: String,
|
||||
pub create_indexes: Vec<SchemaSql>,
|
||||
pub postgres_columns: Vec<PostgresColumn>,
|
||||
pub postgres_constraints: Vec<PostgresConstraint>,
|
||||
pub seeds: Vec<Seed>,
|
||||
pub data_migrations: Vec<DataMigration>,
|
||||
}
|
||||
|
||||
impl Database {
|
||||
pub fn name(&self) -> &str {
|
||||
self.key.connection()
|
||||
}
|
||||
|
||||
pub fn key(&self) -> &DatabaseKey {
|
||||
&self.key
|
||||
}
|
||||
|
||||
pub fn max_page_size(&self) -> usize {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => database.max_page_size(),
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.max_page_size(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn begin(&self) -> Result<Transaction, Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => Ok(Transaction::Sqlite {
|
||||
key: self.key.clone(),
|
||||
lease: Arc::clone(database),
|
||||
transaction: database.begin().await?,
|
||||
}),
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => Ok(Transaction::Postgres {
|
||||
key: self.key.clone(),
|
||||
lease: Arc::clone(database),
|
||||
transaction: database.begin().await?,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(&self, sql: Sql, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => database.execute(sql.sqlite, parameters).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.execute(sql.postgres, parameters).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query(&self, sql: Sql, parameters: Vec<Value>) -> Result<Vec<Row>, Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => database.query(sql.sqlite, parameters).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.query(sql.postgres, parameters).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query_bounded(
|
||||
&self,
|
||||
sql: Sql,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => {
|
||||
database
|
||||
.query_bounded(sql.sqlite, parameters, maximum)
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => {
|
||||
database
|
||||
.query_bounded(sql.postgres, parameters, maximum)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query_owned(
|
||||
&self,
|
||||
sql: OwnedSql,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => database.query_owned(sql.sqlite, parameters).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.query_owned(sql.postgres, parameters).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn stream(
|
||||
&self,
|
||||
sql: Sql,
|
||||
parameters: Vec<Value>,
|
||||
capacity: usize,
|
||||
) -> Result<RowStream<Row>, Error> {
|
||||
let stream = match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => database.stream(sql.sqlite, parameters, capacity).await?,
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => {
|
||||
database.stream(sql.postgres, parameters, capacity).await?
|
||||
}
|
||||
};
|
||||
Ok(stream.hold_database(self.lease_owner()))
|
||||
}
|
||||
|
||||
pub async fn migrate(&self, model: Model) -> Result<(), Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => database.migrate(model).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.migrate(model).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn migrate_foreign_keys(
|
||||
&self,
|
||||
model: String,
|
||||
table: String,
|
||||
constraints: Vec<PostgresConstraint>,
|
||||
) -> Result<(), Error> {
|
||||
let _ = (&model, &table, &constraints);
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(_) => Ok(()),
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => {
|
||||
database
|
||||
.migrate_foreign_keys(model, table, constraints)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn lease_owner(&self) -> Arc<dyn Send + Sync> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => Arc::clone(database) as Arc<dyn Send + Sync>,
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => Arc::clone(database) as Arc<dyn Send + Sync>,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn exclusively_held(&self) -> bool {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => Arc::strong_count(database) == 1,
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => Arc::strong_count(database) == 1,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn prepare_pool(&self) -> Result<(), Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(_) => Ok(()),
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.prepare().await,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn close_pool(&self) -> Result<(), Error> {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => {
|
||||
database.close();
|
||||
Ok(())
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.close().await,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn stop_accepting(&self) {
|
||||
match &self.backend {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Backend::Sqlite(database) => database.close(),
|
||||
#[cfg(feature = "postgres")]
|
||||
Backend::Postgres(database) => database.stop_accepting(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Transaction {
|
||||
pub fn key(&self) -> &DatabaseKey {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite { key, .. } => key,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres { key, .. } => key,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn commit(self) -> Result<(), Error> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite {
|
||||
transaction, lease, ..
|
||||
} => {
|
||||
let result = transaction.commit().await;
|
||||
drop(lease);
|
||||
result
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres {
|
||||
transaction, lease, ..
|
||||
} => {
|
||||
let result = transaction.commit().await;
|
||||
drop(lease);
|
||||
result
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn rollback(self) -> Result<(), Error> {
|
||||
match self {
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite {
|
||||
transaction, lease, ..
|
||||
} => {
|
||||
let result = transaction.rollback().await;
|
||||
drop(lease);
|
||||
result
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres {
|
||||
transaction, lease, ..
|
||||
} => {
|
||||
let result = transaction.rollback().await;
|
||||
drop(lease);
|
||||
result
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Executor<'a> {
|
||||
pub fn new(database: Database, transaction: Option<&'a Transaction>) -> Result<Self, Error> {
|
||||
let Some(transaction) = transaction else {
|
||||
return Ok(Self::Pool(database));
|
||||
};
|
||||
if transaction.key() != database.key() {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Database,
|
||||
format!(
|
||||
"transaction for database key {:?} cannot access database key {:?}",
|
||||
transaction.key(),
|
||||
database.key()
|
||||
),
|
||||
));
|
||||
}
|
||||
match (&database.backend, transaction) {
|
||||
#[cfg(feature = "sqlite")]
|
||||
(Backend::Sqlite(_), Transaction::Sqlite { transaction, .. }) => {
|
||||
Ok(Self::Sqlite(transaction))
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
(Backend::Postgres(_), Transaction::Postgres { transaction, .. }) => {
|
||||
Ok(Self::Postgres(transaction))
|
||||
}
|
||||
#[allow(unreachable_patterns)]
|
||||
_ => Err(Error::new(
|
||||
ErrorKind::Database,
|
||||
"transaction database driver does not match the Model driver",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn max_page_size(&self) -> usize {
|
||||
match self {
|
||||
Self::Pool(database) => database.max_page_size(),
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(transaction) => {
|
||||
crate::sqlite::Executor::Transaction(transaction).max_page_size()
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(transaction) => {
|
||||
crate::postgres::Executor::Transaction(transaction).max_page_size()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(&self, sql: Sql, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.execute(sql, parameters).await,
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(transaction) => {
|
||||
crate::sqlite::Executor::Transaction(transaction)
|
||||
.execute(sql.sqlite, parameters)
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(transaction) => {
|
||||
crate::postgres::Executor::Transaction(transaction)
|
||||
.execute(sql.postgres, parameters)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query(&self, sql: Sql, parameters: Vec<Value>) -> Result<Vec<Row>, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.query(sql, parameters).await,
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(transaction) => {
|
||||
crate::sqlite::Executor::Transaction(transaction)
|
||||
.query(sql.sqlite, parameters)
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(transaction) => {
|
||||
crate::postgres::Executor::Transaction(transaction)
|
||||
.query(sql.postgres, parameters)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query_bounded(
|
||||
&self,
|
||||
sql: Sql,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.query_bounded(sql, parameters, maximum).await,
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(transaction) => {
|
||||
transaction
|
||||
.query_bounded(sql.sqlite, parameters, maximum)
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(transaction) => {
|
||||
transaction
|
||||
.query_bounded(sql.postgres, parameters, maximum)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query_owned(
|
||||
&self,
|
||||
sql: OwnedSql,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.query_owned(sql, parameters).await,
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(transaction) => transaction.query_owned(sql.sqlite, parameters).await,
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(transaction) => transaction.query_owned(sql.postgres, parameters).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn stream(
|
||||
&self,
|
||||
sql: Sql,
|
||||
parameters: Vec<Value>,
|
||||
capacity: usize,
|
||||
) -> Result<RowStream<Row>, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.stream(sql, parameters, capacity).await,
|
||||
#[cfg(feature = "sqlite")]
|
||||
Self::Sqlite(_) => Err(Error::new(
|
||||
ErrorKind::Database,
|
||||
"row streams cannot escape a transaction",
|
||||
)),
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(_) => Err(Error::new(
|
||||
ErrorKind::Database,
|
||||
"row streams cannot escape a transaction",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn database_for(explicit: Option<&str>, package_root: &str) -> Result<Database, Error> {
|
||||
#[cfg(feature = "api")]
|
||||
if let Some(session) = crate::application::current() {
|
||||
return session.databases().database(explicit, package_root);
|
||||
}
|
||||
let (name, config) = crate::config::settings()
|
||||
.resolve_database(explicit, package_root)
|
||||
.map_err(|error| Error::new(ErrorKind::Pool, error))?;
|
||||
let backend = match config {
|
||||
DatabaseConfig::Sqlite { .. } => {
|
||||
#[cfg(feature = "sqlite")]
|
||||
{
|
||||
Backend::Sqlite(crate::sqlite::database(name)?)
|
||||
}
|
||||
#[cfg(not(feature = "sqlite"))]
|
||||
{
|
||||
return Err(Error::new(ErrorKind::Pool, "SQLite support is not bundled"));
|
||||
}
|
||||
}
|
||||
DatabaseConfig::Postgres { .. } => {
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
Backend::Postgres(crate::postgres::database(name)?)
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
{
|
||||
return Err(Error::new(
|
||||
ErrorKind::Pool,
|
||||
"PostgreSQL support is not bundled",
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
Ok(Database {
|
||||
key: DatabaseKey {
|
||||
connection: name.to_owned(),
|
||||
tenant_id: None,
|
||||
},
|
||||
backend,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn tenant_database(
|
||||
connection: &str,
|
||||
tenant_id: i64,
|
||||
create: bool,
|
||||
) -> Result<Database, Error> {
|
||||
let settings = crate::config::current_settings();
|
||||
let config = settings.database(connection).ok_or_else(|| {
|
||||
Error::new(
|
||||
ErrorKind::Pool,
|
||||
format!("database connection '{connection}' is not configured"),
|
||||
)
|
||||
})?;
|
||||
let backend = match config {
|
||||
DatabaseConfig::Sqlite {
|
||||
tenant_directory: Some(directory),
|
||||
max_connections,
|
||||
max_page_size,
|
||||
..
|
||||
} => {
|
||||
#[cfg(not(feature = "sqlite"))]
|
||||
let _ = (directory, max_connections, max_page_size, create);
|
||||
#[cfg(feature = "sqlite")]
|
||||
{
|
||||
if create {
|
||||
std::fs::create_dir_all(directory).map_err(|error| {
|
||||
Error::migration(format!(
|
||||
"cannot create tenant database directory: {error}"
|
||||
))
|
||||
})?;
|
||||
}
|
||||
let path = directory.join(format!("tenant_{tenant_id}.db"));
|
||||
if !create && !path.is_file() {
|
||||
return Err(Error::migration("tenant database is not provisioned"));
|
||||
}
|
||||
Backend::Sqlite(Arc::new(crate::sqlite::Database::open(
|
||||
path,
|
||||
*max_connections,
|
||||
*max_page_size,
|
||||
)?))
|
||||
}
|
||||
#[cfg(not(feature = "sqlite"))]
|
||||
{
|
||||
return Err(Error::new(ErrorKind::Pool, "SQLite support is not bundled"));
|
||||
}
|
||||
}
|
||||
DatabaseConfig::Postgres {
|
||||
url,
|
||||
tenant_database_prefix: Some(prefix),
|
||||
tls,
|
||||
min_connections,
|
||||
max_connections,
|
||||
max_page_size,
|
||||
wait_timeout_ms,
|
||||
io_timeout_ms,
|
||||
..
|
||||
} => {
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
let _ = (
|
||||
url,
|
||||
prefix,
|
||||
tls,
|
||||
min_connections,
|
||||
max_connections,
|
||||
max_page_size,
|
||||
wait_timeout_ms,
|
||||
io_timeout_ms,
|
||||
create,
|
||||
);
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
let database_name = format!("{prefix}_{tenant_id}");
|
||||
if create {
|
||||
crate::postgres::create_database(
|
||||
url,
|
||||
&database_name,
|
||||
*tls,
|
||||
std::time::Duration::from_millis(*wait_timeout_ms),
|
||||
std::time::Duration::from_millis(*io_timeout_ms),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Backend::Postgres(Arc::new(crate::postgres::Database::open(
|
||||
crate::postgres::ConnectionOptions {
|
||||
url,
|
||||
database_name: Some(&database_name),
|
||||
tls: *tls,
|
||||
min_connections: *min_connections,
|
||||
max_connections: *max_connections,
|
||||
max_page_size: *max_page_size,
|
||||
wait_timeout: std::time::Duration::from_millis(*wait_timeout_ms),
|
||||
io_timeout: std::time::Duration::from_millis(*io_timeout_ms),
|
||||
},
|
||||
)?))
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
{
|
||||
return Err(Error::new(
|
||||
ErrorKind::Pool,
|
||||
"PostgreSQL support is not bundled",
|
||||
));
|
||||
}
|
||||
}
|
||||
DatabaseConfig::Sqlite {
|
||||
tenant_directory: None,
|
||||
..
|
||||
} => {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Pool,
|
||||
format!(
|
||||
"database.{connection}.tenant_directory is required for tenant database mode"
|
||||
),
|
||||
));
|
||||
}
|
||||
DatabaseConfig::Postgres {
|
||||
tenant_database_prefix: None,
|
||||
..
|
||||
} => {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Pool,
|
||||
format!(
|
||||
"database.{connection}.tenant_database_prefix is required for tenant database mode"
|
||||
),
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(Database {
|
||||
key: DatabaseKey {
|
||||
connection: connection.to_owned(),
|
||||
tenant_id: Some(tenant_id),
|
||||
},
|
||||
backend,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn prepare() -> Result<(), Error> {
|
||||
#[cfg(feature = "postgres")]
|
||||
crate::postgres::prepare().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn shutdown() -> Result<(), Error> {
|
||||
crate::tenant::shutdown().await?;
|
||||
#[cfg(feature = "sqlite")]
|
||||
crate::sqlite::shutdown();
|
||||
#[cfg(feature = "postgres")]
|
||||
crate::postgres::shutdown().await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
use std::future::Future;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct Endpoint<T> {
|
||||
socket: Mutex<Option<Arc<T>>>,
|
||||
closed: Semaphore,
|
||||
}
|
||||
|
||||
impl<T> Endpoint<T> {
|
||||
pub(crate) fn new(socket: T) -> Self {
|
||||
Self {
|
||||
socket: Mutex::new(Some(Arc::new(socket))),
|
||||
closed: Semaphore::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn get(&self) -> Result<Arc<T>, String> {
|
||||
self.socket
|
||||
.lock()
|
||||
.expect("endpoint lock poisoned")
|
||||
.clone()
|
||||
.ok_or_else(|| "resource is closed".into())
|
||||
}
|
||||
|
||||
pub(crate) fn close(&self) -> Result<(), String> {
|
||||
self.take().map(|_| ())
|
||||
}
|
||||
|
||||
pub(crate) fn take(&self) -> Result<Arc<T>, String> {
|
||||
let socket = self.socket.lock().expect("endpoint lock poisoned").take();
|
||||
self.closed.close();
|
||||
socket.ok_or_else(|| "resource is closed".into())
|
||||
}
|
||||
|
||||
pub(crate) async fn while_open<R>(
|
||||
&self,
|
||||
operation: impl Future<Output = Result<R, String>>,
|
||||
) -> Result<R, String> {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = self.closed.acquire() => Err("resource is closed".into()),
|
||||
result = operation => result,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn closed(&self) {
|
||||
let _closed = self.closed.acquire().await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,701 @@
|
||||
//! Safe extraction primitives for compiler-owned external Worker resources.
|
||||
//!
|
||||
//! The packager supplies verified bytes; this module never downloads, searches
|
||||
//! PATH, follows links, or executes a resource. A caller can therefore use
|
||||
//! the same primitive for an embedded resource table and for a machine-cache
|
||||
//! hit without creating a writable shared execution directory.
|
||||
|
||||
use std::fs::{self, File, OpenOptions};
|
||||
use std::io::{self, Read, Write};
|
||||
use std::path::{Component, Path, PathBuf};
|
||||
|
||||
use serde::Deserialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct Resource<'a> {
|
||||
pub path: &'a str,
|
||||
pub bytes: &'a [u8],
|
||||
pub sha256: &'a str,
|
||||
pub executable: bool,
|
||||
}
|
||||
|
||||
pub(crate) struct PreparedBundle {
|
||||
directory: PathBuf,
|
||||
resources: Vec<Resource<'static>>,
|
||||
}
|
||||
|
||||
pub fn prepare(root: &Path, digest: &str, resources: &[Resource<'static>]) -> Result<(), String> {
|
||||
let session =
|
||||
crate::component::current().ok_or("external resources require a Dever invocation")?;
|
||||
let directory = extract(root, digest, resources)?;
|
||||
session
|
||||
.settings
|
||||
.get_or_init(|| {
|
||||
crate::config::Settings::load_adapter_settings_at(root).map(std::sync::Arc::new)
|
||||
})
|
||||
.as_ref()
|
||||
.map_err(Clone::clone)?;
|
||||
session
|
||||
.resources
|
||||
.set(PreparedBundle {
|
||||
directory,
|
||||
resources: resources.to_vec(),
|
||||
})
|
||||
.map_err(|_| "external resources were initialized twice".into())
|
||||
}
|
||||
|
||||
/// A managed Worker never obtains an interpreter from PATH. Its executable,
|
||||
/// arguments and working directory all come from its verified resource tree.
|
||||
#[derive(Debug)]
|
||||
pub struct WorkerLaunch {
|
||||
pub executable: PathBuf,
|
||||
pub arguments: Vec<std::ffi::OsString>,
|
||||
pub working_directory: PathBuf,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct LaunchManifest {
|
||||
format: String,
|
||||
ecosystem: String,
|
||||
entry: String,
|
||||
executable: String,
|
||||
arguments: Vec<LaunchArgument>,
|
||||
working_directory: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum LaunchArgument {
|
||||
Resource(ResourceArgument),
|
||||
Literal(LiteralArgument),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ResourceArgument {
|
||||
resource: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct LiteralArgument {
|
||||
literal: String,
|
||||
}
|
||||
|
||||
pub fn worker_command(
|
||||
definition: &crate::component::Definition,
|
||||
) -> Result<std::process::Command, String> {
|
||||
sandbox_launch(definition, |launch| launch.command())
|
||||
}
|
||||
|
||||
pub(crate) fn observed_command(
|
||||
definition: &crate::component::Definition,
|
||||
) -> Result<(std::process::Command, dever_sandbox::CommandObservation), String> {
|
||||
sandbox_launch(definition, |launch| launch.observed_command())
|
||||
}
|
||||
|
||||
fn sandbox_launch<T>(
|
||||
definition: &crate::component::Definition,
|
||||
build: impl FnOnce(dever_sandbox::Launch<'_>) -> Result<T, String>,
|
||||
) -> Result<T, String> {
|
||||
let session =
|
||||
crate::component::current().ok_or("external resources require a Dever invocation")?;
|
||||
let bundle = session
|
||||
.resources
|
||||
.get()
|
||||
.ok_or("external Worker resources and sandbox were not prepared")?;
|
||||
let launch = verified_worker_launch(
|
||||
&bundle.directory,
|
||||
&bundle.resources,
|
||||
&definition.entry,
|
||||
&definition.ecosystem,
|
||||
)?;
|
||||
let settings = crate::component::adapter_settings()?
|
||||
.external_grants(&definition.port, &definition.capabilities)?;
|
||||
if definition
|
||||
.capabilities
|
||||
.iter()
|
||||
.any(|capability| capability == "gpu")
|
||||
{
|
||||
return Err("GPU Worker requires a supported device isolation backend".into());
|
||||
}
|
||||
let arguments = launch.arguments;
|
||||
build(dever_sandbox::Launch {
|
||||
assets: &bundle.directory.join("sandbox"),
|
||||
worker: &launch.working_directory,
|
||||
executable: &launch.executable,
|
||||
arguments: &arguments,
|
||||
working_directory: &launch.working_directory,
|
||||
capabilities: dever_sandbox::Capabilities {
|
||||
network: definition
|
||||
.capabilities
|
||||
.iter()
|
||||
.any(|capability| capability == "network"),
|
||||
process: definition
|
||||
.capabilities
|
||||
.iter()
|
||||
.any(|capability| capability == "process"),
|
||||
},
|
||||
grants: &settings,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn verified_worker_launch(
|
||||
directory: &Path,
|
||||
resources: &[Resource<'_>],
|
||||
entry: &str,
|
||||
ecosystem: &str,
|
||||
) -> Result<WorkerLaunch, String> {
|
||||
let (suffix, format, tree) = match ecosystem {
|
||||
"exec" | "pip" | "npm" | "go" => ("worker", "dever-worker-launch-v1", "workers/"),
|
||||
"command" => ("command", "dever-command-launch-v1", "commands/"),
|
||||
_ => return Err("invalid managed Worker ecosystem".into()),
|
||||
};
|
||||
let manifest_path = format!("{entry}.dever-{suffix}.json");
|
||||
let manifest = resources
|
||||
.iter()
|
||||
.find(|resource| resource.path == manifest_path)
|
||||
.ok_or("managed Worker launch manifest is missing")?;
|
||||
if manifest.executable {
|
||||
return Err("managed Worker manifest must not be executable".into());
|
||||
}
|
||||
let text = std::str::from_utf8(manifest.bytes)
|
||||
.map_err(|_| "managed Worker manifest must be UTF-8 JSON")?;
|
||||
crate::wire::parse(text)
|
||||
.map_err(|error| format!("invalid managed Worker manifest: {error}"))?;
|
||||
let manifest: LaunchManifest = serde_json::from_str(text)
|
||||
.map_err(|error| format!("invalid managed Worker manifest: {error}"))?;
|
||||
if manifest.format != format || manifest.ecosystem != ecosystem || manifest.entry != entry {
|
||||
return Err("managed Worker launch identity differs from its checked Adapter".into());
|
||||
}
|
||||
validate_relative_path(&manifest.working_directory)?;
|
||||
if !manifest.working_directory.starts_with(tree) || manifest.arguments.len() > 64 {
|
||||
return Err("invalid managed Worker execution tree or argument count".into());
|
||||
}
|
||||
let prefix = format!("{}/", manifest.working_directory);
|
||||
let resource_path = |path: &str| -> Result<PathBuf, String> {
|
||||
validate_relative_path(path)?;
|
||||
if !path.starts_with(&prefix) || !resources.iter().any(|resource| resource.path == path) {
|
||||
return Err("managed Worker argument escapes its verified execution tree".into());
|
||||
}
|
||||
Ok(directory.join(path))
|
||||
};
|
||||
resource_path(&manifest.executable)?;
|
||||
// Recheck the complete tree once, including executable bits, at every
|
||||
// Worker start/restart before any launch path is exposed to the supervisor.
|
||||
let executable = verified_entry(directory, resources, &manifest.executable)?;
|
||||
let mut total_bytes = 0_usize;
|
||||
let mut arguments = Vec::with_capacity(manifest.arguments.len());
|
||||
for argument in manifest.arguments {
|
||||
let argument = match argument {
|
||||
LaunchArgument::Resource(argument) => {
|
||||
resource_path(&argument.resource)?.into_os_string()
|
||||
}
|
||||
LaunchArgument::Literal(argument) => {
|
||||
if argument.literal.contains('\0') {
|
||||
return Err("managed Worker argument contains NUL".into());
|
||||
}
|
||||
argument.literal.into()
|
||||
}
|
||||
};
|
||||
total_bytes = total_bytes.saturating_add(argument.as_encoded_bytes().len());
|
||||
if total_bytes > 65_536 {
|
||||
return Err("managed Worker arguments exceed the byte budget".into());
|
||||
}
|
||||
arguments.push(argument);
|
||||
}
|
||||
let working_directory = directory.join(manifest.working_directory);
|
||||
if !working_directory.is_dir() {
|
||||
return Err("managed Worker execution directory is missing".into());
|
||||
}
|
||||
reject_symlink_ancestors(&working_directory)?;
|
||||
Ok(WorkerLaunch {
|
||||
executable,
|
||||
arguments,
|
||||
working_directory,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn verified_entry(
|
||||
directory: &Path,
|
||||
resources: &[Resource<'_>],
|
||||
entry: &str,
|
||||
) -> Result<PathBuf, String> {
|
||||
let resource = resources
|
||||
.iter()
|
||||
.find(|resource| resource.path == entry)
|
||||
.ok_or("external Adapter entry is not in the verified manifest")?;
|
||||
if !resource.executable {
|
||||
return Err("external Adapter entry is not executable in the verified manifest".into());
|
||||
}
|
||||
verify_directory(directory, resources)?;
|
||||
let path = directory.join(entry);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
if fs::metadata(&path)
|
||||
.map_err(|error| error.to_string())?
|
||||
.permissions()
|
||||
.mode()
|
||||
& 0o100
|
||||
== 0
|
||||
{
|
||||
return Err("external Adapter entry has no owner execute permission".into());
|
||||
}
|
||||
}
|
||||
path.canonicalize()
|
||||
.map_err(|error| format!("cannot resolve verified external Adapter entry: {error}"))
|
||||
}
|
||||
|
||||
pub fn bundle_digest(resources: &[Resource<'_>]) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
for resource in resources {
|
||||
digest.update(resource.path.as_bytes());
|
||||
digest.update([0]);
|
||||
digest.update(resource.sha256.as_bytes());
|
||||
digest.update([0]);
|
||||
digest.update((resource.bytes.len() as u64).to_le_bytes());
|
||||
digest.update([u8::from(resource.executable)]);
|
||||
}
|
||||
digest
|
||||
.finalize()
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn extract(root: &Path, digest: &str, resources: &[Resource<'_>]) -> Result<PathBuf, String> {
|
||||
validate_digest(digest)?;
|
||||
let mut paths = std::collections::BTreeSet::new();
|
||||
for resource in resources {
|
||||
validate_resource(resource)?;
|
||||
if !paths.insert(resource.path) {
|
||||
return Err(format!(
|
||||
"duplicate external resource path '{}'",
|
||||
resource.path
|
||||
));
|
||||
}
|
||||
}
|
||||
if bundle_digest(resources) != digest {
|
||||
return Err("external resource bundle identity does not match its manifest".into());
|
||||
}
|
||||
ensure_cache_directory(root)?;
|
||||
let target = root.join("data/cache/lib").join(digest);
|
||||
if target.exists() {
|
||||
verify_directory(&target, resources)?;
|
||||
return Ok(target);
|
||||
}
|
||||
let parent = target.parent().ok_or("external resource has no parent")?;
|
||||
let staging = parent.join(format!(
|
||||
".staging-{}-{}",
|
||||
std::process::id(),
|
||||
unique_suffix()
|
||||
));
|
||||
create_private_directory(&staging)?;
|
||||
let result = (|| {
|
||||
let mut paths = std::collections::BTreeSet::new();
|
||||
for resource in resources {
|
||||
validate_resource(resource)?;
|
||||
if !paths.insert(resource.path) {
|
||||
return Err(format!(
|
||||
"duplicate external resource path '{}'",
|
||||
resource.path
|
||||
));
|
||||
}
|
||||
let path = staging.join(resource.path);
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).map_err(|error| error.to_string())?;
|
||||
}
|
||||
let mut options = OpenOptions::new();
|
||||
options.write(true).create_new(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.mode(0o600);
|
||||
}
|
||||
let mut file = options.open(&path).map_err(|error| error.to_string())?;
|
||||
file.write_all(resource.bytes)
|
||||
.map_err(|error| error.to_string())?;
|
||||
file.sync_all().map_err(|error| error.to_string())?;
|
||||
#[cfg(unix)]
|
||||
if resource.executable {
|
||||
set_executable(&file)?;
|
||||
}
|
||||
}
|
||||
sync_directory(&staging)?;
|
||||
match fs::rename(&staging, &target) {
|
||||
Ok(()) => {}
|
||||
Err(error)
|
||||
if matches!(
|
||||
error.kind(),
|
||||
io::ErrorKind::AlreadyExists | io::ErrorKind::DirectoryNotEmpty
|
||||
) =>
|
||||
{
|
||||
verify_directory(&target, resources)?;
|
||||
let _ = fs::remove_dir_all(&staging);
|
||||
}
|
||||
Err(error) => return Err(error.to_string()),
|
||||
}
|
||||
sync_directory(parent)?;
|
||||
Ok(())
|
||||
})();
|
||||
if result.is_err() {
|
||||
let _ = fs::remove_dir_all(&staging);
|
||||
}
|
||||
result
|
||||
.map(|()| target)
|
||||
.map_err(|error: String| format!("cannot publish external resources: {error}"))
|
||||
}
|
||||
|
||||
fn verify_directory(root: &Path, resources: &[Resource<'_>]) -> Result<(), String> {
|
||||
reject_symlink_ancestors(root)?;
|
||||
let metadata = fs::symlink_metadata(root).map_err(|error| error.to_string())?;
|
||||
if !metadata.is_dir() {
|
||||
return Err("external resource cache entry is not a directory".into());
|
||||
}
|
||||
let parent = root
|
||||
.parent()
|
||||
.ok_or("external resource cache has no parent")?;
|
||||
validate_owner(
|
||||
root,
|
||||
&fs::metadata(parent).map_err(|error| error.to_string())?,
|
||||
true,
|
||||
)?;
|
||||
let mut expected = std::collections::BTreeSet::new();
|
||||
for resource in resources {
|
||||
validate_resource(resource)?;
|
||||
if !expected.insert(resource.path) {
|
||||
return Err("duplicate external resource path".into());
|
||||
}
|
||||
let path = root.join(resource.path);
|
||||
reject_symlink_ancestors(&path)?;
|
||||
let metadata = fs::symlink_metadata(&path).map_err(|error| {
|
||||
format!(
|
||||
"cannot inspect external resource '{}': {error}",
|
||||
resource.path
|
||||
)
|
||||
})?;
|
||||
if metadata.file_type().is_symlink() || !metadata.is_file() {
|
||||
return Err(format!(
|
||||
"external resource '{}' is not a regular file",
|
||||
resource.path
|
||||
));
|
||||
}
|
||||
validate_owner(
|
||||
&path,
|
||||
&fs::metadata(root).map_err(|error| error.to_string())?,
|
||||
true,
|
||||
)?;
|
||||
if !verify_resource_file(&path, resource)? {
|
||||
return Err(format!(
|
||||
"external resource '{}' failed integrity verification",
|
||||
resource.path
|
||||
));
|
||||
}
|
||||
}
|
||||
let mut actual = Vec::new();
|
||||
collect_files(root, root, &mut actual)?;
|
||||
if actual.iter().any(|path| !expected.contains(path.as_str())) || actual.len() != expected.len()
|
||||
{
|
||||
return Err("external resource cache contains files outside the verified manifest".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn verify_resource_file(path: &Path, resource: &Resource<'_>) -> Result<bool, String> {
|
||||
let mut file = File::open(path).map_err(|error| error.to_string())?;
|
||||
if file.metadata().map_err(|error| error.to_string())?.len() != resource.bytes.len() as u64 {
|
||||
return Ok(false);
|
||||
}
|
||||
let mut digest = Sha256::new();
|
||||
let mut buffer = [0_u8; 64 * 1024];
|
||||
let mut offset = 0;
|
||||
loop {
|
||||
let read = file.read(&mut buffer).map_err(|error| error.to_string())?;
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
let Some(expected) = resource.bytes.get(offset..offset + read) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if expected != &buffer[..read] {
|
||||
return Ok(false);
|
||||
}
|
||||
digest.update(&buffer[..read]);
|
||||
offset += read;
|
||||
}
|
||||
let actual = digest
|
||||
.finalize()
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect::<String>();
|
||||
Ok(offset == resource.bytes.len() && actual == resource.sha256)
|
||||
}
|
||||
|
||||
fn collect_files(root: &Path, directory: &Path, files: &mut Vec<String>) -> Result<(), String> {
|
||||
for entry in fs::read_dir(directory).map_err(|error| error.to_string())? {
|
||||
let entry = entry.map_err(|error| error.to_string())?;
|
||||
let path = entry.path();
|
||||
let metadata = fs::symlink_metadata(&path).map_err(|error| error.to_string())?;
|
||||
if metadata.file_type().is_symlink() {
|
||||
return Err(format!(
|
||||
"external resource cache contains a symbolic link: {}",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
if metadata.is_dir() {
|
||||
validate_owner(
|
||||
&path,
|
||||
&fs::metadata(root).map_err(|error| error.to_string())?,
|
||||
false,
|
||||
)?;
|
||||
collect_files(root, &path, files)?;
|
||||
} else if metadata.is_file() {
|
||||
let relative = path
|
||||
.strip_prefix(root)
|
||||
.map_err(|error| error.to_string())?
|
||||
.to_str()
|
||||
.ok_or("external resource path is not UTF-8")?;
|
||||
files.push(relative.replace(std::path::MAIN_SEPARATOR, "/"));
|
||||
} else {
|
||||
return Err(format!(
|
||||
"external resource cache contains a non-regular entry: {}",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_cache_directory(root: &Path) -> Result<(), String> {
|
||||
reject_symlink_ancestors(root)?;
|
||||
if !root.exists() {
|
||||
create_private_directory(root)?;
|
||||
}
|
||||
if !fs::symlink_metadata(root)
|
||||
.map_err(|error| error.to_string())?
|
||||
.is_dir()
|
||||
{
|
||||
return Err("external application root must be a directory".into());
|
||||
}
|
||||
let owner = fs::metadata(root).map_err(|error| error.to_string())?;
|
||||
validate_owner(root, &owner, false)?;
|
||||
let data = root.join("data");
|
||||
let cache = data.join("cache");
|
||||
let lib = cache.join("lib");
|
||||
for path in [&data, &cache, &lib] {
|
||||
match fs::symlink_metadata(path) {
|
||||
Ok(metadata) if metadata.file_type().is_symlink() => {
|
||||
return Err(format!(
|
||||
"external resource directory '{}' is a symbolic link",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
Ok(metadata) if !metadata.is_dir() => {
|
||||
return Err(format!(
|
||||
"external resource directory '{}' is not a directory",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(error) if error.kind() == io::ErrorKind::NotFound => {
|
||||
let mut builder = fs::DirBuilder::new();
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::DirBuilderExt;
|
||||
builder.mode(0o700);
|
||||
}
|
||||
match builder.create(path) {
|
||||
Ok(()) => {}
|
||||
Err(create_error) if create_error.kind() == io::ErrorKind::AlreadyExists => {
|
||||
let metadata =
|
||||
fs::symlink_metadata(path).map_err(|error| error.to_string())?;
|
||||
if !metadata.is_dir() || metadata.file_type().is_symlink() {
|
||||
return Err(
|
||||
"external resource directory changed during creation".into()
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(create_error) => {
|
||||
return Err(format!(
|
||||
"cannot create external resource cache '{}': {create_error}",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
fs::set_permissions(path, fs::Permissions::from_mode(0o700))
|
||||
.map_err(|permission_error| permission_error.to_string())?;
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
return Err(format!(
|
||||
"cannot inspect external resource directory '{}': {error}",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
validate_owner(path, &owner, path == &lib)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_resource(resource: &Resource<'_>) -> Result<(), String> {
|
||||
validate_relative_path(resource.path)?;
|
||||
if sha256(resource.bytes) != resource.sha256 {
|
||||
return Err(format!(
|
||||
"external resource '{}' has an invalid digest",
|
||||
resource.path
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_relative_path(value: &str) -> Result<(), String> {
|
||||
let path = Path::new(value);
|
||||
if value.is_empty()
|
||||
|| value.contains(['\\', ':', '\0'])
|
||||
|| path.is_absolute()
|
||||
|| path
|
||||
.components()
|
||||
.any(|component| !matches!(component, Component::Normal(_)))
|
||||
{
|
||||
return Err(format!(
|
||||
"external resource path '{value}' is not safely relative"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reject_symlink_ancestors(path: &Path) -> Result<(), String> {
|
||||
let mut ancestor = PathBuf::new();
|
||||
for component in path.components() {
|
||||
ancestor.push(component);
|
||||
match fs::symlink_metadata(&ancestor) {
|
||||
Ok(metadata) if metadata.file_type().is_symlink() => {
|
||||
return Err("external resource path crosses a symbolic link".into());
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(error) if error.kind() == io::ErrorKind::NotFound => break,
|
||||
Err(error) => return Err(format!("cannot inspect external resource path: {error}")),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_digest(value: &str) -> Result<(), String> {
|
||||
if value.len() != 64
|
||||
|| !value
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase())
|
||||
{
|
||||
return Err("external resource digest must be lowercase SHA-256".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn sha256(bytes: &[u8]) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(bytes);
|
||||
digest
|
||||
.finalize()
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn create_private_directory(path: &Path) -> Result<(), String> {
|
||||
let mut builder = fs::DirBuilder::new();
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::DirBuilderExt;
|
||||
builder.mode(0o700);
|
||||
}
|
||||
builder.create(path).map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
fn validate_owner(path: &Path, owner: &fs::Metadata, private: bool) -> Result<(), String> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::MetadataExt;
|
||||
let metadata = fs::symlink_metadata(path).map_err(|error| error.to_string())?;
|
||||
let prohibited = if private { 0o077 } else { 0o022 };
|
||||
if metadata.uid() != owner.uid() || metadata.mode() & prohibited != 0 {
|
||||
return Err(format!(
|
||||
"external resource path '{}' has an untrusted owner or permissions",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = (path, owner, private);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn sync_directory(path: &Path) -> Result<(), String> {
|
||||
File::open(path)
|
||||
.map_err(|error| error.to_string())?
|
||||
.sync_all()
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn set_executable(file: &File) -> Result<(), String> {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let mut permissions = file
|
||||
.metadata()
|
||||
.map_err(|error| error.to_string())?
|
||||
.permissions();
|
||||
permissions.set_mode(0o700);
|
||||
file.set_permissions(permissions)
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
fn unique_suffix() -> u64 {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
static NEXT: AtomicU64 = AtomicU64::new(0);
|
||||
NEXT.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn extraction_is_verified_and_rejects_path_escape() {
|
||||
let root = std::env::temp_dir().join(format!("dever-external-{}", std::process::id()));
|
||||
let bytes = b"worker";
|
||||
let digest = sha256(bytes);
|
||||
let resources = [Resource {
|
||||
path: "worker.bin",
|
||||
bytes,
|
||||
sha256: &digest,
|
||||
executable: false,
|
||||
}];
|
||||
let bundle = bundle_digest(&resources);
|
||||
let directory = extract(&root, &bundle, &resources).unwrap();
|
||||
assert_eq!(fs::read(directory.join("worker.bin")).unwrap(), bytes);
|
||||
fs::write(directory.join("unexpected.bin"), b"unexpected").unwrap();
|
||||
assert!(extract(&root, &bundle, &resources).is_err());
|
||||
assert!(
|
||||
extract(
|
||||
&root,
|
||||
&bundle,
|
||||
&[Resource {
|
||||
path: "../escape",
|
||||
bytes,
|
||||
sha256: &digest,
|
||||
executable: false
|
||||
}]
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
let _ = fs::remove_dir_all(root);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,548 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use hyper::body::Incoming;
|
||||
use hyper_util::rt::TokioIo;
|
||||
use tokio::sync::{Mutex, OwnedSemaphorePermit, mpsc};
|
||||
use tokio::time::Instant;
|
||||
|
||||
use super::client_body::{RequestBody, download, hold_upload, upload};
|
||||
use super::client_driver::{self, Control, Driver, Entry, Http2Entry, Http2Lease, Http2Pool, Idle};
|
||||
use super::{Budget, Header, Limits, Request, Response, message};
|
||||
use crate::{async_stream::AsyncStream, bytes::Bytes, task, tls::ClientTls};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct PoolLimits {
|
||||
pub idle_ms: i64,
|
||||
pub chunk_bytes: i64,
|
||||
pub read_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct StreamRequest {
|
||||
pub method: String,
|
||||
pub target: String,
|
||||
pub headers: Vec<Header>,
|
||||
pub body: AsyncStream<Result<Bytes, String>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct StreamResponse {
|
||||
pub status: i64,
|
||||
pub headers: Vec<Header>,
|
||||
pub body: AsyncStream<Result<Bytes, String>>,
|
||||
}
|
||||
|
||||
struct State {
|
||||
authority: String,
|
||||
scheme: &'static str,
|
||||
host: String,
|
||||
port: i64,
|
||||
tls: Option<ClientTls>,
|
||||
budget: Budget,
|
||||
chunks: usize,
|
||||
read_timeout: Duration,
|
||||
control: Arc<Control>,
|
||||
idle: Arc<Idle>,
|
||||
http2: Option<Arc<Http2Pool>>,
|
||||
queue: mpsc::Sender<Driver>,
|
||||
driver: Mutex<Option<task::Task<()>>>,
|
||||
}
|
||||
|
||||
impl Drop for State {
|
||||
fn drop(&mut self) {
|
||||
self.control.close();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct HttpClient(Arc<State>);
|
||||
|
||||
impl std::fmt::Debug for HttpClient {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("HttpClient(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn client(
|
||||
origin: &str,
|
||||
tls: ClientTls,
|
||||
limits: Limits,
|
||||
pool: PoolLimits,
|
||||
) -> Result<HttpClient, String> {
|
||||
let budget = limits.check()?;
|
||||
let idle_timeout = task::positive_duration(pool.idle_ms)?;
|
||||
let read_timeout = task::positive_duration(pool.read_ms)?;
|
||||
let chunks = usize::try_from(pool.chunk_bytes)
|
||||
.ok()
|
||||
.filter(|size| *size > 0)
|
||||
.ok_or("HTTP chunk_bytes must be positive")?;
|
||||
let uri: hyper::Uri = origin
|
||||
.parse()
|
||||
.map_err(|error| format!("invalid HTTP origin: {error}"))?;
|
||||
let scheme = match uri.scheme_str() {
|
||||
Some("http") => "http",
|
||||
Some("https") => "https",
|
||||
_ => return Err("HTTP origin must use http or https".into()),
|
||||
};
|
||||
let secure = scheme == "https";
|
||||
let authority = uri
|
||||
.authority()
|
||||
.ok_or("HTTP origin requires a host")?
|
||||
.as_str()
|
||||
.to_owned();
|
||||
if authority.contains('@')
|
||||
|| origin.contains('#')
|
||||
|| uri.query().is_some()
|
||||
|| !matches!(uri.path(), "" | "/")
|
||||
{
|
||||
return Err(
|
||||
"HTTP client origin cannot contain credentials, path, query or fragment".into(),
|
||||
);
|
||||
}
|
||||
let host = uri
|
||||
.host()
|
||||
.ok_or("HTTP origin requires a host")?
|
||||
.trim_matches(['[', ']'])
|
||||
.to_owned();
|
||||
let port = i64::from(uri.port_u16().unwrap_or(if secure { 443 } else { 80 }));
|
||||
let connections = limits.connections as usize;
|
||||
let request_capacity = budget
|
||||
.http2
|
||||
.map(|http2| connections * http2.streams as usize)
|
||||
.unwrap_or(connections);
|
||||
let control = Control::new(request_capacity, connections);
|
||||
let idle = Idle::new(idle_timeout);
|
||||
let http2 = budget
|
||||
.http2
|
||||
.map(|limits| Http2Pool::new(limits.streams as usize, idle_timeout));
|
||||
let (queue, driver) =
|
||||
client_driver::start(control.clone(), idle.clone(), http2.clone(), connections).await?;
|
||||
Ok(HttpClient(Arc::new(State {
|
||||
authority,
|
||||
scheme,
|
||||
host,
|
||||
port,
|
||||
tls: secure.then(|| super::protocol::client_tls(tls, budget.http2)),
|
||||
budget,
|
||||
chunks,
|
||||
read_timeout,
|
||||
control,
|
||||
idle,
|
||||
http2,
|
||||
queue,
|
||||
driver: Mutex::new(Some(driver)),
|
||||
})))
|
||||
}
|
||||
|
||||
impl HttpClient {
|
||||
pub(super) fn control(&self) -> &Arc<Control> {
|
||||
&self.0.control
|
||||
}
|
||||
|
||||
pub(super) fn chunk_limit(&self) -> usize {
|
||||
self.0.chunks
|
||||
}
|
||||
|
||||
pub(super) fn read_timeout(&self) -> Duration {
|
||||
self.0.read_timeout
|
||||
}
|
||||
|
||||
fn encode(
|
||||
&self,
|
||||
mut request: Request,
|
||||
) -> Result<hyper::Request<http_body_util::Full<bytes::Bytes>>, String> {
|
||||
if !request
|
||||
.headers
|
||||
.iter()
|
||||
.any(|header| header.name.eq_ignore_ascii_case("host"))
|
||||
{
|
||||
request.headers.push(Header {
|
||||
name: "host".into(),
|
||||
value: Bytes::from_string(self.0.authority.clone()),
|
||||
});
|
||||
}
|
||||
let mut request = message::encode_request(request, self.0.budget)?;
|
||||
if self.0.http2.is_some() {
|
||||
message::http2_request(&mut request, self.0.scheme)?;
|
||||
}
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
async fn http1_connection(&self) -> Result<Entry, String> {
|
||||
while let Some(mut entry) = self.0.idle.take() {
|
||||
if entry.sender.ready().await.is_ok() {
|
||||
return Ok(entry);
|
||||
}
|
||||
// Discard a closed idle transport before sending any request bytes.
|
||||
}
|
||||
let permit = self
|
||||
.0
|
||||
.control
|
||||
.connections
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| "HTTP client is closed".to_owned())?;
|
||||
let transport =
|
||||
crate::transport::connect(&self.0.host, self.0.port, self.0.tls.as_ref()).await?;
|
||||
super::protocol::check_alpn(&transport, self.0.budget.http2)?;
|
||||
let (sender, connection) = hyper::client::conn::http1::Builder::new()
|
||||
.max_buf_size(self.0.budget.headers)
|
||||
.handshake(TokioIo::new(transport))
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
let (abort, driver) = client_driver::driver(connection, permit);
|
||||
// The connection permit bounds queued plus running drivers together.
|
||||
self.0
|
||||
.queue
|
||||
.try_send(driver)
|
||||
.map_err(|_| "HTTP client driver is closed".to_owned())?;
|
||||
Ok(Entry { sender, abort })
|
||||
}
|
||||
|
||||
async fn new_http2_connection(
|
||||
&self,
|
||||
permit: OwnedSemaphorePermit,
|
||||
) -> Result<Arc<Http2Entry>, String> {
|
||||
let limits = self.0.budget.http2.expect("HTTP/2 client limits");
|
||||
let pool = self.0.http2.as_ref().expect("HTTP/2 client pool");
|
||||
let transport =
|
||||
crate::transport::connect(&self.0.host, self.0.port, self.0.tls.as_ref()).await?;
|
||||
super::protocol::check_alpn(&transport, self.0.budget.http2)?;
|
||||
let (driver, handshake) = client_driver::http2_handshake(
|
||||
transport,
|
||||
permit,
|
||||
limits,
|
||||
self.0.budget.headers,
|
||||
self.0.budget.timeout,
|
||||
Arc::clone(pool),
|
||||
);
|
||||
self.0
|
||||
.queue
|
||||
.try_send(driver)
|
||||
.map_err(|_| "HTTP client driver is closed".to_owned())?;
|
||||
handshake
|
||||
.await
|
||||
.map_err(|_| "HTTP client driver closed during HTTP/2 handshake".to_owned())?
|
||||
}
|
||||
|
||||
async fn http2_connection(&self) -> Result<(Arc<Http2Entry>, OwnedSemaphorePermit), String> {
|
||||
let pool = self.0.http2.as_ref().expect("HTTP/2 client pool");
|
||||
loop {
|
||||
let changed = pool.changed();
|
||||
tokio::pin!(changed);
|
||||
changed.as_mut().enable();
|
||||
if let Some(connection) = pool.take() {
|
||||
return Ok(connection);
|
||||
}
|
||||
let connection = self.0.control.connections.clone().acquire_owned();
|
||||
tokio::pin!(connection);
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = changed => continue,
|
||||
permit = &mut connection => {
|
||||
let permit = permit.map_err(|_| "HTTP client is closed".to_owned())?;
|
||||
let entry = self.new_http2_connection(permit).await?;
|
||||
if let Some(stream) = entry.acquire() {
|
||||
return Ok((entry, stream));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn exchange(
|
||||
&self,
|
||||
request: hyper::Request<RequestBody>,
|
||||
) -> Result<StreamResponse, String> {
|
||||
if self.0.http2.is_some() {
|
||||
self.exchange_http2(request).await
|
||||
} else {
|
||||
self.exchange_http1(request).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn exchange_http1(
|
||||
&self,
|
||||
request: hyper::Request<RequestBody>,
|
||||
) -> Result<StreamResponse, String> {
|
||||
self.0
|
||||
.control
|
||||
.while_open(async {
|
||||
tokio::time::timeout(self.0.budget.timeout, async {
|
||||
let permit = self
|
||||
.0
|
||||
.control
|
||||
.requests
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| "HTTP client is closed".to_owned())?;
|
||||
let mut entry = self.http1_connection().await?;
|
||||
let response = entry
|
||||
.sender
|
||||
.send_request(request)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
let lease = Lease::Http1 {
|
||||
entry: Some(entry),
|
||||
client: self.clone(),
|
||||
_permit: permit,
|
||||
};
|
||||
self.response(response, lease)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "HTTP request timed out".to_owned())?
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn exchange_http2(
|
||||
&self,
|
||||
request: hyper::Request<RequestBody>,
|
||||
) -> Result<StreamResponse, String> {
|
||||
self.0
|
||||
.control
|
||||
.while_open(async {
|
||||
tokio::time::timeout(self.0.budget.timeout, async {
|
||||
let (head, body) = request.into_parts();
|
||||
let (body, upload_lease) = hold_upload(body);
|
||||
let mut pending = Some(hyper::Request::from_parts(head, body));
|
||||
loop {
|
||||
let request_permit = self
|
||||
.0
|
||||
.control
|
||||
.requests
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| "HTTP client is closed".to_owned())?;
|
||||
let (entry, stream_permit) = self.http2_connection().await?;
|
||||
let stream = Http2Lease::new(entry.clone(), stream_permit, request_permit);
|
||||
let mut sender = entry.sender.clone();
|
||||
if sender.ready().await.is_err() {
|
||||
entry.drain();
|
||||
continue;
|
||||
}
|
||||
upload_lease.replace(stream.clone());
|
||||
let request = pending.take().expect("pending HTTP/2 request");
|
||||
let response = match sender.try_send_request(request).await {
|
||||
Ok(response) => response,
|
||||
Err(mut error) => {
|
||||
if sender.is_closed() {
|
||||
entry.drain();
|
||||
}
|
||||
if let Some(request) = error.take_message() {
|
||||
upload_lease.clear();
|
||||
drop(stream);
|
||||
pending = Some(request);
|
||||
continue;
|
||||
}
|
||||
return Err(error.into_error().to_string());
|
||||
}
|
||||
};
|
||||
let lease = Lease::Http2 {
|
||||
client: self.clone(),
|
||||
_stream: stream,
|
||||
};
|
||||
return self.response(response, lease);
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "HTTP request timed out".to_owned())?
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
fn response(
|
||||
&self,
|
||||
response: hyper::Response<Incoming>,
|
||||
lease: Lease,
|
||||
) -> Result<StreamResponse, String> {
|
||||
if response.status() == hyper::StatusCode::SWITCHING_PROTOCOLS {
|
||||
return Err("HTTP client upgrades require the WebSocket API".into());
|
||||
}
|
||||
if message::header_size(response.headers()) + 32 > self.0.budget.headers {
|
||||
return Err("HTTP response headers exceed header_bytes".into());
|
||||
}
|
||||
let (head, body) = response.into_parts();
|
||||
Ok(StreamResponse {
|
||||
status: i64::from(head.status.as_u16()),
|
||||
headers: message::decode_headers(head.headers),
|
||||
body: download(body, lease),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) enum Lease {
|
||||
Http1 {
|
||||
entry: Option<Entry>,
|
||||
client: HttpClient,
|
||||
_permit: OwnedSemaphorePermit,
|
||||
},
|
||||
Http2 {
|
||||
client: HttpClient,
|
||||
_stream: Http2Lease,
|
||||
},
|
||||
}
|
||||
|
||||
impl Lease {
|
||||
pub fn client(&self) -> &HttpClient {
|
||||
match self {
|
||||
Self::Http1 { client, .. } | Self::Http2 { client, .. } => client,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn complete(&mut self) {
|
||||
let Self::Http1 { entry, client, .. } = self else {
|
||||
return;
|
||||
};
|
||||
let mut entry = entry.take().expect("active HTTP lease");
|
||||
if !client.control().closed.is_closed() && entry.sender.ready().await.is_ok() {
|
||||
client.0.idle.put(entry);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn send_http2(
|
||||
address: &str,
|
||||
port: i64,
|
||||
mut request: hyper::Request<http_body_util::Full<bytes::Bytes>>,
|
||||
budget: Budget,
|
||||
) -> Result<Response, String> {
|
||||
use http_body_util::BodyExt;
|
||||
|
||||
message::http2_request(&mut request, "http")?;
|
||||
let deadline = Instant::now() + budget.timeout;
|
||||
let timeout_ms = i64::try_from(budget.timeout.as_millis())
|
||||
.map_err(|_| "HTTP timeout is too large".to_owned())?;
|
||||
let transport = tokio::time::timeout_at(
|
||||
deadline,
|
||||
crate::net::connect_stream_timeout(address, port, timeout_ms),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| "HTTP request timed out".to_owned())??;
|
||||
let (mut owner, spawner, fault) = task::scoped_owner(())?;
|
||||
let limits = budget.http2.expect("HTTP/2 send limits");
|
||||
let builder = limits.client_builder(super::executor::Http2Executor(spawner), budget.headers);
|
||||
let handshake = {
|
||||
let handshake = tokio::time::timeout_at(
|
||||
deadline,
|
||||
owner.enter(builder.handshake(TokioIo::new(transport))),
|
||||
);
|
||||
tokio::pin!(handshake);
|
||||
tokio::select! {
|
||||
result = &mut handshake => match result {
|
||||
Ok(result) => result.map_err(|error| error.to_string()),
|
||||
Err(_) => Err("HTTP request timed out".to_owned()),
|
||||
},
|
||||
error = fault.wait() => Err(error),
|
||||
}
|
||||
};
|
||||
let (mut sender, connection) = match handshake {
|
||||
Ok(connection) => connection,
|
||||
Err(error) => {
|
||||
let _ = owner.stop().await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let request = request.map(|body| body.map_err(|never| match never {}).boxed_unsync());
|
||||
let exchange = async {
|
||||
let response = sender
|
||||
.send_request(request)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
if response.status() == hyper::StatusCode::SWITCHING_PROTOCOLS {
|
||||
return Err("HTTP upgrades are not supported by send".into());
|
||||
}
|
||||
if message::header_size(response.headers()) + 32 > budget.headers {
|
||||
return Err("HTTP response headers exceed header_bytes".into());
|
||||
}
|
||||
let (head, body) = response.into_parts();
|
||||
let body = message::collect_body(body, budget.body)
|
||||
.await
|
||||
.map_err(|status| format!("HTTP response body rejected: {status}"))?;
|
||||
Ok(Response {
|
||||
status: i64::from(head.status.as_u16()),
|
||||
headers: message::decode_headers(head.headers),
|
||||
body,
|
||||
})
|
||||
};
|
||||
let result = {
|
||||
let connection = owner.enter(connection);
|
||||
tokio::pin!(connection, exchange);
|
||||
match tokio::time::timeout_at(deadline, async {
|
||||
tokio::select! {
|
||||
result = &mut exchange => result,
|
||||
result = &mut connection => {
|
||||
result.map_err(|error| error.to_string())?;
|
||||
exchange.await
|
||||
}
|
||||
error = fault.wait() => Err(error),
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(_) => Err("HTTP request timed out".to_owned()),
|
||||
}
|
||||
};
|
||||
let cleanup = owner.stop().await;
|
||||
match result {
|
||||
Ok(response) => {
|
||||
cleanup?;
|
||||
Ok(response)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn open(client: &HttpClient, request: Request) -> Result<StreamResponse, String> {
|
||||
use http_body_util::BodyExt;
|
||||
let request = client
|
||||
.encode(request)?
|
||||
.map(|body| body.map_err(|never| match never {}).boxed_unsync());
|
||||
client.exchange(request).await
|
||||
}
|
||||
|
||||
pub async fn open_stream(
|
||||
client: &HttpClient,
|
||||
request: StreamRequest,
|
||||
) -> Result<StreamResponse, String> {
|
||||
let head = client.encode(Request {
|
||||
method: request.method,
|
||||
target: request.target,
|
||||
headers: request.headers,
|
||||
body: Bytes::new(Vec::new()),
|
||||
})?;
|
||||
let request = head.map(|_| upload(request.body, client.chunk_limit(), client.read_timeout()));
|
||||
client.exchange(request).await
|
||||
}
|
||||
|
||||
pub async fn request(client: &HttpClient, request: Request) -> Result<Response, String> {
|
||||
tokio::time::timeout(client.0.budget.timeout, async {
|
||||
let response = open(client, request).await?;
|
||||
let _owner = response.body.close_on_drop();
|
||||
let mut body = Bytes::new(Vec::new());
|
||||
while let Some(chunk) = response.body.pull().await {
|
||||
let chunk = chunk?;
|
||||
if body.values().len() + chunk.values().len() > client.0.budget.body {
|
||||
return Err("HTTP response body exceeds body_bytes".into());
|
||||
}
|
||||
body = body.concat(&chunk);
|
||||
}
|
||||
Ok(Response {
|
||||
status: response.status,
|
||||
headers: response.headers,
|
||||
body,
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "HTTP request timed out".to_owned())?
|
||||
}
|
||||
|
||||
pub async fn close_client(client: &HttpClient) -> Result<(), String> {
|
||||
client.0.control.close();
|
||||
task::stop_resource(&mut *client.0.driver.lock().await).await
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
use std::io;
|
||||
use std::pin::Pin;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::stream;
|
||||
use http_body_util::{BodyExt, StreamBody, combinators::UnsyncBoxBody};
|
||||
use hyper::body::{Body, Frame, Incoming};
|
||||
|
||||
use super::client::Lease;
|
||||
use super::client_driver::Http2Lease;
|
||||
use crate::{async_stream::AsyncStream, bytes::Bytes, resource::Pull};
|
||||
|
||||
pub(super) type RequestBody = UnsyncBoxBody<bytes::Bytes, io::Error>;
|
||||
|
||||
struct LeasedRequestBody {
|
||||
body: RequestBody,
|
||||
_lease: UploadLease,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct UploadLease(Arc<Mutex<Option<Http2Lease>>>);
|
||||
|
||||
impl UploadLease {
|
||||
pub(super) fn replace(&self, lease: Http2Lease) {
|
||||
*self.0.lock().expect("HTTP/2 upload lease lock") = Some(lease);
|
||||
}
|
||||
|
||||
pub(super) fn clear(&self) {
|
||||
self.0.lock().expect("HTTP/2 upload lease lock").take();
|
||||
}
|
||||
}
|
||||
|
||||
impl Body for LeasedRequestBody {
|
||||
type Data = bytes::Bytes;
|
||||
type Error = io::Error;
|
||||
|
||||
fn poll_frame(
|
||||
mut self: Pin<&mut Self>,
|
||||
context: &mut Context<'_>,
|
||||
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
|
||||
Pin::new(&mut self.body).poll_frame(context)
|
||||
}
|
||||
|
||||
fn is_end_stream(&self) -> bool {
|
||||
self.body.is_end_stream()
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> hyper::body::SizeHint {
|
||||
self.body.size_hint()
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn hold_upload(body: RequestBody) -> (RequestBody, UploadLease) {
|
||||
let lease = UploadLease(Arc::new(Mutex::new(None)));
|
||||
let body = LeasedRequestBody {
|
||||
body,
|
||||
_lease: lease.clone(),
|
||||
}
|
||||
.boxed_unsync();
|
||||
(body, lease)
|
||||
}
|
||||
|
||||
pub(super) fn upload(
|
||||
source: AsyncStream<Result<Bytes, String>>,
|
||||
maximum: usize,
|
||||
timeout: Duration,
|
||||
) -> RequestBody {
|
||||
let owner = source.close_on_drop();
|
||||
StreamBody::new(stream::unfold(
|
||||
(source, owner),
|
||||
move |(source, owner)| async move {
|
||||
let chunk = tokio::time::timeout(timeout, source.pull()).await;
|
||||
let result = match chunk {
|
||||
Ok(Some(Ok(bytes))) if bytes.values().len() <= maximum => {
|
||||
Ok(Frame::data(bytes.into_http()))
|
||||
}
|
||||
Ok(Some(Ok(_))) => Err(io::Error::other("HTTP upload chunk exceeds chunk_bytes")),
|
||||
Ok(Some(Err(error))) => Err(io::Error::other(error)),
|
||||
Ok(None) => return None,
|
||||
Err(_) => Err(io::Error::new(
|
||||
io::ErrorKind::TimedOut,
|
||||
"HTTP upload timed out",
|
||||
)),
|
||||
};
|
||||
if result.is_err() {
|
||||
source.close();
|
||||
}
|
||||
Some((result, (source, owner)))
|
||||
},
|
||||
))
|
||||
.boxed_unsync()
|
||||
}
|
||||
|
||||
struct Download {
|
||||
body: Incoming,
|
||||
pending: bytes::Bytes,
|
||||
lease: Lease,
|
||||
}
|
||||
|
||||
pub(super) fn download(body: Incoming, lease: Lease) -> AsyncStream<Result<Bytes, String>> {
|
||||
AsyncStream::new(stream::unfold(
|
||||
Download {
|
||||
body,
|
||||
pending: bytes::Bytes::new(),
|
||||
lease,
|
||||
},
|
||||
|mut state| async move {
|
||||
let client = state.lease.client().clone();
|
||||
let next = client
|
||||
.control()
|
||||
.while_open(async {
|
||||
tokio::time::timeout(client.read_timeout(), state.next())
|
||||
.await
|
||||
.map_err(|_| "HTTP response read timed out".to_owned())?
|
||||
})
|
||||
.await;
|
||||
match next {
|
||||
Ok(Some((chunk, last))) => {
|
||||
let item = if last {
|
||||
Pull::Last(Ok(chunk))
|
||||
} else {
|
||||
Pull::Item(Ok(chunk))
|
||||
};
|
||||
Some((item, state))
|
||||
}
|
||||
Ok(None) => None,
|
||||
Err(error) => Some((Pull::Last(Err(error)), state)),
|
||||
}
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
impl Download {
|
||||
async fn next(&mut self) -> Result<Option<(Bytes, bool)>, String> {
|
||||
while self.pending.is_empty() {
|
||||
if self.body.is_end_stream() {
|
||||
self.lease.complete().await;
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(frame) = self.body.frame().await else {
|
||||
self.lease.complete().await;
|
||||
return Ok(None);
|
||||
};
|
||||
let frame = frame.map_err(|error| error.to_string())?;
|
||||
match frame.into_data() {
|
||||
Ok(bytes) => self.pending = bytes,
|
||||
Err(frame)
|
||||
if frame
|
||||
.trailers_ref()
|
||||
.is_some_and(|trailers| !trailers.is_empty()) =>
|
||||
{
|
||||
return Err("HTTP response trailers are not supported".into());
|
||||
}
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
let length = self.pending.len().min(self.lease.client().chunk_limit());
|
||||
let chunk = Bytes::from_http(self.pending.split_to(length));
|
||||
let last = self.pending.is_empty() && self.body.is_end_stream();
|
||||
if last {
|
||||
self.lease.complete().await;
|
||||
}
|
||||
Ok(Some((chunk, last)))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::future::{AbortHandle, Abortable};
|
||||
use futures_util::{StreamExt, stream::FuturesUnordered};
|
||||
use hyper_util::rt::TokioIo;
|
||||
use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
|
||||
use tokio::time::Instant;
|
||||
|
||||
use super::client_body::RequestBody;
|
||||
use super::executor::Http2Executor;
|
||||
use super::protocol::Http2Limits;
|
||||
use crate::transport::Transport;
|
||||
|
||||
pub(super) type Driver = Pin<Box<dyn Future<Output = ()> + Send>>;
|
||||
|
||||
pub(super) struct Entry {
|
||||
pub sender: hyper::client::conn::http1::SendRequest<RequestBody>,
|
||||
pub abort: AbortHandle,
|
||||
}
|
||||
|
||||
impl Drop for Entry {
|
||||
fn drop(&mut self) {
|
||||
self.abort.abort();
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct Idle {
|
||||
entries: Mutex<Vec<(Instant, Entry)>>,
|
||||
changed: Notify,
|
||||
timeout: Duration,
|
||||
}
|
||||
|
||||
impl Idle {
|
||||
pub fn new(timeout: Duration) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
entries: Mutex::new(Vec::new()),
|
||||
changed: Notify::new(),
|
||||
timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn take(&self) -> Option<Entry> {
|
||||
let mut entries = self.entries.lock().expect("pool state lock");
|
||||
entries.retain(|(deadline, entry)| *deadline > Instant::now() && !entry.sender.is_closed());
|
||||
entries.pop().map(|(_, entry)| entry)
|
||||
}
|
||||
|
||||
pub fn put(&self, entry: Entry) {
|
||||
self.entries
|
||||
.lock()
|
||||
.expect("pool state lock")
|
||||
.push((Instant::now() + self.timeout, entry));
|
||||
self.changed.notify_one();
|
||||
}
|
||||
|
||||
fn deadline(&self) -> Option<Instant> {
|
||||
let mut entries = self.entries.lock().expect("pool state lock");
|
||||
entries.retain(|(deadline, entry)| *deadline > Instant::now() && !entry.sender.is_closed());
|
||||
entries.iter().map(|(deadline, _)| *deadline).min()
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct Control {
|
||||
pub closed: Semaphore,
|
||||
pub requests: Arc<Semaphore>,
|
||||
pub connections: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl Control {
|
||||
pub fn new(requests: usize, connections: usize) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
closed: Semaphore::new(0),
|
||||
requests: Arc::new(Semaphore::new(requests)),
|
||||
connections: Arc::new(Semaphore::new(connections)),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn close(&self) {
|
||||
self.closed.close();
|
||||
self.requests.close();
|
||||
self.connections.close();
|
||||
}
|
||||
|
||||
pub async fn while_open<T>(
|
||||
&self,
|
||||
future: impl Future<Output = Result<T, String>>,
|
||||
) -> Result<T, String> {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = self.closed.acquire() => Err("HTTP client is closed".into()),
|
||||
result = future => result,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct Http2Entry {
|
||||
pub sender: hyper::client::conn::http2::SendRequest<RequestBody>,
|
||||
streams: Arc<Semaphore>,
|
||||
accepting: AtomicBool,
|
||||
active: AtomicUsize,
|
||||
idle_deadline: Mutex<Option<Instant>>,
|
||||
abort: AbortHandle,
|
||||
changed: Arc<Notify>,
|
||||
idle_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Http2Entry {
|
||||
fn new(
|
||||
sender: hyper::client::conn::http2::SendRequest<RequestBody>,
|
||||
streams: usize,
|
||||
abort: AbortHandle,
|
||||
changed: Arc<Notify>,
|
||||
idle_timeout: Duration,
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
sender,
|
||||
streams: Arc::new(Semaphore::new(streams)),
|
||||
accepting: AtomicBool::new(true),
|
||||
active: AtomicUsize::new(0),
|
||||
idle_deadline: Mutex::new(Some(Instant::now() + idle_timeout)),
|
||||
abort,
|
||||
changed,
|
||||
idle_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn acquire(self: &Arc<Self>) -> Option<OwnedSemaphorePermit> {
|
||||
if !self.accepting.load(Ordering::Acquire) || self.sender.is_closed() {
|
||||
self.drain();
|
||||
return None;
|
||||
}
|
||||
let permit = self.streams.clone().try_acquire_owned().ok()?;
|
||||
if !self.accepting.load(Ordering::Acquire) || self.sender.is_closed() {
|
||||
self.drain();
|
||||
return None;
|
||||
}
|
||||
self.active.fetch_add(1, Ordering::AcqRel);
|
||||
*self.idle_deadline.lock().expect("HTTP/2 pool state lock") = None;
|
||||
Some(permit)
|
||||
}
|
||||
|
||||
fn release(&self) {
|
||||
let previous = self.active.fetch_sub(1, Ordering::AcqRel);
|
||||
debug_assert!(previous > 0, "active HTTP/2 stream count");
|
||||
if previous == 1 && self.accepting.load(Ordering::Acquire) {
|
||||
*self.idle_deadline.lock().expect("HTTP/2 pool state lock") =
|
||||
Some(Instant::now() + self.idle_timeout);
|
||||
}
|
||||
self.changed.notify_waiters();
|
||||
}
|
||||
|
||||
pub fn drain(&self) {
|
||||
if self.accepting.swap(false, Ordering::AcqRel) {
|
||||
self.streams.close();
|
||||
*self.idle_deadline.lock().expect("HTTP/2 pool state lock") = None;
|
||||
self.changed.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
fn retire(&self) {
|
||||
self.drain();
|
||||
self.abort.abort();
|
||||
}
|
||||
|
||||
fn deadline(&self) -> Option<Instant> {
|
||||
*self.idle_deadline.lock().expect("HTTP/2 pool state lock")
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct Http2Pool {
|
||||
entries: Mutex<Vec<Arc<Http2Entry>>>,
|
||||
changed: Arc<Notify>,
|
||||
streams: usize,
|
||||
idle_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Http2Pool {
|
||||
pub fn new(streams: usize, idle_timeout: Duration) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
entries: Mutex::new(Vec::new()),
|
||||
changed: Arc::new(Notify::new()),
|
||||
streams,
|
||||
idle_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn take(&self) -> Option<(Arc<Http2Entry>, OwnedSemaphorePermit)> {
|
||||
let mut entries = self.entries.lock().expect("HTTP/2 pool state lock");
|
||||
entries.retain(|entry| entry.accepting.load(Ordering::Acquire));
|
||||
entries
|
||||
.iter()
|
||||
.find_map(|entry| entry.acquire().map(|permit| (Arc::clone(entry), permit)))
|
||||
}
|
||||
|
||||
pub fn add(
|
||||
&self,
|
||||
sender: hyper::client::conn::http2::SendRequest<RequestBody>,
|
||||
abort: AbortHandle,
|
||||
) -> Arc<Http2Entry> {
|
||||
let entry = Http2Entry::new(
|
||||
sender,
|
||||
self.streams,
|
||||
abort,
|
||||
Arc::clone(&self.changed),
|
||||
self.idle_timeout,
|
||||
);
|
||||
self.entries
|
||||
.lock()
|
||||
.expect("HTTP/2 pool state lock")
|
||||
.push(Arc::clone(&entry));
|
||||
self.changed.notify_waiters();
|
||||
entry
|
||||
}
|
||||
|
||||
pub fn changed(&self) -> tokio::sync::futures::OwnedNotified {
|
||||
self.changed.clone().notified_owned()
|
||||
}
|
||||
|
||||
fn expire_and_deadline(&self) -> Option<Instant> {
|
||||
let now = Instant::now();
|
||||
let mut entries = self.entries.lock().expect("HTTP/2 pool state lock");
|
||||
entries.retain(|entry| {
|
||||
if !entry.accepting.load(Ordering::Acquire) || entry.sender.is_closed() {
|
||||
entry.drain();
|
||||
return false;
|
||||
}
|
||||
if entry.active.load(Ordering::Acquire) == 0
|
||||
&& entry.deadline().is_some_and(|deadline| deadline <= now)
|
||||
{
|
||||
entry.retire();
|
||||
return false;
|
||||
}
|
||||
true
|
||||
});
|
||||
entries.iter().filter_map(|entry| entry.deadline()).min()
|
||||
}
|
||||
|
||||
fn close(&self) {
|
||||
let entries = std::mem::take(&mut *self.entries.lock().expect("HTTP/2 pool state lock"));
|
||||
for entry in entries {
|
||||
entry.retire();
|
||||
}
|
||||
self.changed.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct Http2Lease {
|
||||
_owner: Arc<Http2LeaseOwner>,
|
||||
}
|
||||
|
||||
struct Http2LeaseOwner {
|
||||
entry: Arc<Http2Entry>,
|
||||
stream: Option<OwnedSemaphorePermit>,
|
||||
request: Option<OwnedSemaphorePermit>,
|
||||
}
|
||||
|
||||
impl Http2Lease {
|
||||
pub fn new(
|
||||
entry: Arc<Http2Entry>,
|
||||
stream: OwnedSemaphorePermit,
|
||||
request: OwnedSemaphorePermit,
|
||||
) -> Self {
|
||||
Self {
|
||||
_owner: Arc::new(Http2LeaseOwner {
|
||||
entry,
|
||||
stream: Some(stream),
|
||||
request: Some(request),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Http2LeaseOwner {
|
||||
fn drop(&mut self) {
|
||||
self.stream.take();
|
||||
self.request.take();
|
||||
self.entry.release();
|
||||
}
|
||||
}
|
||||
|
||||
struct DriverOwner {
|
||||
control: Arc<Control>,
|
||||
idle: Arc<Idle>,
|
||||
http2: Option<Arc<Http2Pool>>,
|
||||
}
|
||||
|
||||
impl Drop for DriverOwner {
|
||||
fn drop(&mut self) {
|
||||
self.control.close();
|
||||
self.idle.entries.lock().expect("pool state lock").clear();
|
||||
if let Some(http2) = &self.http2 {
|
||||
http2.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn start(
|
||||
control: Arc<Control>,
|
||||
idle: Arc<Idle>,
|
||||
http2: Option<Arc<Http2Pool>>,
|
||||
capacity: usize,
|
||||
) -> Result<(mpsc::Sender<Driver>, crate::task::Task<()>), String> {
|
||||
let (queue, mut receiver) = mpsc::channel::<Driver>(capacity);
|
||||
// The parent may finish before this future is first polled.
|
||||
let owner = DriverOwner {
|
||||
control: control.clone(),
|
||||
idle: idle.clone(),
|
||||
http2: http2.clone(),
|
||||
};
|
||||
let task = crate::task::run(async move {
|
||||
let _owner = owner;
|
||||
let mut drivers = FuturesUnordered::<Driver>::new();
|
||||
loop {
|
||||
let deadline = [
|
||||
idle.deadline(),
|
||||
http2.as_ref().and_then(|pool| pool.expire_and_deadline()),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.min();
|
||||
let expire = async {
|
||||
match deadline {
|
||||
Some(deadline) => tokio::time::sleep_until(deadline).await,
|
||||
None => std::future::pending().await,
|
||||
}
|
||||
};
|
||||
let http2_changed = async {
|
||||
match &http2 {
|
||||
Some(pool) => pool.changed().await,
|
||||
None => std::future::pending().await,
|
||||
}
|
||||
};
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = control.closed.acquire() => return Ok(()),
|
||||
Some(()) = drivers.next(), if !drivers.is_empty() => {},
|
||||
_ = idle.changed.notified() => {},
|
||||
_ = http2_changed => {},
|
||||
_ = expire => {},
|
||||
next = receiver.recv() => match next {
|
||||
Some(driver) => drivers.push(driver),
|
||||
None => return Ok(()),
|
||||
},
|
||||
}
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
Ok((queue, task))
|
||||
}
|
||||
|
||||
pub(super) fn driver(
|
||||
connection: impl Future<Output = Result<(), hyper::Error>> + Send + 'static,
|
||||
permit: tokio::sync::OwnedSemaphorePermit,
|
||||
) -> (AbortHandle, Driver) {
|
||||
let (abort, registration) = AbortHandle::new_pair();
|
||||
let driver = Box::pin(async move {
|
||||
let _permit = permit;
|
||||
// Protocol/I/O failures are observed through SendRequest/Incoming.
|
||||
let _completion = Abortable::new(connection, registration).await;
|
||||
});
|
||||
(abort, driver)
|
||||
}
|
||||
|
||||
async fn drive_http2<C>(
|
||||
connection: C,
|
||||
mut owner: crate::task::ScopedOwner,
|
||||
fault: crate::task::ScopeFault,
|
||||
entry: Arc<Http2Entry>,
|
||||
registration: futures_util::future::AbortRegistration,
|
||||
) where
|
||||
C: Future<Output = Result<(), hyper::Error>> + Send + 'static,
|
||||
{
|
||||
let completion = {
|
||||
let connection = Abortable::new(owner.enter(connection), registration);
|
||||
tokio::pin!(connection);
|
||||
tokio::select! {
|
||||
result = &mut connection => Some(result),
|
||||
_ = fault.wait() => None,
|
||||
}
|
||||
};
|
||||
entry.drain();
|
||||
match completion {
|
||||
Some(Ok(Ok(()))) => {
|
||||
let _ = owner.finish().await;
|
||||
}
|
||||
Some(Ok(Err(_))) | Some(Err(_)) | None => {
|
||||
let _ = owner.stop().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn http2_handshake(
|
||||
transport: Transport,
|
||||
permit: OwnedSemaphorePermit,
|
||||
limits: Http2Limits,
|
||||
headers: usize,
|
||||
timeout: Duration,
|
||||
pool: Arc<Http2Pool>,
|
||||
) -> (Driver, oneshot::Receiver<Result<Arc<Http2Entry>, String>>) {
|
||||
let (result, receiver) = oneshot::channel();
|
||||
let driver = Box::pin(async move {
|
||||
let (mut owner, spawner, fault) = match crate::task::scoped_owner(permit) {
|
||||
Ok(owner) => owner,
|
||||
Err(error) => {
|
||||
let _ = result.send(Err(error));
|
||||
return;
|
||||
}
|
||||
};
|
||||
let handshake = {
|
||||
let builder = limits.client_builder(Http2Executor(spawner), headers);
|
||||
let handshake = tokio::time::timeout(
|
||||
timeout,
|
||||
owner.enter(builder.handshake(TokioIo::new(transport))),
|
||||
);
|
||||
tokio::pin!(handshake);
|
||||
tokio::select! {
|
||||
result = &mut handshake => match result {
|
||||
Ok(result) => result.map_err(|error| error.to_string()),
|
||||
Err(_) => Err("HTTP/2 handshake timed out".to_owned()),
|
||||
},
|
||||
error = fault.wait() => Err(error),
|
||||
}
|
||||
};
|
||||
let (sender, connection) = match handshake {
|
||||
Ok(connection) => connection,
|
||||
Err(error) => {
|
||||
let _ = result.send(Err(error));
|
||||
let _ = owner.stop().await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let (abort, registration) = AbortHandle::new_pair();
|
||||
let entry = pool.add(sender, abort);
|
||||
let _ = result.send(Ok(Arc::clone(&entry)));
|
||||
drive_http2(connection, owner, fault, entry, registration).await;
|
||||
});
|
||||
(driver, receiver)
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
use std::future::Future;
|
||||
|
||||
use hyper::rt::Executor;
|
||||
|
||||
use crate::task::ScopedSpawner;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct Http2Executor(pub(super) ScopedSpawner);
|
||||
|
||||
impl<F> Executor<F> for Http2Executor
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
fn execute(&self, future: F) {
|
||||
self.0.spawn(future);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
use std::future::Future;
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
|
||||
use hyper::service::service_fn;
|
||||
use hyper_util::rt::TokioIo;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::reply::{HttpReply, LiveLimits, Outgoing, Session};
|
||||
use super::write_timeout::WriteTimeout;
|
||||
use super::{Budget, Limits, Request, RequestContext, read_request, reject};
|
||||
use crate::{net, task};
|
||||
|
||||
/// HTTP/1 每条连接最多一个活跃 handler；升级后仍占用同一个连接和任务名额。
|
||||
pub async fn serve_live<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
limits: Limits,
|
||||
live: LiveLimits,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request, HttpReply) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<(), String>> + Send + 'static,
|
||||
{
|
||||
serve_transport(route, listener, limits, live, None).await
|
||||
}
|
||||
|
||||
pub async fn serve_live_tls<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
limits: Limits,
|
||||
live: LiveLimits,
|
||||
tls: crate::tls::ServerTls,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request, HttpReply) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<(), String>> + Send + 'static,
|
||||
{
|
||||
serve_transport(route, listener, limits, live, Some(tls)).await
|
||||
}
|
||||
|
||||
async fn serve_transport<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
limits: Limits,
|
||||
live: LiveLimits,
|
||||
tls: Option<crate::tls::ServerTls>,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request, HttpReply) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<(), String>> + Send + 'static,
|
||||
{
|
||||
let budget = limits.check()?;
|
||||
live.check()?;
|
||||
let tls = tls.map(|tls| super::protocol::server_tls(tls, budget.http2));
|
||||
task::parallel_each_stream(
|
||||
net::incoming(listener.clone(), true),
|
||||
limits.connections,
|
||||
move |connection| {
|
||||
let route = route.clone();
|
||||
let tls = tls.clone();
|
||||
let listener = listener.clone();
|
||||
async move {
|
||||
let stream = connection?;
|
||||
let stream = match crate::tls::accept(stream, tls.as_ref(), limits.timeout_ms).await
|
||||
{
|
||||
Ok(stream) => stream,
|
||||
Err(_) => return Ok(()),
|
||||
};
|
||||
if super::protocol::check_alpn(&stream, budget.http2).is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(http2) = budget.http2 {
|
||||
return serve_http2(route, listener, stream, budget, live, http2).await;
|
||||
}
|
||||
let handlers = Arc::new(Mutex::new(Some(task::Group::new(1)?)));
|
||||
let upgraded = Arc::new(AtomicBool::new(false));
|
||||
let service = {
|
||||
let handlers = handlers.clone();
|
||||
let upgraded = upgraded.clone();
|
||||
service_fn(move |request| {
|
||||
dispatch(
|
||||
request,
|
||||
route.clone(),
|
||||
budget,
|
||||
live,
|
||||
handlers.clone(),
|
||||
upgraded.clone(),
|
||||
)
|
||||
})
|
||||
};
|
||||
let connection = super::server_builder(budget)
|
||||
.serve_connection(
|
||||
TokioIo::new(WriteTimeout::new(stream, budget.timeout)),
|
||||
service,
|
||||
)
|
||||
.with_upgrades();
|
||||
tokio::pin!(connection);
|
||||
let result = tokio::select! {
|
||||
result = &mut connection => result,
|
||||
_ = listener.closed() => {
|
||||
connection.as_mut().graceful_shutdown();
|
||||
connection.await
|
||||
}
|
||||
};
|
||||
let group = handlers
|
||||
.lock()
|
||||
.await
|
||||
.take()
|
||||
.expect("connection owns handlers");
|
||||
if result.is_ok() && upgraded.load(Ordering::Acquire) {
|
||||
group.wait().await?;
|
||||
} else {
|
||||
group.stop().await?;
|
||||
}
|
||||
super::connection_result(result)
|
||||
}
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn serve_http2<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
stream: crate::transport::Transport,
|
||||
budget: Budget,
|
||||
live: LiveLimits,
|
||||
http2: super::Http2Limits,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request, HttpReply) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<(), String>> + Send + 'static,
|
||||
{
|
||||
let (mut owner, spawner, fault) = task::scoped_owner(())?;
|
||||
let first_request = super::server_h2::FirstRequest::new();
|
||||
let service = {
|
||||
let first_request = first_request.clone();
|
||||
service_fn(move |request| {
|
||||
first_request.mark();
|
||||
let route = route.clone();
|
||||
dispatch_http2(request, route, budget, live)
|
||||
})
|
||||
};
|
||||
let mut builder = http2.server_builder(super::executor::Http2Executor(spawner), budget.headers);
|
||||
super::server_h2::configure_liveness(&mut builder, budget.timeout);
|
||||
let connection = builder.serve_connection(
|
||||
hyper_util::rt::TokioIo::new(WriteTimeout::new(stream, budget.timeout)),
|
||||
service,
|
||||
);
|
||||
let result = owner
|
||||
.enter(super::server_h2::drive(
|
||||
connection,
|
||||
&listener,
|
||||
&first_request,
|
||||
budget.timeout,
|
||||
&fault,
|
||||
|connection| connection.graceful_shutdown(),
|
||||
))
|
||||
.await;
|
||||
super::server_h2::complete(owner, result).await
|
||||
}
|
||||
|
||||
async fn dispatch<F, Fut>(
|
||||
mut request: hyper::Request<hyper::body::Incoming>,
|
||||
route: F,
|
||||
budget: Budget,
|
||||
live: LiveLimits,
|
||||
handlers: Arc<Mutex<Option<task::Group>>>,
|
||||
upgraded: Arc<AtomicBool>,
|
||||
) -> Result<hyper::Response<Outgoing>, std::convert::Infallible>
|
||||
where
|
||||
F: Fn(Request, HttpReply) -> Fut + Send + 'static,
|
||||
Fut: Future<Output = Result<(), String>> + Send + 'static,
|
||||
{
|
||||
let upgrade = if request.headers().contains_key(hyper::header::UPGRADE) {
|
||||
let upgrade = hyper::upgrade::on(&mut request);
|
||||
let mut handshake = hyper::Request::new(());
|
||||
*handshake.method_mut() = request.method().clone();
|
||||
*handshake.uri_mut() = request.uri().clone();
|
||||
*handshake.version_mut() = request.version();
|
||||
*handshake.headers_mut() = request.headers().clone();
|
||||
Some((handshake, upgrade))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let request = match read_request(request, budget).await {
|
||||
Ok(request) => request,
|
||||
Err(status) => return Ok(reject(status, budget).map(Outgoing::Full)),
|
||||
};
|
||||
let context = RequestContext::new("http/1.1", request.target.clone(), None);
|
||||
let (session, head) = Session::new(budget, live, upgrade, upgraded);
|
||||
let reply = session.reply();
|
||||
let state = reply.clone();
|
||||
let handler_context = context.clone();
|
||||
let started = handlers
|
||||
.lock()
|
||||
.await
|
||||
.as_mut()
|
||||
.expect("active connection")
|
||||
.run(async move {
|
||||
let result = session.run(route(request, reply)).await;
|
||||
if let Err(error) = result {
|
||||
handler_context.report(error);
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await;
|
||||
if let Err(error) = started {
|
||||
context.report(error);
|
||||
return Ok(reject(hyper::StatusCode::INTERNAL_SERVER_ERROR, budget).map(Outgoing::Full));
|
||||
}
|
||||
match tokio::time::timeout(budget.timeout, head).await {
|
||||
Ok(Ok(response)) => Ok(response),
|
||||
Ok(Err(_)) => {
|
||||
if state.fault().is_none() {
|
||||
context.report("HTTP live handler ended without a response");
|
||||
}
|
||||
Ok(reject(hyper::StatusCode::INTERNAL_SERVER_ERROR, budget).map(Outgoing::Full))
|
||||
}
|
||||
Err(_) => Ok(reject(hyper::StatusCode::GATEWAY_TIMEOUT, budget).map(Outgoing::Full)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn dispatch_http2<F, Fut>(
|
||||
request: hyper::Request<hyper::body::Incoming>,
|
||||
route: F,
|
||||
budget: Budget,
|
||||
live: LiveLimits,
|
||||
) -> Result<hyper::Response<Outgoing>, std::convert::Infallible>
|
||||
where
|
||||
F: Fn(Request, HttpReply) -> Fut + Send + 'static,
|
||||
Fut: Future<Output = Result<(), String>> + Send + 'static,
|
||||
{
|
||||
if request.method() == hyper::Method::CONNECT
|
||||
|| request.headers().contains_key(hyper::header::UPGRADE)
|
||||
{
|
||||
return Ok(reject(hyper::StatusCode::NOT_IMPLEMENTED, budget).map(Outgoing::Full));
|
||||
}
|
||||
let request = match read_request(request, budget).await {
|
||||
Ok(request) => request,
|
||||
Err(status) => return Ok(reject(status, budget).map(Outgoing::Full)),
|
||||
};
|
||||
let context = RequestContext::new("http/2", request.target.clone(), None);
|
||||
let (session, head) = Session::new(budget, live, None, Arc::new(AtomicBool::new(false)));
|
||||
let reply = session.reply();
|
||||
let state = reply.clone();
|
||||
let mut handler = match task::Group::new(1) {
|
||||
Ok(handler) => handler,
|
||||
Err(error) => {
|
||||
context.report(error);
|
||||
return Ok(reject(hyper::StatusCode::INTERNAL_SERVER_ERROR, budget).map(Outgoing::Full));
|
||||
}
|
||||
};
|
||||
let handler_context = context.clone();
|
||||
let started = handler
|
||||
.run(async move {
|
||||
let result = session.run(route(request, reply)).await;
|
||||
if let Err(error) = result {
|
||||
handler_context.report(error);
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await;
|
||||
if let Err(error) = started {
|
||||
context.report(error);
|
||||
return Ok(reject(hyper::StatusCode::INTERNAL_SERVER_ERROR, budget).map(Outgoing::Full));
|
||||
}
|
||||
match tokio::time::timeout(budget.timeout, head).await {
|
||||
Ok(Ok(response)) => Ok(response.map(|body| body.with_handler(handler))),
|
||||
Ok(Err(_)) => {
|
||||
if let Err(error) = handler.wait().await {
|
||||
context.report(error);
|
||||
} else if state.fault().is_none() {
|
||||
context.report("HTTP live handler ended without a response");
|
||||
}
|
||||
Ok(reject(hyper::StatusCode::INTERNAL_SERVER_ERROR, budget).map(Outgoing::Full))
|
||||
}
|
||||
Err(_) => match handler.stop().await {
|
||||
Ok(()) => Ok(reject(hyper::StatusCode::GATEWAY_TIMEOUT, budget).map(Outgoing::Full)),
|
||||
Err(error) => {
|
||||
context.report(error);
|
||||
Ok(reject(hyper::StatusCode::GATEWAY_TIMEOUT, budget).map(Outgoing::Full))
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
use http_body_util::{BodyExt, Full, LengthLimitError, Limited};
|
||||
use hyper::body::{Body, Incoming};
|
||||
use hyper::header::{CONTENT_LENGTH, HOST, TRANSFER_ENCODING, UPGRADE};
|
||||
use hyper::{HeaderMap, Method, StatusCode, Uri};
|
||||
|
||||
use super::{Budget, Header, Request, Response};
|
||||
use crate::bytes::Bytes;
|
||||
|
||||
#[derive(Clone, Copy, Eq, PartialEq)]
|
||||
enum MessageKind {
|
||||
Request,
|
||||
Response,
|
||||
}
|
||||
|
||||
pub(super) fn header_size(headers: &HeaderMap) -> usize {
|
||||
headers
|
||||
.iter()
|
||||
.map(|(name, value)| name.as_str().len() + value.as_bytes().len() + 4)
|
||||
.sum::<usize>()
|
||||
+ 2
|
||||
}
|
||||
|
||||
pub(super) fn decode_headers(headers: HeaderMap) -> Vec<Header> {
|
||||
headers
|
||||
.iter()
|
||||
.map(|(name, value)| Header {
|
||||
name: name.to_string(),
|
||||
value: Bytes::new(value.as_bytes().to_vec()),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn encode_headers(
|
||||
headers: Vec<Header>,
|
||||
budget: Budget,
|
||||
kind: MessageKind,
|
||||
) -> Result<HeaderMap, String> {
|
||||
let mut result = HeaderMap::new();
|
||||
let mut length = 2;
|
||||
for header in headers {
|
||||
length += header.name.len() + header.value.values().len() + 4;
|
||||
if length > budget.headers {
|
||||
return Err("HTTP headers exceed header_bytes".into());
|
||||
}
|
||||
let name = hyper::header::HeaderName::from_bytes(header.name.as_bytes())
|
||||
.map_err(|error| error.to_string())?;
|
||||
if name == CONTENT_LENGTH || name == TRANSFER_ENCODING {
|
||||
return Err("HTTP framing headers are owned by the engine".into());
|
||||
}
|
||||
if name == UPGRADE {
|
||||
return Err("HTTP upgrades are not supported yet".into());
|
||||
}
|
||||
let value = hyper::header::HeaderValue::from_maybe_shared(header.value.into_http())
|
||||
.map_err(|error| error.to_string())?;
|
||||
if budget.http2.is_some() {
|
||||
check_http2_header(name.as_str(), value.as_bytes(), kind)?;
|
||||
}
|
||||
result
|
||||
.try_append(name, value)
|
||||
.map_err(|error| error.to_string())?;
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn check_http2_header(name: &str, value: &[u8], kind: MessageKind) -> Result<(), String> {
|
||||
if matches!(
|
||||
name,
|
||||
"connection" | "keep-alive" | "proxy-connection" | "transfer-encoding" | "upgrade"
|
||||
) || (name == "te"
|
||||
&& (kind == MessageKind::Response || !value.eq_ignore_ascii_case(b"trailers")))
|
||||
{
|
||||
return Err(format!("HTTP/2 forbids connection-specific header {name}"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Source 继续使用 origin-form；scheme/authority 只在 HTTP/2 线协议边界补齐。
|
||||
pub(super) fn http2_request<B>(
|
||||
request: &mut hyper::Request<B>,
|
||||
scheme: &str,
|
||||
) -> Result<(), String> {
|
||||
let host = request
|
||||
.headers()
|
||||
.get(HOST)
|
||||
.ok_or("HTTP request requires Host")?
|
||||
.to_str()
|
||||
.map_err(|error| error.to_string())?;
|
||||
let authority = host
|
||||
.parse::<hyper::http::uri::Authority>()
|
||||
.map_err(|error| error.to_string())?;
|
||||
if authority.as_str().contains('@') {
|
||||
return Err("HTTP/2 authority cannot contain credentials".into());
|
||||
}
|
||||
let mut uri = request.uri().clone().into_parts();
|
||||
uri.scheme = Some(
|
||||
scheme
|
||||
.parse::<hyper::http::uri::Scheme>()
|
||||
.map_err(|error| error.to_string())?,
|
||||
);
|
||||
uri.authority = Some(authority);
|
||||
*request.uri_mut() = Uri::from_parts(uri).map_err(|error| error.to_string())?;
|
||||
*request.version_mut() = hyper::Version::HTTP_2;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn check_message(
|
||||
headers: &HeaderMap,
|
||||
body: &Bytes,
|
||||
start_bytes: usize,
|
||||
budget: Budget,
|
||||
) -> Result<(), String> {
|
||||
// 预留引擎生成的 Content-Length 和 Date；不截断用户内容。
|
||||
if header_size(headers) + start_bytes + 80 > budget.headers {
|
||||
return Err("HTTP headers exceed header_bytes".into());
|
||||
}
|
||||
if body.values().len() > budget.body {
|
||||
return Err("HTTP body exceeds body_bytes".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn encode_request(
|
||||
request: Request,
|
||||
budget: Budget,
|
||||
) -> Result<hyper::Request<Full<bytes::Bytes>>, String> {
|
||||
let method =
|
||||
Method::from_bytes(request.method.as_bytes()).map_err(|error| error.to_string())?;
|
||||
let uri = request
|
||||
.target
|
||||
.parse::<Uri>()
|
||||
.map_err(|error| error.to_string())?;
|
||||
if method == Method::CONNECT
|
||||
|| uri.scheme().is_some()
|
||||
|| uri.authority().is_some()
|
||||
|| !(request.target.starts_with('/')
|
||||
|| (method == Method::OPTIONS && request.target == "*"))
|
||||
{
|
||||
return Err(
|
||||
"HTTP send requires an origin-form target or OPTIONS *; CONNECT is not supported"
|
||||
.into(),
|
||||
);
|
||||
}
|
||||
let headers = encode_headers(request.headers, budget, MessageKind::Request)?;
|
||||
if headers.get_all(HOST).iter().count() != 1 {
|
||||
return Err("HTTP request requires exactly one Host header".into());
|
||||
}
|
||||
let start_bytes = request.method.len() + request.target.len() + 14;
|
||||
check_message(&headers, &request.body, start_bytes, budget)?;
|
||||
let mut result = hyper::Request::new(Full::new(request.body.into_http()));
|
||||
*result.method_mut() = method;
|
||||
*result.uri_mut() = uri;
|
||||
*result.headers_mut() = headers;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub(super) fn encode_response(
|
||||
response: Response,
|
||||
budget: Budget,
|
||||
) -> Result<hyper::Response<Full<bytes::Bytes>>, String> {
|
||||
let status = u16::try_from(response.status)
|
||||
.ok()
|
||||
.filter(|status| (200..=599).contains(status))
|
||||
.ok_or("HTTP response status must be between 200 and 599")?;
|
||||
let status = StatusCode::from_u16(status).map_err(|error| error.to_string())?;
|
||||
let no_body = matches!(
|
||||
status,
|
||||
StatusCode::NO_CONTENT | StatusCode::RESET_CONTENT | StatusCode::NOT_MODIFIED
|
||||
);
|
||||
if no_body && !response.body.values().is_empty() {
|
||||
return Err("HTTP 204, 205 and 304 responses cannot contain a body".into());
|
||||
}
|
||||
let headers = encode_headers(response.headers, budget, MessageKind::Response)?;
|
||||
check_message(&headers, &response.body, 32, budget)?;
|
||||
let mut result = hyper::Response::new(Full::new(response.body.into_http()));
|
||||
*result.status_mut() = status;
|
||||
*result.headers_mut() = headers;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub(super) async fn collect_body(body: Incoming, limit: usize) -> Result<Bytes, StatusCode> {
|
||||
if body.size_hint().lower() > limit as u64 {
|
||||
return Err(StatusCode::PAYLOAD_TOO_LARGE);
|
||||
}
|
||||
let mut incoming = Limited::new(body, limit);
|
||||
let mut body = Bytes::new(Vec::new());
|
||||
// 单帧直接共享，多帧逐步扩容；不为大量微小帧保留一份元数据队列。
|
||||
while let Some(frame) = incoming.frame().await {
|
||||
let frame = frame.map_err(|error| {
|
||||
if error.is::<LengthLimitError>() {
|
||||
StatusCode::PAYLOAD_TOO_LARGE
|
||||
} else {
|
||||
StatusCode::BAD_REQUEST
|
||||
}
|
||||
})?;
|
||||
match frame.into_data() {
|
||||
Ok(bytes) => body = body.concat(&Bytes::from_http(bytes)),
|
||||
Err(frame)
|
||||
if frame
|
||||
.trailers_ref()
|
||||
.is_some_and(|trailers| !trailers.is_empty()) =>
|
||||
{
|
||||
return Err(StatusCode::NOT_IMPLEMENTED);
|
||||
}
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
Ok(body)
|
||||
}
|
||||
@@ -0,0 +1,550 @@
|
||||
mod client;
|
||||
mod client_body;
|
||||
mod client_driver;
|
||||
mod executor;
|
||||
mod live;
|
||||
mod message;
|
||||
mod protocol;
|
||||
pub(crate) mod reply;
|
||||
mod server_h2;
|
||||
mod write_timeout;
|
||||
|
||||
pub use protocol::Http2Limits;
|
||||
|
||||
pub use client::{
|
||||
HttpClient, PoolLimits, StreamRequest, StreamResponse, client, close_client, open, open_stream,
|
||||
request,
|
||||
};
|
||||
pub use live::{serve_live, serve_live_tls};
|
||||
pub use reply::{HttpReply, LiveLimits, finish, respond, start, write};
|
||||
|
||||
use std::convert::Infallible;
|
||||
use std::future::Future;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use http_body_util::Full;
|
||||
use hyper::body::Incoming;
|
||||
use hyper::service::service_fn;
|
||||
use hyper::{Method, StatusCode};
|
||||
use hyper_util::rt::{TokioIo, TokioTimer};
|
||||
|
||||
use crate::bytes::Bytes;
|
||||
use crate::{net, task};
|
||||
use message::{collect_body, decode_headers, encode_request, encode_response, header_size};
|
||||
use write_timeout::WriteTimeout;
|
||||
|
||||
static NEXT_REQUEST_ID: AtomicU64 = AtomicU64::new(1);
|
||||
tokio::task_local! { static CURRENT_REQUEST: RequestContext; }
|
||||
|
||||
#[cfg(feature = "wire")]
|
||||
pub(crate) fn request_id() -> u64 {
|
||||
CURRENT_REQUEST
|
||||
.try_with(|request| request.id)
|
||||
.unwrap_or_else(|_| NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed))
|
||||
}
|
||||
|
||||
#[cfg(feature = "wire")]
|
||||
pub(crate) fn client_address() -> Option<String> {
|
||||
CURRENT_REQUEST
|
||||
.try_with(|request| request.client_address.clone())
|
||||
.ok()
|
||||
.flatten()
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RequestContext {
|
||||
id: u64,
|
||||
protocol: &'static str,
|
||||
target: String,
|
||||
client_address: Option<String>,
|
||||
}
|
||||
|
||||
impl RequestContext {
|
||||
fn new(protocol: &'static str, target: String, client_address: Option<String>) -> Self {
|
||||
Self {
|
||||
id: NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed),
|
||||
protocol,
|
||||
target,
|
||||
client_address,
|
||||
}
|
||||
}
|
||||
|
||||
fn report(&self, error: impl Into<String>) {
|
||||
crate::log::error(
|
||||
error,
|
||||
vec![
|
||||
crate::log::Field {
|
||||
name: "boundary".into(),
|
||||
value: "http".into(),
|
||||
},
|
||||
crate::log::Field {
|
||||
name: "request_id".into(),
|
||||
value: self.id.to_string(),
|
||||
},
|
||||
crate::log::Field {
|
||||
name: "protocol".into(),
|
||||
value: self.protocol.into(),
|
||||
},
|
||||
crate::log::Field {
|
||||
name: "target".into(),
|
||||
value: self.target.clone(),
|
||||
},
|
||||
],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Header {
|
||||
pub name: String,
|
||||
pub value: Bytes,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Request {
|
||||
pub method: String,
|
||||
pub target: String,
|
||||
pub headers: Vec<Header>,
|
||||
pub body: Bytes,
|
||||
}
|
||||
|
||||
/// Server-only body ownership. Source-level HTTP values remain bounded Request values.
|
||||
pub struct IncomingRequest {
|
||||
pub head: Request,
|
||||
pub body: Incoming,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Response {
|
||||
pub status: i64,
|
||||
pub headers: Vec<Header>,
|
||||
pub body: Bytes,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct Limits {
|
||||
pub header_bytes: i64,
|
||||
pub body_bytes: i64,
|
||||
pub timeout_ms: i64,
|
||||
pub connections: i64,
|
||||
pub http2: Option<Http2Limits>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct Budget {
|
||||
headers: usize,
|
||||
body: usize,
|
||||
timeout: Duration,
|
||||
http2: Option<Http2Limits>,
|
||||
}
|
||||
|
||||
impl Limits {
|
||||
fn check(self) -> Result<Budget, String> {
|
||||
let headers = usize::try_from(self.header_bytes)
|
||||
.ok()
|
||||
.filter(|value| *value >= 8192)
|
||||
.ok_or("HTTP header_bytes must be at least 8192")?;
|
||||
let body =
|
||||
usize::try_from(self.body_bytes).map_err(|_| "HTTP body_bytes must be nonnegative")?;
|
||||
let millis = u64::try_from(self.timeout_ms)
|
||||
.ok()
|
||||
.filter(|value| *value > 0)
|
||||
.ok_or("HTTP timeout_ms must be positive")?;
|
||||
if !(1..=65536).contains(&self.connections) {
|
||||
return Err("HTTP connections must be between 1 and 65536".into());
|
||||
}
|
||||
if let Some(http2) = self.http2 {
|
||||
http2.check(self.connections, headers)?;
|
||||
}
|
||||
Ok(Budget {
|
||||
headers,
|
||||
body,
|
||||
timeout: Duration::from_millis(millis),
|
||||
http2: self.http2,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 连接由现有任务组持有，取消 serve 会沿 Scope 排空连接和 handler 子任务。
|
||||
pub async fn serve<F, Fut>(route: F, listener: net::Listener, limits: Limits) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<Response, String>> + Send + 'static,
|
||||
{
|
||||
serve_buffered(route, listener, limits, None).await
|
||||
}
|
||||
|
||||
pub async fn serve_tls<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
limits: Limits,
|
||||
tls: crate::tls::ServerTls,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<Response, String>> + Send + 'static,
|
||||
{
|
||||
serve_buffered(route, listener, limits, Some(tls)).await
|
||||
}
|
||||
|
||||
async fn serve_buffered<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
limits: Limits,
|
||||
tls: Option<crate::tls::ServerTls>,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(Request) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<Response, String>> + Send + 'static,
|
||||
{
|
||||
let budget = limits.check()?;
|
||||
serve_transport(
|
||||
move |incoming: IncomingRequest| {
|
||||
let route = route.clone();
|
||||
async move {
|
||||
let mut request = incoming.head;
|
||||
request.body = match tokio::time::timeout(
|
||||
budget.timeout,
|
||||
collect_body(incoming.body, budget.body),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(body)) => body,
|
||||
Ok(Err(status)) => {
|
||||
return Ok(Response {
|
||||
status: status.as_u16().into(),
|
||||
headers: Vec::new(),
|
||||
body: Bytes::new(Vec::new()),
|
||||
});
|
||||
}
|
||||
Err(_) => {
|
||||
return Ok(Response {
|
||||
status: 408,
|
||||
headers: Vec::new(),
|
||||
body: Bytes::new(Vec::new()),
|
||||
});
|
||||
}
|
||||
};
|
||||
route(request).await
|
||||
}
|
||||
},
|
||||
listener,
|
||||
limits,
|
||||
tls,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn serve_incoming<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
limits: Limits,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(IncomingRequest) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<Response, String>> + Send + 'static,
|
||||
{
|
||||
serve_transport(route, listener, limits, None).await
|
||||
}
|
||||
|
||||
async fn serve_transport<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
limits: Limits,
|
||||
tls: Option<crate::tls::ServerTls>,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(IncomingRequest) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<Response, String>> + Send + 'static,
|
||||
{
|
||||
let budget = limits.check()?;
|
||||
let tls = tls.map(|tls| protocol::server_tls(tls, budget.http2));
|
||||
task::parallel_each_stream(
|
||||
net::incoming(listener.clone(), true),
|
||||
limits.connections,
|
||||
move |connection| {
|
||||
let route = route.clone();
|
||||
let tls = tls.clone();
|
||||
let listener = listener.clone();
|
||||
async move {
|
||||
let stream = connection?;
|
||||
let client_address = stream
|
||||
.peer_addr()
|
||||
.ok()
|
||||
.map(|address| address.ip().to_string());
|
||||
let stream = match crate::tls::accept(stream, tls.as_ref(), limits.timeout_ms).await
|
||||
{
|
||||
Ok(stream) => stream,
|
||||
Err(_) => return Ok(()), // A peer's TLS failure does not terminate the listener.
|
||||
};
|
||||
if protocol::check_alpn(&stream, budget.http2).is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(http2) = budget.http2 {
|
||||
return serve_http2(route, listener, stream, budget, http2, client_address)
|
||||
.await;
|
||||
}
|
||||
let service = service_fn(move |request| {
|
||||
dispatch(request, route.clone(), budget, client_address.clone())
|
||||
});
|
||||
let connection = server_builder(budget).serve_connection(
|
||||
TokioIo::new(WriteTimeout::new(stream, budget.timeout)),
|
||||
service,
|
||||
);
|
||||
tokio::pin!(connection);
|
||||
let result = tokio::select! {
|
||||
result = &mut connection => result,
|
||||
_ = listener.closed() => {
|
||||
connection.as_mut().graceful_shutdown();
|
||||
connection.await
|
||||
}
|
||||
};
|
||||
connection_result(result)
|
||||
}
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn serve_http2<F, Fut>(
|
||||
route: F,
|
||||
listener: net::Listener,
|
||||
stream: crate::transport::Transport,
|
||||
budget: Budget,
|
||||
http2: Http2Limits,
|
||||
client_address: Option<String>,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(IncomingRequest) -> Fut + Clone + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<Response, String>> + Send + 'static,
|
||||
{
|
||||
let (mut owner, spawner, fault) = task::scoped_owner(())?;
|
||||
let first_request = server_h2::FirstRequest::new();
|
||||
let service = {
|
||||
let first_request = first_request.clone();
|
||||
service_fn(move |request| {
|
||||
first_request.mark();
|
||||
let route = route.clone();
|
||||
dispatch(request, route, budget, client_address.clone())
|
||||
})
|
||||
};
|
||||
let mut builder = http2.server_builder(executor::Http2Executor(spawner), budget.headers);
|
||||
server_h2::configure_liveness(&mut builder, budget.timeout);
|
||||
let connection = builder.serve_connection(
|
||||
TokioIo::new(WriteTimeout::new(stream, budget.timeout)),
|
||||
service,
|
||||
);
|
||||
let result = owner
|
||||
.enter(server_h2::drive(
|
||||
connection,
|
||||
&listener,
|
||||
&first_request,
|
||||
budget.timeout,
|
||||
&fault,
|
||||
|connection| connection.graceful_shutdown(),
|
||||
))
|
||||
.await;
|
||||
server_h2::complete(owner, result).await
|
||||
}
|
||||
|
||||
fn server_builder(budget: Budget) -> hyper::server::conn::http1::Builder {
|
||||
let mut builder = hyper::server::conn::http1::Builder::new();
|
||||
builder
|
||||
.timer(TokioTimer::new())
|
||||
.header_read_timeout(budget.timeout)
|
||||
.max_buf_size(budget.headers);
|
||||
builder
|
||||
}
|
||||
|
||||
fn connection_result(result: Result<(), hyper::Error>) -> Result<(), String> {
|
||||
// 对端断开、畸形报文和 I/O 超时只结束该连接；请求故障已在 dispatch 边界处理。
|
||||
let _ = result;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn dispatch<F, Fut>(
|
||||
request: hyper::Request<Incoming>,
|
||||
route: F,
|
||||
budget: Budget,
|
||||
client_address: Option<String>,
|
||||
) -> Result<hyper::Response<Full<bytes::Bytes>>, Infallible>
|
||||
where
|
||||
F: Fn(IncomingRequest) -> Fut,
|
||||
Fut: Future<Output = Result<Response, String>>,
|
||||
{
|
||||
if request.method() == Method::CONNECT || request.headers().contains_key(hyper::header::UPGRADE)
|
||||
{
|
||||
return Ok(reject(StatusCode::NOT_IMPLEMENTED, budget));
|
||||
}
|
||||
let request = match read_incoming(request, budget) {
|
||||
Ok(request) => request,
|
||||
Err(status) => return Ok(reject(status, budget)),
|
||||
};
|
||||
let context = RequestContext::new(
|
||||
if budget.http2.is_some() {
|
||||
"http/2"
|
||||
} else {
|
||||
"http/1.1"
|
||||
},
|
||||
request.head.target.clone(),
|
||||
client_address,
|
||||
);
|
||||
let response = match tokio::time::timeout(
|
||||
budget.timeout,
|
||||
CURRENT_REQUEST.scope(context.clone(), route(request)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(response)) => response,
|
||||
Ok(Err(error)) => {
|
||||
context.report(error);
|
||||
return Ok(reject(StatusCode::INTERNAL_SERVER_ERROR, budget));
|
||||
}
|
||||
Err(_) => return Ok(reject(StatusCode::GATEWAY_TIMEOUT, budget)),
|
||||
};
|
||||
match encode_response(response, budget) {
|
||||
Ok(response) => Ok(response),
|
||||
Err(error) => {
|
||||
context.report(error);
|
||||
Ok(reject(StatusCode::INTERNAL_SERVER_ERROR, budget))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_request(
|
||||
request: hyper::Request<Incoming>,
|
||||
budget: Budget,
|
||||
) -> Result<Request, StatusCode> {
|
||||
let incoming = read_incoming(request, budget)?;
|
||||
let mut request = incoming.head;
|
||||
request.body = tokio::time::timeout(budget.timeout, collect_body(incoming.body, budget.body))
|
||||
.await
|
||||
.map_err(|_| StatusCode::REQUEST_TIMEOUT)??;
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
fn read_incoming(
|
||||
mut request: hyper::Request<Incoming>,
|
||||
budget: Budget,
|
||||
) -> Result<IncomingRequest, StatusCode> {
|
||||
if request.method() == Method::CONNECT {
|
||||
return Err(StatusCode::NOT_IMPLEMENTED);
|
||||
}
|
||||
let host_count = request
|
||||
.headers()
|
||||
.get_all(hyper::header::HOST)
|
||||
.iter()
|
||||
.count();
|
||||
let target = if request.version() == hyper::Version::HTTP_2 {
|
||||
let authority = request
|
||||
.uri()
|
||||
.authority()
|
||||
.ok_or(StatusCode::BAD_REQUEST)?
|
||||
.as_str()
|
||||
.to_owned();
|
||||
if host_count > 1 {
|
||||
return Err(StatusCode::BAD_REQUEST);
|
||||
}
|
||||
if let Some(host) = request.headers().get(hyper::header::HOST) {
|
||||
if !host.as_bytes().eq_ignore_ascii_case(authority.as_bytes()) {
|
||||
return Err(StatusCode::BAD_REQUEST);
|
||||
}
|
||||
} else {
|
||||
let host = hyper::header::HeaderValue::from_str(&authority)
|
||||
.map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
request.headers_mut().insert(hyper::header::HOST, host);
|
||||
}
|
||||
request
|
||||
.uri()
|
||||
.path_and_query()
|
||||
.map(|value| value.as_str())
|
||||
.unwrap_or("/")
|
||||
.to_owned()
|
||||
} else {
|
||||
if request.version() == hyper::Version::HTTP_11 && host_count != 1 {
|
||||
return Err(StatusCode::BAD_REQUEST);
|
||||
}
|
||||
request.uri().to_string()
|
||||
};
|
||||
let start_bytes = request.method().as_str().len() + target.len() + 14;
|
||||
if header_size(request.headers()) + start_bytes > budget.headers {
|
||||
return Err(StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE);
|
||||
}
|
||||
let (head, body) = request.into_parts();
|
||||
Ok(IncomingRequest {
|
||||
head: Request {
|
||||
method: head.method.to_string(),
|
||||
target,
|
||||
headers: decode_headers(head.headers),
|
||||
body: Bytes::new(Vec::new()),
|
||||
},
|
||||
body,
|
||||
})
|
||||
}
|
||||
|
||||
fn reject(status: StatusCode, budget: Budget) -> hyper::Response<Full<bytes::Bytes>> {
|
||||
let mut response = hyper::Response::new(Full::new(bytes::Bytes::new()));
|
||||
*response.status_mut() = status;
|
||||
if budget.http2.is_none() {
|
||||
response.headers_mut().insert(
|
||||
hyper::header::CONNECTION,
|
||||
hyper::header::HeaderValue::from_static("close"),
|
||||
);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
/// 单次请求独占驱动 future；完成、错误和取消都直接释放连接，不产生后台任务。
|
||||
pub async fn send(
|
||||
address: &str,
|
||||
port: i64,
|
||||
request: Request,
|
||||
limits: Limits,
|
||||
) -> Result<Response, String> {
|
||||
let budget = limits.check()?;
|
||||
let request = encode_request(request, budget)?;
|
||||
if budget.http2.is_some() {
|
||||
return client::send_http2(address, port, request, budget).await;
|
||||
}
|
||||
tokio::time::timeout(budget.timeout, async {
|
||||
let stream = net::connect_stream_timeout(address, port, limits.timeout_ms).await?;
|
||||
let (mut sender, connection) = hyper::client::conn::http1::Builder::new()
|
||||
.max_buf_size(budget.headers)
|
||||
.handshake(TokioIo::new(stream))
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
let exchange = async {
|
||||
let response = sender
|
||||
.send_request(request)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
if response.status() == StatusCode::SWITCHING_PROTOCOLS {
|
||||
return Err("HTTP upgrades are not supported by send".into());
|
||||
}
|
||||
if header_size(response.headers()) + 32 > budget.headers {
|
||||
return Err("HTTP response headers exceed header_bytes".into());
|
||||
}
|
||||
let (head, body) = response.into_parts();
|
||||
let body = collect_body(body, budget.body)
|
||||
.await
|
||||
.map_err(|status| format!("HTTP response body rejected: {status}"))?;
|
||||
Ok(Response {
|
||||
status: i64::from(head.status.as_u16()),
|
||||
headers: decode_headers(head.headers),
|
||||
body,
|
||||
})
|
||||
};
|
||||
tokio::pin!(exchange, connection);
|
||||
tokio::select! {
|
||||
result = &mut exchange => result,
|
||||
result = &mut connection => {
|
||||
result.map_err(|error| error.to_string())?;
|
||||
// Connection: close 的最后一帧可能已入队，继续取完响应。
|
||||
exchange.await
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "HTTP request timed out".to_owned())?
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
use crate::{
|
||||
tls::{ClientTls, ServerTls},
|
||||
transport::Transport,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct Http2Limits {
|
||||
pub streams: i64,
|
||||
pub stream_window_bytes: i64,
|
||||
pub connection_window_bytes: i64,
|
||||
}
|
||||
|
||||
impl Http2Limits {
|
||||
pub(super) fn check(self, connections: i64, headers: usize) -> Result<(), String> {
|
||||
if !(1..=65536).contains(&self.streams) || connections * self.streams > 65536 {
|
||||
return Err("HTTP/2 connections * streams must be between 1 and 65536".into());
|
||||
}
|
||||
if !(1..=i32::MAX as i64).contains(&self.stream_window_bytes) {
|
||||
return Err("HTTP/2 stream_window_bytes must be between 1 and 2147483647".into());
|
||||
}
|
||||
if !(65535..=i32::MAX as i64).contains(&self.connection_window_bytes) {
|
||||
return Err(
|
||||
"HTTP/2 connection_window_bytes must be between 65535 and 2147483647".into(),
|
||||
);
|
||||
}
|
||||
if u32::try_from(headers).is_err() {
|
||||
return Err("HTTP/2 header_bytes must fit an unsigned 32-bit limit".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn client_builder<E: Clone>(
|
||||
self,
|
||||
executor: E,
|
||||
headers: usize,
|
||||
) -> hyper::client::conn::http2::Builder<E> {
|
||||
let mut builder = hyper::client::conn::http2::Builder::new(executor);
|
||||
builder
|
||||
.timer(hyper_util::rt::TokioTimer::new())
|
||||
.adaptive_window(false)
|
||||
.initial_stream_window_size(self.stream_window_bytes as u32)
|
||||
.initial_connection_window_size(self.connection_window_bytes as u32)
|
||||
.initial_max_send_streams(self.streams as usize)
|
||||
.max_concurrent_streams(0)
|
||||
.max_header_list_size(headers as u32)
|
||||
.max_frame_size(FRAME_BYTES)
|
||||
.max_send_buf_size(FRAME_BYTES as usize);
|
||||
builder
|
||||
}
|
||||
|
||||
pub(super) fn server_builder<E>(
|
||||
self,
|
||||
executor: E,
|
||||
headers: usize,
|
||||
) -> hyper::server::conn::http2::Builder<E> {
|
||||
let mut builder = hyper::server::conn::http2::Builder::new(executor);
|
||||
builder
|
||||
.timer(hyper_util::rt::TokioTimer::new())
|
||||
.adaptive_window(false)
|
||||
.initial_stream_window_size(self.stream_window_bytes as u32)
|
||||
.initial_connection_window_size(self.connection_window_bytes as u32)
|
||||
.max_concurrent_streams(self.streams as u32)
|
||||
.max_header_list_size(headers as u32)
|
||||
.max_frame_size(FRAME_BYTES)
|
||||
.max_send_buf_size(FRAME_BYTES as usize);
|
||||
builder
|
||||
}
|
||||
}
|
||||
|
||||
const FRAME_BYTES: u32 = 16384;
|
||||
|
||||
// HTTP 单独派生 ALPN；不能修改 WSS 和其他 TLS 使用者共享的配置。
|
||||
pub(super) fn client_tls(tls: ClientTls, http2: Option<Http2Limits>) -> ClientTls {
|
||||
let mut config = (*tls.0).clone();
|
||||
config.alpn_protocols = vec![alpn(http2).to_vec()];
|
||||
ClientTls(std::sync::Arc::new(config))
|
||||
}
|
||||
|
||||
pub(super) fn server_tls(tls: ServerTls, http2: Option<Http2Limits>) -> ServerTls {
|
||||
let mut config = (*tls.0).clone();
|
||||
config.alpn_protocols = vec![alpn(http2).to_vec()];
|
||||
ServerTls(std::sync::Arc::new(config))
|
||||
}
|
||||
|
||||
fn alpn(http2: Option<Http2Limits>) -> &'static [u8] {
|
||||
if http2.is_some() { b"h2" } else { b"http/1.1" }
|
||||
}
|
||||
|
||||
pub(super) fn check_alpn(stream: &Transport, http2: Option<Http2Limits>) -> Result<(), String> {
|
||||
if let Transport::Tls(stream) = stream
|
||||
&& http2.is_some()
|
||||
&& stream.get_ref().1.alpn_protocol() != Some(b"h2")
|
||||
{
|
||||
return Err("HTTP/2 TLS peer did not negotiate h2".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,456 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use http_body_util::Full;
|
||||
use hyper::body::{Body, Frame, SizeHint};
|
||||
use tokio::sync::{Mutex as AsyncMutex, Semaphore, mpsc, oneshot};
|
||||
use tokio::time::Sleep;
|
||||
|
||||
use super::{Budget, Header, Response, message::encode_response};
|
||||
use crate::{bytes::Bytes, task, websocket};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct LiveLimits {
|
||||
pub chunk_bytes: i64,
|
||||
pub idle_ms: i64,
|
||||
pub heartbeat_ms: i64,
|
||||
}
|
||||
|
||||
impl LiveLimits {
|
||||
pub(super) fn check(self) -> Result<(), String> {
|
||||
if self.chunk_bytes <= 0 || usize::try_from(self.chunk_bytes).is_err() {
|
||||
return Err("HTTP chunk_bytes must be positive".into());
|
||||
}
|
||||
if self.idle_ms <= 0 || self.heartbeat_ms <= 0 || self.heartbeat_ms >= self.idle_ms {
|
||||
return Err("HTTP live limits require 0 < heartbeat_ms < idle_ms".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
type Head = hyper::Response<Outgoing>;
|
||||
enum Mode {
|
||||
Pending,
|
||||
Streaming {
|
||||
sender: mpsc::Sender<bytes::Bytes>,
|
||||
sse: bool,
|
||||
},
|
||||
Finished,
|
||||
Upgraded(Option<websocket::WebSocket>),
|
||||
}
|
||||
|
||||
struct State {
|
||||
mode: Mode,
|
||||
head: Option<oneshot::Sender<Head>>,
|
||||
upgrade: Option<(hyper::Request<()>, hyper::upgrade::OnUpgrade)>,
|
||||
fault: Option<String>,
|
||||
body_cancelled: bool,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
state: Mutex<State>,
|
||||
writer: AsyncMutex<()>,
|
||||
closed: Semaphore,
|
||||
budget: Budget,
|
||||
live: LiveLimits,
|
||||
upgraded: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl Shared {
|
||||
fn close_body(&self, cancelled: bool) {
|
||||
if cancelled {
|
||||
self.state
|
||||
.lock()
|
||||
.expect("reply lock poisoned")
|
||||
.body_cancelled = true;
|
||||
}
|
||||
self.closed.close();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct HttpReply(Arc<Shared>);
|
||||
|
||||
impl std::fmt::Debug for HttpReply {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("HttpReply(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpReply {
|
||||
pub(crate) fn chunk_limit(&self) -> usize {
|
||||
self.0.live.chunk_bytes as usize
|
||||
}
|
||||
|
||||
pub(crate) fn fault(&self) -> Option<String> {
|
||||
self.0
|
||||
.state
|
||||
.lock()
|
||||
.expect("reply lock poisoned")
|
||||
.fault
|
||||
.clone()
|
||||
}
|
||||
|
||||
fn publish(&self, prepare: impl FnOnce() -> (Head, Mode)) -> Result<(), String> {
|
||||
let mut state = self.0.state.lock().expect("reply lock poisoned");
|
||||
if !matches!(state.mode, Mode::Pending) {
|
||||
return Err("HTTP response has already started".into());
|
||||
}
|
||||
let head = state.head.take().ok_or("HTTP response is closed")?;
|
||||
// 只有真正提交的正文才能拥有取消信号；拒绝重复提交时不能关闭已有响应。
|
||||
let (response, mode) = prepare();
|
||||
state.upgrade.take();
|
||||
state.mode = mode;
|
||||
// SendError 会析构响应正文；先释放状态锁，正文 Drop 才能登记传输取消。
|
||||
drop(state);
|
||||
if head.send(response).is_err() {
|
||||
// 空正文从创建起就是 end-stream，不能只依赖 BufferedBody::drop 判断取消。
|
||||
self.0.close_body(true);
|
||||
return Err("HTTP peer is closed".into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn upgrade(
|
||||
&self,
|
||||
limits: websocket::Limits,
|
||||
) -> Result<websocket::WebSocket, String> {
|
||||
limits.check()?;
|
||||
let (head, response, upgrade) = {
|
||||
let mut state = self.0.state.lock().expect("reply lock poisoned");
|
||||
if !matches!(state.mode, Mode::Pending) {
|
||||
return Err("HTTP response has already started".into());
|
||||
}
|
||||
let (request, upgrade) = state.upgrade.take().ok_or("HTTP upgrade is unavailable")?;
|
||||
let response =
|
||||
tokio_tungstenite::tungstenite::handshake::server::create_response(&request)
|
||||
.map_err(|error| error.to_string())?;
|
||||
let response = response.map(|()| Outgoing::Full(Full::new(bytes::Bytes::new())));
|
||||
state.mode = Mode::Upgraded(None);
|
||||
self.0.upgraded.store(true, Ordering::Release);
|
||||
let head = state.head.take().ok_or("HTTP response is closed")?;
|
||||
(head, response, upgrade)
|
||||
};
|
||||
if head.send(response).is_err() {
|
||||
self.0.close_body(true);
|
||||
return Err("HTTP peer is closed".into());
|
||||
}
|
||||
let stream = tokio::time::timeout(self.0.budget.timeout, upgrade)
|
||||
.await
|
||||
.map_err(|_| "WebSocket upgrade timed out")?
|
||||
.map_err(|error| error.to_string())?;
|
||||
// 取回已知的底层 I/O 和预读帧；升级后不再继承 HTTP 写入期限。
|
||||
let parts = stream
|
||||
.downcast::<hyper_util::rt::TokioIo<super::write_timeout::WriteTimeout>>()
|
||||
.map_err(|_| "WebSocket upgrade transport does not match the HTTP server")?;
|
||||
let socket = websocket::from_server(
|
||||
parts.io.into_inner().into_inner(),
|
||||
parts.read_buf.to_vec(),
|
||||
limits,
|
||||
)
|
||||
.await?;
|
||||
let mut state = self.0.state.lock().expect("reply lock poisoned");
|
||||
state.mode = Mode::Upgraded(Some(socket.clone()));
|
||||
Ok(socket)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct Session(HttpReply);
|
||||
|
||||
impl Session {
|
||||
pub(super) fn new(
|
||||
budget: Budget,
|
||||
live: LiveLimits,
|
||||
upgrade: Option<(hyper::Request<()>, hyper::upgrade::OnUpgrade)>,
|
||||
upgraded: Arc<AtomicBool>,
|
||||
) -> (Self, oneshot::Receiver<Head>) {
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
(
|
||||
Self(HttpReply(Arc::new(Shared {
|
||||
state: Mutex::new(State {
|
||||
mode: Mode::Pending,
|
||||
head: Some(sender),
|
||||
upgrade,
|
||||
fault: None,
|
||||
body_cancelled: false,
|
||||
}),
|
||||
writer: AsyncMutex::new(()),
|
||||
closed: Semaphore::new(0),
|
||||
budget,
|
||||
live,
|
||||
upgraded,
|
||||
}))),
|
||||
receiver,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn reply(&self) -> HttpReply {
|
||||
self.0.clone()
|
||||
}
|
||||
|
||||
pub(super) async fn run(
|
||||
self,
|
||||
route: impl Future<Output = Result<(), String>>,
|
||||
) -> Result<(), String> {
|
||||
let result = tokio::select! {
|
||||
biased;
|
||||
result = route => result,
|
||||
_ = self.0.0.closed.acquire() => Ok(()),
|
||||
};
|
||||
if let Err(error) = &result {
|
||||
let mut state = self.0.0.state.lock().expect("reply lock poisoned");
|
||||
// 传输侧先取消正文时，随后被唤醒的 write 失败属于该流，不是程序故障。
|
||||
if state.body_cancelled {
|
||||
return Ok(());
|
||||
}
|
||||
state.fault = Some(error.clone());
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Session {
|
||||
fn drop(&mut self) {
|
||||
let mut state = self.0.0.state.lock().expect("reply lock poisoned");
|
||||
// 函数返回或取消都释放正文发送端/升级连接；外部别名不能延长会话。
|
||||
if let Mode::Upgraded(Some(socket)) = &state.mode {
|
||||
socket.abort();
|
||||
}
|
||||
state.mode = Mode::Finished;
|
||||
state.upgrade.take();
|
||||
state.head.take();
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn respond(reply: &HttpReply, response: Response) -> Result<(), String> {
|
||||
let _writer = reply.0.writer.lock().await;
|
||||
let response = encode_response(response, reply.0.budget)?;
|
||||
reply.publish(|| {
|
||||
(
|
||||
response.map(|body| {
|
||||
Outgoing::Buffered(BufferedBody {
|
||||
body,
|
||||
reply: reply.clone(),
|
||||
})
|
||||
}),
|
||||
Mode::Finished,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn start(reply: &HttpReply, status: i64, headers: Vec<Header>) -> Result<(), String> {
|
||||
start_stream(reply, status, headers, false).await
|
||||
}
|
||||
|
||||
pub(crate) async fn start_stream(
|
||||
reply: &HttpReply,
|
||||
status: i64,
|
||||
headers: Vec<Header>,
|
||||
sse: bool,
|
||||
) -> Result<(), String> {
|
||||
if matches!(status, 204 | 205 | 304) {
|
||||
return Err("bodyless status requires http.respond".into());
|
||||
}
|
||||
let _writer = reply.0.writer.lock().await;
|
||||
let response = encode_response(
|
||||
Response {
|
||||
status,
|
||||
headers,
|
||||
body: Bytes::new(Vec::new()),
|
||||
},
|
||||
reply.0.budget,
|
||||
)?;
|
||||
reply.publish(|| {
|
||||
let (sender, receiver) = mpsc::channel(1);
|
||||
let millis = if sse {
|
||||
reply.0.live.heartbeat_ms
|
||||
} else {
|
||||
reply.0.live.idle_ms
|
||||
};
|
||||
let interval = Duration::from_millis(millis as u64);
|
||||
let body = StreamBody {
|
||||
receiver,
|
||||
reply: reply.clone(),
|
||||
interval,
|
||||
deadline: Box::pin(tokio::time::sleep(interval)),
|
||||
sse,
|
||||
ended: false,
|
||||
};
|
||||
(
|
||||
response.map(|_| Outgoing::Stream(body)),
|
||||
Mode::Streaming { sender, sse },
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn write(reply: &HttpReply, bytes: Bytes) -> Result<(), String> {
|
||||
write_chunk(reply, bytes, false).await
|
||||
}
|
||||
|
||||
pub(crate) async fn write_chunk(reply: &HttpReply, bytes: Bytes, sse: bool) -> Result<(), String> {
|
||||
if bytes.values().len() > reply.chunk_limit() {
|
||||
return Err("HTTP chunk exceeds chunk_bytes".into());
|
||||
}
|
||||
tokio::time::timeout(reply.0.budget.timeout, async {
|
||||
let _writer = reply.0.writer.lock().await;
|
||||
let sender = {
|
||||
let state = reply.0.state.lock().expect("reply lock poisoned");
|
||||
match &state.mode {
|
||||
Mode::Streaming { sender, sse: mode } if *mode == sse => sender.clone(),
|
||||
_ => return Err("HTTP reply is not an open stream of this kind".into()),
|
||||
}
|
||||
};
|
||||
sender
|
||||
.send(bytes.into_http())
|
||||
.await
|
||||
.map_err(|_| "HTTP peer is closed".into())
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "HTTP stream write timed out".to_owned())?
|
||||
}
|
||||
|
||||
pub async fn finish(reply: &HttpReply) -> Result<(), String> {
|
||||
let _writer = reply.0.writer.lock().await;
|
||||
let mut state = reply.0.state.lock().expect("reply lock poisoned");
|
||||
if !matches!(state.mode, Mode::Streaming { .. }) {
|
||||
return Err("HTTP reply is not an open stream".into());
|
||||
}
|
||||
state.mode = Mode::Finished;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) enum Outgoing {
|
||||
Full(Full<bytes::Bytes>),
|
||||
Buffered(BufferedBody),
|
||||
Stream(StreamBody),
|
||||
Handler(HandlerBody),
|
||||
}
|
||||
|
||||
impl Outgoing {
|
||||
pub(super) fn with_handler(self, handler: task::Group) -> Self {
|
||||
Self::Handler(HandlerBody {
|
||||
body: Box::new(self),
|
||||
_handler: handler,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct HandlerBody {
|
||||
body: Box<Outgoing>,
|
||||
_handler: task::Group,
|
||||
}
|
||||
|
||||
pub(super) struct BufferedBody {
|
||||
body: Full<bytes::Bytes>,
|
||||
reply: HttpReply,
|
||||
}
|
||||
|
||||
impl Drop for BufferedBody {
|
||||
fn drop(&mut self) {
|
||||
self.reply.0.close_body(!self.body.is_end_stream());
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct StreamBody {
|
||||
receiver: mpsc::Receiver<bytes::Bytes>,
|
||||
reply: HttpReply,
|
||||
interval: Duration,
|
||||
deadline: Pin<Box<Sleep>>,
|
||||
sse: bool,
|
||||
ended: bool,
|
||||
}
|
||||
|
||||
impl Drop for StreamBody {
|
||||
fn drop(&mut self) {
|
||||
self.reply.0.close_body(!self.ended);
|
||||
}
|
||||
}
|
||||
|
||||
impl Body for Outgoing {
|
||||
type Data = bytes::Bytes;
|
||||
type Error = std::io::Error;
|
||||
|
||||
fn poll_frame(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
|
||||
let body = match self.get_mut() {
|
||||
Self::Stream(body) => body,
|
||||
Self::Full(body) => {
|
||||
return Pin::new(body)
|
||||
.poll_frame(cx)
|
||||
.map(|frame| frame.map(|frame| frame.map_err(|never| match never {})));
|
||||
}
|
||||
Self::Buffered(body) => {
|
||||
return Pin::new(&mut body.body)
|
||||
.poll_frame(cx)
|
||||
.map(|frame| frame.map(|frame| frame.map_err(|never| match never {})));
|
||||
}
|
||||
Self::Handler(body) => return Pin::new(body.body.as_mut()).poll_frame(cx),
|
||||
};
|
||||
if body.ended {
|
||||
return Poll::Ready(None);
|
||||
}
|
||||
match body.receiver.poll_recv(cx) {
|
||||
Poll::Ready(Some(bytes)) => {
|
||||
body.deadline
|
||||
.as_mut()
|
||||
.reset(tokio::time::Instant::now() + body.interval);
|
||||
return Poll::Ready(Some(Ok(Frame::data(bytes))));
|
||||
}
|
||||
Poll::Ready(None) => {
|
||||
body.ended = true;
|
||||
let fault = body
|
||||
.reply
|
||||
.0
|
||||
.state
|
||||
.lock()
|
||||
.expect("reply lock poisoned")
|
||||
.fault
|
||||
.clone();
|
||||
return Poll::Ready(fault.map(|fault| Err(std::io::Error::other(fault))));
|
||||
}
|
||||
Poll::Pending => {}
|
||||
}
|
||||
if body.deadline.as_mut().poll(cx).is_pending() {
|
||||
return Poll::Pending;
|
||||
}
|
||||
body.deadline
|
||||
.as_mut()
|
||||
.reset(tokio::time::Instant::now() + body.interval);
|
||||
if body.sse {
|
||||
Poll::Ready(Some(Ok(Frame::data(bytes::Bytes::from_static(b":\n\n")))))
|
||||
} else {
|
||||
body.ended = true;
|
||||
body.reply.0.close_body(true);
|
||||
Poll::Ready(Some(Err(std::io::Error::new(
|
||||
std::io::ErrorKind::TimedOut,
|
||||
"HTTP stream idle timeout",
|
||||
))))
|
||||
}
|
||||
}
|
||||
|
||||
fn is_end_stream(&self) -> bool {
|
||||
match self {
|
||||
Self::Full(body) => body.is_end_stream(),
|
||||
Self::Buffered(body) => body.body.is_end_stream(),
|
||||
Self::Stream(body) => body.ended,
|
||||
Self::Handler(body) => body.body.is_end_stream(),
|
||||
}
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> SizeHint {
|
||||
match self {
|
||||
Self::Full(body) => body.size_hint(),
|
||||
Self::Buffered(body) => body.body.size_hint(),
|
||||
Self::Stream(_) => SizeHint::default(),
|
||||
Self::Handler(body) => body.body.size_hint(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use crate::{
|
||||
net,
|
||||
task::{ScopeFault, ScopedOwner},
|
||||
};
|
||||
|
||||
pub(super) fn configure_liveness<E>(
|
||||
builder: &mut hyper::server::conn::http2::Builder<E>,
|
||||
timeout: Duration,
|
||||
) {
|
||||
builder
|
||||
.keep_alive_interval(Some(timeout))
|
||||
.keep_alive_timeout(timeout);
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct FirstRequest(Arc<FirstRequestState>);
|
||||
|
||||
struct FirstRequestState {
|
||||
seen: AtomicBool,
|
||||
changed: Notify,
|
||||
}
|
||||
|
||||
impl FirstRequest {
|
||||
pub(super) fn new() -> Self {
|
||||
Self(Arc::new(FirstRequestState {
|
||||
seen: AtomicBool::new(false),
|
||||
changed: Notify::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
pub(super) fn mark(&self) {
|
||||
if !self.0.seen.swap(true, Ordering::AcqRel) {
|
||||
self.0.changed.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait(&self) {
|
||||
loop {
|
||||
let changed = self.0.changed.notified();
|
||||
if self.0.seen.load(Ordering::Acquire) {
|
||||
return;
|
||||
}
|
||||
changed.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn drive<C, Shutdown>(
|
||||
connection: C,
|
||||
listener: &net::Listener,
|
||||
first_request: &FirstRequest,
|
||||
timeout: Duration,
|
||||
fault: &ScopeFault,
|
||||
mut shutdown: Shutdown,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
C: Future<Output = Result<(), hyper::Error>>,
|
||||
Shutdown: for<'a> FnMut(Pin<&'a mut C>),
|
||||
{
|
||||
tokio::pin!(connection);
|
||||
tokio::select! {
|
||||
biased;
|
||||
error = fault.wait() => return Err(error),
|
||||
result = &mut connection => return super::connection_result(result),
|
||||
_ = listener.closed() => {
|
||||
shutdown(connection.as_mut());
|
||||
return drain(connection.as_mut(), fault).await;
|
||||
}
|
||||
_ = first_request.wait() => {}
|
||||
_ = tokio::time::sleep(timeout) => return Ok(()),
|
||||
}
|
||||
tokio::select! {
|
||||
biased;
|
||||
error = fault.wait() => Err(error),
|
||||
result = &mut connection => super::connection_result(result),
|
||||
_ = listener.closed() => {
|
||||
shutdown(connection.as_mut());
|
||||
drain(connection.as_mut(), fault).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn drain<C>(mut connection: Pin<&mut C>, fault: &ScopeFault) -> Result<(), String>
|
||||
where
|
||||
C: Future<Output = Result<(), hyper::Error>>,
|
||||
{
|
||||
tokio::select! {
|
||||
biased;
|
||||
error = fault.wait() => Err(error),
|
||||
result = connection.as_mut() => super::connection_result(result),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn complete(owner: ScopedOwner, result: Result<(), String>) -> Result<(), String> {
|
||||
match result {
|
||||
Ok(()) => owner.finish().await,
|
||||
Err(error) => {
|
||||
let _ = owner.stop().await;
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
use std::future::Future;
|
||||
use std::io::{self, IoSlice};
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::transport::Transport;
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
use tokio::time::Sleep;
|
||||
|
||||
// Header/body 的期限由协议层持有；这里仅限制无法取得进展的写操作。
|
||||
pub(super) struct WriteTimeout {
|
||||
stream: Transport,
|
||||
timeout: Duration,
|
||||
deadline: Option<Pin<Box<Sleep>>>,
|
||||
}
|
||||
|
||||
impl WriteTimeout {
|
||||
pub(super) fn into_inner(self) -> Transport {
|
||||
self.stream
|
||||
}
|
||||
pub fn new(stream: Transport, timeout: Duration) -> Self {
|
||||
Self {
|
||||
stream,
|
||||
timeout,
|
||||
deadline: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn output<T>(
|
||||
&mut self,
|
||||
context: &mut Context<'_>,
|
||||
write: impl FnOnce(Pin<&mut Transport>, &mut Context<'_>) -> Poll<io::Result<T>>,
|
||||
) -> Poll<io::Result<T>> {
|
||||
match write(Pin::new(&mut self.stream), context) {
|
||||
Poll::Ready(result) => {
|
||||
self.deadline = None;
|
||||
Poll::Ready(result)
|
||||
}
|
||||
Poll::Pending => {
|
||||
let deadline = self
|
||||
.deadline
|
||||
.get_or_insert_with(|| Box::pin(tokio::time::sleep(self.timeout)));
|
||||
match deadline.as_mut().poll(context) {
|
||||
Poll::Ready(()) => Poll::Ready(Err(io::Error::new(
|
||||
io::ErrorKind::TimedOut,
|
||||
"HTTP write timed out",
|
||||
))),
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncRead for WriteTimeout {
|
||||
fn poll_read(
|
||||
mut self: Pin<&mut Self>,
|
||||
context: &mut Context<'_>,
|
||||
buffer: &mut ReadBuf<'_>,
|
||||
) -> Poll<io::Result<()>> {
|
||||
Pin::new(&mut self.stream).poll_read(context, buffer)
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for WriteTimeout {
|
||||
fn poll_write(
|
||||
mut self: Pin<&mut Self>,
|
||||
context: &mut Context<'_>,
|
||||
bytes: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
self.output(context, |stream, context| stream.poll_write(context, bytes))
|
||||
}
|
||||
|
||||
fn poll_write_vectored(
|
||||
mut self: Pin<&mut Self>,
|
||||
context: &mut Context<'_>,
|
||||
buffers: &[IoSlice<'_>],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
self.output(context, |stream, context| {
|
||||
stream.poll_write_vectored(context, buffers)
|
||||
})
|
||||
}
|
||||
|
||||
fn is_write_vectored(&self) -> bool {
|
||||
self.stream.is_write_vectored()
|
||||
}
|
||||
|
||||
fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
self.output(context, |stream, context| stream.poll_flush(context))
|
||||
}
|
||||
|
||||
fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
self.output(context, |stream, context| stream.poll_shutdown(context))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,343 @@
|
||||
//! Compiler-owned durable queue bridge. All persistence uses the ordinary database executor.
|
||||
use crate::database::{Database, Executor, Sql, Transaction};
|
||||
use crate::orm::{self, Error, Row, Value};
|
||||
use crate::wire::Encoded;
|
||||
use std::cell::RefCell;
|
||||
use std::future::Future;
|
||||
|
||||
pub mod resources;
|
||||
pub mod store;
|
||||
pub mod worker;
|
||||
|
||||
pub use crate::time::Clock;
|
||||
|
||||
tokio::task_local! { static EXECUTION_CONTEXT: ExecutionContext; }
|
||||
|
||||
struct ExecutionContext {
|
||||
identity: ExecutionIdentity,
|
||||
permission_key: RefCell<Option<String>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum ExecutionIdentity {
|
||||
System {
|
||||
tenant_id: Option<i64>,
|
||||
},
|
||||
User {
|
||||
tenant_id: Option<i64>,
|
||||
provider: String,
|
||||
site: String,
|
||||
subject: String,
|
||||
session: String,
|
||||
tenant_claim: Option<String>,
|
||||
permission_key: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl ExecutionIdentity {
|
||||
pub fn tenant_id(&self) -> Option<i64> {
|
||||
match self {
|
||||
Self::System { tenant_id } | Self::User { tenant_id, .. } => *tenant_id,
|
||||
}
|
||||
}
|
||||
|
||||
fn validate(&self) -> Result<(), Error> {
|
||||
if self.tenant_id().is_some_and(|tenant_id| tenant_id <= 0) {
|
||||
return Err(Error::invalid_data(
|
||||
"Job execution tenant id must be positive",
|
||||
));
|
||||
}
|
||||
if let Self::User {
|
||||
provider,
|
||||
site,
|
||||
subject,
|
||||
session,
|
||||
tenant_claim,
|
||||
permission_key,
|
||||
..
|
||||
} = self
|
||||
{
|
||||
for (value, label) in [
|
||||
(provider.as_str(), "provider"),
|
||||
(site.as_str(), "site"),
|
||||
(subject.as_str(), "subject"),
|
||||
(session.as_str(), "session"),
|
||||
] {
|
||||
bounded(value, 256, &format!("Job execution {label}"))?;
|
||||
}
|
||||
if let Some(tenant_claim) = tenant_claim {
|
||||
bounded(tenant_claim, 256, "Job execution tenant claim")?;
|
||||
}
|
||||
if !valid_permission_key(permission_key) {
|
||||
return Err(Error::invalid_data(
|
||||
"Job execution permission key is invalid",
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn scope_execution<T>(
|
||||
identity: ExecutionIdentity,
|
||||
future: impl Future<Output = T>,
|
||||
) -> impl Future<Output = T> {
|
||||
// Pin before constructing the context wrapper so its poll frame does not
|
||||
// carry another copy of a large business Future.
|
||||
let future = Box::pin(future);
|
||||
async move {
|
||||
let permission_key = match &identity {
|
||||
ExecutionIdentity::User { permission_key, .. } if !permission_key.is_empty() => {
|
||||
Some(permission_key.clone())
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
EXECUTION_CONTEXT
|
||||
.scope(
|
||||
ExecutionContext {
|
||||
identity,
|
||||
permission_key: RefCell::new(permission_key),
|
||||
},
|
||||
future,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub fn scope_system<T>(
|
||||
tenant_id: Option<i64>,
|
||||
future: impl Future<Output = T>,
|
||||
) -> impl Future<Output = T> {
|
||||
scope_execution(ExecutionIdentity::System { tenant_id }, future)
|
||||
}
|
||||
|
||||
pub fn execution_identity() -> Result<ExecutionIdentity, Error> {
|
||||
EXECUTION_CONTEXT
|
||||
.try_with(|context| {
|
||||
let mut identity = context.identity.clone();
|
||||
if let ExecutionIdentity::User { permission_key, .. } = &mut identity {
|
||||
*permission_key = context.permission_key.borrow().clone().unwrap_or_default();
|
||||
}
|
||||
identity
|
||||
})
|
||||
.map_err(|_| Error::invalid_data("trusted Job execution identity is required"))
|
||||
}
|
||||
|
||||
/// Compiled component ownership does not enable tenant isolation; deployment settings do.
|
||||
pub async fn require_components(
|
||||
identity: &ExecutionIdentity,
|
||||
components: &[&str],
|
||||
manifest: &[&str],
|
||||
) -> Result<(), store::Outcome> {
|
||||
if components.is_empty() || crate::config::current_settings().tenant().is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
let tenant_id = identity
|
||||
.tenant_id()
|
||||
.ok_or(store::Outcome::InvalidIdentity)?;
|
||||
crate::tenant::require_components(tenant_id, components, manifest)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
if error.kind() == orm::ErrorKind::InvalidData {
|
||||
store::Outcome::ComponentDisabled
|
||||
} else {
|
||||
store::Outcome::Retry
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "api")]
|
||||
pub fn bind_permission(permission_key: &str) -> Result<(), Error> {
|
||||
if !valid_permission_key(permission_key) {
|
||||
return Err(Error::invalid_data(
|
||||
"Job execution permission key is invalid",
|
||||
));
|
||||
}
|
||||
EXECUTION_CONTEXT
|
||||
.try_with(|context| match &context.identity {
|
||||
ExecutionIdentity::User { .. } => {
|
||||
let mut current = context.permission_key.borrow_mut();
|
||||
match current.as_deref() {
|
||||
Some(value) if value != permission_key => Err(Error::invalid_data(
|
||||
"Job execution permission cannot change within one request",
|
||||
)),
|
||||
Some(_) => Ok(()),
|
||||
None => {
|
||||
*current = Some(permission_key.to_owned());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
ExecutionIdentity::System { .. } => Ok(()),
|
||||
})
|
||||
.map_err(|_| Error::invalid_data("trusted Job execution identity is required"))?
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Spec {
|
||||
pub target: &'static str,
|
||||
pub schema: &'static str,
|
||||
pub attempts: u32,
|
||||
pub timeout_ms: u32,
|
||||
}
|
||||
|
||||
impl Spec {
|
||||
pub fn validate(&self) -> Result<(), Error> {
|
||||
bounded(self.target, 1024, "Job target")?;
|
||||
bounded(self.schema, 65536, "Job schema")?;
|
||||
if !(1..=100).contains(&self.attempts)
|
||||
|| !(1..=crate::config::MAX_JOB_TIMEOUT_MS).contains(&self.timeout_ms)
|
||||
{
|
||||
return Err(Error::invalid_data("invalid Job retry or timeout policy"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn now() -> Result<i64, Error> {
|
||||
crate::time::now().map_err(Error::invalid_data)
|
||||
}
|
||||
|
||||
fn bounded(value: &str, maximum: usize, label: &str) -> Result<(), Error> {
|
||||
if value.is_empty() || value.len() > maximum {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"{label} must contain 1..={maximum} bytes"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn valid_permission_key(permission_key: &str) -> bool {
|
||||
let parts = permission_key.split('.').collect::<Vec<_>>();
|
||||
parts.len() == 4
|
||||
&& parts.iter().all(|part| {
|
||||
part.starts_with(|character: char| character.is_ascii_lowercase())
|
||||
&& part.split('_').all(|word| {
|
||||
!word.is_empty()
|
||||
&& word.chars().all(|character| {
|
||||
character.is_ascii_lowercase() || character.is_ascii_digit()
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn text(row: &Row, index: usize) -> Result<String, Error> {
|
||||
orm::text(row.get(index)?.clone())
|
||||
}
|
||||
fn int(row: &Row, index: usize) -> Result<i64, Error> {
|
||||
orm::int(row.get(index)?.clone())
|
||||
}
|
||||
fn t(value: impl Into<String>) -> Value {
|
||||
Value::Text(value.into())
|
||||
}
|
||||
fn n(value: i64) -> Value {
|
||||
Value::Int(value)
|
||||
}
|
||||
|
||||
async fn finish_transaction<T>(
|
||||
transaction: Transaction,
|
||||
result: Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
match result {
|
||||
Ok(value) => {
|
||||
transaction.commit().await?;
|
||||
Ok(value)
|
||||
}
|
||||
Err(error) => match transaction.rollback().await {
|
||||
Ok(()) => Err(error),
|
||||
Err(rollback) => Err(error.cause(rollback)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn enqueue(
|
||||
database: Database,
|
||||
ambient: Option<&Transaction>,
|
||||
spec: &Spec,
|
||||
payload: Encoded,
|
||||
key: &str,
|
||||
run_at: i64,
|
||||
) -> Result<crate::Id, Error> {
|
||||
spec.validate()?;
|
||||
bounded(key, 1024, "Job idempotency key")?;
|
||||
// Business keys and scheduler keys occupy disjoint namespaces.
|
||||
let key = format!("business:{key}");
|
||||
let tenant_id = database.key().tenant_id().unwrap_or(0);
|
||||
let identity = execution_identity()?;
|
||||
identity.validate()?;
|
||||
if tenant_id > 0 && identity.tenant_id() != Some(tenant_id) {
|
||||
return Err(Error::invalid_data(
|
||||
"Job execution tenant does not match tenant storage",
|
||||
));
|
||||
}
|
||||
if let Some(transaction) = ambient {
|
||||
let executor = Executor::new(database, Some(transaction))?;
|
||||
return insert(
|
||||
&executor,
|
||||
tenant_id,
|
||||
&identity,
|
||||
spec,
|
||||
payload.as_str(),
|
||||
&key,
|
||||
run_at,
|
||||
)
|
||||
.await
|
||||
.map(crate::Id);
|
||||
}
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = insert(
|
||||
&executor,
|
||||
tenant_id,
|
||||
&identity,
|
||||
spec,
|
||||
payload.as_str(),
|
||||
&key,
|
||||
run_at,
|
||||
)
|
||||
.await;
|
||||
finish_transaction(transaction, result).await.map(crate::Id)
|
||||
}
|
||||
|
||||
async fn insert(
|
||||
executor: &Executor<'_>,
|
||||
tenant_id: i64,
|
||||
identity: &ExecutionIdentity,
|
||||
spec: &Spec,
|
||||
payload: &str,
|
||||
key: &str,
|
||||
run_at: i64,
|
||||
) -> Result<String, Error> {
|
||||
let id = orm::Uuid::new_v7()?.to_string();
|
||||
let (kind, scope_tenant_id, provider, site, subject, session, tenant_claim, permission_key) =
|
||||
match identity {
|
||||
ExecutionIdentity::System { tenant_id } => {
|
||||
("system", tenant_id.unwrap_or(0), "", "", "", "", "", "")
|
||||
}
|
||||
ExecutionIdentity::User {
|
||||
tenant_id,
|
||||
provider,
|
||||
site,
|
||||
subject,
|
||||
session,
|
||||
tenant_claim,
|
||||
permission_key,
|
||||
} => (
|
||||
"user",
|
||||
tenant_id.unwrap_or(0),
|
||||
provider.as_str(),
|
||||
site.as_str(),
|
||||
subject.as_str(),
|
||||
session.as_str(),
|
||||
tenant_claim.as_deref().unwrap_or(""),
|
||||
permission_key.as_str(),
|
||||
),
|
||||
};
|
||||
// RETURNING holds the conflicting row through this transaction. A separate
|
||||
// SELECT could miss it if a worker completed the active job in between.
|
||||
let rows = executor.query(Sql {
|
||||
sqlite: "INSERT INTO _dever_jobs (id,tenant_id,scope_tenant_id,execution_kind,auth_provider,auth_site,auth_subject,auth_session,auth_tenant,auth_permission,target,schema,payload,dedupe_key,state,run_at,attempt,max_attempts,timeout_ms,claim_token,lease_until,error,updated_at) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,'pending',?15,0,?16,?17,'',0,'',?15) ON CONFLICT(tenant_id,scope_tenant_id,target,dedupe_key) WHERE state IN ('pending','running') DO UPDATE SET dedupe_key=excluded.dedupe_key RETURNING id",
|
||||
postgres: "INSERT INTO _dever_jobs (id,tenant_id,scope_tenant_id,execution_kind,auth_provider,auth_site,auth_subject,auth_session,auth_tenant,auth_permission,target,schema,payload,dedupe_key,state,run_at,attempt,max_attempts,timeout_ms,claim_token,lease_until,error,updated_at) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,'pending',$15,0,$16,$17,'',0,'',$15) ON CONFLICT(tenant_id,scope_tenant_id,target,dedupe_key) WHERE state IN ('pending','running') DO UPDATE SET dedupe_key=excluded.dedupe_key RETURNING id",
|
||||
}, vec![t(id), n(tenant_id), n(scope_tenant_id), t(kind), t(provider), t(site), t(subject), t(session), t(tenant_claim), t(permission_key), t(spec.target), t(spec.schema), t(payload), t(key), n(run_at), n(i64::from(spec.attempts)), n(i64::from(spec.timeout_ms))]).await?;
|
||||
text(&orm::required_row(rows, "Job enqueue")?, 0)
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
//! Application-owned Job bindings; storage and scheduling remain in their original owners.
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use super::{Clock, store, worker::Binding};
|
||||
use crate::{database::Database, orm::Error};
|
||||
|
||||
pub struct Resources {
|
||||
pub bindings: Vec<Binding>,
|
||||
pub clock: Clock,
|
||||
}
|
||||
|
||||
impl Resources {
|
||||
pub fn validate(&self, worker: bool) -> Result<(), String> {
|
||||
let settings = crate::config::current_settings();
|
||||
for binding in &self.bindings {
|
||||
binding.spec.validate().map_err(|error| error.to_string())?;
|
||||
if worker && let Some((explicit, root)) = binding.selector() {
|
||||
settings.validate_job_connection(explicit, root, binding.spec.timeout_ms)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn initialize(&self) -> Result<(), Error> {
|
||||
let mut initialized = BTreeSet::new();
|
||||
for binding in &self.bindings {
|
||||
if binding.uses_tenant_storage() {
|
||||
continue;
|
||||
}
|
||||
let database = binding.database(None).await?;
|
||||
if initialized.insert(database.key().clone()) {
|
||||
store::initialize(database.clone()).await?;
|
||||
}
|
||||
Self::validate_schedule(binding, database).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn has_tenant_jobs(&self) -> bool {
|
||||
self.bindings.iter().any(Binding::uses_tenant_storage)
|
||||
}
|
||||
|
||||
pub async fn migrate_tenant(
|
||||
&self,
|
||||
tenant_id: i64,
|
||||
fingerprint: &str,
|
||||
databases: &mut BTreeMap<String, Database>,
|
||||
) -> Result<(), Error> {
|
||||
let settings = crate::config::current_settings();
|
||||
let mut initialized = BTreeSet::new();
|
||||
for binding in self
|
||||
.bindings
|
||||
.iter()
|
||||
.filter(|binding| binding.uses_tenant_storage())
|
||||
{
|
||||
let (explicit, root) = binding
|
||||
.selector()
|
||||
.ok_or_else(|| Error::migration("tenant Job requires scoped storage"))?;
|
||||
let connection = settings
|
||||
.resolve_database(explicit, root)
|
||||
.map_err(Error::migration)?
|
||||
.0
|
||||
.to_owned();
|
||||
if !databases.contains_key(&connection) {
|
||||
let database =
|
||||
crate::tenant::begin_migration(&connection, tenant_id, fingerprint).await?;
|
||||
databases.insert(connection.clone(), database);
|
||||
}
|
||||
let database = databases[&connection].clone();
|
||||
if initialized.insert(connection) {
|
||||
store::initialize(database.clone()).await?;
|
||||
}
|
||||
Self::validate_schedule(binding, database).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn validate_schedule(binding: &Binding, database: Database) -> Result<(), Error> {
|
||||
if let Some(schedule) = &binding.schedule {
|
||||
store::validate_schedule(database, &binding.spec, schedule).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,556 @@
|
||||
use super::*;
|
||||
|
||||
/// A fencing token belongs to exactly one claim, including reclaim by the same process.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Claim {
|
||||
pub id: String,
|
||||
/// Physical queue database tenant, not the execution scope.
|
||||
pub tenant_id: Option<i64>,
|
||||
pub execution: ExecutionIdentity,
|
||||
pub target: String,
|
||||
pub schema: String,
|
||||
pub payload: String,
|
||||
pub attempt: i64,
|
||||
pub max_attempts: i64,
|
||||
pub timeout_ms: i64,
|
||||
pub token: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub enum Outcome {
|
||||
Success,
|
||||
Retry,
|
||||
Timeout,
|
||||
UnknownTarget,
|
||||
SchemaMismatch,
|
||||
InvalidPayload,
|
||||
InvalidPolicy,
|
||||
IdentityRejected,
|
||||
InvalidIdentity,
|
||||
ComponentDisabled,
|
||||
}
|
||||
|
||||
impl Outcome {
|
||||
fn error(self) -> &'static str {
|
||||
match self {
|
||||
Self::Success => "",
|
||||
Self::Retry => "handler_failed",
|
||||
Self::Timeout => "handler_timeout",
|
||||
Self::UnknownTarget => "unknown_target",
|
||||
Self::SchemaMismatch => "schema_mismatch",
|
||||
Self::InvalidPayload => "invalid_payload",
|
||||
Self::InvalidPolicy => "invalid_policy",
|
||||
Self::IdentityRejected => "identity_rejected",
|
||||
Self::InvalidIdentity => "invalid_identity",
|
||||
Self::ComponentDisabled => "component_disabled",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn initialize(database: Database) -> Result<(), Error> {
|
||||
let tenant_id = database.key().tenant_id();
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = initialize_in(&executor, tenant_id).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
async fn initialize_in(executor: &Executor<'_>, tenant_id: Option<i64>) -> Result<(), Error> {
|
||||
// Serialize private schema installation independently of Model migrations.
|
||||
executor
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT 1",
|
||||
postgres: "SELECT pg_advisory_xact_lock(1684371045, 1785684595) IS NULL",
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
executor.execute(Sql {
|
||||
sqlite: "CREATE TABLE IF NOT EXISTS _dever_job_version (singleton INTEGER PRIMARY KEY CHECK(singleton=1), version BIGINT NOT NULL)",
|
||||
postgres: "CREATE TABLE IF NOT EXISTS _dever_job_version (singleton BIGINT PRIMARY KEY CHECK(singleton=1), version BIGINT NOT NULL)",
|
||||
}, vec![]).await?;
|
||||
let versions = executor
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT version FROM _dever_job_version",
|
||||
postgres: "SELECT version FROM _dever_job_version",
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
if let Some(version) = versions.first() {
|
||||
if versions.len() != 1 {
|
||||
return Err(Error::migration("invalid private Job schema version state"));
|
||||
}
|
||||
return match int(version, 0)? {
|
||||
5 => Ok(()),
|
||||
4 => migrate_v4(executor).await,
|
||||
3 => {
|
||||
migrate_v3(executor).await?;
|
||||
migrate_v4(executor).await
|
||||
}
|
||||
2 => {
|
||||
migrate_v2(executor).await?;
|
||||
migrate_v3(executor).await?;
|
||||
migrate_v4(executor).await
|
||||
}
|
||||
1 => {
|
||||
migrate_v1(executor, tenant_id).await?;
|
||||
migrate_v2(executor).await?;
|
||||
migrate_v3(executor).await?;
|
||||
migrate_v4(executor).await
|
||||
}
|
||||
_ => Err(Error::migration("unsupported private Job schema version")),
|
||||
};
|
||||
}
|
||||
const JOBS: &str = "CREATE TABLE _dever_jobs (id TEXT PRIMARY KEY, tenant_id BIGINT NOT NULL, scope_tenant_id BIGINT NOT NULL, execution_kind TEXT NOT NULL CHECK(execution_kind IN ('system','user','legacy')), auth_provider TEXT NOT NULL, auth_site TEXT NOT NULL, auth_subject TEXT NOT NULL, auth_session TEXT NOT NULL, auth_tenant TEXT NOT NULL, auth_permission TEXT NOT NULL, target TEXT NOT NULL, schema TEXT NOT NULL, payload TEXT NOT NULL, dedupe_key TEXT NOT NULL, state TEXT NOT NULL CHECK(state IN ('pending','running','succeeded','dead','blocked')), run_at BIGINT NOT NULL, attempt BIGINT NOT NULL, max_attempts BIGINT NOT NULL, timeout_ms BIGINT NOT NULL, claim_token TEXT NOT NULL, lease_until BIGINT NOT NULL, error TEXT NOT NULL, updated_at BIGINT NOT NULL)";
|
||||
const ACTIVE: &str = "CREATE UNIQUE INDEX _dever_jobs_active ON _dever_jobs(tenant_id,scope_tenant_id,target,dedupe_key) WHERE state IN ('pending','running')";
|
||||
const DUE: &str = "CREATE INDEX _dever_jobs_due ON _dever_jobs(state,run_at,lease_until)";
|
||||
const SCHEDULE: &str = "CREATE TABLE _dever_job_schedule (identity TEXT PRIMARY KEY, fingerprint TEXT NOT NULL, last_minute BIGINT NOT NULL)";
|
||||
for sql in [
|
||||
JOBS,
|
||||
ACTIVE,
|
||||
DUE,
|
||||
SCHEDULE,
|
||||
"INSERT INTO _dever_job_version(singleton,version) VALUES (1,5)",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn migrate_v4(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
for sql in [
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN auth_permission TEXT NOT NULL DEFAULT ''",
|
||||
"UPDATE _dever_jobs SET state='blocked',error='authorization_contract_changed',claim_token='',lease_until=0 WHERE execution_kind='user' AND state IN ('pending','running')",
|
||||
"ALTER TABLE _dever_jobs DROP COLUMN auth_capabilities",
|
||||
"UPDATE _dever_job_version SET version=5 WHERE singleton=1",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn migrate_v3(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
for sql in [
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN auth_capabilities TEXT NOT NULL DEFAULT '[]'",
|
||||
"UPDATE _dever_jobs SET state='blocked',error='missing_execution_capabilities',claim_token='',lease_until=0 WHERE execution_kind='user' AND state IN ('pending','running')",
|
||||
"UPDATE _dever_job_version SET version=4 WHERE singleton=1",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn migrate_v2(executor: &Executor<'_>) -> Result<(), Error> {
|
||||
for sql in [
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN scope_tenant_id BIGINT NOT NULL DEFAULT 0",
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN execution_kind TEXT NOT NULL DEFAULT 'legacy'",
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN auth_provider TEXT NOT NULL DEFAULT ''",
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN auth_site TEXT NOT NULL DEFAULT ''",
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN auth_subject TEXT NOT NULL DEFAULT ''",
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN auth_session TEXT NOT NULL DEFAULT ''",
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN auth_tenant TEXT NOT NULL DEFAULT ''",
|
||||
"UPDATE _dever_jobs SET state='blocked',error='missing_execution_identity',claim_token='',lease_until=0 WHERE state IN ('pending','running')",
|
||||
"DROP INDEX IF EXISTS _dever_jobs_active",
|
||||
"CREATE UNIQUE INDEX _dever_jobs_active ON _dever_jobs(tenant_id,scope_tenant_id,target,dedupe_key) WHERE state IN ('pending','running')",
|
||||
"UPDATE _dever_job_version SET version=3 WHERE singleton=1",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn migrate_v1(executor: &Executor<'_>, tenant_id: Option<i64>) -> Result<(), Error> {
|
||||
if tenant_id.is_some() {
|
||||
let rows = executor
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT COUNT(*) FROM _dever_jobs",
|
||||
postgres: "SELECT COUNT(*) FROM _dever_jobs",
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
if int(&orm::required_row(rows, "legacy Job row count")?, 0)? != 0 {
|
||||
return Err(Error::migration(
|
||||
"legacy Job rows have no tenant ownership; drain or explicitly migrate them before enabling tenant Jobs",
|
||||
));
|
||||
}
|
||||
}
|
||||
for sql in [
|
||||
"ALTER TABLE _dever_jobs ADD COLUMN tenant_id BIGINT NOT NULL DEFAULT 0",
|
||||
"DROP INDEX IF EXISTS _dever_jobs_active",
|
||||
"CREATE UNIQUE INDEX _dever_jobs_active ON _dever_jobs(tenant_id,target,dedupe_key) WHERE state IN ('pending','running')",
|
||||
"UPDATE _dever_job_version SET version=2 WHERE singleton=1",
|
||||
] {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: sql,
|
||||
postgres: sql,
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn claim(
|
||||
database: Database,
|
||||
clock: &Clock,
|
||||
lease_ms: i64,
|
||||
available: usize,
|
||||
) -> Result<Vec<Claim>, Error> {
|
||||
if !(1..=256).contains(&available)
|
||||
|| !(100..=crate::config::MAX_JOB_LEASE_MS).contains(&lease_ms)
|
||||
{
|
||||
return Err(Error::invalid_data("invalid Job claim capacity or lease"));
|
||||
}
|
||||
let tenant_id = database.key().tenant_id().unwrap_or(0);
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = claim_in(&executor, tenant_id, clock, lease_ms, available).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
async fn claim_in(
|
||||
executor: &Executor<'_>,
|
||||
expected_tenant: i64,
|
||||
clock: &Clock,
|
||||
lease_ms: i64,
|
||||
available: usize,
|
||||
) -> Result<Vec<Claim>, Error> {
|
||||
let now = clock.now()?;
|
||||
// Exhausted rows consume bounded cleanup work, never an execution slot. In
|
||||
// particular, an empty claim must not hide eligible jobs behind dead rows.
|
||||
executor.execute(Sql {
|
||||
sqlite: "UPDATE _dever_jobs SET state='dead',error='attempts_exhausted',updated_at=?1 WHERE id IN (SELECT id FROM _dever_jobs WHERE attempt>=max_attempts AND ((state='pending' AND run_at<=?1) OR (state='running' AND lease_until<=?1)) ORDER BY run_at,id LIMIT ?2)",
|
||||
postgres: "UPDATE _dever_jobs SET state='dead',error='attempts_exhausted',updated_at=$1 WHERE id IN (SELECT id FROM _dever_jobs WHERE attempt>=max_attempts AND ((state='pending' AND run_at<=$1) OR (state='running' AND lease_until<=$1)) ORDER BY run_at,id LIMIT $2 FOR UPDATE SKIP LOCKED)",
|
||||
}, vec![n(now), n(available as i64)]).await?;
|
||||
let rows = executor.query(Sql {
|
||||
sqlite: "SELECT id,tenant_id,scope_tenant_id,execution_kind,auth_provider,auth_site,auth_subject,auth_session,auth_tenant,auth_permission,target,schema,payload,attempt,max_attempts,timeout_ms FROM _dever_jobs WHERE attempt<max_attempts AND ((state='pending' AND run_at<=?1) OR (state='running' AND lease_until<=?1)) ORDER BY run_at,id LIMIT ?2",
|
||||
postgres: "SELECT id,tenant_id,scope_tenant_id,execution_kind,auth_provider,auth_site,auth_subject,auth_session,auth_tenant,auth_permission,target,schema,payload,attempt,max_attempts,timeout_ms FROM _dever_jobs WHERE attempt<max_attempts AND ((state='pending' AND run_at<=$1) OR (state='running' AND lease_until<=$1)) ORDER BY run_at,id LIMIT $2 FOR UPDATE SKIP LOCKED",
|
||||
}, vec![n(now), n(available as i64)]).await?;
|
||||
// BEGIN/row locking may wait. Start the lease only after ownership is locked.
|
||||
let now = clock.now()?;
|
||||
let lease_until = now
|
||||
.checked_add(lease_ms)
|
||||
.ok_or_else(|| Error::invalid_data("Job lease overflow"))?;
|
||||
let mut claims = Vec::with_capacity(rows.len());
|
||||
for row in rows {
|
||||
let id = text(&row, 0)?;
|
||||
let tenant = int(&row, 1)?;
|
||||
if tenant != expected_tenant {
|
||||
return Err(Error::migration(
|
||||
"Job tenant ownership does not match its physical database",
|
||||
));
|
||||
}
|
||||
let scope_tenant_id = int(&row, 2)?;
|
||||
if scope_tenant_id < 0 {
|
||||
return Err(Error::migration("Job execution tenant id is invalid"));
|
||||
}
|
||||
let execution = match text(&row, 3)?.as_str() {
|
||||
"system" => ExecutionIdentity::System {
|
||||
tenant_id: (scope_tenant_id > 0).then_some(scope_tenant_id),
|
||||
},
|
||||
"user" => ExecutionIdentity::User {
|
||||
tenant_id: (scope_tenant_id > 0).then_some(scope_tenant_id),
|
||||
provider: text(&row, 4)?,
|
||||
site: text(&row, 5)?,
|
||||
subject: text(&row, 6)?,
|
||||
session: text(&row, 7)?,
|
||||
tenant_claim: match text(&row, 8)? {
|
||||
value if value.is_empty() => None,
|
||||
value => Some(value),
|
||||
},
|
||||
permission_key: text(&row, 9)?,
|
||||
},
|
||||
_ => {
|
||||
return Err(Error::migration(
|
||||
"active Job has no trusted execution identity",
|
||||
));
|
||||
}
|
||||
};
|
||||
execution.validate()?;
|
||||
if tenant > 0 && execution.tenant_id() != Some(tenant) {
|
||||
return Err(Error::migration(
|
||||
"Job execution tenant does not match tenant storage",
|
||||
));
|
||||
}
|
||||
let token = orm::Uuid::new_v7()?.to_string();
|
||||
let changed = executor.execute(Sql {
|
||||
sqlite: "UPDATE _dever_jobs SET state='running',attempt=attempt+1,claim_token=?2,lease_until=?3,updated_at=?4 WHERE id=?1 AND attempt<max_attempts AND ((state='pending' AND run_at<=?4) OR (state='running' AND lease_until<=?4))",
|
||||
postgres: "UPDATE _dever_jobs SET state='running',attempt=attempt+1,claim_token=$2,lease_until=$3,updated_at=$4 WHERE id=$1 AND attempt<max_attempts AND ((state='pending' AND run_at<=$4) OR (state='running' AND lease_until<=$4))",
|
||||
}, vec![t(&id), t(&token), n(lease_until), n(now)]).await?;
|
||||
if changed == 1 {
|
||||
claims.push(Claim {
|
||||
id,
|
||||
tenant_id: (tenant > 0).then_some(tenant),
|
||||
execution,
|
||||
target: text(&row, 10)?,
|
||||
schema: text(&row, 11)?,
|
||||
payload: text(&row, 12)?,
|
||||
attempt: int(&row, 13)? + 1,
|
||||
max_attempts: int(&row, 14)?,
|
||||
timeout_ms: int(&row, 15)?,
|
||||
token,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(claims)
|
||||
}
|
||||
|
||||
pub async fn renew(
|
||||
database: Database,
|
||||
claim: &Claim,
|
||||
clock: &Clock,
|
||||
lease_ms: i64,
|
||||
) -> Result<bool, Error> {
|
||||
if !(100..=crate::config::MAX_JOB_LEASE_MS).contains(&lease_ms) {
|
||||
return Err(Error::invalid_data("invalid Job lease"));
|
||||
}
|
||||
let tenant_id = database.key().tenant_id().unwrap_or(0);
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = renew_in(&executor, tenant_id, claim, clock, lease_ms).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
async fn lock_claim(
|
||||
executor: &Executor<'_>,
|
||||
expected_tenant: i64,
|
||||
claim: &Claim,
|
||||
) -> Result<bool, Error> {
|
||||
let tenant_id = claim.tenant_id.unwrap_or(0);
|
||||
if expected_tenant != tenant_id {
|
||||
return Err(Error::migration(
|
||||
"Job claim tenant does not match its physical database",
|
||||
));
|
||||
}
|
||||
// PostgreSQL needs the same ordering as SQLite's BEGIN IMMEDIATE: obtain the
|
||||
// write lock before reading time, so lock wait cannot revive an expired lease.
|
||||
executor.query(Sql {
|
||||
sqlite: "SELECT id FROM _dever_jobs WHERE id=?1 AND claim_token=?2 AND tenant_id=?3 AND state='running'",
|
||||
postgres: "SELECT id FROM _dever_jobs WHERE id=$1 AND claim_token=$2 AND tenant_id=$3 AND state='running' FOR UPDATE",
|
||||
}, vec![t(&claim.id), t(&claim.token), n(tenant_id)]).await.map(|rows| !rows.is_empty())
|
||||
}
|
||||
|
||||
async fn renew_in(
|
||||
executor: &Executor<'_>,
|
||||
expected_tenant: i64,
|
||||
claim: &Claim,
|
||||
clock: &Clock,
|
||||
lease_ms: i64,
|
||||
) -> Result<bool, Error> {
|
||||
if !lock_claim(executor, expected_tenant, claim).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
let now = clock.now()?;
|
||||
let until = now
|
||||
.checked_add(lease_ms)
|
||||
.ok_or_else(|| Error::invalid_data("Job lease overflow"))?;
|
||||
executor.execute(Sql {
|
||||
sqlite: "UPDATE _dever_jobs SET lease_until=?4,updated_at=?3 WHERE id=?1 AND claim_token=?2 AND state='running' AND lease_until>?3",
|
||||
postgres: "UPDATE _dever_jobs SET lease_until=$4,updated_at=$3 WHERE id=$1 AND claim_token=$2 AND state='running' AND lease_until>$3",
|
||||
}, vec![t(&claim.id), t(&claim.token), n(now), n(until)]).await.map(|count| count == 1)
|
||||
}
|
||||
|
||||
pub async fn complete(
|
||||
database: Database,
|
||||
claim: &Claim,
|
||||
outcome: Outcome,
|
||||
clock: &Clock,
|
||||
retry_base_ms: i64,
|
||||
retry_max_ms: i64,
|
||||
) -> Result<bool, Error> {
|
||||
if retry_base_ms < 1 || retry_max_ms < retry_base_ms || retry_max_ms > 86_400_000 {
|
||||
return Err(Error::invalid_data("invalid Job retry bounds"));
|
||||
}
|
||||
let tenant_id = database.key().tenant_id().unwrap_or(0);
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = complete_in(
|
||||
&executor,
|
||||
tenant_id,
|
||||
claim,
|
||||
outcome,
|
||||
clock,
|
||||
retry_base_ms,
|
||||
retry_max_ms,
|
||||
)
|
||||
.await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
async fn complete_in(
|
||||
executor: &Executor<'_>,
|
||||
expected_tenant: i64,
|
||||
claim: &Claim,
|
||||
outcome: Outcome,
|
||||
clock: &Clock,
|
||||
retry_base_ms: i64,
|
||||
retry_max_ms: i64,
|
||||
) -> Result<bool, Error> {
|
||||
if !lock_claim(executor, expected_tenant, claim).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
let now = clock.now()?;
|
||||
let state = match outcome {
|
||||
Outcome::Success => "succeeded",
|
||||
Outcome::UnknownTarget
|
||||
| Outcome::SchemaMismatch
|
||||
| Outcome::InvalidPayload
|
||||
| Outcome::InvalidPolicy
|
||||
| Outcome::IdentityRejected
|
||||
| Outcome::InvalidIdentity
|
||||
| Outcome::ComponentDisabled => "blocked",
|
||||
_ if claim.attempt >= claim.max_attempts => "dead",
|
||||
_ => "pending",
|
||||
};
|
||||
let exponent = u32::try_from(claim.attempt.saturating_sub(1))
|
||||
.unwrap_or(u32::MAX)
|
||||
.min(62);
|
||||
let delay = retry_base_ms
|
||||
.saturating_mul(1i64 << exponent)
|
||||
.min(retry_max_ms);
|
||||
let run_at = now
|
||||
.checked_add(delay)
|
||||
.ok_or_else(|| Error::invalid_data("Job retry time overflow"))?;
|
||||
executor.execute(Sql {
|
||||
sqlite: "UPDATE _dever_jobs SET state=?4,error=?5,run_at=?6,lease_until=0,claim_token='',updated_at=?3 WHERE id=?1 AND claim_token=?2 AND state='running' AND lease_until>?3",
|
||||
postgres: "UPDATE _dever_jobs SET state=$4,error=$5,run_at=$6,lease_until=0,claim_token='',updated_at=$3 WHERE id=$1 AND claim_token=$2 AND state='running' AND lease_until>$3",
|
||||
}, vec![t(&claim.id), t(&claim.token), n(now), t(state), t(outcome.error()), n(run_at)]).await.map(|count| count == 1)
|
||||
}
|
||||
|
||||
pub async fn materialize(
|
||||
database: Database,
|
||||
spec: &Spec,
|
||||
cron: &crate::cron::Cron,
|
||||
now: i64,
|
||||
) -> Result<usize, Error> {
|
||||
spec.validate()?;
|
||||
let tenant_id = database.key().tenant_id().unwrap_or(0);
|
||||
let transaction = database.begin().await?;
|
||||
let executor = Executor::new(database, Some(&transaction))?;
|
||||
let result = materialize_in(&executor, tenant_id, spec, cron, now.div_euclid(60_000)).await;
|
||||
finish_transaction(transaction, result).await
|
||||
}
|
||||
|
||||
/// Validate persisted schedule identity before any application listener starts.
|
||||
pub async fn validate_schedule(
|
||||
database: Database,
|
||||
spec: &Spec,
|
||||
cron: &crate::cron::Cron,
|
||||
) -> Result<(), Error> {
|
||||
let rows = database
|
||||
.query(
|
||||
Sql {
|
||||
sqlite: "SELECT fingerprint FROM _dever_job_schedule WHERE identity=?1",
|
||||
postgres: "SELECT fingerprint FROM _dever_job_schedule WHERE identity=$1",
|
||||
},
|
||||
vec![t(spec.target)],
|
||||
)
|
||||
.await?;
|
||||
if let Some(row) = rows.first() {
|
||||
validate_fingerprint(row, &cron.fingerprint())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_fingerprint(cursor: &Row, expected: &str) -> Result<(), Error> {
|
||||
if text(cursor, 0)? != expected {
|
||||
return Err(Error::migration("Job schedule fingerprint changed"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn materialize_in(
|
||||
executor: &Executor<'_>,
|
||||
tenant_id: i64,
|
||||
spec: &Spec,
|
||||
cron: &crate::cron::Cron,
|
||||
minute: i64,
|
||||
) -> Result<usize, Error> {
|
||||
let fingerprint = cron.fingerprint();
|
||||
executor.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_job_schedule(identity,fingerprint,last_minute) VALUES (?1,?2,?3) ON CONFLICT DO NOTHING",
|
||||
postgres: "INSERT INTO _dever_job_schedule(identity,fingerprint,last_minute) VALUES ($1,$2,$3) ON CONFLICT DO NOTHING",
|
||||
}, vec![t(spec.target), t(&fingerprint), n(minute)]).await?;
|
||||
let rows = executor.query(Sql {
|
||||
sqlite: "SELECT fingerprint,last_minute FROM _dever_job_schedule WHERE identity=?1",
|
||||
postgres: "SELECT fingerprint,last_minute FROM _dever_job_schedule WHERE identity=$1 FOR UPDATE",
|
||||
}, vec![t(spec.target)]).await?;
|
||||
let cursor = orm::required_row(rows, "Job schedule cursor")?;
|
||||
validate_fingerprint(&cursor, &fingerprint)?;
|
||||
let previous = int(&cursor, 1)?;
|
||||
let end = minute.min(previous.saturating_add(256));
|
||||
let mut inserted = 0;
|
||||
for slot in previous.saturating_add(1)..=end {
|
||||
if cron.matches_minute(slot) {
|
||||
let run_at = slot
|
||||
.checked_mul(60_000)
|
||||
.ok_or_else(|| Error::invalid_data("Job schedule time overflow"))?;
|
||||
let identity = ExecutionIdentity::System {
|
||||
tenant_id: (tenant_id > 0).then_some(tenant_id),
|
||||
};
|
||||
insert(
|
||||
executor,
|
||||
tenant_id,
|
||||
&identity,
|
||||
spec,
|
||||
"null",
|
||||
&format!("schedule:{}:{slot}", spec.target),
|
||||
run_at,
|
||||
)
|
||||
.await?;
|
||||
inserted += 1;
|
||||
}
|
||||
}
|
||||
if end > previous {
|
||||
executor
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: "UPDATE _dever_job_schedule SET last_minute=?2 WHERE identity=?1",
|
||||
postgres: "UPDATE _dever_job_schedule SET last_minute=$2 WHERE identity=$1",
|
||||
},
|
||||
vec![t(spec.target), n(end)],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(inserted)
|
||||
}
|
||||
@@ -0,0 +1,470 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use futures_util::{StreamExt, stream::FuturesUnordered};
|
||||
|
||||
use super::{Clock, Spec, store};
|
||||
use crate::{
|
||||
config::JobSettings,
|
||||
database::{Database, ModelScope, StorageBinding},
|
||||
orm::ErrorKind,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
enum BindingStorage {
|
||||
Fixed(Database),
|
||||
Scoped(StorageBinding),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Binding {
|
||||
storage: BindingStorage,
|
||||
tenant: bool,
|
||||
pub spec: Spec,
|
||||
pub schedule: Option<crate::cron::Cron>,
|
||||
}
|
||||
|
||||
impl Binding {
|
||||
pub fn fixed(
|
||||
database: Database,
|
||||
spec: &'static Spec,
|
||||
schedule: Option<crate::cron::Cron>,
|
||||
) -> Self {
|
||||
Self {
|
||||
storage: BindingStorage::Fixed(database),
|
||||
tenant: false,
|
||||
spec: spec.clone(),
|
||||
schedule,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn scoped(
|
||||
storage: &'static StorageBinding,
|
||||
scope: ModelScope,
|
||||
spec: &'static Spec,
|
||||
schedule: Option<crate::cron::Cron>,
|
||||
) -> Self {
|
||||
Self::owned(*storage, scope, spec.clone(), schedule)
|
||||
}
|
||||
|
||||
pub fn owned(
|
||||
storage: StorageBinding,
|
||||
scope: ModelScope,
|
||||
spec: Spec,
|
||||
schedule: Option<crate::cron::Cron>,
|
||||
) -> Self {
|
||||
Self {
|
||||
storage: BindingStorage::Scoped(storage),
|
||||
tenant: scope == ModelScope::Tenant,
|
||||
spec,
|
||||
schedule,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn uses_tenant_storage(&self) -> bool {
|
||||
self.tenant && crate::config::current_settings().tenant().is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn selector(&self) -> Option<(Option<&str>, &str)> {
|
||||
match &self.storage {
|
||||
BindingStorage::Fixed(_) => None,
|
||||
BindingStorage::Scoped(storage) => Some(storage.selector()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn database(
|
||||
&self,
|
||||
tenant_id: Option<i64>,
|
||||
) -> Result<Database, crate::orm::Error> {
|
||||
match &self.storage {
|
||||
BindingStorage::Fixed(database) => Ok(database.clone()),
|
||||
BindingStorage::Scoped(storage) => match tenant_id {
|
||||
Some(tenant_id) => crate::tenant::scope(tenant_id, storage.database()).await,
|
||||
None => storage.database().await,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn global_connections(bindings: &[Binding]) -> Result<Vec<Database>, String> {
|
||||
let mut databases = BTreeMap::new();
|
||||
for binding in bindings
|
||||
.iter()
|
||||
.filter(|binding| !binding.uses_tenant_storage())
|
||||
{
|
||||
let database = binding
|
||||
.database(None)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
databases.insert(database.key().clone(), database);
|
||||
}
|
||||
Ok(databases.into_values().collect())
|
||||
}
|
||||
|
||||
async fn tenant_connections(bindings: &[Binding], tenant_id: i64) -> Result<Vec<Database>, String> {
|
||||
let mut databases = BTreeMap::new();
|
||||
for binding in bindings
|
||||
.iter()
|
||||
.filter(|binding| binding.uses_tenant_storage())
|
||||
{
|
||||
let database = match binding.database(Some(tenant_id)).await {
|
||||
Ok(database) => database,
|
||||
Err(error) if error.kind() == ErrorKind::PoolExhausted => continue,
|
||||
Err(error) => return Err(error.to_string()),
|
||||
};
|
||||
databases.insert(database.key().clone(), database);
|
||||
}
|
||||
Ok(databases.into_values().collect())
|
||||
}
|
||||
|
||||
async fn schedule_global(bindings: &[Binding], clock: &Clock) -> Result<(), String> {
|
||||
let now = clock.now().map_err(|error| error.to_string())?;
|
||||
for binding in bindings
|
||||
.iter()
|
||||
.filter(|binding| !binding.uses_tenant_storage())
|
||||
{
|
||||
if let Some(cron) = &binding.schedule {
|
||||
let database = binding
|
||||
.database(None)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
store::materialize(database, &binding.spec, cron, now)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn schedule_tenant(
|
||||
bindings: &[Binding],
|
||||
clock: &Clock,
|
||||
tenant_id: i64,
|
||||
) -> Result<(), String> {
|
||||
let now = clock.now().map_err(|error| error.to_string())?;
|
||||
for binding in bindings
|
||||
.iter()
|
||||
.filter(|binding| binding.uses_tenant_storage())
|
||||
{
|
||||
let Some(cron) = &binding.schedule else {
|
||||
continue;
|
||||
};
|
||||
let database = match binding.database(Some(tenant_id)).await {
|
||||
Ok(database) => database,
|
||||
Err(error) if error.kind() == ErrorKind::PoolExhausted => continue,
|
||||
Err(error) => return Err(error.to_string()),
|
||||
};
|
||||
store::materialize(database, &binding.spec, cron, now)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn execute<F, Fut>(
|
||||
database: Database,
|
||||
claim: store::Claim,
|
||||
settings: JobSettings,
|
||||
clock: Clock,
|
||||
dispatch: F,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(store::Claim) -> Fut + Send + 'static,
|
||||
Fut: Future<Output = store::Outcome> + Send + 'static,
|
||||
{
|
||||
if database.key().tenant_id() != claim.tenant_id {
|
||||
return Err("Job claim tenant does not match its physical database".into());
|
||||
}
|
||||
let current_settings = crate::config::current_settings();
|
||||
let connection = current_settings
|
||||
.database(database.name())
|
||||
.ok_or("Job database is not configured")?;
|
||||
let outcome = if settings
|
||||
.validate_lease(connection, claim.timeout_ms)
|
||||
.is_err()
|
||||
|| !(1..=100).contains(&claim.max_attempts)
|
||||
{
|
||||
store::Outcome::InvalidPolicy
|
||||
} else {
|
||||
let identity = claim.execution.clone();
|
||||
let tenant_id = identity.tenant_id();
|
||||
let input = claim.clone();
|
||||
let task = crate::task::run(async move {
|
||||
let execution = async move {
|
||||
match tenant_id {
|
||||
Some(tenant_id) => crate::tenant::scope(tenant_id, dispatch(input)).await,
|
||||
None => dispatch(input).await,
|
||||
}
|
||||
};
|
||||
let outcome = crate::job::scope_execution(identity, execution).await;
|
||||
Ok(outcome)
|
||||
})
|
||||
.await?;
|
||||
match crate::task::wait_lease(task, claim.timeout_ms, settings.lease_ms / 3, || async {
|
||||
store::renew(database.clone(), &claim, &clock, settings.lease_ms)
|
||||
.await
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.await?
|
||||
{
|
||||
crate::task::LeaseResult::Completed(outcome) => outcome,
|
||||
crate::task::LeaseResult::Timeout => store::Outcome::Timeout,
|
||||
crate::task::LeaseResult::Lost => return Ok(()),
|
||||
}
|
||||
};
|
||||
store::complete(
|
||||
database,
|
||||
&claim,
|
||||
outcome,
|
||||
&clock,
|
||||
settings.retry_base_ms,
|
||||
settings.retry_max_ms,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drain_connections<F, Fut>(
|
||||
databases: Vec<Database>,
|
||||
settings: JobSettings,
|
||||
clock: &Clock,
|
||||
limit: i64,
|
||||
count: &mut i64,
|
||||
dispatch: &F,
|
||||
) -> Result<bool, String>
|
||||
where
|
||||
F: Fn(store::Claim) -> Fut + Clone + Send + 'static,
|
||||
Fut: Future<Output = store::Outcome> + Send + 'static,
|
||||
{
|
||||
let mut progress = false;
|
||||
for database in databases {
|
||||
if *count == limit {
|
||||
return Ok(progress);
|
||||
}
|
||||
for claim in store::claim(database.clone(), clock, settings.lease_ms, 1)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?
|
||||
{
|
||||
execute(
|
||||
database.clone(),
|
||||
claim,
|
||||
settings,
|
||||
clock.clone(),
|
||||
dispatch.clone(),
|
||||
)
|
||||
.await?;
|
||||
*count += 1;
|
||||
progress = true;
|
||||
}
|
||||
}
|
||||
Ok(progress)
|
||||
}
|
||||
|
||||
pub async fn drain<F, Fut>(
|
||||
bindings: Vec<Binding>,
|
||||
clock: Clock,
|
||||
limit: i64,
|
||||
dispatch: F,
|
||||
) -> Result<i64, String>
|
||||
where
|
||||
F: Fn(store::Claim) -> Fut + Clone + Send + 'static,
|
||||
Fut: Future<Output = store::Outcome> + Send + 'static,
|
||||
{
|
||||
if !(1..=10_000).contains(&limit) {
|
||||
return Err("test Job drain limit must be between 1 and 10000".into());
|
||||
}
|
||||
let settings = JobSettings::for_test();
|
||||
schedule_global(&bindings, &clock).await?;
|
||||
let tenant_id = if bindings.iter().any(Binding::uses_tenant_storage) {
|
||||
let tenant_id = crate::tenant::current_id().map_err(|error| error.to_string())?;
|
||||
schedule_tenant(&bindings, &clock, tenant_id).await?;
|
||||
Some(tenant_id)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let mut count = 0;
|
||||
loop {
|
||||
let mut progress = drain_connections(
|
||||
global_connections(&bindings).await?,
|
||||
settings,
|
||||
&clock,
|
||||
limit,
|
||||
&mut count,
|
||||
&dispatch,
|
||||
)
|
||||
.await?;
|
||||
if count == limit {
|
||||
return Ok(count);
|
||||
}
|
||||
if let Some(tenant_id) = tenant_id {
|
||||
progress |= drain_connections(
|
||||
tenant_connections(&bindings, tenant_id).await?,
|
||||
settings,
|
||||
&clock,
|
||||
limit,
|
||||
&mut count,
|
||||
&dispatch,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
if !progress {
|
||||
return Ok(count);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type ActiveJob = Pin<Box<dyn Future<Output = Result<(), String>> + Send>>;
|
||||
|
||||
async fn claim_into<F, Fut>(
|
||||
databases: Vec<Database>,
|
||||
settings: JobSettings,
|
||||
clock: &Clock,
|
||||
dispatch: &F,
|
||||
active: &mut FuturesUnordered<ActiveJob>,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(store::Claim) -> Fut + Clone + Send + 'static,
|
||||
Fut: Future<Output = store::Outcome> + Send + 'static,
|
||||
{
|
||||
for database in databases {
|
||||
if crate::lifecycle::stopping() {
|
||||
break;
|
||||
}
|
||||
let available = settings.workers.saturating_sub(active.len());
|
||||
if available == 0 {
|
||||
break;
|
||||
}
|
||||
let claims = store::claim(database.clone(), clock, settings.lease_ms, available)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
if crate::lifecycle::stopping() {
|
||||
break;
|
||||
}
|
||||
for claim in claims {
|
||||
let task = crate::task::run(execute(
|
||||
database.clone(),
|
||||
claim,
|
||||
settings,
|
||||
clock.clone(),
|
||||
dispatch.clone(),
|
||||
))
|
||||
.await?;
|
||||
active.push(Box::pin(crate::task::wait(task)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn claim_ready_tenants<F, Fut>(
|
||||
bindings: &[Binding],
|
||||
settings: JobSettings,
|
||||
clock: &Clock,
|
||||
dispatch: &F,
|
||||
active: &mut FuturesUnordered<ActiveJob>,
|
||||
after: &mut i64,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(store::Claim) -> Fut + Clone + Send + 'static,
|
||||
Fut: Future<Output = store::Outcome> + Send + 'static,
|
||||
{
|
||||
if active.len() >= settings.workers {
|
||||
return Ok(());
|
||||
}
|
||||
let tenant_ids = crate::tenant::ready_ids_after(*after, 256)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
if tenant_ids.is_empty() {
|
||||
*after = 0;
|
||||
return Ok(());
|
||||
}
|
||||
for tenant_id in tenant_ids {
|
||||
if crate::lifecycle::stopping() || active.len() >= settings.workers {
|
||||
break;
|
||||
}
|
||||
schedule_tenant(bindings, clock, tenant_id).await?;
|
||||
claim_into(
|
||||
tenant_connections(bindings, tenant_id).await?,
|
||||
settings,
|
||||
clock,
|
||||
dispatch,
|
||||
active,
|
||||
)
|
||||
.await?;
|
||||
*after = tenant_id;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn serve<F, Fut>(
|
||||
bindings: Vec<Binding>,
|
||||
clock: Clock,
|
||||
settings: JobSettings,
|
||||
dispatch: F,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
F: Fn(store::Claim) -> Fut + Clone + Send + 'static,
|
||||
Fut: Future<Output = store::Outcome> + Send + 'static,
|
||||
{
|
||||
let settings = settings.validate()?;
|
||||
let has_tenant_jobs = bindings.iter().any(Binding::uses_tenant_storage);
|
||||
let mut tenant_after = 0;
|
||||
let mut prefer_tenant = has_tenant_jobs;
|
||||
let mut active = FuturesUnordered::new();
|
||||
loop {
|
||||
if crate::lifecycle::stopping() {
|
||||
break;
|
||||
}
|
||||
schedule_global(&bindings, &clock).await?;
|
||||
if prefer_tenant {
|
||||
claim_ready_tenants(
|
||||
&bindings,
|
||||
settings,
|
||||
&clock,
|
||||
&dispatch,
|
||||
&mut active,
|
||||
&mut tenant_after,
|
||||
)
|
||||
.await?;
|
||||
claim_into(
|
||||
global_connections(&bindings).await?,
|
||||
settings,
|
||||
&clock,
|
||||
&dispatch,
|
||||
&mut active,
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
claim_into(
|
||||
global_connections(&bindings).await?,
|
||||
settings,
|
||||
&clock,
|
||||
&dispatch,
|
||||
&mut active,
|
||||
)
|
||||
.await?;
|
||||
if has_tenant_jobs {
|
||||
claim_ready_tenants(
|
||||
&bindings,
|
||||
settings,
|
||||
&clock,
|
||||
&dispatch,
|
||||
&mut active,
|
||||
&mut tenant_after,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
prefer_tenant = has_tenant_jobs && !prefer_tenant;
|
||||
tokio::select! {
|
||||
_ = crate::lifecycle::shutdown_requested() => break,
|
||||
Some(result) = active.next(), if !active.is_empty() => result?,
|
||||
_ = tokio::time::sleep(std::time::Duration::from_millis(settings.poll_ms as u64)) => {}
|
||||
}
|
||||
}
|
||||
while let Some(result) = active.next().await {
|
||||
result?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
use std::io::{self, Write};
|
||||
|
||||
pub mod abi;
|
||||
#[cfg(feature = "wire")]
|
||||
pub mod api;
|
||||
#[cfg(feature = "api")]
|
||||
pub mod application;
|
||||
pub mod async_stream;
|
||||
#[cfg(feature = "api")]
|
||||
pub mod auth;
|
||||
pub mod bytes;
|
||||
pub mod channel;
|
||||
pub mod collections;
|
||||
#[cfg(feature = "wire")]
|
||||
pub mod component;
|
||||
pub mod concurrent;
|
||||
#[cfg(any(feature = "database", feature = "wire"))]
|
||||
pub mod config;
|
||||
pub mod cron;
|
||||
#[cfg(feature = "crypto")]
|
||||
pub mod crypto;
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub mod database;
|
||||
mod endpoint;
|
||||
#[cfg(feature = "external")]
|
||||
pub mod external;
|
||||
pub mod http;
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub mod job;
|
||||
#[cfg(any(feature = "database", feature = "wire"))]
|
||||
pub mod lifecycle;
|
||||
pub mod log;
|
||||
pub mod net;
|
||||
pub mod number;
|
||||
pub mod orm;
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
mod pool;
|
||||
#[cfg(feature = "postgres")]
|
||||
pub mod postgres;
|
||||
pub mod process;
|
||||
pub mod render;
|
||||
pub mod resource;
|
||||
pub mod secret;
|
||||
#[cfg(feature = "sqlite")]
|
||||
pub mod sqlite;
|
||||
pub mod sse;
|
||||
pub mod task;
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub mod tenant;
|
||||
pub mod text;
|
||||
pub mod time;
|
||||
pub mod tls;
|
||||
mod transport;
|
||||
pub mod websocket;
|
||||
#[cfg(feature = "wire")]
|
||||
pub mod wire;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct Id(pub String);
|
||||
|
||||
/// 将故障格式化移出正常计算路径，避免每个数值操作携带重复格式化代码。
|
||||
#[cold]
|
||||
#[inline(never)]
|
||||
pub fn located_error(location: &str, error: impl std::fmt::Display) -> String {
|
||||
format!("{location}: {error}")
|
||||
}
|
||||
|
||||
pub fn println(text: &str) -> io::Result<()> {
|
||||
write_line(&mut io::stdout().lock(), text)
|
||||
}
|
||||
|
||||
pub fn write_line(output: &mut impl Write, text: &str) -> io::Result<()> {
|
||||
writeln!(output, "{text}")?;
|
||||
output.flush()
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
//! One application owner broadcasts stop-admission and joins its structured root.
|
||||
use crate::config::{RuntimeMode, RuntimeSettings};
|
||||
use std::future::Future;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
pub(crate) struct Lifecycle {
|
||||
settings: RuntimeSettings,
|
||||
stop: tokio::sync::watch::Sender<bool>,
|
||||
}
|
||||
|
||||
static LIFECYCLE: OnceLock<Arc<Lifecycle>> = OnceLock::new();
|
||||
|
||||
fn current() -> Option<Arc<Lifecycle>> {
|
||||
#[cfg(feature = "api")]
|
||||
if let Some(session) = crate::application::current() {
|
||||
return session.lifecycle();
|
||||
}
|
||||
LIFECYCLE.get().cloned()
|
||||
}
|
||||
|
||||
pub fn configure(api: bool, worker: bool) -> Result<(), String> {
|
||||
let settings = crate::config::current_settings().services(api, worker)?;
|
||||
let (stop, _) = tokio::sync::watch::channel(false);
|
||||
let state = Arc::new(Lifecycle { settings, stop });
|
||||
#[cfg(feature = "api")]
|
||||
if let Some(session) = crate::application::current() {
|
||||
return session.set_lifecycle(state);
|
||||
}
|
||||
LIFECYCLE
|
||||
.set(state)
|
||||
.map_err(|_| "application lifecycle initialized twice".into())
|
||||
}
|
||||
|
||||
pub fn api_enabled() -> bool {
|
||||
current().is_none_or(|state| matches!(state.settings.mode, RuntimeMode::Api | RuntimeMode::All))
|
||||
}
|
||||
|
||||
pub fn worker_enabled() -> bool {
|
||||
current()
|
||||
.is_some_and(|state| matches!(state.settings.mode, RuntimeMode::Worker | RuntimeMode::All))
|
||||
}
|
||||
|
||||
pub fn stopping() -> bool {
|
||||
current().is_some_and(|state| *state.stop.borrow())
|
||||
}
|
||||
|
||||
pub async fn shutdown_requested() {
|
||||
let Some(state) = current() else {
|
||||
return std::future::pending().await;
|
||||
};
|
||||
let mut receiver = state.stop.subscribe();
|
||||
while !*receiver.borrow_and_update() {
|
||||
if receiver.changed().await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run<F, T, E>(future: F) -> Result<T, E>
|
||||
where
|
||||
F: Future<Output = Result<T, E>> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
E: From<String> + Send + 'static,
|
||||
{
|
||||
let state =
|
||||
current().ok_or_else(|| E::from("application lifecycle is not configured".into()))?;
|
||||
#[cfg(unix)]
|
||||
let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
||||
.map_err(|error| E::from(format!("cannot register SIGTERM: {error}")))?;
|
||||
#[cfg(unix)]
|
||||
let mut interrupt = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
|
||||
.map_err(|error| E::from(format!("cannot register SIGINT: {error}")))?;
|
||||
let signal = async {
|
||||
#[cfg(unix)]
|
||||
tokio::select! { _ = terminate.recv() => {}, _ = interrupt.recv() => {} }
|
||||
#[cfg(not(unix))]
|
||||
tokio::signal::ctrl_c()
|
||||
.await
|
||||
.map_err(|error| format!("cannot receive Ctrl-C: {error}"))?;
|
||||
state.stop.send_replace(true);
|
||||
Ok(())
|
||||
};
|
||||
let root = crate::task::run_typed(future).await.map_err(E::from)?;
|
||||
crate::task::wait_shutdown(root, signal, state.settings.shutdown_ms).await
|
||||
}
|
||||
|
||||
pub async fn serve_both<A, W>(api: A, worker: W) -> Result<(), String>
|
||||
where
|
||||
A: Future<Output = Result<(), String>>,
|
||||
W: Future<Output = Result<(), String>>,
|
||||
{
|
||||
tokio::try_join!(api, worker)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::io::{self, BufWriter, Write};
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Condvar, Mutex, MutexGuard, OnceLock};
|
||||
|
||||
const QUEUE_CAPACITY: usize = 128;
|
||||
const WRITER_STACK_BYTES: usize = 128 * 1024;
|
||||
const MAX_MESSAGE_BYTES: usize = 4 * 1024;
|
||||
const MAX_FIELDS: usize = 16;
|
||||
const MAX_FIELD_NAME_BYTES: usize = 128;
|
||||
const MAX_FIELD_VALUE_BYTES: usize = 1024;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum Level {
|
||||
Debug,
|
||||
Info,
|
||||
Warn,
|
||||
Error,
|
||||
}
|
||||
|
||||
impl Level {
|
||||
fn name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Debug => "debug",
|
||||
Self::Info => "info",
|
||||
Self::Warn => "warn",
|
||||
Self::Error => "error",
|
||||
}
|
||||
}
|
||||
|
||||
fn drops_when_full(self) -> bool {
|
||||
matches!(self, Self::Debug | Self::Info)
|
||||
}
|
||||
|
||||
#[cfg(feature = "wire")]
|
||||
fn enabled(self, configured: crate::config::LogLevel) -> bool {
|
||||
let actual = match self {
|
||||
Self::Debug => 0,
|
||||
Self::Info => 1,
|
||||
Self::Warn => 2,
|
||||
Self::Error => 3,
|
||||
};
|
||||
let minimum = match configured {
|
||||
crate::config::LogLevel::Debug => 0,
|
||||
crate::config::LogLevel::Info => 1,
|
||||
crate::config::LogLevel::Warn => 2,
|
||||
crate::config::LogLevel::Error => 3,
|
||||
};
|
||||
actual >= minimum
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct Field {
|
||||
pub name: String,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
struct Record {
|
||||
level: Level,
|
||||
message: String,
|
||||
fields: Vec<Field>,
|
||||
dropped: u64,
|
||||
}
|
||||
|
||||
enum Command {
|
||||
Record(Record),
|
||||
Flush,
|
||||
}
|
||||
|
||||
struct Logger {
|
||||
queue: Mutex<Queue>,
|
||||
available: Condvar,
|
||||
space: Condvar,
|
||||
flushed: Condvar,
|
||||
flush_call: Mutex<()>,
|
||||
dropped: AtomicU64,
|
||||
}
|
||||
|
||||
struct Queue {
|
||||
commands: VecDeque<Command>,
|
||||
closed: bool,
|
||||
flushed: bool,
|
||||
}
|
||||
|
||||
struct WriterGuard<'a>(&'a Logger);
|
||||
|
||||
impl Drop for WriterGuard<'_> {
|
||||
fn drop(&mut self) {
|
||||
self.0.close();
|
||||
}
|
||||
}
|
||||
|
||||
static LOGGER: OnceLock<Result<Arc<Logger>, String>> = OnceLock::new();
|
||||
|
||||
pub fn debug(message: impl Into<String>, fields: Vec<Field>) {
|
||||
emit(Level::Debug, message, fields);
|
||||
}
|
||||
|
||||
pub fn info(message: impl Into<String>, fields: Vec<Field>) {
|
||||
emit(Level::Info, message, fields);
|
||||
}
|
||||
|
||||
pub fn warn(message: impl Into<String>, fields: Vec<Field>) {
|
||||
emit(Level::Warn, message, fields);
|
||||
}
|
||||
|
||||
pub fn error(message: impl Into<String>, fields: Vec<Field>) {
|
||||
emit(Level::Error, message, fields);
|
||||
}
|
||||
|
||||
pub fn emit(level: Level, message: impl Into<String>, fields: Vec<Field>) {
|
||||
#[cfg(feature = "wire")]
|
||||
if crate::config::try_current_settings()
|
||||
.is_some_and(|settings| !level.enabled(settings.log().level))
|
||||
{
|
||||
return;
|
||||
}
|
||||
#[cfg(feature = "api")]
|
||||
let fields = {
|
||||
let mut fields = fields;
|
||||
for field in crate::auth::log_fields() {
|
||||
if !fields.iter().any(|existing| existing.name == field.name) {
|
||||
fields.push(field);
|
||||
}
|
||||
}
|
||||
fields
|
||||
};
|
||||
let record = Record {
|
||||
level,
|
||||
message: truncate(message.into(), MAX_MESSAGE_BYTES),
|
||||
fields: fields
|
||||
.into_iter()
|
||||
.take(MAX_FIELDS)
|
||||
.map(|field| Field {
|
||||
name: truncate(field.name, MAX_FIELD_NAME_BYTES),
|
||||
value: truncate(field.value, MAX_FIELD_VALUE_BYTES),
|
||||
})
|
||||
.collect(),
|
||||
dropped: 0,
|
||||
};
|
||||
match logger() {
|
||||
Ok(logger) => logger.submit(record),
|
||||
Err(error) => emergency(&Record {
|
||||
level: Level::Error,
|
||||
message: error.clone(),
|
||||
fields: Vec::new(),
|
||||
dropped: 0,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn flush() {
|
||||
if let Ok(logger) = logger() {
|
||||
logger.flush();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dropped_events() -> u64 {
|
||||
logger()
|
||||
.as_ref()
|
||||
.map(|logger| logger.dropped.load(Ordering::Relaxed))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn logger() -> &'static Result<Arc<Logger>, String> {
|
||||
LOGGER.get_or_init(|| {
|
||||
let logger = Arc::new(Logger::new());
|
||||
let writer = Arc::clone(&logger);
|
||||
std::thread::Builder::new()
|
||||
.name("dever-log".into())
|
||||
.stack_size(WRITER_STACK_BYTES)
|
||||
.spawn(move || {
|
||||
let mut output = BufWriter::new(io::stderr().lock());
|
||||
writer.write(&mut output);
|
||||
})
|
||||
.map_err(|error| format!("cannot start Dever log writer: {error}"))?;
|
||||
Ok(logger)
|
||||
})
|
||||
}
|
||||
|
||||
impl Logger {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
queue: Mutex::new(Queue {
|
||||
commands: VecDeque::with_capacity(QUEUE_CAPACITY),
|
||||
closed: false,
|
||||
flushed: false,
|
||||
}),
|
||||
available: Condvar::new(),
|
||||
space: Condvar::new(),
|
||||
flushed: Condvar::new(),
|
||||
flush_call: Mutex::new(()),
|
||||
dropped: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn queue(&self) -> MutexGuard<'_, Queue> {
|
||||
self.queue.lock().unwrap_or_else(|error| error.into_inner())
|
||||
}
|
||||
|
||||
fn wait_for_space<'a>(&self, mut queue: MutexGuard<'a, Queue>) -> MutexGuard<'a, Queue> {
|
||||
while !queue.closed && queue.commands.len() == QUEUE_CAPACITY {
|
||||
queue = self
|
||||
.space
|
||||
.wait(queue)
|
||||
.unwrap_or_else(|error| error.into_inner());
|
||||
}
|
||||
queue
|
||||
}
|
||||
|
||||
fn submit(&self, mut record: Record) {
|
||||
record.dropped = self.dropped.swap(0, Ordering::Relaxed);
|
||||
let mut queue = self.queue();
|
||||
if !queue.closed && queue.commands.len() == QUEUE_CAPACITY && record.level.drops_when_full()
|
||||
{
|
||||
self.dropped
|
||||
.fetch_add(record.dropped.saturating_add(1), Ordering::Relaxed);
|
||||
return;
|
||||
}
|
||||
queue = self.wait_for_space(queue);
|
||||
if queue.closed {
|
||||
drop(queue);
|
||||
emergency(&record);
|
||||
return;
|
||||
}
|
||||
queue.commands.push_back(Command::Record(record));
|
||||
self.available.notify_one();
|
||||
}
|
||||
|
||||
fn flush(&self) {
|
||||
// A single pending flush pairs the fixed acknowledgement with its
|
||||
// caller. Queue storage and waiting state never allocate after startup.
|
||||
let _caller = self
|
||||
.flush_call
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner());
|
||||
let mut queue = self.wait_for_space(self.queue());
|
||||
if queue.closed {
|
||||
return;
|
||||
}
|
||||
queue.flushed = false;
|
||||
queue.commands.push_back(Command::Flush);
|
||||
self.available.notify_one();
|
||||
while !queue.flushed && !queue.closed {
|
||||
queue = self
|
||||
.flushed
|
||||
.wait(queue)
|
||||
.unwrap_or_else(|error| error.into_inner());
|
||||
}
|
||||
}
|
||||
|
||||
fn receive(&self) -> Option<Command> {
|
||||
let mut queue = self.queue();
|
||||
while queue.commands.is_empty() && !queue.closed {
|
||||
queue = self
|
||||
.available
|
||||
.wait(queue)
|
||||
.unwrap_or_else(|error| error.into_inner());
|
||||
}
|
||||
let command = queue.commands.pop_front();
|
||||
self.space.notify_one();
|
||||
command
|
||||
}
|
||||
|
||||
fn write(&self, output: &mut impl Write) {
|
||||
let _guard = WriterGuard(self);
|
||||
while let Some(command) = self.receive() {
|
||||
match command {
|
||||
Command::Record(record) => {
|
||||
if write_record(output, &record).is_err() {
|
||||
emergency(&record);
|
||||
}
|
||||
}
|
||||
Command::Flush => {
|
||||
let _ = output.flush();
|
||||
let mut queue = self.queue();
|
||||
queue.flushed = true;
|
||||
self.flushed.notify_one();
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = output.flush();
|
||||
}
|
||||
|
||||
fn close(&self) {
|
||||
let mut queue = self.queue();
|
||||
queue.closed = true;
|
||||
queue.commands.clear();
|
||||
self.available.notify_all();
|
||||
self.space.notify_all();
|
||||
self.flushed.notify_all();
|
||||
}
|
||||
}
|
||||
|
||||
fn write_record(output: &mut impl Write, record: &Record) -> io::Result<()> {
|
||||
write!(
|
||||
output,
|
||||
"{{\"level\":\"{}\",\"message\":",
|
||||
record.level.name()
|
||||
)?;
|
||||
write_json_string(output, &record.message)?;
|
||||
if record.dropped > 0 {
|
||||
write!(output, ",\"dropped_events\":{}", record.dropped)?;
|
||||
}
|
||||
write!(output, ",\"fields\":{{")?;
|
||||
for (index, field) in record.fields.iter().enumerate() {
|
||||
if index > 0 {
|
||||
output.write_all(b",")?;
|
||||
}
|
||||
write_json_string(output, &field.name)?;
|
||||
output.write_all(b":")?;
|
||||
write_json_string(output, &field.value)?;
|
||||
}
|
||||
output.write_all(b"}}\n")?;
|
||||
if record.level == Level::Error {
|
||||
output.flush()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_json_string(output: &mut impl Write, value: &str) -> io::Result<()> {
|
||||
output.write_all(b"\"")?;
|
||||
for character in value.chars() {
|
||||
match character {
|
||||
'"' => output.write_all(b"\\\"")?,
|
||||
'\\' => output.write_all(b"\\\\")?,
|
||||
'\n' => output.write_all(b"\\n")?,
|
||||
'\r' => output.write_all(b"\\r")?,
|
||||
'\t' => output.write_all(b"\\t")?,
|
||||
character if character.is_control() => write!(output, "\\u{:04x}", character as u32)?,
|
||||
character => write!(output, "{character}")?,
|
||||
}
|
||||
}
|
||||
output.write_all(b"\"")
|
||||
}
|
||||
|
||||
fn emergency(record: &Record) {
|
||||
let mut output = io::stderr().lock();
|
||||
let _ = write_record(&mut output, record);
|
||||
}
|
||||
|
||||
fn truncate(mut value: String, limit: usize) -> String {
|
||||
if value.len() <= limit {
|
||||
return value;
|
||||
}
|
||||
let end = value
|
||||
.char_indices()
|
||||
.map(|(index, _)| index)
|
||||
.take_while(|index| *index <= limit)
|
||||
.last()
|
||||
.unwrap_or(0);
|
||||
value.truncate(end);
|
||||
value
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../../../test/runtime_logging.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,297 @@
|
||||
use std::future::Future;
|
||||
use std::io::ErrorKind;
|
||||
use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::stream;
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio::sync::Mutex as AsyncMutex;
|
||||
|
||||
use crate::async_stream::AsyncStream;
|
||||
use crate::bytes::Bytes;
|
||||
use crate::endpoint::Endpoint;
|
||||
use crate::resource::{Pull, read_limit};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Listener(Arc<Endpoint<TcpListener>>);
|
||||
|
||||
impl Listener {
|
||||
pub fn close(&self) -> Result<(), String> {
|
||||
self.0.close()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct SocketState {
|
||||
endpoint: Endpoint<TcpStream>,
|
||||
reader: AsyncMutex<()>,
|
||||
writer: AsyncMutex<()>,
|
||||
timeout_millis: AtomicU64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Socket(Arc<SocketState>);
|
||||
|
||||
impl Socket {
|
||||
fn new(stream: TcpStream) -> Self {
|
||||
Self(Arc::new(SocketState {
|
||||
endpoint: Endpoint::new(stream),
|
||||
reader: AsyncMutex::new(()),
|
||||
writer: AsyncMutex::new(()),
|
||||
timeout_millis: AtomicU64::new(0),
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn close(&self) -> Result<(), String> {
|
||||
self.0.endpoint.close()
|
||||
}
|
||||
|
||||
async fn operation<R>(
|
||||
&self,
|
||||
operation: impl Future<Output = Result<R, String>>,
|
||||
) -> Result<R, String> {
|
||||
let operation = self.0.endpoint.while_open(operation);
|
||||
match self.0.timeout_millis.load(Ordering::Relaxed) {
|
||||
0 => operation.await,
|
||||
millis => timed(Duration::from_millis(millis), operation).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn connect(host: &str, port: i64) -> Result<Socket, String> {
|
||||
connect_stream(host, port).await.map(Socket::new)
|
||||
}
|
||||
|
||||
pub(crate) async fn connect_stream(host: &str, port: i64) -> Result<TcpStream, String> {
|
||||
let addresses = addresses(host, port).await?;
|
||||
let stream = TcpStream::connect(addresses.as_slice())
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
prepare_stream(stream)
|
||||
}
|
||||
|
||||
pub async fn connect_timeout(address: &str, port: i64, millis: i64) -> Result<Socket, String> {
|
||||
connect_stream_timeout(address, port, millis)
|
||||
.await
|
||||
.map(Socket::new)
|
||||
}
|
||||
|
||||
pub(crate) async fn connect_stream_timeout(
|
||||
address: &str,
|
||||
port: i64,
|
||||
millis: i64,
|
||||
) -> Result<TcpStream, String> {
|
||||
let address = address
|
||||
.parse::<IpAddr>()
|
||||
.map_err(|_| "timed connect requires a numeric IP address")?;
|
||||
let address = SocketAddr::new(address, port_number(port)?);
|
||||
timed(timeout_duration(millis)?, async {
|
||||
let stream = TcpStream::connect(address)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
prepare_stream(stream)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn listen(host: &str, port: i64) -> Result<Listener, String> {
|
||||
let addresses = addresses(host, port).await?;
|
||||
let listener = TcpListener::bind(addresses.as_slice())
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(Listener(Arc::new(Endpoint::new(listener))))
|
||||
}
|
||||
|
||||
pub async fn accept(listener: &Listener) -> Result<Socket, String> {
|
||||
accept_stream(listener).await.map(Socket::new)
|
||||
}
|
||||
|
||||
async fn accept_stream(listener: &Listener) -> Result<TcpStream, String> {
|
||||
let socket = listener.0.get()?;
|
||||
listener
|
||||
.0
|
||||
.while_open(async {
|
||||
let (stream, _) = socket.accept().await.map_err(|error| error.to_string())?;
|
||||
prepare_stream(stream)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn port(listener: &Listener) -> Result<i64, String> {
|
||||
listener
|
||||
.0
|
||||
.get()?
|
||||
.local_addr()
|
||||
.map(|address| i64::from(address.port()))
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
fn prepare_stream(stream: TcpStream) -> Result<TcpStream, String> {
|
||||
stream
|
||||
.set_nodelay(true)
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(stream)
|
||||
}
|
||||
|
||||
pub fn timeout(socket: &Socket, millis: i64) -> Result<(), String> {
|
||||
let duration = timeout_duration(millis)?;
|
||||
socket.0.endpoint.get()?;
|
||||
socket
|
||||
.0
|
||||
.timeout_millis
|
||||
.store(duration.as_millis() as u64, Ordering::Relaxed);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn read(socket: &Socket, limit: i64) -> Result<Option<Bytes>, String> {
|
||||
let limit = read_limit(limit)?;
|
||||
let stream = socket.0.endpoint.get()?;
|
||||
socket
|
||||
.operation(async {
|
||||
let _reader = socket.0.reader.lock().await;
|
||||
// 等待读锁和可读事件后再分配缓冲区，避免等待者各占一份内存。
|
||||
stream.readable().await.map_err(|error| error.to_string())?;
|
||||
let mut bytes = Vec::new();
|
||||
bytes
|
||||
.try_reserve_exact(limit)
|
||||
.map_err(|error| error.to_string())?;
|
||||
bytes.resize(limit, 0);
|
||||
loop {
|
||||
match stream.try_read(&mut bytes) {
|
||||
Ok(length) => {
|
||||
bytes.truncate(length);
|
||||
return Ok((length != 0).then(|| Bytes::new(bytes)));
|
||||
}
|
||||
Err(error) if error.kind() == ErrorKind::WouldBlock => {
|
||||
stream.readable().await.map_err(|error| error.to_string())?;
|
||||
}
|
||||
Err(error) => return Err(error.to_string()),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
// 部分写入中断时先关闭连接，再释放写锁，阻止下一次写入接在残缺内容后。
|
||||
struct WriteProgress<'a> {
|
||||
socket: &'a Socket,
|
||||
partial: bool,
|
||||
}
|
||||
|
||||
impl Drop for WriteProgress<'_> {
|
||||
fn drop(&mut self) {
|
||||
if self.partial {
|
||||
let _ = self.socket.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn write(socket: &Socket, bytes: &Bytes) -> Result<(), String> {
|
||||
let stream = socket.0.endpoint.get()?;
|
||||
socket
|
||||
.operation(async {
|
||||
let _writer = socket.0.writer.lock().await;
|
||||
let mut progress = WriteProgress {
|
||||
socket,
|
||||
partial: false,
|
||||
};
|
||||
let mut offset = 0;
|
||||
while offset < bytes.values().len() {
|
||||
stream.writable().await.map_err(|error| error.to_string())?;
|
||||
match stream.try_write(&bytes.values()[offset..]) {
|
||||
Ok(0) => return Err("socket write returned zero bytes".into()),
|
||||
Ok(length) => {
|
||||
offset += length;
|
||||
progress.partial = offset < bytes.values().len();
|
||||
}
|
||||
Err(error) if error.kind() == ErrorKind::WouldBlock => {}
|
||||
Err(error) => return Err(error.to_string()),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn connections(listener: Listener) -> AsyncStream<Result<Socket, String>> {
|
||||
incoming(listener, false).map(|connection| connection.map(Socket::new))
|
||||
}
|
||||
|
||||
pub(crate) fn incoming(
|
||||
listener: Listener,
|
||||
end_on_close: bool,
|
||||
) -> AsyncStream<Result<TcpStream, String>> {
|
||||
AsyncStream::new(stream::unfold(listener, move |listener| async move {
|
||||
let result = tokio::select! {
|
||||
biased;
|
||||
_ = listener.closed(), if end_on_close => return None,
|
||||
result = accept_stream(&listener) => result,
|
||||
};
|
||||
let event = match result {
|
||||
Ok(socket) => Pull::Item(Ok(socket)),
|
||||
Err(error) => Pull::Last(Err(error)),
|
||||
};
|
||||
Some((event, listener))
|
||||
}))
|
||||
}
|
||||
|
||||
impl Listener {
|
||||
pub(crate) async fn closed(&self) {
|
||||
self.0.closed().await;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn chunks(socket: Socket, limit: i64) -> Result<AsyncStream<Result<Bytes, String>>, String> {
|
||||
read_limit(limit)?;
|
||||
socket.0.endpoint.get()?;
|
||||
Ok(AsyncStream::new(stream::unfold(
|
||||
socket,
|
||||
move |socket| async move {
|
||||
let event = match read(&socket, limit).await {
|
||||
Ok(Some(bytes)) => Pull::Item(Ok(bytes)),
|
||||
Ok(None) => return None,
|
||||
Err(error) => Pull::Last(Err(error)),
|
||||
};
|
||||
Some((event, socket))
|
||||
},
|
||||
)))
|
||||
}
|
||||
|
||||
async fn addresses(host: &str, port: i64) -> Result<Vec<SocketAddr>, String> {
|
||||
let port = port_number(port)?;
|
||||
if let Ok(address) = host.parse::<IpAddr>() {
|
||||
return Ok(vec![SocketAddr::new(address, port)]);
|
||||
}
|
||||
let host = host.to_owned();
|
||||
// DNS 复用有界 blocking 通道及其结构化清理。
|
||||
crate::task::blocking(move || {
|
||||
(host.as_str(), port)
|
||||
.to_socket_addrs()
|
||||
.map(Iterator::collect)
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
fn port_number(port: i64) -> Result<u16, String> {
|
||||
u16::try_from(port).map_err(|_| "port must be between 0 and 65535".into())
|
||||
}
|
||||
|
||||
fn timeout_duration(millis: i64) -> Result<Duration, String> {
|
||||
let millis = u64::try_from(millis)
|
||||
.ok()
|
||||
.filter(|millis| *millis > 0)
|
||||
.ok_or("timeout must be a positive Int")?;
|
||||
Ok(Duration::from_millis(millis))
|
||||
}
|
||||
|
||||
async fn timed<T>(
|
||||
duration: Duration,
|
||||
operation: impl Future<Output = Result<T, String>>,
|
||||
) -> Result<T, String> {
|
||||
tokio::time::timeout(duration, operation)
|
||||
.await
|
||||
.map_err(|_| "network operation timed out".to_owned())?
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
use std::cell::RefCell;
|
||||
use std::cmp::Ordering;
|
||||
use std::fmt;
|
||||
|
||||
use dec::{Context, Decimal128, Rounding};
|
||||
|
||||
thread_local! {
|
||||
static CONTEXT: RefCell<Context<Decimal128>> = RefCell::new(Context::default());
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct DecimalValue(Decimal128);
|
||||
|
||||
impl Eq for DecimalValue {}
|
||||
|
||||
impl PartialOrd for DecimalValue {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
|
||||
impl Ord for DecimalValue {
|
||||
fn cmp(&self, other: &Self) -> Ordering {
|
||||
self.0
|
||||
.partial_cmp(&other.0)
|
||||
.expect("Decimal values are finite")
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for DecimalValue {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let reduced = Context::<Decimal128>::default().reduce(self.0);
|
||||
if reduced.is_zero() {
|
||||
formatter.write_str("0")
|
||||
} else {
|
||||
write!(formatter, "{reduced}")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DecimalValue {
|
||||
pub const ZERO: Self = Self(Decimal128::ZERO);
|
||||
|
||||
pub fn parse(text: &str) -> Result<Self, &'static str> {
|
||||
CONTEXT.with_borrow_mut(|context| {
|
||||
context.clear_status();
|
||||
let value = context.parse(text).map_err(|_| "invalid Decimal literal")?;
|
||||
let value = finite_result(context, value)?;
|
||||
if context.status().inexact() {
|
||||
return Err("Decimal literal is not exactly representable in 34 digits");
|
||||
}
|
||||
Ok(value)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn from_int(value: i64) -> Self {
|
||||
Self(Decimal128::from(value))
|
||||
}
|
||||
|
||||
pub fn from_bytes(bytes: [u8; 16]) -> Result<Self, &'static str> {
|
||||
let value = Decimal128::from_le_bytes(bytes);
|
||||
if value.is_finite() {
|
||||
Ok(Self(value))
|
||||
} else {
|
||||
Err("Decimal must be finite")
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_bytes(self) -> [u8; 16] {
|
||||
self.0.to_le_bytes()
|
||||
}
|
||||
|
||||
pub fn checked_add(self, other: Self) -> Result<Self, &'static str> {
|
||||
operation(|context| context.add(self.0, other.0))
|
||||
}
|
||||
|
||||
pub fn checked_sub(self, other: Self) -> Result<Self, &'static str> {
|
||||
operation(|context| context.sub(self.0, other.0))
|
||||
}
|
||||
|
||||
pub fn checked_mul(self, other: Self) -> Result<Self, &'static str> {
|
||||
operation(|context| context.mul(self.0, other.0))
|
||||
}
|
||||
|
||||
pub fn checked_div(self, other: Self) -> Result<Self, &'static str> {
|
||||
operation(|context| context.div(self.0, other.0))
|
||||
}
|
||||
|
||||
pub fn checked_neg(self) -> Result<Self, &'static str> {
|
||||
operation(|context| context.minus(self.0))
|
||||
}
|
||||
|
||||
pub fn round(self, places: i64) -> Result<Self, &'static str> {
|
||||
let exponent = places
|
||||
.checked_neg()
|
||||
.and_then(|value| i32::try_from(value).ok())
|
||||
.ok_or("Decimal rounding precision is out of range")?;
|
||||
if !(-6176..=6111).contains(&exponent) {
|
||||
return Err("Decimal rounding precision is out of range");
|
||||
}
|
||||
operation(|context| {
|
||||
let mut quantum = Decimal128::ONE;
|
||||
context.set_exponent(&mut quantum, exponent);
|
||||
context.quantize(self.0, quantum)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn next_up(self) -> Option<Self> {
|
||||
let value = Context::<Decimal128>::default().next_plus(self.0);
|
||||
value.is_finite().then_some(Self(value))
|
||||
}
|
||||
|
||||
pub fn next_down(self) -> Option<Self> {
|
||||
let value = Context::<Decimal128>::default().next_minus(self.0);
|
||||
value.is_finite().then_some(Self(value))
|
||||
}
|
||||
|
||||
pub fn minimum() -> Self {
|
||||
Self::parse("-9.999999999999999999999999999999999E+6144").expect("Decimal128 minimum")
|
||||
}
|
||||
|
||||
pub fn maximum() -> Self {
|
||||
Self::parse("9.999999999999999999999999999999999E+6144").expect("Decimal128 maximum")
|
||||
}
|
||||
}
|
||||
|
||||
fn operation(
|
||||
apply: impl FnOnce(&mut Context<Decimal128>) -> Decimal128,
|
||||
) -> Result<DecimalValue, &'static str> {
|
||||
CONTEXT.with_borrow_mut(|context| {
|
||||
context.clear_status();
|
||||
context.set_rounding(Rounding::HalfEven);
|
||||
let value = apply(context);
|
||||
finite_result(context, value)
|
||||
})
|
||||
}
|
||||
|
||||
fn finite_result(
|
||||
context: &Context<Decimal128>,
|
||||
value: Decimal128,
|
||||
) -> Result<DecimalValue, &'static str> {
|
||||
let status = context.status();
|
||||
if status.division_by_zero() || status.division_undefined() {
|
||||
Err("Decimal division by zero")
|
||||
} else if status.overflow() {
|
||||
Err("Decimal overflow")
|
||||
} else if status.underflow() {
|
||||
Err("Decimal underflow")
|
||||
} else if status.invalid_operation() || status.invalid_context() || !value.is_finite() {
|
||||
Err("invalid Decimal operation")
|
||||
} else {
|
||||
Ok(DecimalValue(value))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn int_add(left: i64, right: i64) -> Result<i64, &'static str> {
|
||||
left.checked_add(right).ok_or("Int overflow")
|
||||
}
|
||||
|
||||
pub fn parse_int(text: &str) -> Option<i64> {
|
||||
text.parse().ok()
|
||||
}
|
||||
|
||||
pub fn parse_decimal(text: &str) -> Option<DecimalValue> {
|
||||
decimal_text(text)
|
||||
.then(|| DecimalValue::parse(text).ok())
|
||||
.flatten()
|
||||
}
|
||||
|
||||
pub fn parse_float(text: &str) -> Option<f64> {
|
||||
if !decimal_text(text) && !matches!(text, "NaN" | "inf" | "-inf") {
|
||||
return None;
|
||||
}
|
||||
text.parse().ok()
|
||||
}
|
||||
|
||||
fn decimal_text(text: &str) -> bool {
|
||||
text.bytes().any(|byte| byte.is_ascii_digit())
|
||||
&& text
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit() || matches!(byte, b'+' | b'-' | b'.' | b'e' | b'E'))
|
||||
}
|
||||
|
||||
pub fn int_sub(left: i64, right: i64) -> Result<i64, &'static str> {
|
||||
left.checked_sub(right).ok_or("Int overflow")
|
||||
}
|
||||
|
||||
pub fn int_mul(left: i64, right: i64) -> Result<i64, &'static str> {
|
||||
left.checked_mul(right).ok_or("Int overflow")
|
||||
}
|
||||
|
||||
pub fn int_neg(value: i64) -> Result<i64, &'static str> {
|
||||
value.checked_neg().ok_or("Int overflow")
|
||||
}
|
||||
|
||||
pub fn int_div(left: i64, right: i64) -> Result<i64, &'static str> {
|
||||
if right == 0 {
|
||||
return Err("Int division by zero");
|
||||
}
|
||||
left.checked_div(right).ok_or("Int overflow")
|
||||
}
|
||||
|
||||
pub fn int_rem(left: i64, right: i64) -> Result<i64, &'static str> {
|
||||
if right == 0 {
|
||||
return Err("Int division by zero");
|
||||
}
|
||||
left.checked_rem(right).ok_or("Int overflow")
|
||||
}
|
||||
@@ -0,0 +1,695 @@
|
||||
use std::fmt;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
|
||||
use crate::number::DecimalValue;
|
||||
use futures_util::{Stream, StreamExt, stream};
|
||||
use tokio::sync::{Mutex, Semaphore};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Value {
|
||||
Null,
|
||||
Bool(bool),
|
||||
Int(i64),
|
||||
Float(f64),
|
||||
Decimal(String),
|
||||
Text(String),
|
||||
Bytes(Vec<u8>),
|
||||
Uuid(Uuid),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Related<T> {
|
||||
Unloaded,
|
||||
Loaded(Box<T>),
|
||||
}
|
||||
|
||||
impl<T: crate::render::Render> crate::render::Render for Related<T> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
match self {
|
||||
Self::Unloaded => output.push_str("Related.Unloaded"),
|
||||
Self::Loaded(value) => {
|
||||
output.push_str("Related.Loaded(");
|
||||
value.render_to(output);
|
||||
output.push(')');
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
|
||||
pub struct Uuid([u8; 16]);
|
||||
|
||||
impl Uuid {
|
||||
pub const fn from_bytes(value: [u8; 16]) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
|
||||
pub const fn into_bytes(self) -> [u8; 16] {
|
||||
self.0
|
||||
}
|
||||
|
||||
pub fn parse(text: &str) -> Result<Self, Error> {
|
||||
if text.len() != 36
|
||||
|| ![8, 13, 18, 23]
|
||||
.into_iter()
|
||||
.all(|index| text.as_bytes()[index] == b'-')
|
||||
{
|
||||
return Err(Error::invalid_data(
|
||||
"Uuid must use the canonical 8-4-4-4-12 form",
|
||||
));
|
||||
}
|
||||
let mut output = [0; 16];
|
||||
let mut source = text.bytes().filter(|byte| *byte != b'-');
|
||||
for byte in &mut output {
|
||||
let high = hex(source.next().expect("validated Uuid length"))?;
|
||||
let low = hex(source.next().expect("validated Uuid length"))?;
|
||||
*byte = (high << 4) | low;
|
||||
}
|
||||
Ok(Self(output))
|
||||
}
|
||||
|
||||
#[cfg(feature = "database")]
|
||||
pub fn new_v7() -> Result<Self, Error> {
|
||||
let millis = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map_err(|error| Error::new(ErrorKind::Database, error.to_string()))?
|
||||
.as_millis();
|
||||
if millis > 0xffff_ffff_ffff {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Database,
|
||||
"system time exceeds the UUIDv7 timestamp range",
|
||||
));
|
||||
}
|
||||
let mut value = [0; 16];
|
||||
getrandom::getrandom(&mut value[6..])
|
||||
.map_err(|error| Error::new(ErrorKind::Database, error.to_string()))?;
|
||||
let timestamp = (millis as u64).to_be_bytes();
|
||||
value[..6].copy_from_slice(×tamp[2..]);
|
||||
value[6] = (value[6] & 0x0f) | 0x70;
|
||||
value[8] = (value[8] & 0x3f) | 0x80;
|
||||
Ok(Self(value))
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for Uuid {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
for (index, byte) in self.0.iter().enumerate() {
|
||||
if matches!(index, 4 | 6 | 8 | 10) {
|
||||
formatter.write_str("-")?;
|
||||
}
|
||||
write!(formatter, "{byte:02x}")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Row {
|
||||
values: Vec<Value>,
|
||||
}
|
||||
|
||||
type RowProducer<T> = Pin<Box<dyn Stream<Item = Result<T, Error>> + Send>>;
|
||||
|
||||
struct RowStreamState<T> {
|
||||
producer: Mutex<Option<RowProducer<T>>>,
|
||||
closed: Semaphore,
|
||||
worker: StdMutex<Option<crate::task::ScopedOwner>>,
|
||||
worker_fault: Option<crate::task::ScopeFault>,
|
||||
_database_lease: Option<Arc<dyn Send + Sync>>,
|
||||
}
|
||||
|
||||
/// Internal database stream used by generated Model code. Dever exposes it as AsyncStream<T>.
|
||||
pub struct RowStream<T>(Arc<RowStreamState<T>>);
|
||||
|
||||
impl<T> Clone for RowStream<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self(Arc::clone(&self.0))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> fmt::Debug for RowStream<T> {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("RowStream(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> RowStream<T> {
|
||||
pub(crate) fn new(producer: impl Stream<Item = Result<T, Error>> + Send + 'static) -> Self {
|
||||
Self(Arc::new(RowStreamState {
|
||||
producer: Mutex::new(Some(Box::pin(producer))),
|
||||
closed: Semaphore::new(0),
|
||||
worker: StdMutex::new(None),
|
||||
worker_fault: None,
|
||||
_database_lease: None,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub(crate) fn hold_database(mut self, lease: Arc<dyn Send + Sync>) -> Self {
|
||||
Arc::get_mut(&mut self.0)
|
||||
.expect("new database stream has one owner")
|
||||
._database_lease = Some(lease);
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn pull(&self) -> Result<Option<T>, Error> {
|
||||
let mut producer = self.0.producer.lock().await;
|
||||
let next = match producer.as_mut() {
|
||||
Some(producer) => match &self.0.worker_fault {
|
||||
Some(fault) => tokio::select! {
|
||||
biased;
|
||||
_ = self.0.closed.acquire() => None,
|
||||
message = fault.wait() => Some(Err(Error::new(ErrorKind::Database, message))),
|
||||
value = producer.next() => value,
|
||||
},
|
||||
None => tokio::select! {
|
||||
biased;
|
||||
_ = self.0.closed.acquire() => None,
|
||||
value = producer.next() => value,
|
||||
},
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
let result = match next {
|
||||
Some(Ok(value)) => Ok(Some(value)),
|
||||
Some(Err(error)) => {
|
||||
producer.take();
|
||||
Err(error)
|
||||
}
|
||||
None => {
|
||||
producer.take();
|
||||
Ok(None)
|
||||
}
|
||||
};
|
||||
if result.as_ref().is_err() || matches!(&result, Ok(None)) {
|
||||
drop(producer);
|
||||
self.0.finish_worker();
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
pub fn close(&self) {
|
||||
self.0.closed.close();
|
||||
if let Ok(mut producer) = self.0.producer.try_lock() {
|
||||
producer.take();
|
||||
}
|
||||
self.0.stop_worker();
|
||||
}
|
||||
|
||||
pub fn close_on_drop(&self) -> CloseRowStreamOnDrop<T> {
|
||||
CloseRowStreamOnDrop(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Send + 'static> RowStream<T> {
|
||||
#[cfg(any(feature = "sqlite", feature = "postgres"))]
|
||||
pub(crate) fn from_receiver(
|
||||
receiver: tokio::sync::mpsc::Receiver<Result<T, Error>>,
|
||||
worker: crate::task::ScopedOwner,
|
||||
worker_fault: crate::task::ScopeFault,
|
||||
) -> Self {
|
||||
Self(Arc::new(RowStreamState {
|
||||
producer: Mutex::new(Some(Box::pin(stream::unfold(
|
||||
receiver,
|
||||
|mut receiver| async move { receiver.recv().await.map(|value| (value, receiver)) },
|
||||
)))),
|
||||
closed: Semaphore::new(0),
|
||||
worker: StdMutex::new(Some(worker)),
|
||||
worker_fault: Some(worker_fault),
|
||||
_database_lease: None,
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn map<U: Send + 'static>(
|
||||
self,
|
||||
transform: impl FnMut(T) -> Result<U, Error> + Send + 'static,
|
||||
) -> RowStream<U> {
|
||||
let owner = self.close_on_drop();
|
||||
RowStream::new(stream::unfold(
|
||||
(owner, transform),
|
||||
|(owner, mut transform)| async move {
|
||||
owner.0.pull().await.transpose().map(|value| {
|
||||
let value = value.and_then(&mut transform);
|
||||
(value, (owner, transform))
|
||||
})
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> RowStreamState<T> {
|
||||
fn finish_worker(&self) {
|
||||
let owner = self
|
||||
.worker
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.take();
|
||||
if let Some(owner) = owner {
|
||||
owner.finish_in_background();
|
||||
}
|
||||
}
|
||||
|
||||
fn stop_worker(&self) {
|
||||
let owner = self
|
||||
.worker
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.take();
|
||||
if let Some(owner) = owner {
|
||||
owner.stop_in_background();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Drop for RowStreamState<T> {
|
||||
fn drop(&mut self) {
|
||||
self.closed.close();
|
||||
self.producer.get_mut().take();
|
||||
let owner = self
|
||||
.worker
|
||||
.get_mut()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.take();
|
||||
if let Some(owner) = owner {
|
||||
owner.stop_in_background();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CloseRowStreamOnDrop<T>(RowStream<T>);
|
||||
|
||||
impl<T> Drop for CloseRowStreamOnDrop<T> {
|
||||
fn drop(&mut self) {
|
||||
self.0.close();
|
||||
}
|
||||
}
|
||||
|
||||
impl Row {
|
||||
pub fn new(values: Vec<Value>) -> Self {
|
||||
Self { values }
|
||||
}
|
||||
|
||||
pub fn get(&self, index: usize) -> Result<&Value, Error> {
|
||||
self.values.get(index).ok_or_else(|| {
|
||||
Error::invalid_data(format!("database row has no column at index {index}"))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.values.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.values.is_empty()
|
||||
}
|
||||
|
||||
pub fn into_values(self) -> std::vec::IntoIter<Value> {
|
||||
self.values.into_iter()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn required_row(rows: Vec<Row>, operation: &str) -> Result<Row, Error> {
|
||||
let mut rows = rows.into_iter();
|
||||
let row = rows.next().ok_or_else(|| {
|
||||
Error::invalid_data(format!("SQL operation '{operation}' returned no rows"))
|
||||
})?;
|
||||
if rows.next().is_some() {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"SQL operation '{operation}' returned more than one row"
|
||||
)));
|
||||
}
|
||||
Ok(row)
|
||||
}
|
||||
|
||||
pub fn optional_row(rows: Vec<Row>, operation: &str) -> Result<Option<Row>, Error> {
|
||||
let mut rows = rows.into_iter();
|
||||
let row = rows.next();
|
||||
if rows.next().is_some() {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"SQL operation '{operation}' returned more than one row"
|
||||
)));
|
||||
}
|
||||
Ok(row)
|
||||
}
|
||||
|
||||
pub fn bounded_rows(rows: Vec<Row>, maximum: usize) -> Result<Vec<Row>, Error> {
|
||||
if rows.len() > maximum {
|
||||
let message = if maximum == 1 {
|
||||
"database query returned more than one row".into()
|
||||
} else {
|
||||
format!("database query returned more than {maximum} rows")
|
||||
};
|
||||
Err(Error::invalid_data(message))
|
||||
} else {
|
||||
Ok(rows)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum ErrorKind {
|
||||
Pool,
|
||||
PoolExhausted,
|
||||
Connection,
|
||||
Timeout,
|
||||
Cancelled,
|
||||
Database,
|
||||
Constraint,
|
||||
NotFound,
|
||||
InvalidData,
|
||||
Migration,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct Error {
|
||||
kind: ErrorKind,
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
message: message.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn invalid_data(message: impl Into<String>) -> Self {
|
||||
Self::new(ErrorKind::InvalidData, message)
|
||||
}
|
||||
|
||||
pub fn not_found(model: &str) -> Self {
|
||||
Self::new(ErrorKind::NotFound, format!("{model} was not found"))
|
||||
}
|
||||
|
||||
pub fn migration(message: impl Into<String>) -> Self {
|
||||
Self::new(ErrorKind::Migration, message)
|
||||
}
|
||||
|
||||
pub fn cause(mut self, error: impl fmt::Display) -> Self {
|
||||
self.message.push_str("; caused by: ");
|
||||
self.message.push_str(&error.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn kind(&self) -> ErrorKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
pub(crate) fn keeps_connection(&self) -> bool {
|
||||
matches!(
|
||||
self.kind,
|
||||
ErrorKind::Constraint
|
||||
| ErrorKind::NotFound
|
||||
| ErrorKind::InvalidData
|
||||
| ErrorKind::Migration
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for Error {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str(&self.message)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for Error {}
|
||||
|
||||
pub fn boolean(value: Value) -> Result<bool, Error> {
|
||||
match value {
|
||||
Value::Bool(value) => Ok(value),
|
||||
Value::Int(0) => Ok(false),
|
||||
Value::Int(1) => Ok(true),
|
||||
_ => Err(Error::invalid_data("expected a database Bool")),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn int(value: Value) -> Result<i64, Error> {
|
||||
match value {
|
||||
Value::Int(value) => Ok(value),
|
||||
_ => Err(Error::invalid_data("expected a database Int")),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn float(value: Value) -> Result<f64, Error> {
|
||||
match value {
|
||||
Value::Float(value) => Ok(value),
|
||||
Value::Int(value) => Ok(value as f64),
|
||||
_ => Err(Error::invalid_data("expected a database Float")),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decimal(value: Value) -> Result<DecimalValue, Error> {
|
||||
match value {
|
||||
Value::Decimal(value) | Value::Text(value) => decimal_text(&value)
|
||||
.and_then(|value| DecimalValue::parse(&value).map_err(Error::invalid_data)),
|
||||
_ => Err(Error::invalid_data("expected a database Decimal")),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decimal_text(value: &str) -> Result<String, Error> {
|
||||
if value.starts_with('D') {
|
||||
decode_decimal_storage(value).map(|value| value.to_string())
|
||||
} else {
|
||||
DecimalValue::parse(value)
|
||||
.map(|value| value.to_string())
|
||||
.map_err(Error::invalid_data)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decimal_value(value: DecimalValue, precision: u8, scale: u8) -> Result<Value, Error> {
|
||||
decimal_storage(&value.to_string(), precision, scale).map(Value::Decimal)
|
||||
}
|
||||
|
||||
pub fn decimal_storage(text: &str, precision: u8, scale: u8) -> Result<String, Error> {
|
||||
let (negative, mut digits) = scaled_decimal(text, precision, scale)?;
|
||||
if negative {
|
||||
for digit in &mut digits {
|
||||
*digit = b'9' - (*digit - b'0');
|
||||
}
|
||||
}
|
||||
Ok(format!(
|
||||
"D{}{}:{scale:02}",
|
||||
if negative { '0' } else { '1' },
|
||||
String::from_utf8(digits).expect("ASCII decimal digits")
|
||||
))
|
||||
}
|
||||
|
||||
pub fn uuid(value: Value) -> Result<Uuid, Error> {
|
||||
match value {
|
||||
Value::Uuid(value) => Ok(value),
|
||||
Value::Bytes(value) => <[u8; 16]>::try_from(value)
|
||||
.map(Uuid::from_bytes)
|
||||
.map_err(|_| Error::invalid_data("database Uuid must contain 16 bytes")),
|
||||
_ => Err(Error::invalid_data("expected a database Uuid")),
|
||||
}
|
||||
}
|
||||
|
||||
fn scaled_decimal(text: &str, precision: u8, scale: u8) -> Result<(bool, Vec<u8>), Error> {
|
||||
let (negative, text) = match text.as_bytes().first() {
|
||||
Some(b'-') => (true, &text[1..]),
|
||||
Some(b'+') => (false, &text[1..]),
|
||||
_ => (false, text),
|
||||
};
|
||||
let (mantissa, exponent) = text
|
||||
.find(['e', 'E'])
|
||||
.map(|index| {
|
||||
text[index + 1..]
|
||||
.parse::<i32>()
|
||||
.map(|exponent| (&text[..index], exponent))
|
||||
.map_err(|_| Error::invalid_data("invalid Decimal exponent"))
|
||||
})
|
||||
.transpose()?
|
||||
.unwrap_or((text, 0));
|
||||
let mut point = None;
|
||||
let mut digits = Vec::with_capacity(mantissa.len() + usize::from(scale));
|
||||
for byte in mantissa.bytes() {
|
||||
if byte == b'.' && point.is_none() {
|
||||
point = Some(digits.len());
|
||||
} else if byte.is_ascii_digit() {
|
||||
digits.push(byte);
|
||||
} else {
|
||||
return Err(Error::invalid_data("invalid Decimal value"));
|
||||
}
|
||||
}
|
||||
if digits.is_empty() {
|
||||
return Err(Error::invalid_data("invalid Decimal value"));
|
||||
}
|
||||
let integer_digits = i64::try_from(point.unwrap_or(digits.len()))
|
||||
.map_err(|_| Error::invalid_data("Decimal value is too long"))?
|
||||
.checked_add(i64::from(exponent))
|
||||
.ok_or_else(|| Error::invalid_data("Decimal exponent is out of range"))?;
|
||||
let scaled_length = integer_digits
|
||||
.checked_add(i64::from(scale))
|
||||
.ok_or_else(|| Error::invalid_data("Decimal scale is out of range"))?;
|
||||
let first_nonzero = digits.iter().position(|digit| *digit != b'0');
|
||||
let Some(first_nonzero) = first_nonzero else {
|
||||
return Ok((false, vec![b'0'; usize::from(precision)]));
|
||||
};
|
||||
if scaled_length <= i64::try_from(first_nonzero).expect("index fits i64") {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"Decimal has more than {scale} fractional digits"
|
||||
)));
|
||||
}
|
||||
let significant_length = scaled_length - i64::try_from(first_nonzero).expect("index fits i64");
|
||||
if significant_length > i64::from(precision) {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"Decimal exceeds precision {precision} with scale {scale}"
|
||||
)));
|
||||
}
|
||||
let scaled_length = usize::try_from(scaled_length)
|
||||
.map_err(|_| Error::invalid_data("Decimal scale is out of range"))?;
|
||||
if scaled_length < digits.len() {
|
||||
if digits[scaled_length..].iter().any(|digit| *digit != b'0') {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"Decimal has more than {scale} fractional digits"
|
||||
)));
|
||||
}
|
||||
digits.truncate(scaled_length);
|
||||
} else {
|
||||
digits.resize(scaled_length, b'0');
|
||||
}
|
||||
let mut magnitude = digits.split_off(first_nonzero);
|
||||
let mut padded = vec![b'0'; usize::from(precision) - magnitude.len()];
|
||||
padded.append(&mut magnitude);
|
||||
Ok((negative, padded))
|
||||
}
|
||||
|
||||
fn decode_decimal_storage(text: &str) -> Result<DecimalValue, Error> {
|
||||
let Some(encoded) = text.strip_prefix('D') else {
|
||||
return Err(Error::invalid_data("invalid stored Decimal"));
|
||||
};
|
||||
let (sign, encoded) = encoded
|
||||
.split_at_checked(1)
|
||||
.ok_or_else(|| Error::invalid_data("invalid stored Decimal"))?;
|
||||
let (digits, scale) = encoded
|
||||
.rsplit_once(':')
|
||||
.ok_or_else(|| Error::invalid_data("invalid stored Decimal"))?;
|
||||
let scale = scale
|
||||
.parse::<usize>()
|
||||
.map_err(|_| Error::invalid_data("invalid stored Decimal scale"))?;
|
||||
if !matches!(sign, "0" | "1")
|
||||
|| digits.len() < scale
|
||||
|| !digits.bytes().all(|byte| byte.is_ascii_digit())
|
||||
{
|
||||
return Err(Error::invalid_data("invalid stored Decimal"));
|
||||
}
|
||||
let digits = if sign == "0" {
|
||||
digits
|
||||
.bytes()
|
||||
.map(|digit| char::from(b'9' - (digit - b'0')))
|
||||
.collect::<String>()
|
||||
} else {
|
||||
digits.to_owned()
|
||||
};
|
||||
let point = digits.len() - scale;
|
||||
let mut decimal = if scale == 0 {
|
||||
digits
|
||||
} else {
|
||||
format!("{}.{}", &digits[..point], &digits[point..])
|
||||
};
|
||||
if sign == "0" {
|
||||
decimal.insert(0, '-');
|
||||
}
|
||||
DecimalValue::parse(&decimal).map_err(Error::invalid_data)
|
||||
}
|
||||
|
||||
fn hex(value: u8) -> Result<u8, Error> {
|
||||
match value {
|
||||
b'0'..=b'9' => Ok(value - b'0'),
|
||||
b'a'..=b'f' => Ok(value - b'a' + 10),
|
||||
b'A'..=b'F' => Ok(value - b'A' + 10),
|
||||
_ => Err(Error::invalid_data("Uuid contains a non-hexadecimal digit")),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn text(value: Value) -> Result<String, Error> {
|
||||
match value {
|
||||
Value::Text(value) => Ok(value),
|
||||
_ => Err(Error::invalid_data("expected database Text")),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn bytes(value: Value) -> Result<crate::bytes::Bytes, Error> {
|
||||
match value {
|
||||
Value::Bytes(value) => Ok(crate::bytes::Bytes::new(value)),
|
||||
_ => Err(Error::invalid_data("expected database Bytes")),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn nullable<T>(
|
||||
value: Value,
|
||||
decode: impl FnOnce(Value) -> Result<T, Error>,
|
||||
) -> Result<Option<T>, Error> {
|
||||
match value {
|
||||
Value::Null => Ok(None),
|
||||
value => decode(value).map(Some),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn column(values: &mut std::vec::IntoIter<Value>, name: &str) -> Result<Value, Error> {
|
||||
values
|
||||
.next()
|
||||
.ok_or_else(|| Error::invalid_data(format!("database row is missing column '{name}'")))
|
||||
}
|
||||
|
||||
pub fn scalar(row: Row) -> Result<Value, Error> {
|
||||
let mut values = row.into_values();
|
||||
let value = column(&mut values, "scalar")?;
|
||||
if values.next().is_some() {
|
||||
return Err(Error::invalid_data(
|
||||
"database scalar query returned more than one column",
|
||||
));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub fn pagination(page: i64, size: i64, maximum: usize) -> Result<(i64, i64, i64), Error> {
|
||||
let maximum = i64::try_from(maximum)
|
||||
.map_err(|_| Error::invalid_data("database page limit is too large"))?;
|
||||
if page < 1 {
|
||||
return Err(Error::invalid_data("page must be at least 1"));
|
||||
}
|
||||
if size < 1 || size > maximum {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"page size must be between 1 and {maximum}"
|
||||
)));
|
||||
}
|
||||
let offset = page
|
||||
.checked_sub(1)
|
||||
.and_then(|page| page.checked_mul(size))
|
||||
.ok_or_else(|| Error::invalid_data("page offset is too large"))?;
|
||||
Ok((page, size, offset))
|
||||
}
|
||||
|
||||
pub fn cursor_size(size: i64, maximum: usize) -> Result<i64, Error> {
|
||||
let maximum = i64::try_from(maximum)
|
||||
.map_err(|_| Error::invalid_data("database cursor limit is too large"))?;
|
||||
if size < 1 || size > maximum || size == i64::MAX {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"cursor size must be between 1 and {maximum}"
|
||||
)));
|
||||
}
|
||||
Ok(size)
|
||||
}
|
||||
|
||||
pub fn relation_limit(maximum: usize) -> Result<i64, Error> {
|
||||
i64::try_from(maximum).map_err(|_| Error::invalid_data("database relation limit is too large"))
|
||||
}
|
||||
|
||||
pub fn stream_capacity(size: i64, maximum: usize) -> Result<usize, Error> {
|
||||
let maximum = i64::try_from(maximum)
|
||||
.map_err(|_| Error::invalid_data("database stream limit is too large"))?;
|
||||
if size < 1 || size > maximum {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"stream buffer size must be between 1 and {maximum}"
|
||||
)));
|
||||
}
|
||||
usize::try_from(size).map_err(|_| Error::invalid_data("stream buffer size is too large"))
|
||||
}
|
||||
|
||||
pub fn affected(value: usize) -> Result<i64, Error> {
|
||||
i64::try_from(value).map_err(|_| Error::invalid_data("affected row count is too large"))
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
#[cfg(feature = "sqlite")]
|
||||
pub(crate) const WAIT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
|
||||
pub(crate) const CREATE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
|
||||
pub(crate) const RECYCLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
|
||||
|
||||
pub(crate) struct PoolObjectGuard<T> {
|
||||
object: Option<T>,
|
||||
detach: fn(T),
|
||||
}
|
||||
|
||||
impl<T> PoolObjectGuard<T> {
|
||||
pub(crate) fn new(object: T, detach: fn(T)) -> Self {
|
||||
Self {
|
||||
object: Some(object),
|
||||
detach,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn object(&self) -> &T {
|
||||
self.object.as_ref().expect("active pooled object guard")
|
||||
}
|
||||
|
||||
pub(crate) fn release(mut self) -> T {
|
||||
self.object.take().expect("active pooled object guard")
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Drop for PoolObjectGuard<T> {
|
||||
fn drop(&mut self) {
|
||||
if let Some(object) = self.object.take() {
|
||||
(self.detach)(object);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,922 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::error::Error as _;
|
||||
use std::future::Future;
|
||||
use std::str::FromStr;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use deadpool_postgres::{
|
||||
Manager, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime, Timeouts,
|
||||
};
|
||||
use futures_util::TryStreamExt;
|
||||
use tokio_postgres::types::{FromSql, IsNull, ToSql, Type, to_sql_checked};
|
||||
|
||||
use crate::config::{Database as DatabaseConfig, PostgresTls, Settings};
|
||||
use crate::orm::{Error, ErrorKind, Row, RowStream, Uuid, Value};
|
||||
use crate::pool::{CREATE_TIMEOUT, PoolObjectGuard, RECYCLE_TIMEOUT};
|
||||
|
||||
mod migration;
|
||||
|
||||
const STATEMENT_CACHE_CAPACITY: usize = 64;
|
||||
|
||||
static DATABASES: OnceLock<BTreeMap<String, Arc<Database>>> = OnceLock::new();
|
||||
|
||||
pub struct Database {
|
||||
pool: Pool,
|
||||
min_connections: usize,
|
||||
max_page_size: usize,
|
||||
io_timeout: Duration,
|
||||
}
|
||||
|
||||
pub struct ConnectionOptions<'a> {
|
||||
pub url: &'a str,
|
||||
pub database_name: Option<&'a str>,
|
||||
pub tls: PostgresTls,
|
||||
pub min_connections: usize,
|
||||
pub max_connections: usize,
|
||||
pub max_page_size: usize,
|
||||
pub wait_timeout: Duration,
|
||||
pub io_timeout: Duration,
|
||||
}
|
||||
|
||||
pub struct Transaction {
|
||||
connection: Option<deadpool_postgres::Object>,
|
||||
max_page_size: usize,
|
||||
io_timeout: Duration,
|
||||
}
|
||||
|
||||
pub enum Executor<'a> {
|
||||
Pool(&'a Database),
|
||||
Transaction(&'a Transaction),
|
||||
}
|
||||
|
||||
struct Parameter(Value);
|
||||
struct NumericText(String);
|
||||
#[derive(Debug)]
|
||||
struct UuidBytes(Uuid);
|
||||
|
||||
impl Database {
|
||||
pub fn open(options: ConnectionOptions<'_>) -> Result<Self, Error> {
|
||||
let ConnectionOptions {
|
||||
url,
|
||||
database_name,
|
||||
tls,
|
||||
min_connections,
|
||||
max_connections,
|
||||
max_page_size,
|
||||
wait_timeout,
|
||||
io_timeout,
|
||||
} = options;
|
||||
let mut config = tokio_postgres::Config::from_str(url)
|
||||
.map_err(|error| Error::new(ErrorKind::Pool, error.to_string()))?;
|
||||
if let Some(database_name) = database_name {
|
||||
config.dbname(database_name);
|
||||
}
|
||||
config.connect_timeout(CREATE_TIMEOUT);
|
||||
let pool = match tls {
|
||||
PostgresTls::Disabled => {
|
||||
config.ssl_mode(tokio_postgres::config::SslMode::Disable);
|
||||
build_pool(config, tokio_postgres::NoTls, max_connections, wait_timeout)?
|
||||
}
|
||||
PostgresTls::System => {
|
||||
config.ssl_mode(tokio_postgres::config::SslMode::Require);
|
||||
let (tls, _) = tokio_postgres_rustls::MakeRustlsConnect::with_native_certs()
|
||||
.map_err(|errors| {
|
||||
Error::new(
|
||||
ErrorKind::Pool,
|
||||
errors
|
||||
.into_iter()
|
||||
.map(|error| error.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join("; "),
|
||||
)
|
||||
})?;
|
||||
build_pool(config, tls, max_connections, wait_timeout)?
|
||||
}
|
||||
};
|
||||
Ok(Self {
|
||||
pool,
|
||||
min_connections,
|
||||
max_page_size,
|
||||
io_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn max_page_size(&self) -> usize {
|
||||
self.max_page_size
|
||||
}
|
||||
|
||||
pub(crate) async fn prepare(&self) -> Result<(), Error> {
|
||||
let mut connections = Vec::with_capacity(self.min_connections);
|
||||
for _ in 0..self.min_connections {
|
||||
connections.push(self.connection().await?);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn execute(&self, sql: &'static str, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = io(
|
||||
self.io_timeout,
|
||||
execute_on(connection.object(), sql, parameters),
|
||||
)
|
||||
.await;
|
||||
finish(connection, result)
|
||||
}
|
||||
|
||||
pub async fn query(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = io(
|
||||
self.io_timeout,
|
||||
query_on(connection.object(), sql, parameters),
|
||||
)
|
||||
.await;
|
||||
finish(connection, result)
|
||||
}
|
||||
|
||||
pub async fn query_bounded(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = io(
|
||||
self.io_timeout,
|
||||
query_bounded_on(connection.object(), sql, parameters, maximum),
|
||||
)
|
||||
.await;
|
||||
finish(connection, result)
|
||||
}
|
||||
|
||||
pub async fn query_owned(
|
||||
&self,
|
||||
sql: String,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = io(
|
||||
self.io_timeout,
|
||||
query_owned_on(connection.object(), &sql, parameters),
|
||||
)
|
||||
.await;
|
||||
finish(connection, result)
|
||||
}
|
||||
|
||||
pub async fn stream(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
capacity: usize,
|
||||
) -> Result<RowStream<Row>, Error> {
|
||||
if capacity == 0 || capacity > self.max_page_size {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"row stream capacity must be between 1 and {}",
|
||||
self.max_page_size
|
||||
)));
|
||||
}
|
||||
let connection = self.connection().await?;
|
||||
let io_timeout = self.io_timeout;
|
||||
let (rows, receiver) = tokio::sync::mpsc::channel(capacity);
|
||||
let (ready, initialized) = tokio::sync::oneshot::channel();
|
||||
let (worker, spawner, worker_fault) =
|
||||
crate::task::scoped_owner(()).map_err(|error| Error::new(ErrorKind::Pool, error))?;
|
||||
spawner.spawn(async move {
|
||||
let connection = guarded(connection);
|
||||
let mut ready = Some(ready);
|
||||
let result = async {
|
||||
let statement = io(io_timeout, prepare_cached(connection.object(), sql)).await?;
|
||||
let parameters = parameters.into_iter().map(Parameter).collect::<Vec<_>>();
|
||||
let references = parameters
|
||||
.iter()
|
||||
.map(|value| value as &(dyn ToSql + Sync))
|
||||
.collect::<Vec<_>>();
|
||||
let source = io(io_timeout, async {
|
||||
connection
|
||||
.object()
|
||||
.query_raw(&statement, references)
|
||||
.await
|
||||
.map_err(database_error)
|
||||
})
|
||||
.await?;
|
||||
tokio::pin!(source);
|
||||
ready
|
||||
.take()
|
||||
.expect("pending row stream initialization")
|
||||
.send(Ok(()))
|
||||
.ok();
|
||||
while let Some(row) = io(io_timeout, async {
|
||||
source.as_mut().try_next().await.map_err(database_error)
|
||||
})
|
||||
.await?
|
||||
{
|
||||
if rows.send(decode_row(row)).await.is_err() {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Cancelled,
|
||||
"PostgreSQL row stream consumer closed",
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok::<_, Error>(())
|
||||
}
|
||||
.await;
|
||||
match result {
|
||||
Ok(()) => drop(connection.release()),
|
||||
Err(error) => {
|
||||
if let Some(ready) = ready.take() {
|
||||
let _ = ready.send(Err(error));
|
||||
} else {
|
||||
let _ = rows.send(Err(error)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
initialized
|
||||
.await
|
||||
.map_err(|_| Error::new(ErrorKind::Pool, "PostgreSQL row stream worker stopped"))??;
|
||||
Ok(RowStream::from_receiver(receiver, worker, worker_fault))
|
||||
}
|
||||
|
||||
pub async fn begin(&self) -> Result<Transaction, Error> {
|
||||
let connection = guarded(self.connection().await?);
|
||||
io(self.io_timeout, async {
|
||||
connection
|
||||
.object()
|
||||
.batch_execute("BEGIN")
|
||||
.await
|
||||
.map_err(database_error)
|
||||
})
|
||||
.await?;
|
||||
Ok(Transaction {
|
||||
connection: Some(connection.release()),
|
||||
max_page_size: self.max_page_size,
|
||||
io_timeout: self.io_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn migrate(&self, model: crate::database::Model) -> Result<(), Error> {
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = io(
|
||||
self.io_timeout,
|
||||
migration::apply(connection.object(), model),
|
||||
)
|
||||
.await;
|
||||
finish(connection, result)
|
||||
}
|
||||
|
||||
pub async fn migrate_foreign_keys(
|
||||
&self,
|
||||
model: String,
|
||||
table: String,
|
||||
constraints: Vec<crate::database::PostgresConstraint>,
|
||||
) -> Result<(), Error> {
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = io(
|
||||
self.io_timeout,
|
||||
migration::apply_foreign_keys(connection.object(), &model, &table, &constraints),
|
||||
)
|
||||
.await;
|
||||
finish(connection, result)
|
||||
}
|
||||
|
||||
async fn connection(&self) -> Result<deadpool_postgres::Object, Error> {
|
||||
self.pool.get().await.map_err(pool_error)
|
||||
}
|
||||
|
||||
pub async fn close(&self) -> Result<(), Error> {
|
||||
self.stop_accepting();
|
||||
self.drain().await
|
||||
}
|
||||
|
||||
pub(crate) fn stop_accepting(&self) {
|
||||
self.pool.close();
|
||||
}
|
||||
|
||||
async fn drain(&self) -> Result<(), Error> {
|
||||
io(self.io_timeout, async {
|
||||
while self.pool.status().size != 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_database(
|
||||
url: &str,
|
||||
database_name: &str,
|
||||
tls: PostgresTls,
|
||||
wait_timeout: Duration,
|
||||
io_timeout: Duration,
|
||||
) -> Result<(), Error> {
|
||||
if database_name.is_empty()
|
||||
|| database_name.len() > 63
|
||||
|| !database_name
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_')
|
||||
{
|
||||
return Err(Error::migration("invalid PostgreSQL tenant database name"));
|
||||
}
|
||||
let database = Database::open(ConnectionOptions {
|
||||
url,
|
||||
database_name: None,
|
||||
tls,
|
||||
min_connections: 0,
|
||||
max_connections: 1,
|
||||
max_page_size: 1,
|
||||
wait_timeout,
|
||||
io_timeout,
|
||||
})?;
|
||||
let connection = guarded(database.connection().await?);
|
||||
let sql = format!("CREATE DATABASE \"{database_name}\"");
|
||||
let result = io(io_timeout, async {
|
||||
match connection.object().batch_execute(&sql).await {
|
||||
Ok(()) => Ok(()),
|
||||
Err(error) if error.code().is_some_and(|code| code.code() == "42P04") => Ok(()),
|
||||
Err(error) => Err(database_error(error)),
|
||||
}
|
||||
})
|
||||
.await;
|
||||
let result = finish(connection, result);
|
||||
database.close().await?;
|
||||
result
|
||||
}
|
||||
|
||||
impl Transaction {
|
||||
pub async fn execute(&self, sql: &'static str, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
io(
|
||||
self.io_timeout,
|
||||
execute_on(self.connection(), sql, parameters),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn query(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
io(
|
||||
self.io_timeout,
|
||||
query_on(self.connection(), sql, parameters),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn query_bounded(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
io(
|
||||
self.io_timeout,
|
||||
query_bounded_on(self.connection(), sql, parameters, maximum),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn query_owned(
|
||||
&self,
|
||||
sql: String,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
io(
|
||||
self.io_timeout,
|
||||
query_owned_on(self.connection(), &sql, parameters),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn commit(mut self) -> Result<(), Error> {
|
||||
self.finish("COMMIT").await
|
||||
}
|
||||
|
||||
pub async fn rollback(mut self) -> Result<(), Error> {
|
||||
self.finish("ROLLBACK").await
|
||||
}
|
||||
|
||||
fn connection(&self) -> &deadpool_postgres::Object {
|
||||
self.connection
|
||||
.as_ref()
|
||||
.expect("open PostgreSQL transaction")
|
||||
}
|
||||
|
||||
async fn finish(&mut self, sql: &'static str) -> Result<(), Error> {
|
||||
let connection = guarded(self.connection.take().expect("open PostgreSQL transaction"));
|
||||
let result = io(self.io_timeout, async {
|
||||
connection
|
||||
.object()
|
||||
.batch_execute(sql)
|
||||
.await
|
||||
.map_err(database_error)
|
||||
})
|
||||
.await;
|
||||
finish(connection, result)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Transaction {
|
||||
fn drop(&mut self) {
|
||||
if let Some(connection) = self.connection.take() {
|
||||
drop(deadpool_postgres::Object::take(connection));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn guarded(connection: deadpool_postgres::Object) -> PoolObjectGuard<deadpool_postgres::Object> {
|
||||
PoolObjectGuard::new(connection, detach)
|
||||
}
|
||||
|
||||
fn detach(connection: deadpool_postgres::Object) {
|
||||
drop(deadpool_postgres::Object::take(connection));
|
||||
}
|
||||
|
||||
impl Executor<'_> {
|
||||
pub fn max_page_size(&self) -> usize {
|
||||
match self {
|
||||
Self::Pool(database) => database.max_page_size(),
|
||||
Self::Transaction(transaction) => transaction.max_page_size,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(&self, sql: &'static str, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.execute(sql, parameters).await,
|
||||
Self::Transaction(transaction) => transaction.execute(sql, parameters).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.query(sql, parameters).await,
|
||||
Self::Transaction(transaction) => transaction.query(sql, parameters).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn initialize(settings: &Settings) -> Result<(), String> {
|
||||
DATABASES
|
||||
.set(configured_databases(settings)?)
|
||||
.map_err(|_| "PostgreSQL databases were initialized more than once".to_owned())
|
||||
}
|
||||
|
||||
pub(crate) fn configured_databases(
|
||||
settings: &Settings,
|
||||
) -> Result<BTreeMap<String, Arc<Database>>, String> {
|
||||
let mut databases = BTreeMap::new();
|
||||
for (name, config) in settings.databases() {
|
||||
if let DatabaseConfig::Postgres {
|
||||
url,
|
||||
tls,
|
||||
min_connections,
|
||||
max_connections,
|
||||
max_page_size,
|
||||
wait_timeout_ms,
|
||||
io_timeout_ms,
|
||||
..
|
||||
} = config
|
||||
{
|
||||
let database = Database::open(ConnectionOptions {
|
||||
url,
|
||||
database_name: None,
|
||||
tls: *tls,
|
||||
min_connections: *min_connections,
|
||||
max_connections: *max_connections,
|
||||
max_page_size: *max_page_size,
|
||||
wait_timeout: Duration::from_millis(*wait_timeout_ms),
|
||||
io_timeout: Duration::from_millis(*io_timeout_ms),
|
||||
})
|
||||
.map_err(|error| format!("cannot initialize database.{name}: {error}"))?;
|
||||
databases.insert(name.clone(), Arc::new(database));
|
||||
}
|
||||
}
|
||||
Ok(databases)
|
||||
}
|
||||
|
||||
pub async fn prepare() -> Result<(), Error> {
|
||||
for database in DATABASES
|
||||
.get()
|
||||
.expect("PostgreSQL databases are initialized before prepare")
|
||||
.values()
|
||||
{
|
||||
database.prepare().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn shutdown() -> Result<(), Error> {
|
||||
let databases = DATABASES
|
||||
.get()
|
||||
.expect("PostgreSQL databases are initialized before shutdown")
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
for database in &databases {
|
||||
database.stop_accepting();
|
||||
}
|
||||
let mut first_error = None;
|
||||
for database in databases {
|
||||
if let Err(error) = database.drain().await
|
||||
&& first_error.is_none()
|
||||
{
|
||||
first_error = Some(error);
|
||||
}
|
||||
}
|
||||
first_error.map_or(Ok(()), Err)
|
||||
}
|
||||
|
||||
pub fn database(name: &str) -> Result<Arc<Database>, Error> {
|
||||
DATABASES
|
||||
.get()
|
||||
.and_then(|databases| databases.get(name))
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
Error::new(
|
||||
ErrorKind::Pool,
|
||||
format!("PostgreSQL database '{name}' is not initialized"),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn build_pool<T>(
|
||||
config: tokio_postgres::Config,
|
||||
tls: T,
|
||||
max_connections: usize,
|
||||
wait_timeout: Duration,
|
||||
) -> Result<Pool, Error>
|
||||
where
|
||||
T: tokio_postgres::tls::MakeTlsConnect<tokio_postgres::Socket> + Clone + Send + Sync + 'static,
|
||||
T::Stream: Send + Sync,
|
||||
T::TlsConnect: Send + Sync,
|
||||
<T::TlsConnect as tokio_postgres::tls::TlsConnect<tokio_postgres::Socket>>::Future: Send,
|
||||
{
|
||||
let manager = Manager::from_config(
|
||||
config,
|
||||
tls,
|
||||
ManagerConfig {
|
||||
recycling_method: RecyclingMethod::Fast,
|
||||
},
|
||||
);
|
||||
let mut pool_config = PoolConfig::new(max_connections);
|
||||
pool_config.timeouts = Timeouts {
|
||||
wait: Some(wait_timeout),
|
||||
create: Some(CREATE_TIMEOUT),
|
||||
recycle: Some(RECYCLE_TIMEOUT),
|
||||
};
|
||||
Pool::builder(manager)
|
||||
.config(pool_config)
|
||||
.runtime(Runtime::Tokio1)
|
||||
.build()
|
||||
.map_err(|error| Error::new(ErrorKind::Pool, error.to_string()))
|
||||
}
|
||||
|
||||
async fn execute_on(
|
||||
connection: &deadpool_postgres::Object,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<usize, Error> {
|
||||
let statement = prepare_cached(connection, sql).await?;
|
||||
let parameters = parameters.into_iter().map(Parameter).collect::<Vec<_>>();
|
||||
let references = parameters
|
||||
.iter()
|
||||
.map(|value| value as &(dyn ToSql + Sync))
|
||||
.collect::<Vec<_>>();
|
||||
connection
|
||||
.execute(&statement, &references)
|
||||
.await
|
||||
.map(|affected| affected as usize)
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
async fn query_on(
|
||||
connection: &deadpool_postgres::Object,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
let statement = prepare_cached(connection, sql).await?;
|
||||
let parameters = parameters.into_iter().map(Parameter).collect::<Vec<_>>();
|
||||
let references = parameters
|
||||
.iter()
|
||||
.map(|value| value as &(dyn ToSql + Sync))
|
||||
.collect::<Vec<_>>();
|
||||
connection
|
||||
.query(&statement, &references)
|
||||
.await
|
||||
.map_err(database_error)?
|
||||
.into_iter()
|
||||
.map(decode_row)
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn query_bounded_on(
|
||||
connection: &deadpool_postgres::Object,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
let statement = prepare_cached(connection, sql).await?;
|
||||
let parameters = parameters.into_iter().map(Parameter).collect::<Vec<_>>();
|
||||
let references = parameters
|
||||
.iter()
|
||||
.map(|value| value as &(dyn ToSql + Sync))
|
||||
.collect::<Vec<_>>();
|
||||
let rows = connection
|
||||
.query_raw(&statement, references)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
tokio::pin!(rows);
|
||||
let mut output = Vec::with_capacity(maximum.saturating_add(1));
|
||||
while output.len() <= maximum {
|
||||
let Some(row) = rows.try_next().await.map_err(database_error)? else {
|
||||
break;
|
||||
};
|
||||
output.push(decode_row(row)?);
|
||||
}
|
||||
crate::orm::bounded_rows(output, maximum)
|
||||
}
|
||||
|
||||
async fn query_owned_on(
|
||||
connection: &deadpool_postgres::Object,
|
||||
sql: &str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
let statement = connection.prepare(sql).await.map_err(database_error)?;
|
||||
let parameters = parameters.into_iter().map(Parameter).collect::<Vec<_>>();
|
||||
let references = parameters
|
||||
.iter()
|
||||
.map(|value| value as &(dyn ToSql + Sync))
|
||||
.collect::<Vec<_>>();
|
||||
connection
|
||||
.query(&statement, &references)
|
||||
.await
|
||||
.map_err(database_error)?
|
||||
.into_iter()
|
||||
.map(decode_row)
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn prepare_cached(
|
||||
connection: &deadpool_postgres::Object,
|
||||
sql: &'static str,
|
||||
) -> Result<tokio_postgres::Statement, Error> {
|
||||
let statement = connection
|
||||
.prepare_cached(sql)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
if connection.statement_cache.size() > STATEMENT_CACHE_CAPACITY {
|
||||
connection.statement_cache.remove(sql, &[]);
|
||||
}
|
||||
Ok(statement)
|
||||
}
|
||||
|
||||
fn decode_row(row: tokio_postgres::Row) -> Result<Row, Error> {
|
||||
let mut values = Vec::with_capacity(row.len());
|
||||
for (index, column) in row.columns().iter().enumerate() {
|
||||
let value = match *column.type_() {
|
||||
Type::BOOL => optional(&row, index, |value: bool| Value::Bool(value))?,
|
||||
Type::INT2 => optional(&row, index, |value: i16| Value::Int(i64::from(value)))?,
|
||||
Type::INT4 => optional(&row, index, |value: i32| Value::Int(i64::from(value)))?,
|
||||
Type::INT8 => optional(&row, index, Value::Int)?,
|
||||
Type::FLOAT4 => optional(&row, index, |value: f32| Value::Float(f64::from(value)))?,
|
||||
Type::FLOAT8 => optional(&row, index, Value::Float)?,
|
||||
Type::BYTEA => optional(&row, index, Value::Bytes)?,
|
||||
Type::NUMERIC => optional(&row, index, |value: NumericText| Value::Decimal(value.0))?,
|
||||
Type::UUID => optional(&row, index, |value: UuidBytes| Value::Uuid(value.0))?,
|
||||
Type::TEXT | Type::VARCHAR | Type::BPCHAR | Type::NAME => {
|
||||
optional(&row, index, Value::Text)?
|
||||
}
|
||||
ref ty => {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"unsupported PostgreSQL column type '{ty}'"
|
||||
)));
|
||||
}
|
||||
};
|
||||
values.push(value);
|
||||
}
|
||||
Ok(Row::new(values))
|
||||
}
|
||||
|
||||
fn optional<T>(
|
||||
row: &tokio_postgres::Row,
|
||||
index: usize,
|
||||
convert: impl FnOnce(T) -> Value,
|
||||
) -> Result<Value, Error>
|
||||
where
|
||||
T: for<'a> tokio_postgres::types::FromSql<'a>,
|
||||
{
|
||||
row.try_get::<_, Option<T>>(index)
|
||||
.map(|value| value.map(convert).unwrap_or(Value::Null))
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
impl ToSql for Parameter {
|
||||
fn to_sql(
|
||||
&self,
|
||||
ty: &Type,
|
||||
output: &mut bytes::BytesMut,
|
||||
) -> Result<IsNull, Box<dyn std::error::Error + Send + Sync>> {
|
||||
match &self.0 {
|
||||
Value::Null => Ok(IsNull::Yes),
|
||||
Value::Bool(value) => value.to_sql(ty, output),
|
||||
Value::Int(value) => value.to_sql(ty, output),
|
||||
Value::Float(value) => value.to_sql(ty, output),
|
||||
Value::Decimal(value) => crate::orm::decimal_text(value)?.to_sql(ty, output),
|
||||
Value::Text(value) => value.to_sql(ty, output),
|
||||
Value::Bytes(value) => value.to_sql(ty, output),
|
||||
Value::Uuid(value) => UuidBytes(*value).to_sql(ty, output),
|
||||
}
|
||||
}
|
||||
|
||||
fn accepts(_: &Type) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
to_sql_checked!();
|
||||
}
|
||||
|
||||
impl<'a> FromSql<'a> for NumericText {
|
||||
fn from_sql(_: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
|
||||
decode_numeric(raw)
|
||||
.map(Self)
|
||||
.map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidData, message).into())
|
||||
}
|
||||
|
||||
fn accepts(ty: &Type) -> bool {
|
||||
*ty == Type::NUMERIC
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> FromSql<'a> for UuidBytes {
|
||||
fn from_sql(_: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
|
||||
raw.try_into().map(Uuid::from_bytes).map(Self).map_err(|_| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"PostgreSQL UUID is not 16 bytes",
|
||||
)
|
||||
.into()
|
||||
})
|
||||
}
|
||||
|
||||
fn accepts(ty: &Type) -> bool {
|
||||
*ty == Type::UUID
|
||||
}
|
||||
}
|
||||
|
||||
impl ToSql for UuidBytes {
|
||||
fn to_sql(
|
||||
&self,
|
||||
_: &Type,
|
||||
output: &mut bytes::BytesMut,
|
||||
) -> Result<IsNull, Box<dyn std::error::Error + Send + Sync>> {
|
||||
output.extend_from_slice(&self.0.into_bytes());
|
||||
Ok(IsNull::No)
|
||||
}
|
||||
|
||||
fn accepts(ty: &Type) -> bool {
|
||||
*ty == Type::UUID
|
||||
}
|
||||
|
||||
to_sql_checked!();
|
||||
}
|
||||
|
||||
fn decode_numeric(raw: &[u8]) -> Result<String, &'static str> {
|
||||
if raw.len() < 8 || !(raw.len() - 8).is_multiple_of(2) {
|
||||
return Err("invalid PostgreSQL NUMERIC payload");
|
||||
}
|
||||
let read = |offset: usize| i16::from_be_bytes([raw[offset], raw[offset + 1]]);
|
||||
let digit_count = read(0);
|
||||
let weight = read(2);
|
||||
let sign = u16::from_be_bytes([raw[4], raw[5]]);
|
||||
let scale = read(6);
|
||||
if digit_count < 0
|
||||
|| scale < 0
|
||||
|| raw.len() != 8 + usize::try_from(digit_count).unwrap_or(usize::MAX) * 2
|
||||
|| !matches!(sign, 0x0000 | 0x4000)
|
||||
{
|
||||
return Err("unsupported PostgreSQL NUMERIC payload");
|
||||
}
|
||||
let digits = raw[8..]
|
||||
.as_chunks::<2>()
|
||||
.0
|
||||
.iter()
|
||||
.map(|digit| u16::from_be_bytes([digit[0], digit[1]]))
|
||||
.collect::<Vec<_>>();
|
||||
if digits.iter().any(|digit| *digit >= 10_000) {
|
||||
return Err("invalid PostgreSQL NUMERIC digit");
|
||||
}
|
||||
let digit_at = |exponent: i16| {
|
||||
let index = i32::from(weight) - i32::from(exponent);
|
||||
usize::try_from(index)
|
||||
.ok()
|
||||
.and_then(|index| digits.get(index))
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
};
|
||||
let mut text = String::new();
|
||||
if sign == 0x4000 && digits.iter().any(|digit| *digit != 0) {
|
||||
text.push('-');
|
||||
}
|
||||
if weight >= 0 {
|
||||
for exponent in (0..=weight).rev() {
|
||||
let digit = digit_at(exponent);
|
||||
if exponent == weight {
|
||||
text.push_str(&digit.to_string());
|
||||
} else {
|
||||
use std::fmt::Write as _;
|
||||
write!(text, "{digit:04}").expect("string formatting");
|
||||
}
|
||||
}
|
||||
} else {
|
||||
text.push('0');
|
||||
}
|
||||
let scale = usize::try_from(scale).expect("validated NUMERIC scale");
|
||||
if scale > 0 {
|
||||
text.push('.');
|
||||
let groups = scale.div_ceil(4);
|
||||
for group in 1..=groups {
|
||||
use std::fmt::Write as _;
|
||||
write!(text, "{:04}", digit_at(-(group as i16))).expect("string formatting");
|
||||
}
|
||||
text.truncate(text.len() - (groups * 4 - scale));
|
||||
}
|
||||
Ok(text)
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Parameter {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("Parameter(<value>)")
|
||||
}
|
||||
}
|
||||
|
||||
fn database_error(error: tokio_postgres::Error) -> Error {
|
||||
let kind = match error.code() {
|
||||
Some(
|
||||
&tokio_postgres::error::SqlState::UNIQUE_VIOLATION
|
||||
| &tokio_postgres::error::SqlState::FOREIGN_KEY_VIOLATION
|
||||
| &tokio_postgres::error::SqlState::CHECK_VIOLATION
|
||||
| &tokio_postgres::error::SqlState::NOT_NULL_VIOLATION,
|
||||
) => ErrorKind::Constraint,
|
||||
Some(code) if code == &tokio_postgres::error::SqlState::QUERY_CANCELED => {
|
||||
ErrorKind::Cancelled
|
||||
}
|
||||
Some(
|
||||
&tokio_postgres::error::SqlState::ADMIN_SHUTDOWN
|
||||
| &tokio_postgres::error::SqlState::CRASH_SHUTDOWN
|
||||
| &tokio_postgres::error::SqlState::CANNOT_CONNECT_NOW
|
||||
| &tokio_postgres::error::SqlState::DATABASE_DROPPED,
|
||||
) => ErrorKind::Connection,
|
||||
Some(code) if code.code().starts_with("08") => ErrorKind::Connection,
|
||||
_ if error.is_closed() => ErrorKind::Connection,
|
||||
_ => ErrorKind::Database,
|
||||
};
|
||||
let mut result = Error::new(kind, error.to_string());
|
||||
if let Some(source) = error.source() {
|
||||
result = result.cause(source);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
async fn io<T>(
|
||||
timeout: Duration,
|
||||
operation: impl Future<Output = Result<T, Error>>,
|
||||
) -> Result<T, Error> {
|
||||
tokio::time::timeout(timeout, operation)
|
||||
.await
|
||||
.map_err(|_| Error::new(ErrorKind::Timeout, "PostgreSQL I/O timed out"))?
|
||||
}
|
||||
|
||||
fn finish<T>(
|
||||
connection: PoolObjectGuard<deadpool_postgres::Object>,
|
||||
result: Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
if match &result {
|
||||
Ok(_) => true,
|
||||
Err(error) => error.keeps_connection(),
|
||||
} {
|
||||
drop(connection.release());
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn pool_error(error: deadpool_postgres::PoolError) -> Error {
|
||||
let kind = match &error {
|
||||
deadpool_postgres::PoolError::Timeout(deadpool_postgres::TimeoutType::Wait) => {
|
||||
ErrorKind::PoolExhausted
|
||||
}
|
||||
deadpool_postgres::PoolError::Timeout(_) => ErrorKind::Timeout,
|
||||
deadpool_postgres::PoolError::Backend(_) => ErrorKind::Connection,
|
||||
deadpool_postgres::PoolError::Closed => ErrorKind::Connection,
|
||||
_ => ErrorKind::Pool,
|
||||
};
|
||||
Error::new(kind, error.to_string())
|
||||
}
|
||||
@@ -0,0 +1,660 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use tokio_postgres::types::ToSql;
|
||||
|
||||
use crate::database::{MigrationPhase, Model, PostgresConstraint, Schema};
|
||||
use crate::orm::Error;
|
||||
|
||||
use super::{database_error, execute_on};
|
||||
|
||||
const CREATE_HISTORY: &str = "\
|
||||
CREATE TABLE IF NOT EXISTS _dever_history (\
|
||||
model TEXT NOT NULL,\
|
||||
kind TEXT NOT NULL,\
|
||||
name TEXT NOT NULL,\
|
||||
revision TEXT NOT NULL,\
|
||||
definition TEXT NOT NULL,\
|
||||
applied_at BIGINT NOT NULL DEFAULT ((extract(epoch from clock_timestamp()) * 1000)::bigint),\
|
||||
PRIMARY KEY (model, kind, name)\
|
||||
)";
|
||||
|
||||
pub(super) async fn apply(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: Model,
|
||||
) -> Result<(), Error> {
|
||||
connection
|
||||
.batch_execute("BEGIN")
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
let result = apply_locked(connection, &model).await;
|
||||
match result {
|
||||
Ok(()) => connection
|
||||
.batch_execute("COMMIT")
|
||||
.await
|
||||
.map_err(database_error),
|
||||
Err(error) => match connection.batch_execute("ROLLBACK").await {
|
||||
Ok(()) => Err(error),
|
||||
Err(rollback) => Err(error.cause(rollback)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn apply_foreign_keys(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &str,
|
||||
table: &str,
|
||||
constraints: &[PostgresConstraint],
|
||||
) -> Result<(), Error> {
|
||||
connection
|
||||
.batch_execute("BEGIN")
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
let result = apply_foreign_keys_locked(connection, model, table, constraints).await;
|
||||
match result {
|
||||
Ok(()) => connection
|
||||
.batch_execute("COMMIT")
|
||||
.await
|
||||
.map_err(database_error),
|
||||
Err(error) => match connection.batch_execute("ROLLBACK").await {
|
||||
Ok(()) => Err(error),
|
||||
Err(rollback) => Err(error.cause(rollback)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn apply_foreign_keys_locked(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &str,
|
||||
table: &str,
|
||||
constraints: &[PostgresConstraint],
|
||||
) -> Result<(), Error> {
|
||||
connection
|
||||
.query_one(
|
||||
"SELECT pg_advisory_xact_lock(hashtextextended($1, 0))",
|
||||
&[&model],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
for constraint in constraints {
|
||||
let exists = connection
|
||||
.query_one(
|
||||
"SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_constraint con JOIN pg_catalog.pg_class rel ON rel.oid = con.conrelid JOIN pg_catalog.pg_namespace nsp ON nsp.oid = rel.relnamespace WHERE nsp.nspname = current_schema() AND rel.relname = $1 AND con.conname = $2 AND con.contype = 'f')",
|
||||
&[&table, &constraint.name],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?
|
||||
.get::<_, bool>(0);
|
||||
if !exists {
|
||||
alter_table(
|
||||
connection,
|
||||
table,
|
||||
&format!(
|
||||
"ADD CONSTRAINT {} {}",
|
||||
quote(&constraint.name),
|
||||
constraint.definition
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
store_catalog_history(connection, model, table).await
|
||||
}
|
||||
|
||||
async fn apply_locked(connection: &deadpool_postgres::Object, model: &Model) -> Result<(), Error> {
|
||||
connection
|
||||
.query_one(
|
||||
"SELECT pg_advisory_xact_lock(hashtextextended($1, 0))",
|
||||
&[&model.schema.model],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
connection
|
||||
.batch_execute(CREATE_HISTORY)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
let table_exists = table_exists(connection, &model.schema.table).await?;
|
||||
let previous = load_schema(connection, &model.schema.model).await?;
|
||||
let mut pending = BTreeSet::new();
|
||||
match (table_exists, previous.as_ref()) {
|
||||
(false, None) => {
|
||||
connection
|
||||
.batch_execute(&model.create_table.postgres)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
}
|
||||
(true, Some(previous)) => {
|
||||
validate_migrations(connection, &model.schema).await?;
|
||||
crate::database::validate_migration_order(previous, &model.schema)?;
|
||||
validate_catalog(connection, previous).await?;
|
||||
let applied = load_migrations(connection, &model.schema.model).await?;
|
||||
pending = model
|
||||
.schema
|
||||
.migrations
|
||||
.iter()
|
||||
.filter(|migration| !applied.contains_key(&migration.name))
|
||||
.map(|migration| migration.name.as_str())
|
||||
.collect();
|
||||
apply_data(connection, model, &pending, MigrationPhase::Before).await?;
|
||||
if previous.revision != model.schema.revision {
|
||||
migrate_schema(connection, previous, model).await?;
|
||||
}
|
||||
}
|
||||
(true, None) => {
|
||||
return Err(Error::migration(format!(
|
||||
"table '{}' exists without Dever schema history",
|
||||
model.schema.table
|
||||
)));
|
||||
}
|
||||
(false, Some(_)) => {
|
||||
return Err(Error::migration(format!(
|
||||
"Dever schema history exists but table '{}' is missing",
|
||||
model.schema.table
|
||||
)));
|
||||
}
|
||||
}
|
||||
for index in &model.create_indexes {
|
||||
connection
|
||||
.batch_execute(&index.postgres)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
}
|
||||
apply_data(connection, model, &pending, MigrationPhase::After).await?;
|
||||
if previous
|
||||
.as_ref()
|
||||
.is_none_or(|schema| schema.seed_revision != model.schema.seed_revision)
|
||||
{
|
||||
for seed in &model.seeds {
|
||||
execute_on(connection, seed.sql.postgres, seed.parameters.clone()).await?;
|
||||
}
|
||||
}
|
||||
store_history(connection, &model.schema).await
|
||||
}
|
||||
|
||||
async fn apply_data(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &Model,
|
||||
pending: &BTreeSet<&str>,
|
||||
phase: MigrationPhase,
|
||||
) -> Result<(), Error> {
|
||||
for statement in &model.data_migrations {
|
||||
if statement.phase == phase && pending.contains(statement.name.as_str()) {
|
||||
execute_on(
|
||||
connection,
|
||||
statement.sql.postgres,
|
||||
statement.parameters.clone(),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn migrate_schema(
|
||||
connection: &deadpool_postgres::Object,
|
||||
previous: &Schema,
|
||||
model: &Model,
|
||||
) -> Result<(), Error> {
|
||||
let applied = load_migrations(connection, &model.schema.model).await?;
|
||||
let allowed_drops = model
|
||||
.schema
|
||||
.migrations
|
||||
.iter()
|
||||
.filter(|migration| !applied.contains_key(&migration.name))
|
||||
.flat_map(|migration| migration.drops.iter().cloned())
|
||||
.collect::<BTreeSet<_>>();
|
||||
let previous_fields = previous
|
||||
.fields
|
||||
.iter()
|
||||
.map(|field| (field.name.as_str(), field))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let columns = model
|
||||
.postgres_columns
|
||||
.iter()
|
||||
.map(|column| (column.name.as_str(), column))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let mut mapped = Vec::new();
|
||||
let mut consumed = BTreeSet::new();
|
||||
for target in &model.schema.fields {
|
||||
let source = previous_fields
|
||||
.get(target.name.as_str())
|
||||
.copied()
|
||||
.or_else(|| {
|
||||
target
|
||||
.rename_from
|
||||
.as_deref()
|
||||
.and_then(|name| previous_fields.get(name).copied())
|
||||
});
|
||||
if let Some(source) = source {
|
||||
if !compatible_types(&source.ty, &target.ty) {
|
||||
return Err(Error::migration(format!(
|
||||
"field '{}.{}' cannot change type from '{}' to '{}' automatically",
|
||||
model.schema.model, target.name, source.ty, target.ty
|
||||
)));
|
||||
}
|
||||
if !consumed.insert(source.name.as_str()) {
|
||||
return Err(Error::migration(format!(
|
||||
"field '{}.{}' reuses migration source '{}'",
|
||||
model.schema.model, target.name, source.name
|
||||
)));
|
||||
}
|
||||
}
|
||||
mapped.push((target, source));
|
||||
}
|
||||
for source in &previous.fields {
|
||||
if !consumed.contains(source.name.as_str()) && !allowed_drops.contains(&source.name) {
|
||||
return Err(Error::migration(format!(
|
||||
"field '{}.{}' was removed without a named migrate drop",
|
||||
model.schema.model, source.name
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
drop_dever_constraints(connection, &model.schema.table).await?;
|
||||
for (target, source) in &mapped {
|
||||
if let Some(source) = source
|
||||
&& source.name != target.name
|
||||
{
|
||||
alter_table(
|
||||
connection,
|
||||
&model.schema.table,
|
||||
&format!(
|
||||
"RENAME COLUMN {} TO {}",
|
||||
quote(&source.name),
|
||||
quote(&target.name)
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
for source in &previous.fields {
|
||||
if !consumed.contains(source.name.as_str()) {
|
||||
alter_table(
|
||||
connection,
|
||||
&model.schema.table,
|
||||
&format!("DROP COLUMN {}", quote(&source.name)),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
for (target, source) in mapped {
|
||||
let column = columns
|
||||
.get(target.name.as_str())
|
||||
.expect("generated PostgreSQL column");
|
||||
let Some(source) = source else {
|
||||
alter_table(
|
||||
connection,
|
||||
&model.schema.table,
|
||||
&format!("ADD COLUMN {}", column.definition),
|
||||
)
|
||||
.await?;
|
||||
continue;
|
||||
};
|
||||
let previous_type = postgres_type(&source.ty);
|
||||
if previous_type != column.sql_type {
|
||||
alter_table(
|
||||
connection,
|
||||
&model.schema.table,
|
||||
&format!(
|
||||
"ALTER COLUMN {} TYPE {} USING {}::{}",
|
||||
quote(&target.name),
|
||||
column.sql_type,
|
||||
quote(&target.name),
|
||||
column.sql_type,
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
if source.nullable != target.nullable {
|
||||
alter_table(
|
||||
connection,
|
||||
&model.schema.table,
|
||||
&format!(
|
||||
"ALTER COLUMN {} {} NOT NULL",
|
||||
quote(&target.name),
|
||||
if target.nullable { "DROP" } else { "SET" },
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
if source.default != target.default || source.generated != target.generated {
|
||||
let action = column
|
||||
.default
|
||||
.as_ref()
|
||||
.map(|default| format!("SET DEFAULT {default}"))
|
||||
.unwrap_or_else(|| "DROP DEFAULT".into());
|
||||
alter_table(
|
||||
connection,
|
||||
&model.schema.table,
|
||||
&format!("ALTER COLUMN {} {action}", quote(&target.name)),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
for constraint in model
|
||||
.postgres_constraints
|
||||
.iter()
|
||||
.filter(|constraint| !constraint.foreign_key)
|
||||
{
|
||||
alter_table(
|
||||
connection,
|
||||
&model.schema.table,
|
||||
&format!(
|
||||
"ADD CONSTRAINT {} {}",
|
||||
quote(&constraint.name),
|
||||
constraint.definition
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
for index in &previous.indexes {
|
||||
connection
|
||||
.batch_execute(&format!("DROP INDEX IF EXISTS {}", quote(&index.name)))
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_dever_constraints(
|
||||
connection: &deadpool_postgres::Object,
|
||||
table: &str,
|
||||
) -> Result<(), Error> {
|
||||
let constraints = connection
|
||||
.query(
|
||||
"SELECT con.conname FROM pg_catalog.pg_constraint con JOIN pg_catalog.pg_class rel ON rel.oid = con.conrelid JOIN pg_catalog.pg_namespace nsp ON nsp.oid = rel.relnamespace WHERE nsp.nspname = current_schema() AND rel.relname = $1 AND left(con.conname, 7) = '_dever_'",
|
||||
&[&table],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
for row in constraints {
|
||||
alter_table(
|
||||
connection,
|
||||
table,
|
||||
&format!("DROP CONSTRAINT {}", quote(row.get::<_, &str>(0))),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn alter_table(
|
||||
connection: &deadpool_postgres::Object,
|
||||
table: &str,
|
||||
operation: &str,
|
||||
) -> Result<(), Error> {
|
||||
connection
|
||||
.batch_execute(&format!("ALTER TABLE {} {operation}", quote(table)))
|
||||
.await
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
fn compatible_types(previous: &str, desired: &str) -> bool {
|
||||
previous == desired
|
||||
|| (previous.starts_with("Text") && desired.starts_with("Text"))
|
||||
|| (previous.starts_with("Bytes") && desired.starts_with("Bytes"))
|
||||
|| (previous == "Int" && desired == "Float")
|
||||
}
|
||||
|
||||
fn postgres_type(logical: &str) -> String {
|
||||
if logical == "Bool" {
|
||||
"BOOLEAN".into()
|
||||
} else if logical == "Float" {
|
||||
"DOUBLE PRECISION".into()
|
||||
} else if logical.starts_with("Bytes") {
|
||||
"BYTEA".into()
|
||||
} else if let Some(decimal) = logical
|
||||
.strip_prefix("Decimal(")
|
||||
.and_then(|value| value.strip_suffix(')'))
|
||||
{
|
||||
format!("NUMERIC({decimal})")
|
||||
} else if logical == "Uuid" {
|
||||
"UUID".into()
|
||||
} else if logical == "Int"
|
||||
|| logical == "DateTime"
|
||||
|| logical == "Date"
|
||||
|| logical == "Time"
|
||||
|| logical == "Duration"
|
||||
|| logical.starts_with("id:")
|
||||
{
|
||||
"BIGINT".into()
|
||||
} else {
|
||||
"TEXT".into()
|
||||
}
|
||||
}
|
||||
|
||||
fn quote(identifier: &str) -> String {
|
||||
format!("\"{}\"", identifier.replace('"', "\"\""))
|
||||
}
|
||||
|
||||
async fn table_exists(connection: &deadpool_postgres::Object, table: &str) -> Result<bool, Error> {
|
||||
connection
|
||||
.query_one(
|
||||
"SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = current_schema() AND c.relname = $1 AND c.relkind = 'r')",
|
||||
&[&table],
|
||||
)
|
||||
.await
|
||||
.map(|row| row.get(0))
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
async fn load_schema(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &str,
|
||||
) -> Result<Option<Schema>, Error> {
|
||||
let row = connection
|
||||
.query_opt(
|
||||
"SELECT definition FROM _dever_history WHERE model = $1 AND kind = 'schema' AND name = 'current'",
|
||||
&[&model],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?;
|
||||
row.map(|row| {
|
||||
serde_json::from_str(row.get(0))
|
||||
.map_err(|error| Error::migration(format!("invalid Dever schema history: {error}")))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn validate_catalog(
|
||||
connection: &deadpool_postgres::Object,
|
||||
schema: &Schema,
|
||||
) -> Result<(), Error> {
|
||||
let recorded = load_catalog(connection, &schema.model)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
Error::migration(format!(
|
||||
"Model '{}' has no recorded database catalog",
|
||||
schema.model
|
||||
))
|
||||
})?;
|
||||
if catalog(connection, &schema.table).await? != recorded {
|
||||
return Err(Error::migration(format!(
|
||||
"table '{}' does not match its recorded Dever schema",
|
||||
schema.table
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn load_catalog(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &str,
|
||||
) -> Result<Option<String>, Error> {
|
||||
connection
|
||||
.query_opt(
|
||||
"SELECT definition FROM _dever_history WHERE model = $1 AND kind = 'catalog' AND name = 'current'",
|
||||
&[&model],
|
||||
)
|
||||
.await
|
||||
.map(|row| row.map(|row| row.get(0)))
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
async fn catalog(connection: &deadpool_postgres::Object, table: &str) -> Result<String, Error> {
|
||||
let columns = connection
|
||||
.query(
|
||||
"SELECT attribute.attname, pg_catalog.format_type(attribute.atttypid, attribute.atttypmod), attribute.attnotnull, attribute.attidentity::text, pg_catalog.pg_get_expr(default_value.adbin, default_value.adrelid) FROM pg_catalog.pg_attribute attribute JOIN pg_catalog.pg_class relation ON relation.oid = attribute.attrelid JOIN pg_catalog.pg_namespace namespace ON namespace.oid = relation.relnamespace LEFT JOIN pg_catalog.pg_attrdef default_value ON default_value.adrelid = relation.oid AND default_value.adnum = attribute.attnum WHERE namespace.nspname = current_schema() AND relation.relname = $1 AND attribute.attnum > 0 AND NOT attribute.attisdropped ORDER BY attribute.attnum",
|
||||
&[&table],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
(
|
||||
row.get::<_, String>(0),
|
||||
row.get::<_, String>(1),
|
||||
row.get::<_, bool>(2),
|
||||
row.get::<_, String>(3),
|
||||
row.get::<_, Option<String>>(4),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let constraints = connection
|
||||
.query(
|
||||
"SELECT constraint_.conname, constraint_.contype::text, pg_catalog.pg_get_constraintdef(constraint_.oid, false) FROM pg_catalog.pg_constraint constraint_ JOIN pg_catalog.pg_class relation ON relation.oid = constraint_.conrelid JOIN pg_catalog.pg_namespace namespace ON namespace.oid = relation.relnamespace WHERE namespace.nspname = current_schema() AND relation.relname = $1 ORDER BY constraint_.conname",
|
||||
&[&table],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
(
|
||||
row.get::<_, String>(0),
|
||||
row.get::<_, String>(1),
|
||||
row.get::<_, String>(2),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let indexes = connection
|
||||
.query(
|
||||
"SELECT index_relation.relname, pg_catalog.pg_get_indexdef(index_relation.oid) FROM pg_catalog.pg_index index_ JOIN pg_catalog.pg_class table_relation ON table_relation.oid = index_.indrelid JOIN pg_catalog.pg_class index_relation ON index_relation.oid = index_.indexrelid JOIN pg_catalog.pg_namespace namespace ON namespace.oid = table_relation.relnamespace WHERE namespace.nspname = current_schema() AND table_relation.relname = $1 AND left(index_relation.relname, 11) = '_dever_idx_' ORDER BY index_relation.relname",
|
||||
&[&table],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)?
|
||||
.into_iter()
|
||||
.map(|row| (row.get::<_, String>(0), row.get::<_, String>(1)))
|
||||
.collect::<Vec<_>>();
|
||||
serde_json::to_string(&(columns, constraints, indexes))
|
||||
.map_err(|error| Error::migration(format!("cannot encode PostgreSQL catalog: {error}")))
|
||||
}
|
||||
|
||||
async fn validate_migrations(
|
||||
connection: &deadpool_postgres::Object,
|
||||
schema: &Schema,
|
||||
) -> Result<(), Error> {
|
||||
let applied = load_migrations(connection, &schema.model).await?;
|
||||
let desired = schema
|
||||
.migrations
|
||||
.iter()
|
||||
.map(|migration| (migration.name.as_str(), migration))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
for (name, revision) in applied {
|
||||
let Some(migration) = desired.get(name.as_str()) else {
|
||||
return Err(Error::migration(format!(
|
||||
"applied migration '{}.{}' was removed from source",
|
||||
schema.model, name
|
||||
)));
|
||||
};
|
||||
if migration.revision != revision {
|
||||
return Err(Error::migration(format!(
|
||||
"applied migration '{}.{}' was modified",
|
||||
schema.model, name
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn load_migrations(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &str,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
connection
|
||||
.query(
|
||||
"SELECT name, revision FROM _dever_history WHERE model = $1 AND kind = 'migration'",
|
||||
&[&model],
|
||||
)
|
||||
.await
|
||||
.map_err(database_error)
|
||||
.map(|rows| {
|
||||
rows.into_iter()
|
||||
.map(|row| (row.get(0), row.get(1)))
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
|
||||
async fn store_history(
|
||||
connection: &deadpool_postgres::Object,
|
||||
schema: &Schema,
|
||||
) -> Result<(), Error> {
|
||||
let definition = serde_json::to_string(schema).map_err(|error| {
|
||||
Error::migration(format!("cannot encode Dever schema history: {error}"))
|
||||
})?;
|
||||
store_history_row(
|
||||
connection,
|
||||
&schema.model,
|
||||
"schema",
|
||||
"current",
|
||||
&schema.revision,
|
||||
&definition,
|
||||
)
|
||||
.await?;
|
||||
for migration in &schema.migrations {
|
||||
store_history_row(
|
||||
connection,
|
||||
&schema.model,
|
||||
"migration",
|
||||
&migration.name,
|
||||
&migration.revision,
|
||||
"",
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
store_history_row(
|
||||
connection,
|
||||
&schema.model,
|
||||
"seed",
|
||||
"current",
|
||||
&schema.seed_revision,
|
||||
"",
|
||||
)
|
||||
.await?;
|
||||
store_catalog_history(connection, &schema.model, &schema.table).await
|
||||
}
|
||||
|
||||
async fn store_catalog_history(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &str,
|
||||
table: &str,
|
||||
) -> Result<(), Error> {
|
||||
store_history_row(
|
||||
connection,
|
||||
model,
|
||||
"catalog",
|
||||
"current",
|
||||
"actual",
|
||||
&catalog(connection, table).await?,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn store_history_row(
|
||||
connection: &deadpool_postgres::Object,
|
||||
model: &str,
|
||||
kind: &str,
|
||||
name: &str,
|
||||
revision: &str,
|
||||
definition: &str,
|
||||
) -> Result<(), Error> {
|
||||
let parameters: [&(dyn ToSql + Sync); 5] = [&model, &kind, &name, &revision, &definition];
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO _dever_history (model, kind, name, revision, definition) VALUES ($1, $2, $3, $4, $5) \
|
||||
ON CONFLICT (model, kind, name) DO UPDATE SET revision = EXCLUDED.revision, definition = EXCLUDED.definition, applied_at = (extract(epoch from clock_timestamp()) * 1000)::bigint WHERE _dever_history.kind <> 'migration'",
|
||||
¶meters,
|
||||
)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(database_error)
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
use crate::collections::List;
|
||||
|
||||
pub fn arguments() -> Result<List<String>, String> {
|
||||
// argv[0] 是宿主可执行文件路径，不属于应用参数。
|
||||
std::env::args_os()
|
||||
.skip(1)
|
||||
.map(|argument| {
|
||||
argument
|
||||
.into_string()
|
||||
.map_err(|_| "process argument is not valid UTF-8".into())
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(List::new)
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
use std::fmt::{Display, Write};
|
||||
|
||||
use crate::Id;
|
||||
use crate::bytes::Bytes;
|
||||
use crate::collections::{List, Map, MapEntry};
|
||||
use crate::number::DecimalValue;
|
||||
use crate::resource::{Resource, Stream};
|
||||
|
||||
pub trait Render {
|
||||
fn render_to(&self, output: &mut String);
|
||||
|
||||
fn render(&self) -> String {
|
||||
let mut output = String::new();
|
||||
self.render_to(&mut output);
|
||||
output
|
||||
}
|
||||
}
|
||||
|
||||
fn write_display(output: &mut String, value: impl Display) {
|
||||
write!(output, "{value}").expect("writing to a String cannot fail");
|
||||
}
|
||||
|
||||
impl Render for String {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str(self);
|
||||
}
|
||||
}
|
||||
impl Render for bool {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
write_display(output, self);
|
||||
}
|
||||
}
|
||||
impl Render for i64 {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
write_display(output, self);
|
||||
}
|
||||
}
|
||||
impl Render for f64 {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
write_display(output, self);
|
||||
}
|
||||
}
|
||||
impl Render for DecimalValue {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
write_display(output, self);
|
||||
}
|
||||
}
|
||||
|
||||
impl Render for crate::orm::Uuid {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
use std::fmt::Write as _;
|
||||
write!(output, "{self}").expect("string formatting");
|
||||
}
|
||||
}
|
||||
impl Render for Id {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str(&self.0);
|
||||
}
|
||||
}
|
||||
impl Render for Bytes {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
write!(output, "{:?}", self.values()).expect("writing to a String cannot fail");
|
||||
}
|
||||
}
|
||||
impl<T> Render for Resource<T> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("<resource>");
|
||||
}
|
||||
}
|
||||
impl<T> Render for Stream<T> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("<stream>");
|
||||
}
|
||||
}
|
||||
impl<T> Render for crate::async_stream::AsyncStream<T> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("<async-stream>");
|
||||
}
|
||||
}
|
||||
impl Render for crate::net::Socket {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("<resource>");
|
||||
}
|
||||
}
|
||||
impl Render for crate::net::Listener {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("<resource>");
|
||||
}
|
||||
}
|
||||
|
||||
impl Render for crate::http::HttpReply {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("HttpReply(<resource>)");
|
||||
}
|
||||
}
|
||||
impl Render for crate::websocket::WebSocket {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("WebSocket(<resource>)");
|
||||
}
|
||||
}
|
||||
|
||||
impl Render for crate::tls::ClientTls {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("ClientTls(<resource>)");
|
||||
}
|
||||
}
|
||||
impl Render for crate::tls::ServerTls {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("ServerTls(<resource>)");
|
||||
}
|
||||
}
|
||||
impl Render for crate::http::HttpClient {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("HttpClient(<resource>)");
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Render for crate::channel::Channel<T> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("Channel(<resource>)");
|
||||
}
|
||||
}
|
||||
impl<T: Render> Render for Option<T> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
match self {
|
||||
Some(value) => value.render_to(output),
|
||||
None => output.push_str("null"),
|
||||
}
|
||||
}
|
||||
}
|
||||
impl<T: Render> Render for List<T> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push('[');
|
||||
for (index, value) in self.values().iter().enumerate() {
|
||||
if index != 0 {
|
||||
output.push_str(", ");
|
||||
}
|
||||
value.render_to(output);
|
||||
}
|
||||
output.push(']');
|
||||
}
|
||||
}
|
||||
impl<K: Render, V: Render> Render for Map<K, V> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push('{');
|
||||
for (index, (key, value)) in self.pairs().enumerate() {
|
||||
if index != 0 {
|
||||
output.push_str(", ");
|
||||
}
|
||||
key.render_to(output);
|
||||
output.push_str(" = ");
|
||||
value.render_to(output);
|
||||
}
|
||||
output.push('}');
|
||||
}
|
||||
}
|
||||
impl<K: Render, V: Render> Render for MapEntry<K, V> {
|
||||
fn render_to(&self, output: &mut String) {
|
||||
output.push_str("{key = ");
|
||||
self.key.render_to(output);
|
||||
output.push_str(", value = ");
|
||||
self.value.render_to(output);
|
||||
output.push('}');
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{Read, Write};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::bytes::Bytes;
|
||||
|
||||
pub enum Pull<T> {
|
||||
Item(T),
|
||||
Last(T),
|
||||
End,
|
||||
}
|
||||
|
||||
struct StreamState<T> {
|
||||
producer: Option<Box<dyn FnMut() -> Pull<T>>>,
|
||||
}
|
||||
|
||||
pub struct Stream<T>(Arc<Mutex<StreamState<T>>>);
|
||||
|
||||
impl<T> Clone for Stream<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self(Arc::clone(&self.0))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> std::fmt::Debug for Stream<T> {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("Stream(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: 'static> Stream<T> {
|
||||
pub fn new(producer: impl FnMut() -> Pull<T> + 'static) -> Self {
|
||||
Self(Arc::new(Mutex::new(StreamState {
|
||||
producer: Some(Box::new(producer)),
|
||||
})))
|
||||
}
|
||||
|
||||
fn pull_next(&self) -> Pull<T> {
|
||||
let mut state = self.0.lock().expect("stream lock poisoned");
|
||||
let Some(producer) = state.producer.as_mut() else {
|
||||
return Pull::End;
|
||||
};
|
||||
let result = producer.as_mut()();
|
||||
if matches!(&result, Pull::Last(_) | Pull::End) {
|
||||
state.producer = None;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
pub fn pull(&self) -> Option<T> {
|
||||
match self.pull_next() {
|
||||
Pull::Item(value) | Pull::Last(value) => Some(value),
|
||||
Pull::End => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn close(&self) {
|
||||
self.0.lock().expect("stream lock poisoned").producer = None;
|
||||
}
|
||||
|
||||
pub fn map<U: 'static>(self, mut transform: impl FnMut(T) -> U + 'static) -> Stream<U> {
|
||||
Stream::new(move || match self.pull_next() {
|
||||
Pull::Item(value) => Pull::Item(transform(value)),
|
||||
Pull::Last(value) => Pull::Last(transform(value)),
|
||||
Pull::End => Pull::End,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Resource<T>(Arc<Mutex<Option<T>>>);
|
||||
|
||||
impl<T> Clone for Resource<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self(Arc::clone(&self.0))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Resource<T> {
|
||||
pub fn new(value: T) -> Self {
|
||||
Self(Arc::new(Mutex::new(Some(value))))
|
||||
}
|
||||
|
||||
pub fn close(&self) -> Result<(), String> {
|
||||
self.0
|
||||
.lock()
|
||||
.map_err(|_| "resource lock failed".to_owned())?
|
||||
.take()
|
||||
.map(|_| ())
|
||||
.ok_or_else(|| "resource is closed".to_owned())
|
||||
}
|
||||
|
||||
pub fn is_open(&self) -> Result<bool, String> {
|
||||
self.0
|
||||
.lock()
|
||||
.map(|resource| resource.is_some())
|
||||
.map_err(|_| "resource lock failed".to_owned())
|
||||
}
|
||||
|
||||
fn with<R>(&self, operation: impl FnOnce(&mut T) -> std::io::Result<R>) -> Result<R, String> {
|
||||
let mut guard = self
|
||||
.0
|
||||
.lock()
|
||||
.map_err(|_| "resource lock failed".to_owned())?;
|
||||
let resource = guard
|
||||
.as_mut()
|
||||
.ok_or_else(|| "resource is closed".to_owned())?;
|
||||
operation(resource).map_err(|error| error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub type FileHandle = Resource<File>;
|
||||
|
||||
pub fn file_open(path: &str) -> Result<FileHandle, String> {
|
||||
File::open(path)
|
||||
.map(Resource::new)
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
pub fn file_create(path: &str) -> Result<FileHandle, String> {
|
||||
OpenOptions::new()
|
||||
.write(true)
|
||||
.create_new(true)
|
||||
.open(path)
|
||||
.map(Resource::new)
|
||||
.map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
pub fn read<T: Read>(resource: &Resource<T>, limit: i64) -> Result<Option<Bytes>, String> {
|
||||
let limit = read_limit(limit)?;
|
||||
resource.with(|source| {
|
||||
let mut bytes = Vec::new();
|
||||
bytes
|
||||
.try_reserve_exact(limit)
|
||||
.map_err(std::io::Error::other)?;
|
||||
bytes.resize(limit, 0);
|
||||
let length = source.read(&mut bytes)?;
|
||||
bytes.truncate(length);
|
||||
Ok((length != 0).then(|| Bytes::new(bytes)))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn read_chunks<T: Read + 'static>(
|
||||
resource: Resource<T>,
|
||||
limit: i64,
|
||||
) -> Result<Stream<Result<Bytes, String>>, String> {
|
||||
read_limit(limit)?;
|
||||
if !resource.is_open()? {
|
||||
return Err("resource is closed".to_owned());
|
||||
}
|
||||
Ok(Stream::new(move || match read(&resource, limit) {
|
||||
Ok(Some(bytes)) => Pull::Item(Ok(bytes)),
|
||||
Ok(None) => Pull::End,
|
||||
Err(message) => Pull::Last(Err(message)),
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn read_limit(limit: i64) -> Result<usize, String> {
|
||||
usize::try_from(limit)
|
||||
.ok()
|
||||
.filter(|limit| *limit > 0)
|
||||
.ok_or_else(|| "read size must be a positive Int".to_owned())
|
||||
}
|
||||
|
||||
pub fn write<T: Write>(resource: &Resource<T>, bytes: &Bytes) -> Result<(), String> {
|
||||
resource.with(|output| {
|
||||
output.write_all(bytes.values())?;
|
||||
output.flush()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn stdout_write(text: &str) -> Result<(), String> {
|
||||
let mut output = std::io::stdout().lock();
|
||||
output
|
||||
.write_all(text.as_bytes())
|
||||
.and_then(|_| output.flush())
|
||||
.map_err(|error| format!("standard output failed: {error}"))
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
use std::fmt;
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
/// Opaque source value. Only approved runtime sinks may inspect its bytes.
|
||||
#[derive(Clone)]
|
||||
// The payload is consumed only when an approved crypto sink is linked.
|
||||
#[cfg_attr(not(feature = "crypto"), allow(dead_code))]
|
||||
pub struct Secret(pub(crate) Zeroizing<Vec<u8>>);
|
||||
|
||||
impl Secret {
|
||||
/// Compiler-owned input decoding and test construction; no source-level reveal.
|
||||
pub fn from_input(bytes: Vec<u8>) -> Self {
|
||||
Self(Zeroizing::new(bytes))
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for Secret {
|
||||
fn fmt(&self, output: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
output.write_str("Secret([REDACTED])")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,578 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use deadpool_sqlite::rusqlite::types::{Value as SqliteValue, ValueRef};
|
||||
use deadpool_sqlite::{Config, Hook, HookError, Pool, PoolConfig, Runtime, Timeouts};
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
use crate::config::{Database as DatabaseConfig, Settings};
|
||||
use crate::orm::{Error, ErrorKind, Row, RowStream, Value};
|
||||
use crate::pool::{CREATE_TIMEOUT, PoolObjectGuard, RECYCLE_TIMEOUT, WAIT_TIMEOUT};
|
||||
|
||||
mod migration;
|
||||
|
||||
const CONNECTION_SETUP: &str = "\
|
||||
PRAGMA foreign_keys = ON;\n\
|
||||
PRAGMA synchronous = NORMAL;\n\
|
||||
PRAGMA busy_timeout = 5000;";
|
||||
|
||||
static DATABASES: OnceLock<BTreeMap<String, Arc<Database>>> = OnceLock::new();
|
||||
|
||||
pub struct Database {
|
||||
pool: Pool,
|
||||
max_page_size: usize,
|
||||
initialized: OnceCell<()>,
|
||||
}
|
||||
|
||||
pub struct Transaction {
|
||||
connection: Option<deadpool_sqlite::Object>,
|
||||
max_page_size: usize,
|
||||
}
|
||||
|
||||
pub enum Executor<'a> {
|
||||
Pool(&'a Database),
|
||||
Transaction(&'a Transaction),
|
||||
}
|
||||
|
||||
impl Database {
|
||||
pub fn open(
|
||||
path: impl Into<std::path::PathBuf>,
|
||||
max_connections: usize,
|
||||
max_page_size: usize,
|
||||
) -> Result<Self, Error> {
|
||||
let config = Config::new(path);
|
||||
let manager = deadpool_sqlite::Manager::from_config(&config, Runtime::Tokio1);
|
||||
let mut pool_config = PoolConfig::new(max_connections);
|
||||
pool_config.timeouts = Timeouts {
|
||||
wait: Some(WAIT_TIMEOUT),
|
||||
create: Some(CREATE_TIMEOUT),
|
||||
recycle: Some(RECYCLE_TIMEOUT),
|
||||
};
|
||||
let pool = Pool::builder(manager)
|
||||
.config(pool_config)
|
||||
.runtime(Runtime::Tokio1)
|
||||
.post_create(Hook::async_fn(|connection, _| {
|
||||
Box::pin(async move {
|
||||
connection
|
||||
.interact(|connection| connection.execute_batch(CONNECTION_SETUP))
|
||||
.await
|
||||
.map_err(|error| HookError::message(error.to_string()))?
|
||||
.map_err(HookError::Backend)
|
||||
})
|
||||
}))
|
||||
.build()
|
||||
.map_err(|error| Error::new(ErrorKind::Pool, error.to_string()))?;
|
||||
Ok(Self {
|
||||
pool,
|
||||
max_page_size,
|
||||
initialized: OnceCell::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn max_page_size(&self) -> usize {
|
||||
self.max_page_size
|
||||
}
|
||||
|
||||
pub(crate) fn close(&self) {
|
||||
self.pool.close();
|
||||
}
|
||||
|
||||
pub async fn execute(&self, sql: &'static str, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
self.prepare().await?;
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = execute_on(connection.object(), sql, parameters).await;
|
||||
drop(connection.release());
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn query(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
self.prepare().await?;
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = query_on(connection.object(), sql, parameters).await;
|
||||
drop(connection.release());
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn query_bounded(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
self.prepare().await?;
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = query_bounded_on(connection.object(), sql, parameters, maximum).await;
|
||||
drop(connection.release());
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn query_owned(
|
||||
&self,
|
||||
sql: String,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
self.prepare().await?;
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = query_owned_on(connection.object(), sql, parameters).await;
|
||||
drop(connection.release());
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn stream(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
capacity: usize,
|
||||
) -> Result<RowStream<Row>, Error> {
|
||||
self.prepare().await?;
|
||||
if capacity == 0 || capacity > self.max_page_size {
|
||||
return Err(Error::invalid_data(format!(
|
||||
"row stream capacity must be between 1 and {}",
|
||||
self.max_page_size
|
||||
)));
|
||||
}
|
||||
let connection = self.connection().await?;
|
||||
let (rows, receiver) = tokio::sync::mpsc::channel(capacity);
|
||||
let (ready, initialized) = tokio::sync::oneshot::channel();
|
||||
let (worker, spawner, worker_fault) =
|
||||
crate::task::scoped_owner(()).map_err(|error| Error::new(ErrorKind::Pool, error))?;
|
||||
spawner.spawn(async move {
|
||||
let connection = guarded(connection);
|
||||
let result = connection
|
||||
.object()
|
||||
.interact(move |connection| stream_rows(connection, sql, parameters, rows, ready))
|
||||
.await;
|
||||
if matches!(result, Ok(true)) {
|
||||
drop(connection.release());
|
||||
}
|
||||
});
|
||||
initialized
|
||||
.await
|
||||
.map_err(|_| Error::new(ErrorKind::Pool, "SQLite row stream worker stopped"))??;
|
||||
Ok(RowStream::from_receiver(receiver, worker, worker_fault))
|
||||
}
|
||||
|
||||
pub async fn execute_batch(&self, sql: String) -> Result<(), Error> {
|
||||
self.prepare().await?;
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = connection
|
||||
.object()
|
||||
.interact(move |connection| connection.execute_batch(&sql))
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
.map_err(database_error);
|
||||
drop(connection.release());
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn begin(&self) -> Result<Transaction, Error> {
|
||||
self.begin_with("BEGIN IMMEDIATE").await
|
||||
}
|
||||
|
||||
async fn begin_with(&self, statement: &'static str) -> Result<Transaction, Error> {
|
||||
self.prepare().await?;
|
||||
let connection = guarded(self.connection().await?);
|
||||
batch_on(connection.object(), statement).await?;
|
||||
Ok(Transaction {
|
||||
connection: Some(connection.release()),
|
||||
max_page_size: self.max_page_size,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn migrate(&self, model: crate::database::Model) -> Result<(), Error> {
|
||||
self.prepare().await?;
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = migration::apply(connection.object(), model).await;
|
||||
if result.is_ok() {
|
||||
drop(connection.release());
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
async fn prepare(&self) -> Result<(), Error> {
|
||||
self.initialized
|
||||
.get_or_try_init(|| async {
|
||||
let connection = guarded(self.connection().await?);
|
||||
let result = connection
|
||||
.object()
|
||||
.interact(|connection| connection.pragma_update(None, "journal_mode", "WAL"))
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
.map_err(database_error);
|
||||
if result.is_ok() {
|
||||
drop(connection.release());
|
||||
}
|
||||
result
|
||||
})
|
||||
.await
|
||||
.copied()
|
||||
}
|
||||
|
||||
async fn connection(&self) -> Result<deadpool_sqlite::Object, Error> {
|
||||
self.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|error| Error::new(ErrorKind::Pool, error.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
impl Transaction {
|
||||
pub async fn execute(&self, sql: &'static str, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
execute_on(self.connection(), sql, parameters).await
|
||||
}
|
||||
|
||||
pub async fn query(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
query_on(self.connection(), sql, parameters).await
|
||||
}
|
||||
|
||||
pub async fn query_bounded(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
query_bounded_on(self.connection(), sql, parameters, maximum).await
|
||||
}
|
||||
|
||||
pub async fn query_owned(
|
||||
&self,
|
||||
sql: String,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
query_owned_on(self.connection(), sql, parameters).await
|
||||
}
|
||||
|
||||
pub async fn commit(mut self) -> Result<(), Error> {
|
||||
self.finish("COMMIT").await
|
||||
}
|
||||
|
||||
pub async fn rollback(mut self) -> Result<(), Error> {
|
||||
self.finish("ROLLBACK").await
|
||||
}
|
||||
|
||||
fn connection(&self) -> &deadpool_sqlite::Object {
|
||||
self.connection.as_ref().expect("open SQLite transaction")
|
||||
}
|
||||
|
||||
async fn finish(&mut self, sql: &'static str) -> Result<(), Error> {
|
||||
let connection = guarded(self.connection.take().expect("open SQLite transaction"));
|
||||
let result = batch_on(connection.object(), sql).await;
|
||||
if result.is_ok() {
|
||||
drop(connection.release());
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Transaction {
|
||||
fn drop(&mut self) {
|
||||
if let Some(connection) = self.connection.take() {
|
||||
drop(deadpool_sqlite::Object::take(connection));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn guarded(connection: deadpool_sqlite::Object) -> PoolObjectGuard<deadpool_sqlite::Object> {
|
||||
PoolObjectGuard::new(connection, detach)
|
||||
}
|
||||
|
||||
fn detach(connection: deadpool_sqlite::Object) {
|
||||
drop(deadpool_sqlite::Object::take(connection));
|
||||
}
|
||||
|
||||
impl Executor<'_> {
|
||||
pub fn max_page_size(&self) -> usize {
|
||||
match self {
|
||||
Self::Pool(database) => database.max_page_size(),
|
||||
Self::Transaction(transaction) => transaction.max_page_size,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(&self, sql: &'static str, parameters: Vec<Value>) -> Result<usize, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.execute(sql, parameters).await,
|
||||
Self::Transaction(transaction) => transaction.execute(sql, parameters).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query(
|
||||
&self,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
match self {
|
||||
Self::Pool(database) => database.query(sql, parameters).await,
|
||||
Self::Transaction(transaction) => transaction.query(sql, parameters).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn initialize(settings: &Settings) -> Result<(), String> {
|
||||
DATABASES
|
||||
.set(configured_databases(settings)?)
|
||||
.map_err(|_| "SQLite databases were initialized more than once".to_owned())
|
||||
}
|
||||
|
||||
pub(crate) fn configured_databases(
|
||||
settings: &Settings,
|
||||
) -> Result<BTreeMap<String, Arc<Database>>, String> {
|
||||
let mut databases = BTreeMap::new();
|
||||
for (name, config) in settings.databases() {
|
||||
if let DatabaseConfig::Sqlite {
|
||||
path,
|
||||
max_connections,
|
||||
max_page_size,
|
||||
..
|
||||
} = config
|
||||
{
|
||||
let database = Database::open(path, *max_connections, *max_page_size)
|
||||
.map_err(|error| format!("cannot initialize database.{name}: {error}"))?;
|
||||
databases.insert(name.clone(), Arc::new(database));
|
||||
}
|
||||
}
|
||||
Ok(databases)
|
||||
}
|
||||
|
||||
pub fn database(name: &str) -> Result<Arc<Database>, Error> {
|
||||
DATABASES
|
||||
.get()
|
||||
.and_then(|databases| databases.get(name))
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
Error::new(
|
||||
ErrorKind::Pool,
|
||||
format!("SQLite database '{name}' is not initialized"),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn shutdown() {
|
||||
if let Some(databases) = DATABASES.get() {
|
||||
for database in databases.values() {
|
||||
database.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn database_for(explicit: Option<&str>, package_root: &str) -> Result<Arc<Database>, Error> {
|
||||
let (name, config) = crate::config::settings()
|
||||
.resolve_database(explicit, package_root)
|
||||
.map_err(|error| Error::new(ErrorKind::Pool, error))?;
|
||||
if !matches!(config, DatabaseConfig::Sqlite { .. }) {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Pool,
|
||||
format!("database connection '{name}' is not SQLite"),
|
||||
));
|
||||
}
|
||||
database(name)
|
||||
}
|
||||
|
||||
fn sqlite_values(parameters: Vec<Value>) -> Vec<SqliteValue> {
|
||||
parameters
|
||||
.into_iter()
|
||||
.map(|value| match value {
|
||||
Value::Null => SqliteValue::Null,
|
||||
Value::Bool(value) => SqliteValue::Integer(i64::from(value)),
|
||||
Value::Int(value) => SqliteValue::Integer(value),
|
||||
Value::Float(value) => SqliteValue::Real(value),
|
||||
Value::Decimal(value) => SqliteValue::Text(value),
|
||||
Value::Text(value) => SqliteValue::Text(value),
|
||||
Value::Bytes(value) => SqliteValue::Blob(value),
|
||||
Value::Uuid(value) => SqliteValue::Blob(value.into_bytes().to_vec()),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn execute_on(
|
||||
connection: &deadpool_sqlite::Object,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<usize, Error> {
|
||||
connection
|
||||
.interact(move |connection| {
|
||||
let parameters = sqlite_values(parameters);
|
||||
connection.execute(sql, deadpool_sqlite::rusqlite::params_from_iter(parameters))
|
||||
})
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
async fn query_on(
|
||||
connection: &deadpool_sqlite::Object,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
connection
|
||||
.interact(move |connection| {
|
||||
let parameters = sqlite_values(parameters);
|
||||
let mut statement = connection.prepare_cached(sql)?;
|
||||
let column_count = statement.column_count();
|
||||
let mut rows =
|
||||
statement.query(deadpool_sqlite::rusqlite::params_from_iter(parameters))?;
|
||||
let mut output = Vec::new();
|
||||
while let Some(row) = rows.next()? {
|
||||
let values = (0..column_count)
|
||||
.map(|index| value(index, row.get_ref(index)?))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
output.push(Row::new(values));
|
||||
}
|
||||
Ok(output)
|
||||
})
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
async fn query_bounded_on(
|
||||
connection: &deadpool_sqlite::Object,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
maximum: usize,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
let rows = connection
|
||||
.interact(move |connection| {
|
||||
let parameters = sqlite_values(parameters);
|
||||
let mut statement = connection.prepare_cached(sql)?;
|
||||
let column_count = statement.column_count();
|
||||
let mut rows =
|
||||
statement.query(deadpool_sqlite::rusqlite::params_from_iter(parameters))?;
|
||||
let mut output = Vec::with_capacity(maximum.saturating_add(1));
|
||||
while output.len() <= maximum {
|
||||
let Some(row) = rows.next()? else {
|
||||
break;
|
||||
};
|
||||
let values = (0..column_count)
|
||||
.map(|index| value(index, row.get_ref(index)?))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
output.push(Row::new(values));
|
||||
}
|
||||
Ok(output)
|
||||
})
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
.map_err(database_error)?;
|
||||
crate::orm::bounded_rows(rows, maximum)
|
||||
}
|
||||
|
||||
async fn query_owned_on(
|
||||
connection: &deadpool_sqlite::Object,
|
||||
sql: String,
|
||||
parameters: Vec<Value>,
|
||||
) -> Result<Vec<Row>, Error> {
|
||||
connection
|
||||
.interact(move |connection| {
|
||||
let parameters = sqlite_values(parameters);
|
||||
let mut statement = connection.prepare(&sql)?;
|
||||
let column_count = statement.column_count();
|
||||
let mut rows =
|
||||
statement.query(deadpool_sqlite::rusqlite::params_from_iter(parameters))?;
|
||||
let mut output = Vec::new();
|
||||
while let Some(row) = rows.next()? {
|
||||
let values = (0..column_count)
|
||||
.map(|index| value(index, row.get_ref(index)?))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
output.push(Row::new(values));
|
||||
}
|
||||
Ok(output)
|
||||
})
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
fn stream_rows(
|
||||
connection: &mut deadpool_sqlite::rusqlite::Connection,
|
||||
sql: &'static str,
|
||||
parameters: Vec<Value>,
|
||||
rows: tokio::sync::mpsc::Sender<Result<Row, Error>>,
|
||||
ready: tokio::sync::oneshot::Sender<Result<(), Error>>,
|
||||
) -> bool {
|
||||
let parameters = sqlite_values(parameters);
|
||||
let mut statement = match connection.prepare_cached(sql) {
|
||||
Ok(statement) => statement,
|
||||
Err(error) => {
|
||||
let _ = ready.send(Err(database_error(error)));
|
||||
return false;
|
||||
}
|
||||
};
|
||||
let column_count = statement.column_count();
|
||||
let mut source = match statement.query(deadpool_sqlite::rusqlite::params_from_iter(parameters))
|
||||
{
|
||||
Ok(source) => source,
|
||||
Err(error) => {
|
||||
let _ = ready.send(Err(database_error(error)));
|
||||
return false;
|
||||
}
|
||||
};
|
||||
if ready.send(Ok(())).is_err() {
|
||||
return true;
|
||||
}
|
||||
loop {
|
||||
let row = match source.next() {
|
||||
Ok(Some(row)) => row,
|
||||
Ok(None) => return true,
|
||||
Err(error) => {
|
||||
let _ = rows.blocking_send(Err(database_error(error)));
|
||||
return false;
|
||||
}
|
||||
};
|
||||
let values = (0..column_count)
|
||||
.map(|index| value(index, row.get_ref(index)?))
|
||||
.collect::<Result<Vec<_>, deadpool_sqlite::rusqlite::Error>>();
|
||||
let item = values.map(Row::new).map_err(database_error);
|
||||
let healthy = item.is_ok();
|
||||
if rows.blocking_send(item).is_err() {
|
||||
return true;
|
||||
}
|
||||
if !healthy {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn batch_on(connection: &deadpool_sqlite::Object, sql: &'static str) -> Result<(), Error> {
|
||||
connection
|
||||
.interact(move |connection| connection.execute_batch(sql))
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
fn value(column: usize, value: ValueRef<'_>) -> Result<Value, deadpool_sqlite::rusqlite::Error> {
|
||||
Ok(match value {
|
||||
ValueRef::Null => Value::Null,
|
||||
ValueRef::Integer(value) => Value::Int(value),
|
||||
ValueRef::Real(value) => Value::Float(value),
|
||||
ValueRef::Text(value) => Value::Text(
|
||||
std::str::from_utf8(value)
|
||||
.map_err(|error| deadpool_sqlite::rusqlite::Error::Utf8Error(column, error))?
|
||||
.to_owned(),
|
||||
),
|
||||
ValueRef::Blob(value) => Value::Bytes(value.to_vec()),
|
||||
})
|
||||
}
|
||||
|
||||
fn interact_error(error: deadpool_sqlite::InteractError) -> Error {
|
||||
Error::new(ErrorKind::Pool, error.to_string())
|
||||
}
|
||||
|
||||
fn database_error(error: deadpool_sqlite::rusqlite::Error) -> Error {
|
||||
let kind = match &error {
|
||||
deadpool_sqlite::rusqlite::Error::SqliteFailure(failure, _)
|
||||
if failure.code == deadpool_sqlite::rusqlite::ErrorCode::ConstraintViolation =>
|
||||
{
|
||||
ErrorKind::Constraint
|
||||
}
|
||||
_ => ErrorKind::Database,
|
||||
};
|
||||
Error::new(kind, error.to_string())
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use deadpool_sqlite::rusqlite::{Connection, OptionalExtension, params};
|
||||
|
||||
use crate::database::{Field, Migration, MigrationPhase, Model, Schema};
|
||||
use crate::orm::Error;
|
||||
|
||||
use super::{database_error, interact_error, sqlite_values};
|
||||
|
||||
const CREATE_HISTORY: &str = "\
|
||||
CREATE TABLE IF NOT EXISTS _dever_history (\
|
||||
model TEXT NOT NULL,\
|
||||
kind TEXT NOT NULL,\
|
||||
name TEXT NOT NULL,\
|
||||
revision TEXT NOT NULL,\
|
||||
definition TEXT NOT NULL,\
|
||||
applied_at INTEGER NOT NULL DEFAULT (unixepoch()),\
|
||||
PRIMARY KEY (model, kind, name)\
|
||||
)";
|
||||
|
||||
pub(super) async fn apply(connection: &deadpool_sqlite::Object, model: Model) -> Result<(), Error> {
|
||||
connection
|
||||
.interact(move |connection| transaction(connection, model))
|
||||
.await
|
||||
.map_err(interact_error)?
|
||||
}
|
||||
|
||||
fn transaction(connection: &mut Connection, model: Model) -> Result<(), Error> {
|
||||
connection
|
||||
.execute_batch("BEGIN IMMEDIATE; PRAGMA defer_foreign_keys = ON;")
|
||||
.map_err(database_error)?;
|
||||
let result = apply_locked(connection, &model);
|
||||
let result = match result {
|
||||
Ok(()) => match connection.execute_batch("COMMIT") {
|
||||
Ok(()) => Ok(()),
|
||||
Err(commit) => rollback(connection, database_error(commit)),
|
||||
},
|
||||
Err(error) => rollback(connection, error),
|
||||
};
|
||||
connection.flush_prepared_statement_cache();
|
||||
result
|
||||
}
|
||||
|
||||
fn rollback(connection: &Connection, error: Error) -> Result<(), Error> {
|
||||
match connection.execute_batch("ROLLBACK") {
|
||||
Ok(()) => Err(error),
|
||||
Err(rollback) => Err(error.cause(rollback)),
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_locked(connection: &Connection, model: &Model) -> Result<(), Error> {
|
||||
connection
|
||||
.execute_batch(CREATE_HISTORY)
|
||||
.map_err(database_error)?;
|
||||
let table_exists = table_exists(connection, &model.schema.table)?;
|
||||
let previous = load_schema(connection, &model.schema.model)?;
|
||||
match (table_exists, previous.as_ref()) {
|
||||
(false, None) => create_model(connection, model)?,
|
||||
(true, Some(previous)) => update_model(connection, model, previous)?,
|
||||
(true, None) => {
|
||||
return Err(Error::migration(format!(
|
||||
"table '{}' exists without Dever schema history",
|
||||
model.schema.table
|
||||
)));
|
||||
}
|
||||
(false, Some(_)) => {
|
||||
return Err(Error::migration(format!(
|
||||
"Dever schema history exists but table '{}' is missing",
|
||||
model.schema.table
|
||||
)));
|
||||
}
|
||||
}
|
||||
apply_seed(connection, model, previous.as_ref())?;
|
||||
store_history(connection, &model.schema)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_model(connection: &Connection, model: &Model) -> Result<(), Error> {
|
||||
connection
|
||||
.execute_batch(&model.create_table.sqlite)
|
||||
.map_err(database_error)?;
|
||||
create_indexes(connection, &model.create_indexes)
|
||||
}
|
||||
|
||||
fn update_model(connection: &Connection, model: &Model, previous: &Schema) -> Result<(), Error> {
|
||||
if previous.table != model.schema.table {
|
||||
return Err(Error::migration(format!(
|
||||
"Model '{}' cannot rename table '{}' to '{}' automatically",
|
||||
model.schema.model, previous.table, model.schema.table
|
||||
)));
|
||||
}
|
||||
validate_migrations(connection, &model.schema)?;
|
||||
crate::database::validate_migration_order(previous, &model.schema)?;
|
||||
validate_catalog(connection, previous)?;
|
||||
let pending = unapplied_migrations(connection, &model.schema)?;
|
||||
apply_data(connection, model, &pending, MigrationPhase::Before)?;
|
||||
if previous.revision == model.schema.revision {
|
||||
create_indexes(connection, &model.create_indexes)?;
|
||||
} else if fields_match(&previous.fields, &model.schema.fields) {
|
||||
synchronize_indexes(connection, previous, model)?;
|
||||
} else {
|
||||
rebuild_table(connection, previous, model, &pending)?;
|
||||
}
|
||||
apply_data(connection, model, &pending, MigrationPhase::After)
|
||||
}
|
||||
|
||||
fn apply_data(
|
||||
connection: &Connection,
|
||||
model: &Model,
|
||||
pending: &[Migration],
|
||||
phase: MigrationPhase,
|
||||
) -> Result<(), Error> {
|
||||
for statement in &model.data_migrations {
|
||||
if statement.phase != phase
|
||||
|| !pending
|
||||
.iter()
|
||||
.any(|migration| migration.name == statement.name)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let parameters = sqlite_values(statement.parameters.clone());
|
||||
connection
|
||||
.execute(
|
||||
statement.sql.sqlite,
|
||||
deadpool_sqlite::rusqlite::params_from_iter(parameters),
|
||||
)
|
||||
.map_err(database_error)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn rebuild_table(
|
||||
connection: &Connection,
|
||||
previous: &Schema,
|
||||
model: &Model,
|
||||
new_migrations: &[Migration],
|
||||
) -> Result<(), Error> {
|
||||
let allowed_drops = new_migrations
|
||||
.iter()
|
||||
.flat_map(|migration| migration.drops.iter().cloned())
|
||||
.collect::<BTreeSet<_>>();
|
||||
let previous_fields = previous
|
||||
.fields
|
||||
.iter()
|
||||
.map(|field| (field.name.as_str(), field))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let mut copied = Vec::new();
|
||||
let mut consumed = BTreeSet::new();
|
||||
for target in &model.schema.fields {
|
||||
let source = previous_fields
|
||||
.get(target.name.as_str())
|
||||
.copied()
|
||||
.or_else(|| {
|
||||
target
|
||||
.rename_from
|
||||
.as_deref()
|
||||
.and_then(|name| previous_fields.get(name).copied())
|
||||
});
|
||||
let Some(source) = source else {
|
||||
continue;
|
||||
};
|
||||
if !compatible_types(&source.ty, &target.ty) {
|
||||
return Err(Error::migration(format!(
|
||||
"field '{}.{}' cannot change type from '{}' to '{}' automatically",
|
||||
model.schema.model, target.name, source.ty, target.ty
|
||||
)));
|
||||
}
|
||||
if !consumed.insert(source.name.as_str()) {
|
||||
return Err(Error::migration(format!(
|
||||
"field '{}.{}' reuses migration source '{}'",
|
||||
model.schema.model, target.name, source.name
|
||||
)));
|
||||
}
|
||||
copied.push((target.name.as_str(), source.name.as_str()));
|
||||
}
|
||||
for source in &previous.fields {
|
||||
if !consumed.contains(source.name.as_str()) && !allowed_drops.contains(&source.name) {
|
||||
return Err(Error::migration(format!(
|
||||
"field '{}.{}' was removed without a named migrate drop",
|
||||
model.schema.model, source.name
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
connection
|
||||
.execute_batch(&model.create_temporary_table)
|
||||
.map_err(database_error)?;
|
||||
let targets = copied
|
||||
.iter()
|
||||
.map(|(target, _)| quote(target))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let sources = copied
|
||||
.iter()
|
||||
.map(|(_, source)| quote(source))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let temporary = temporary_table(&model.schema);
|
||||
connection
|
||||
.execute_batch(&format!(
|
||||
"INSERT INTO {} ({targets}) SELECT {sources} FROM {};\nDROP TABLE {};\nALTER TABLE {} RENAME TO {};",
|
||||
quote(&temporary),
|
||||
quote(&model.schema.table),
|
||||
quote(&model.schema.table),
|
||||
quote(&temporary),
|
||||
quote(&model.schema.table),
|
||||
))
|
||||
.map_err(database_error)?;
|
||||
create_indexes(connection, &model.create_indexes)?;
|
||||
let mut foreign_keys = connection
|
||||
.prepare("PRAGMA foreign_key_check")
|
||||
.map_err(database_error)?;
|
||||
if foreign_keys
|
||||
.query([])
|
||||
.map_err(database_error)?
|
||||
.next()
|
||||
.map_err(database_error)?
|
||||
.is_some()
|
||||
{
|
||||
return Err(Error::migration(format!(
|
||||
"Model '{}' migration would leave invalid foreign keys",
|
||||
model.schema.model
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn synchronize_indexes(
|
||||
connection: &Connection,
|
||||
previous: &Schema,
|
||||
model: &Model,
|
||||
) -> Result<(), Error> {
|
||||
for index in &previous.indexes {
|
||||
connection
|
||||
.execute_batch(&format!("DROP INDEX IF EXISTS {}", quote(&index.name)))
|
||||
.map_err(database_error)?;
|
||||
}
|
||||
create_indexes(connection, &model.create_indexes)
|
||||
}
|
||||
|
||||
fn create_indexes(
|
||||
connection: &Connection,
|
||||
indexes: &[crate::database::SchemaSql],
|
||||
) -> Result<(), Error> {
|
||||
for index in indexes {
|
||||
connection
|
||||
.execute_batch(&index.sqlite)
|
||||
.map_err(database_error)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_seed(
|
||||
connection: &Connection,
|
||||
model: &Model,
|
||||
previous: Option<&Schema>,
|
||||
) -> Result<(), Error> {
|
||||
if previous.is_some_and(|schema| schema.seed_revision == model.schema.seed_revision) {
|
||||
return Ok(());
|
||||
}
|
||||
for seed in &model.seeds {
|
||||
let parameters = sqlite_values(seed.parameters.clone());
|
||||
connection
|
||||
.execute(
|
||||
seed.sql.sqlite,
|
||||
deadpool_sqlite::rusqlite::params_from_iter(parameters),
|
||||
)
|
||||
.map_err(database_error)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_catalog(connection: &Connection, schema: &Schema) -> Result<(), Error> {
|
||||
let recorded = load_catalog(connection, &schema.model)?.ok_or_else(|| {
|
||||
Error::migration(format!(
|
||||
"Model '{}' has no recorded database catalog",
|
||||
schema.model
|
||||
))
|
||||
})?;
|
||||
if catalog(connection, &schema.table)? != recorded {
|
||||
return Err(Error::migration(format!(
|
||||
"table '{}' does not match its recorded Dever schema",
|
||||
schema.table
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_migrations(connection: &Connection, schema: &Schema) -> Result<(), Error> {
|
||||
let applied = load_migrations(connection, &schema.model)?;
|
||||
let desired = schema
|
||||
.migrations
|
||||
.iter()
|
||||
.map(|migration| (migration.name.as_str(), migration))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
for (name, revision) in applied {
|
||||
let Some(migration) = desired.get(name.as_str()) else {
|
||||
return Err(Error::migration(format!(
|
||||
"applied migration '{}.{}' was removed from source",
|
||||
schema.model, name
|
||||
)));
|
||||
};
|
||||
if migration.revision != revision {
|
||||
return Err(Error::migration(format!(
|
||||
"applied migration '{}.{}' was modified",
|
||||
schema.model, name
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn unapplied_migrations(connection: &Connection, schema: &Schema) -> Result<Vec<Migration>, Error> {
|
||||
let applied = load_migrations(connection, &schema.model)?;
|
||||
Ok(schema
|
||||
.migrations
|
||||
.iter()
|
||||
.filter(|migration| !applied.contains_key(&migration.name))
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn table_exists(connection: &Connection, table: &str) -> Result<bool, Error> {
|
||||
connection
|
||||
.query_row(
|
||||
"SELECT 1 FROM sqlite_schema WHERE type = 'table' AND name = ?1",
|
||||
[table],
|
||||
|_| Ok(()),
|
||||
)
|
||||
.optional()
|
||||
.map(|value| value.is_some())
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
fn load_schema(connection: &Connection, model: &str) -> Result<Option<Schema>, Error> {
|
||||
let definition = connection
|
||||
.query_row(
|
||||
"SELECT definition FROM _dever_history WHERE model = ?1 AND kind = 'schema' AND name = 'current'",
|
||||
[model],
|
||||
|row| row.get::<_, String>(0),
|
||||
)
|
||||
.optional()
|
||||
.map_err(database_error)?;
|
||||
definition
|
||||
.map(|definition| {
|
||||
serde_json::from_str(&definition)
|
||||
.map_err(|error| Error::migration(format!("invalid Dever schema history: {error}")))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn load_catalog(connection: &Connection, model: &str) -> Result<Option<String>, Error> {
|
||||
connection
|
||||
.query_row(
|
||||
"SELECT definition FROM _dever_history WHERE model = ?1 AND kind = 'catalog' AND name = 'current'",
|
||||
[model],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.optional()
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
fn catalog(connection: &Connection, table: &str) -> Result<String, Error> {
|
||||
let table_sql = connection
|
||||
.query_row(
|
||||
"SELECT sql FROM sqlite_schema WHERE type = 'table' AND name = ?1",
|
||||
[table],
|
||||
|row| row.get::<_, String>(0),
|
||||
)
|
||||
.map_err(database_error)?;
|
||||
let mut statement = connection
|
||||
.prepare(
|
||||
"SELECT name, sql FROM sqlite_schema WHERE type = 'index' AND tbl_name = ?1 AND substr(name, 1, 11) = '_dever_idx_' ORDER BY name",
|
||||
)
|
||||
.map_err(database_error)?;
|
||||
let indexes = statement
|
||||
.query_map([table], |row| {
|
||||
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
|
||||
})
|
||||
.map_err(database_error)?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(database_error)?;
|
||||
serde_json::to_string(&(table_sql, indexes))
|
||||
.map_err(|error| Error::migration(format!("cannot encode SQLite catalog: {error}")))
|
||||
}
|
||||
|
||||
fn load_migrations(
|
||||
connection: &Connection,
|
||||
model: &str,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
let mut statement = connection
|
||||
.prepare(
|
||||
"SELECT name, revision FROM _dever_history WHERE model = ?1 AND kind = 'migration'",
|
||||
)
|
||||
.map_err(database_error)?;
|
||||
statement
|
||||
.query_map([model], |row| Ok((row.get(0)?, row.get(1)?)))
|
||||
.map_err(database_error)?
|
||||
.collect::<Result<BTreeMap<_, _>, _>>()
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
fn store_history(connection: &Connection, schema: &Schema) -> Result<(), Error> {
|
||||
let definition = serde_json::to_string(schema).map_err(|error| {
|
||||
Error::migration(format!("cannot encode Dever schema history: {error}"))
|
||||
})?;
|
||||
store_history_row(
|
||||
connection,
|
||||
&schema.model,
|
||||
"schema",
|
||||
"current",
|
||||
&schema.revision,
|
||||
&definition,
|
||||
)?;
|
||||
for migration in &schema.migrations {
|
||||
store_history_row(
|
||||
connection,
|
||||
&schema.model,
|
||||
"migration",
|
||||
&migration.name,
|
||||
&migration.revision,
|
||||
"",
|
||||
)?;
|
||||
}
|
||||
store_history_row(
|
||||
connection,
|
||||
&schema.model,
|
||||
"seed",
|
||||
"current",
|
||||
&schema.seed_revision,
|
||||
"",
|
||||
)?;
|
||||
store_history_row(
|
||||
connection,
|
||||
&schema.model,
|
||||
"catalog",
|
||||
"current",
|
||||
&schema.revision,
|
||||
&catalog(connection, &schema.table)?,
|
||||
)
|
||||
}
|
||||
|
||||
fn store_history_row(
|
||||
connection: &Connection,
|
||||
model: &str,
|
||||
kind: &str,
|
||||
name: &str,
|
||||
revision: &str,
|
||||
definition: &str,
|
||||
) -> Result<(), Error> {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO _dever_history (model, kind, name, revision, definition) VALUES (?1, ?2, ?3, ?4, ?5) \
|
||||
ON CONFLICT (model, kind, name) DO UPDATE SET revision = excluded.revision, definition = excluded.definition, applied_at = unixepoch() WHERE _dever_history.kind <> 'migration'",
|
||||
params![model, kind, name, revision, definition],
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(database_error)
|
||||
}
|
||||
|
||||
fn fields_match(previous: &[Field], desired: &[Field]) -> bool {
|
||||
previous == desired
|
||||
}
|
||||
|
||||
fn compatible_types(previous: &str, desired: &str) -> bool {
|
||||
previous == desired
|
||||
|| (previous.starts_with("Text") && desired.starts_with("Text"))
|
||||
|| (previous.starts_with("Bytes") && desired.starts_with("Bytes"))
|
||||
|| (previous == "Int" && desired == "Float")
|
||||
|| (sqlite_integer_reference(previous) && sqlite_integer_reference(desired))
|
||||
}
|
||||
|
||||
fn sqlite_integer_reference(ty: &str) -> bool {
|
||||
ty == "Int" || ty.starts_with("id:")
|
||||
}
|
||||
|
||||
pub fn temporary_table(schema: &Schema) -> String {
|
||||
format!("_dever_{}_{}", schema.table, &schema.revision[..8])
|
||||
}
|
||||
|
||||
fn quote(identifier: &str) -> String {
|
||||
format!("\"{}\"", identifier.replace('"', "\"\""))
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
use crate::bytes::Bytes;
|
||||
use crate::http::{Header, HttpReply};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Event {
|
||||
pub event: String,
|
||||
pub data: String,
|
||||
pub id: Option<String>,
|
||||
pub retry_ms: Option<i64>,
|
||||
}
|
||||
|
||||
pub async fn start(reply: &HttpReply, mut headers: Vec<Header>) -> Result<(), String> {
|
||||
if reply.chunk_limit() < 3 {
|
||||
return Err("SSE chunk_bytes must fit a three-byte heartbeat".into());
|
||||
}
|
||||
if headers
|
||||
.iter()
|
||||
.any(|header| header.name.eq_ignore_ascii_case("content-type"))
|
||||
{
|
||||
return Err("SSE owns the Content-Type header".into());
|
||||
}
|
||||
headers.push(Header {
|
||||
name: "content-type".into(),
|
||||
value: Bytes::from_string("text/event-stream; charset=utf-8".into()),
|
||||
});
|
||||
if !headers
|
||||
.iter()
|
||||
.any(|header| header.name.eq_ignore_ascii_case("cache-control"))
|
||||
{
|
||||
headers.push(Header {
|
||||
name: "cache-control".into(),
|
||||
value: Bytes::from_string("no-cache".into()),
|
||||
});
|
||||
}
|
||||
crate::http::reply::start_stream(reply, 200, headers, true).await
|
||||
}
|
||||
|
||||
pub async fn send(reply: &HttpReply, event: Event) -> Result<(), String> {
|
||||
let bytes = encode(event, reply.chunk_limit())?;
|
||||
crate::http::reply::write_chunk(reply, bytes, true).await
|
||||
}
|
||||
|
||||
// 按 EventSource 的行语义编码；先计算展开长度，避免超过单块上限后才分配。
|
||||
pub fn encode(event: Event, limit: usize) -> Result<Bytes, String> {
|
||||
if event.event.contains(['\r', '\n']) {
|
||||
return Err("SSE event cannot contain a line break".into());
|
||||
}
|
||||
if event
|
||||
.id
|
||||
.as_ref()
|
||||
.is_some_and(|id| id.contains(['\r', '\n', '\0']))
|
||||
{
|
||||
return Err("SSE id cannot contain a line break or NUL".into());
|
||||
}
|
||||
if event.retry_ms.is_some_and(|retry| retry < 0) {
|
||||
return Err("SSE retry_ms must be nonnegative".into());
|
||||
}
|
||||
let retry = event.retry_ms.map(|retry| retry.to_string());
|
||||
// split_terminator 丢弃末尾空行，会改变 data 尾部换行；这里保留每一行。
|
||||
let newlines = event.data.bytes().filter(|byte| *byte == b'\n').count();
|
||||
let lines = || {
|
||||
event
|
||||
.data
|
||||
.split('\n')
|
||||
.enumerate()
|
||||
.flat_map(|(index, line)| {
|
||||
let line = if index < newlines {
|
||||
line.strip_suffix('\r').unwrap_or(line)
|
||||
} else {
|
||||
line
|
||||
};
|
||||
line.split('\r')
|
||||
})
|
||||
};
|
||||
let length = 1
|
||||
+ lines().map(|line| 7 + line.len()).sum::<usize>()
|
||||
+ if event.event.is_empty() {
|
||||
0
|
||||
} else {
|
||||
8 + event.event.len()
|
||||
}
|
||||
+ event.id.as_ref().map_or(0, |id| 5 + id.len())
|
||||
+ retry.as_ref().map_or(0, |retry| 8 + retry.len());
|
||||
if length > limit {
|
||||
return Err("SSE event exceeds chunk_bytes".into());
|
||||
}
|
||||
let mut encoded = String::with_capacity(length);
|
||||
let mut field = |name: &str, value: &str| {
|
||||
encoded.push_str(name);
|
||||
encoded.push_str(": ");
|
||||
encoded.push_str(value);
|
||||
encoded.push('\n');
|
||||
};
|
||||
if !event.event.is_empty() {
|
||||
field("event", &event.event);
|
||||
}
|
||||
if let Some(id) = event.id {
|
||||
field("id", &id);
|
||||
}
|
||||
if let Some(retry) = retry {
|
||||
field("retry", &retry);
|
||||
}
|
||||
for line in lines() {
|
||||
field("data", line);
|
||||
}
|
||||
encoded.push('\n');
|
||||
Ok(Bytes::from_string(encoded))
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,270 @@
|
||||
use std::future::Future;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use tokio::sync::{Notify, oneshot};
|
||||
|
||||
use super::{
|
||||
CURRENT_SCOPE, CompleteOnDrop, Completion, Registration, Scope, TaskOutcome, start_task,
|
||||
supervise_with,
|
||||
};
|
||||
|
||||
/// Completion of an externally supervised resource still belongs to the
|
||||
/// requesting task: cancellation must drain it without aborting its cleanup.
|
||||
#[cfg(feature = "wire")]
|
||||
pub(crate) fn scoped_completion() -> Result<ScopedCompletion, String> {
|
||||
let scope = super::current_scope()?;
|
||||
let completion = Completion::new();
|
||||
let id = scope.register(None, completion.clone());
|
||||
Ok(ScopedCompletion {
|
||||
registration: Some(Registration { scope, id }),
|
||||
completion,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "wire")]
|
||||
pub(crate) struct ScopedCompletion {
|
||||
registration: Option<Registration>,
|
||||
completion: Completion,
|
||||
}
|
||||
|
||||
#[cfg(feature = "wire")]
|
||||
impl Drop for ScopedCompletion {
|
||||
fn drop(&mut self) {
|
||||
if let Some(registration) = self.registration.take() {
|
||||
registration.finish();
|
||||
}
|
||||
self.completion.complete();
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn scoped_owner<G>(guard: G) -> Result<(ScopedOwner, ScopedSpawner, ScopeFault), String>
|
||||
where
|
||||
G: Send + 'static,
|
||||
{
|
||||
let parent = super::current_scope()?;
|
||||
Ok(scoped_owner_in(parent, guard))
|
||||
}
|
||||
|
||||
fn scoped_owner_in<G>(parent: Scope, guard: G) -> (ScopedOwner, ScopedSpawner, ScopeFault)
|
||||
where
|
||||
G: Send + 'static,
|
||||
{
|
||||
let scope = parent.child();
|
||||
let completion = Completion::new();
|
||||
let registration = Registration {
|
||||
id: parent.register(None, completion.clone()),
|
||||
scope: parent,
|
||||
};
|
||||
let admission = Arc::new(Mutex::new(true));
|
||||
let fault = ScopeFault::new();
|
||||
let task_slots = Arc::clone(&scope.limits.task_slots);
|
||||
(
|
||||
ScopedOwner {
|
||||
scope: scope.clone(),
|
||||
admission: Arc::clone(&admission),
|
||||
fault: fault.clone(),
|
||||
registration: Some(registration),
|
||||
completion: Some(completion),
|
||||
guard: Some(Box::new(guard)),
|
||||
},
|
||||
ScopedSpawner {
|
||||
scope,
|
||||
admission,
|
||||
fault: fault.clone(),
|
||||
task_slots,
|
||||
},
|
||||
fault,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) struct ScopedOwner {
|
||||
scope: Scope,
|
||||
admission: Arc<Mutex<bool>>,
|
||||
fault: ScopeFault,
|
||||
registration: Option<Registration>,
|
||||
completion: Option<Completion>,
|
||||
guard: Option<Box<dyn Send>>,
|
||||
}
|
||||
|
||||
impl ScopedOwner {
|
||||
pub(crate) async fn enter<F>(&mut self, future: F) -> F::Output
|
||||
where
|
||||
F: Future,
|
||||
{
|
||||
CURRENT_SCOPE.scope(self.scope.clone(), future).await
|
||||
}
|
||||
|
||||
pub(crate) async fn finish(mut self) -> Result<(), String> {
|
||||
self.start_cleanup(false, true)
|
||||
.expect("scoped owner cleaned once")
|
||||
.await
|
||||
.unwrap_or_else(|_| Err("scoped task cleanup stopped".into()))
|
||||
}
|
||||
|
||||
pub(crate) async fn stop(mut self) -> Result<(), String> {
|
||||
self.start_cleanup(true, true)
|
||||
.expect("scoped owner cleaned once")
|
||||
.await
|
||||
.unwrap_or_else(|_| Err("scoped task cleanup stopped".into()))
|
||||
}
|
||||
|
||||
pub(crate) fn finish_in_background(mut self) {
|
||||
let _ = self.start_cleanup(false, false);
|
||||
}
|
||||
|
||||
pub(crate) fn stop_in_background(mut self) {
|
||||
let _ = self.start_cleanup(true, false);
|
||||
}
|
||||
|
||||
fn close_admission(&self) {
|
||||
*self
|
||||
.admission
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner()) = false;
|
||||
}
|
||||
|
||||
fn start_cleanup(
|
||||
&mut self,
|
||||
stopping: bool,
|
||||
cancel_if_unobserved: bool,
|
||||
) -> Option<oneshot::Receiver<Result<(), String>>> {
|
||||
self.close_admission();
|
||||
let completion = self.completion.take()?;
|
||||
let (mut sender, receiver) = oneshot::channel();
|
||||
let scope = self.scope.clone();
|
||||
let registration = self.registration.take();
|
||||
let guard = self.guard.take();
|
||||
let fault = self.fault.clone();
|
||||
tokio::spawn(async move {
|
||||
let complete = CompleteOnDrop(completion);
|
||||
if stopping {
|
||||
scope.cancel_and_wait().await;
|
||||
} else if cancel_if_unobserved {
|
||||
tokio::select! {
|
||||
_ = scope.wait() => {}
|
||||
_ = sender.closed() => scope.cancel_and_wait().await,
|
||||
}
|
||||
} else {
|
||||
scope.wait().await;
|
||||
}
|
||||
drop(guard);
|
||||
if let Some(registration) = registration {
|
||||
registration.finish();
|
||||
}
|
||||
let result = fault.result();
|
||||
drop(complete);
|
||||
let _ = sender.send(result);
|
||||
});
|
||||
Some(receiver)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ScopedOwner {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.start_cleanup(true, true);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ScopedSpawner {
|
||||
scope: Scope,
|
||||
admission: Arc<Mutex<bool>>,
|
||||
fault: ScopeFault,
|
||||
task_slots: Arc<tokio::sync::Semaphore>,
|
||||
}
|
||||
|
||||
impl ScopedSpawner {
|
||||
pub(crate) fn spawn<F>(&self, future: F)
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
let accepting = self
|
||||
.admission
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner());
|
||||
if !*accepting {
|
||||
return;
|
||||
}
|
||||
let permit = match Arc::clone(&self.task_slots)
|
||||
.try_acquire_owned()
|
||||
.map_err(|_| "async task capacity is exhausted".to_owned())
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(error) => {
|
||||
self.fault.report(error);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let started = start_task(
|
||||
&self.scope,
|
||||
async move {
|
||||
future.await;
|
||||
Ok(())
|
||||
},
|
||||
permit,
|
||||
);
|
||||
let scope = self.scope.clone();
|
||||
let fault = self.fault.clone();
|
||||
tokio::spawn(supervise_with(started, move |id, outcome| {
|
||||
match outcome {
|
||||
TaskOutcome::Finished(Err(error)) | TaskOutcome::Panicked(error) => {
|
||||
fault.report(error.clone());
|
||||
}
|
||||
TaskOutcome::Finished(Ok(())) | TaskOutcome::Stopped => {}
|
||||
}
|
||||
scope.unregister(id);
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ScopeFault(Arc<FaultState>);
|
||||
|
||||
struct FaultState {
|
||||
first: Mutex<Option<String>>,
|
||||
changed: Notify,
|
||||
}
|
||||
|
||||
impl ScopeFault {
|
||||
fn new() -> Self {
|
||||
Self(Arc::new(FaultState {
|
||||
first: Mutex::new(None),
|
||||
changed: Notify::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn report(&self, error: String) {
|
||||
let mut first = self
|
||||
.0
|
||||
.first
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner());
|
||||
if first.is_none() {
|
||||
*first = Some(error);
|
||||
drop(first);
|
||||
self.0.changed.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn wait(&self) -> String {
|
||||
loop {
|
||||
let changed = self.0.changed.notified();
|
||||
if let Some(error) = self.current() {
|
||||
return error;
|
||||
}
|
||||
changed.await;
|
||||
}
|
||||
}
|
||||
|
||||
fn result(&self) -> Result<(), String> {
|
||||
self.current().map_or(Ok(()), Err)
|
||||
}
|
||||
|
||||
fn current(&self) -> Option<String> {
|
||||
self.0
|
||||
.first
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.clone()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::future::Future;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::database::{Database, OwnedSql, Sql};
|
||||
use crate::orm::{self, Error, ErrorKind, Value};
|
||||
|
||||
tokio::task_local! {
|
||||
static CURRENT_TENANT: i64;
|
||||
}
|
||||
|
||||
static MANAGER: OnceLock<Arc<Manager>> = OnceLock::new();
|
||||
|
||||
pub(crate) struct Manager {
|
||||
control: Database,
|
||||
schema_fingerprint: String,
|
||||
max_pools: usize,
|
||||
idle_timeout: Duration,
|
||||
state: Mutex<State>,
|
||||
}
|
||||
|
||||
struct State {
|
||||
accepting: bool,
|
||||
pools: BTreeMap<(String, i64), PoolEntry>,
|
||||
}
|
||||
|
||||
struct PoolEntry {
|
||||
database: Database,
|
||||
last_used: Instant,
|
||||
ready: bool,
|
||||
}
|
||||
|
||||
impl State {
|
||||
async fn retire_pool(&mut self, key: &(String, i64)) -> Result<(), Error> {
|
||||
let Some(entry) = self.pools.get_mut(key) else {
|
||||
return Ok(());
|
||||
};
|
||||
entry.ready = false;
|
||||
let result = entry.database.close_pool().await;
|
||||
self.pools.remove(key);
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn scope<T>(tenant_id: i64, future: impl Future<Output = T>) -> T {
|
||||
CURRENT_TENANT.scope(tenant_id, future).await
|
||||
}
|
||||
|
||||
pub fn current_id() -> Result<i64, Error> {
|
||||
CURRENT_TENANT
|
||||
.try_with(|tenant_id| *tenant_id)
|
||||
.map_err(|_| Error::new(ErrorKind::InvalidData, "tenant context is required"))
|
||||
.and_then(|tenant_id| {
|
||||
if tenant_id > 0 {
|
||||
Ok(tenant_id)
|
||||
} else {
|
||||
Err(Error::new(
|
||||
ErrorKind::InvalidData,
|
||||
"tenant id must be positive",
|
||||
))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub fn schema_fingerprint() -> Result<String, Error> {
|
||||
Ok(manager()?.schema_fingerprint.clone())
|
||||
}
|
||||
|
||||
pub async fn initialize(schema_fingerprint: &str) -> Result<(), Error> {
|
||||
if schema_fingerprint.is_empty() || schema_fingerprint.len() > 65_536 {
|
||||
return Err(Error::migration("invalid tenant schema fingerprint"));
|
||||
}
|
||||
let configuration = crate::config::current_settings();
|
||||
let settings = configuration
|
||||
.tenant()
|
||||
.ok_or_else(|| Error::new(ErrorKind::Pool, "tenant settings are not configured"))?;
|
||||
let control = crate::database::database_for(Some(settings.database()), "")?;
|
||||
control.execute(Sql {
|
||||
sqlite: "CREATE TABLE IF NOT EXISTS _dever_tenant_storage (tenant_id INTEGER PRIMARY KEY, state TEXT NOT NULL, schema_fingerprint TEXT NOT NULL, updated_at INTEGER NOT NULL)",
|
||||
postgres: "CREATE TABLE IF NOT EXISTS _dever_tenant_storage (tenant_id BIGINT PRIMARY KEY, state TEXT NOT NULL, schema_fingerprint TEXT NOT NULL, updated_at BIGINT NOT NULL)",
|
||||
}, Vec::new()).await?;
|
||||
control.execute(Sql {
|
||||
sqlite: "CREATE TABLE IF NOT EXISTS _dever_tenant_component (tenant_id INTEGER NOT NULL, component TEXT NOT NULL, disabled_at INTEGER NOT NULL, PRIMARY KEY (tenant_id, component), FOREIGN KEY (tenant_id) REFERENCES _dever_tenant_storage (tenant_id) ON DELETE CASCADE)",
|
||||
postgres: "CREATE TABLE IF NOT EXISTS _dever_tenant_component (tenant_id BIGINT NOT NULL, component TEXT NOT NULL, disabled_at BIGINT NOT NULL, PRIMARY KEY (tenant_id, component), FOREIGN KEY (tenant_id) REFERENCES _dever_tenant_storage (tenant_id) ON DELETE CASCADE)",
|
||||
}, Vec::new()).await?;
|
||||
let manager = Arc::new(Manager {
|
||||
control,
|
||||
schema_fingerprint: schema_fingerprint.to_owned(),
|
||||
max_pools: settings.max_pools(),
|
||||
idle_timeout: Duration::from_millis(settings.idle_timeout_ms()),
|
||||
state: Mutex::new(State {
|
||||
accepting: true,
|
||||
pools: BTreeMap::new(),
|
||||
}),
|
||||
});
|
||||
#[cfg(feature = "api")]
|
||||
if let Some(session) = crate::application::current() {
|
||||
return session.set_tenant(manager);
|
||||
}
|
||||
MANAGER
|
||||
.set(manager)
|
||||
.map_err(|_| Error::migration("tenant storage was initialized more than once"))
|
||||
}
|
||||
|
||||
pub async fn database(connection: &str) -> Result<Database, Error> {
|
||||
let tenant_id = current_id()?;
|
||||
let manager = manager()?;
|
||||
manager.require_ready(tenant_id).await?;
|
||||
let key = (connection.to_owned(), tenant_id);
|
||||
let mut state = manager.state.lock().await;
|
||||
if !state.accepting {
|
||||
return Err(Error::new(
|
||||
ErrorKind::Pool,
|
||||
"tenant storage is shutting down",
|
||||
));
|
||||
}
|
||||
if let Some(entry) = state.pools.get_mut(&key)
|
||||
&& entry.ready
|
||||
{
|
||||
entry.last_used = Instant::now();
|
||||
return Ok(entry.database.clone());
|
||||
}
|
||||
// An interrupted preparation remains owned, but must never become a usable pool.
|
||||
state.retire_pool(&key).await?;
|
||||
let now = Instant::now();
|
||||
let expired = state
|
||||
.pools
|
||||
.iter()
|
||||
.filter(|(_, entry)| {
|
||||
entry.database.exclusively_held()
|
||||
&& now.duration_since(entry.last_used) >= manager.idle_timeout
|
||||
})
|
||||
.map(|(key, _)| key.clone())
|
||||
.collect::<Vec<_>>();
|
||||
for expired_key in expired {
|
||||
state.retire_pool(&expired_key).await?;
|
||||
}
|
||||
let mut reusable = state
|
||||
.pools
|
||||
.iter()
|
||||
.filter(|(_, entry)| entry.database.exclusively_held())
|
||||
.map(|(key, entry)| (key.clone(), entry.last_used))
|
||||
.collect::<Vec<_>>();
|
||||
reusable.sort_by_key(|(_, last_used)| *last_used);
|
||||
for (reusable_key, _) in reusable {
|
||||
if state.pools.len() < manager.max_pools {
|
||||
break;
|
||||
}
|
||||
state.retire_pool(&reusable_key).await?;
|
||||
}
|
||||
if state.pools.len() >= manager.max_pools {
|
||||
return Err(Error::new(
|
||||
ErrorKind::PoolExhausted,
|
||||
"tenant database pool limit reached",
|
||||
));
|
||||
}
|
||||
let database = crate::database::tenant_database(connection, tenant_id, false).await?;
|
||||
state.pools.insert(
|
||||
key.clone(),
|
||||
PoolEntry {
|
||||
database: database.clone(),
|
||||
last_used: now,
|
||||
ready: false,
|
||||
},
|
||||
);
|
||||
let preparation = async {
|
||||
database.prepare_pool().await?;
|
||||
manager.require_local_marker(&database, tenant_id).await
|
||||
}
|
||||
.await;
|
||||
if let Err(error) = preparation {
|
||||
let cleanup = state.retire_pool(&key).await;
|
||||
return Err(match cleanup {
|
||||
Ok(()) => error,
|
||||
Err(cleanup) => error.cause(cleanup),
|
||||
});
|
||||
}
|
||||
state
|
||||
.pools
|
||||
.get_mut(&key)
|
||||
.expect("registered tenant pool")
|
||||
.ready = true;
|
||||
Ok(database)
|
||||
}
|
||||
|
||||
pub async fn ready_ids_after(after: i64, limit: usize) -> Result<Vec<i64>, Error> {
|
||||
if after < 0 || !(1..=256).contains(&limit) {
|
||||
return Err(Error::invalid_data("invalid tenant worker page"));
|
||||
}
|
||||
let manager = manager()?;
|
||||
let rows = manager.control.query(Sql {
|
||||
sqlite: "SELECT tenant_id FROM _dever_tenant_storage WHERE state='ready' AND schema_fingerprint=?1 AND tenant_id>?2 ORDER BY tenant_id LIMIT ?3",
|
||||
postgres: "SELECT tenant_id FROM _dever_tenant_storage WHERE state='ready' AND schema_fingerprint=$1 AND tenant_id>$2 ORDER BY tenant_id LIMIT $3",
|
||||
}, vec![
|
||||
Value::Text(manager.schema_fingerprint.clone()),
|
||||
Value::Int(after),
|
||||
Value::Int(limit as i64),
|
||||
]).await?;
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
let tenant_id = orm::int(row.get(0)?.clone())?;
|
||||
if tenant_id <= 0 {
|
||||
return Err(Error::migration(
|
||||
"tenant registry contains an invalid tenant id",
|
||||
));
|
||||
}
|
||||
Ok(tenant_id)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn disable_component(
|
||||
tenant_id: i64,
|
||||
component: &str,
|
||||
tenant_components: &[&str],
|
||||
) -> Result<(), Error> {
|
||||
require_tenant_component(component, tenant_components)?;
|
||||
let manager = manager()?;
|
||||
manager.require_ready(tenant_id).await?;
|
||||
let disabled_at = crate::time::now().map_err(Error::migration)?;
|
||||
manager.control.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_tenant_component (tenant_id,component,disabled_at) VALUES (?1,?2,?3) ON CONFLICT(tenant_id,component) DO UPDATE SET disabled_at=excluded.disabled_at",
|
||||
postgres: "INSERT INTO _dever_tenant_component (tenant_id,component,disabled_at) VALUES ($1,$2,$3) ON CONFLICT(tenant_id,component) DO UPDATE SET disabled_at=excluded.disabled_at",
|
||||
}, vec![
|
||||
Value::Int(tenant_id),
|
||||
Value::Text(component.to_owned()),
|
||||
Value::Int(disabled_at),
|
||||
]).await.map(|_| ())
|
||||
}
|
||||
|
||||
pub async fn enable_component(
|
||||
tenant_id: i64,
|
||||
component: &str,
|
||||
tenant_components: &[&str],
|
||||
) -> Result<(), Error> {
|
||||
require_tenant_component(component, tenant_components)?;
|
||||
let manager = manager()?;
|
||||
manager.require_ready(tenant_id).await?;
|
||||
manager
|
||||
.control
|
||||
.execute(
|
||||
Sql {
|
||||
sqlite: "DELETE FROM _dever_tenant_component WHERE tenant_id=?1 AND component=?2",
|
||||
postgres: "DELETE FROM _dever_tenant_component WHERE tenant_id=$1 AND component=$2",
|
||||
},
|
||||
vec![Value::Int(tenant_id), Value::Text(component.to_owned())],
|
||||
)
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
pub async fn require_components(
|
||||
tenant_id: i64,
|
||||
components: &[&str],
|
||||
tenant_components: &[&str],
|
||||
) -> Result<(), Error> {
|
||||
let components = components.iter().copied().collect::<BTreeSet<_>>();
|
||||
for component in &components {
|
||||
require_tenant_component(component, tenant_components)?;
|
||||
}
|
||||
if components.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let manager = manager()?;
|
||||
manager.require_ready(tenant_id).await?;
|
||||
let sqlite_placeholders = (2..=components.len() + 1)
|
||||
.map(|index| format!("?{index}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let postgres_placeholders = (2..=components.len() + 1)
|
||||
.map(|index| format!("${index}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let mut parameters = Vec::with_capacity(components.len() + 1);
|
||||
parameters.push(Value::Int(tenant_id));
|
||||
parameters.extend(
|
||||
components
|
||||
.iter()
|
||||
.map(|component| Value::Text((*component).to_owned())),
|
||||
);
|
||||
let rows = manager.control.query_owned(OwnedSql {
|
||||
sqlite: format!("SELECT component FROM _dever_tenant_component WHERE tenant_id=?1 AND component IN ({sqlite_placeholders}) ORDER BY component LIMIT 1"),
|
||||
postgres: format!("SELECT component FROM _dever_tenant_component WHERE tenant_id=$1 AND component IN ({postgres_placeholders}) ORDER BY component LIMIT 1"),
|
||||
}, parameters).await?;
|
||||
if let Some(row) = orm::optional_row(rows, "tenant component lookup")? {
|
||||
let component = orm::text(row.get(0)?.clone())?;
|
||||
return Err(Error::invalid_data(format!(
|
||||
"tenant component '{component}' is disabled"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn begin_migration(
|
||||
connection: &str,
|
||||
tenant_id: i64,
|
||||
schema_fingerprint: &str,
|
||||
) -> Result<Database, Error> {
|
||||
if tenant_id <= 0 {
|
||||
return Err(Error::migration("tenant id must be positive"));
|
||||
}
|
||||
let manager = manager()?;
|
||||
if schema_fingerprint != manager.schema_fingerprint {
|
||||
return Err(Error::migration(
|
||||
"tenant migration schema fingerprint does not match this application",
|
||||
));
|
||||
}
|
||||
manager
|
||||
.write_state(tenant_id, "migrating", schema_fingerprint)
|
||||
.await?;
|
||||
let database = crate::database::tenant_database(connection, tenant_id, true).await?;
|
||||
#[cfg(feature = "api")]
|
||||
if let Some(session) = crate::application::current() {
|
||||
session.own_migration_database(database.clone());
|
||||
}
|
||||
if let Err(error) = database.prepare_pool().await {
|
||||
return Err(match database.close_pool().await {
|
||||
Ok(()) => error,
|
||||
Err(cleanup) => error.cause(cleanup),
|
||||
});
|
||||
}
|
||||
Ok(database)
|
||||
}
|
||||
|
||||
pub async fn mark_database_ready(
|
||||
database: &Database,
|
||||
tenant_id: i64,
|
||||
schema_fingerprint: &str,
|
||||
) -> Result<(), Error> {
|
||||
let manager = manager()?;
|
||||
if schema_fingerprint != manager.schema_fingerprint {
|
||||
return Err(Error::migration(
|
||||
"tenant migration schema fingerprint does not match this application",
|
||||
));
|
||||
}
|
||||
if database.key().tenant_id() != Some(tenant_id) {
|
||||
return Err(Error::migration(
|
||||
"tenant migration database key does not match tenant id",
|
||||
));
|
||||
}
|
||||
database.execute(Sql {
|
||||
sqlite: "CREATE TABLE IF NOT EXISTS _dever_tenant_ready (marker INTEGER PRIMARY KEY CHECK (marker = 1), tenant_id INTEGER NOT NULL, schema_fingerprint TEXT NOT NULL)",
|
||||
postgres: "CREATE TABLE IF NOT EXISTS _dever_tenant_ready (marker BIGINT PRIMARY KEY CHECK (marker = 1), tenant_id BIGINT NOT NULL, schema_fingerprint TEXT NOT NULL)",
|
||||
}, Vec::new()).await?;
|
||||
database.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_tenant_ready (marker,tenant_id,schema_fingerprint) VALUES (1,?1,?2) ON CONFLICT(marker) DO UPDATE SET tenant_id=excluded.tenant_id,schema_fingerprint=excluded.schema_fingerprint",
|
||||
postgres: "INSERT INTO _dever_tenant_ready (marker,tenant_id,schema_fingerprint) VALUES (1,$1,$2) ON CONFLICT(marker) DO UPDATE SET tenant_id=excluded.tenant_id,schema_fingerprint=excluded.schema_fingerprint",
|
||||
}, vec![Value::Int(tenant_id), Value::Text(schema_fingerprint.to_owned())]).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn finish_migration(tenant_id: i64, schema_fingerprint: &str) -> Result<(), Error> {
|
||||
let manager = manager()?;
|
||||
if schema_fingerprint != manager.schema_fingerprint {
|
||||
return Err(Error::migration(
|
||||
"tenant migration schema fingerprint does not match this application",
|
||||
));
|
||||
}
|
||||
manager
|
||||
.write_state(tenant_id, "ready", schema_fingerprint)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn shutdown() -> Result<(), Error> {
|
||||
let Some(manager) = MANAGER.get() else {
|
||||
return Ok(());
|
||||
};
|
||||
manager.close().await
|
||||
}
|
||||
|
||||
impl Manager {
|
||||
pub(crate) async fn close(&self) -> Result<(), Error> {
|
||||
let databases = {
|
||||
let mut state = self.state.lock().await;
|
||||
state.accepting = false;
|
||||
std::mem::take(&mut state.pools)
|
||||
.into_values()
|
||||
.map(|entry| entry.database)
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
for database in &databases {
|
||||
database.stop_accepting();
|
||||
}
|
||||
let mut first_error = None;
|
||||
for database in databases {
|
||||
if let Err(error) = database.close_pool().await
|
||||
&& first_error.is_none()
|
||||
{
|
||||
first_error = Some(error);
|
||||
}
|
||||
}
|
||||
first_error.map_or(Ok(()), Err)
|
||||
}
|
||||
}
|
||||
|
||||
fn manager() -> Result<Arc<Manager>, Error> {
|
||||
#[cfg(feature = "api")]
|
||||
if let Some(session) = crate::application::current() {
|
||||
return session
|
||||
.tenant()
|
||||
.ok_or_else(|| Error::new(ErrorKind::Pool, "tenant storage is not initialized"));
|
||||
}
|
||||
MANAGER
|
||||
.get()
|
||||
.cloned()
|
||||
.ok_or_else(|| Error::new(ErrorKind::Pool, "tenant storage is not initialized"))
|
||||
}
|
||||
|
||||
fn require_tenant_component(component: &str, tenant_components: &[&str]) -> Result<(), Error> {
|
||||
if tenant_components.contains(&component) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(Error::invalid_data(format!(
|
||||
"component '{component}' is not a tenant component"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
impl Manager {
|
||||
async fn require_ready(&self, tenant_id: i64) -> Result<(), Error> {
|
||||
let rows = self.control.query(Sql {
|
||||
sqlite: "SELECT state, schema_fingerprint FROM _dever_tenant_storage WHERE tenant_id = ?1",
|
||||
postgres: "SELECT state, schema_fingerprint FROM _dever_tenant_storage WHERE tenant_id = $1",
|
||||
}, vec![Value::Int(tenant_id)]).await?;
|
||||
let row = orm::optional_row(rows, "tenant storage lookup")?
|
||||
.ok_or_else(|| Error::migration("tenant storage is not provisioned"))?;
|
||||
let state = orm::text(row.get(0)?.clone())?;
|
||||
let fingerprint = orm::text(row.get(1)?.clone())?;
|
||||
if state != "ready" {
|
||||
return Err(Error::migration("tenant storage is not ready"));
|
||||
}
|
||||
if fingerprint != self.schema_fingerprint {
|
||||
return Err(Error::migration(
|
||||
"tenant storage schema fingerprint is stale",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_state(
|
||||
&self,
|
||||
tenant_id: i64,
|
||||
state: &str,
|
||||
fingerprint: &str,
|
||||
) -> Result<(), Error> {
|
||||
let updated_at = crate::time::now().map_err(Error::migration)?;
|
||||
self.control.execute(Sql {
|
||||
sqlite: "INSERT INTO _dever_tenant_storage (tenant_id,state,schema_fingerprint,updated_at) VALUES (?1,?2,?3,?4) ON CONFLICT(tenant_id) DO UPDATE SET state=excluded.state,schema_fingerprint=excluded.schema_fingerprint,updated_at=excluded.updated_at",
|
||||
postgres: "INSERT INTO _dever_tenant_storage (tenant_id,state,schema_fingerprint,updated_at) VALUES ($1,$2,$3,$4) ON CONFLICT(tenant_id) DO UPDATE SET state=excluded.state,schema_fingerprint=excluded.schema_fingerprint,updated_at=excluded.updated_at",
|
||||
}, vec![
|
||||
Value::Int(tenant_id),
|
||||
Value::Text(state.to_owned()),
|
||||
Value::Text(fingerprint.to_owned()),
|
||||
Value::Int(updated_at),
|
||||
]).await.map(|_| ())
|
||||
}
|
||||
|
||||
async fn require_local_marker(&self, database: &Database, tenant_id: i64) -> Result<(), Error> {
|
||||
let rows = database.query(Sql {
|
||||
sqlite: "SELECT tenant_id, schema_fingerprint FROM _dever_tenant_ready WHERE marker = 1",
|
||||
postgres: "SELECT tenant_id, schema_fingerprint FROM _dever_tenant_ready WHERE marker = 1",
|
||||
}, Vec::new()).await.map_err(|error| {
|
||||
Error::migration(format!("tenant database ready marker is missing: {error}"))
|
||||
})?;
|
||||
let row = orm::optional_row(rows, "tenant database ready marker")?
|
||||
.ok_or_else(|| Error::migration("tenant database ready marker is missing"))?;
|
||||
if orm::int(row.get(0)?.clone())? != tenant_id {
|
||||
return Err(Error::migration(
|
||||
"tenant database ready marker has the wrong tenant id",
|
||||
));
|
||||
}
|
||||
if orm::text(row.get(1)?.clone())? != self.schema_fingerprint {
|
||||
return Err(Error::migration(
|
||||
"tenant database ready marker has a stale schema fingerprint",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
use crate::collections::List;
|
||||
|
||||
pub fn codepoint(value: &str) -> Option<i64> {
|
||||
let mut chars = value.chars();
|
||||
let first = chars.next()?;
|
||||
chars
|
||||
.next()
|
||||
.is_none()
|
||||
.then_some(i64::from(u32::from(first)))
|
||||
}
|
||||
|
||||
pub fn from_codepoint(value: i64) -> Option<String> {
|
||||
char::from_u32(u32::try_from(value).ok()?).map(String::from)
|
||||
}
|
||||
|
||||
pub fn at(value: &str, index: i64) -> Option<String> {
|
||||
value
|
||||
.chars()
|
||||
.nth(usize::try_from(index).ok()?)
|
||||
.map(String::from)
|
||||
}
|
||||
|
||||
/// Text 的位置按 Unicode scalar 计数；Bytes API 单独使用字节偏移。
|
||||
pub fn slice(value: &str, start: i64, end: i64) -> Option<String> {
|
||||
let start = usize::try_from(start).ok()?;
|
||||
let end = usize::try_from(end).ok()?;
|
||||
if start > end {
|
||||
return None;
|
||||
}
|
||||
let mut boundaries = value
|
||||
.char_indices()
|
||||
.map(|(index, _)| index)
|
||||
.chain([value.len()]);
|
||||
let first = boundaries.nth(start)?;
|
||||
let last = if start == end {
|
||||
first
|
||||
} else {
|
||||
boundaries.nth(end - start - 1)?
|
||||
};
|
||||
Some(value[first..last].to_owned())
|
||||
}
|
||||
|
||||
pub fn index_of(value: &str, pattern: &str) -> Option<i64> {
|
||||
value
|
||||
.find(pattern)
|
||||
.map(|offset| value[..offset].chars().count() as i64)
|
||||
}
|
||||
|
||||
pub fn split(value: &str, separator: &str) -> List<String> {
|
||||
List::new(pieces(value, separator).collect())
|
||||
}
|
||||
|
||||
pub fn pieces<'a>(value: &'a str, separator: &'a str) -> impl Iterator<Item = String> + 'a {
|
||||
value
|
||||
.split(separator)
|
||||
.filter(move |piece| !separator.is_empty() || !piece.is_empty())
|
||||
.map(str::to_owned)
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
use chrono::{Datelike, Timelike};
|
||||
use std::sync::OnceLock;
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
const MILLIS_PER_DAY: i64 = 86_400_000;
|
||||
|
||||
/// A worker and its application test share this clock; advancing it never sleeps.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct Clock(Option<std::sync::Arc<std::sync::atomic::AtomicI64>>);
|
||||
|
||||
impl Clock {
|
||||
pub fn for_test() -> Self {
|
||||
Self(Some(std::sync::Arc::new(
|
||||
std::sync::atomic::AtomicI64::new(0),
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn now(&self) -> Result<i64, crate::orm::Error> {
|
||||
match &self.0 {
|
||||
Some(clock) => Ok(clock.load(std::sync::atomic::Ordering::SeqCst)),
|
||||
None => now().map_err(crate::orm::Error::invalid_data),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn advance(&self, millis: i64) -> Result<(), String> {
|
||||
if millis < 0 {
|
||||
return Err("test clock cannot move backwards".into());
|
||||
}
|
||||
let clock = self
|
||||
.0
|
||||
.as_ref()
|
||||
.ok_or("production clock cannot be advanced")?;
|
||||
clock
|
||||
.fetch_update(
|
||||
std::sync::atomic::Ordering::SeqCst,
|
||||
std::sync::atomic::Ordering::SeqCst,
|
||||
|current| current.checked_add(millis),
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(|_| "test clock overflow".into())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn now() -> Result<i64, String> {
|
||||
let value = unix_millis()?;
|
||||
datetime(value)?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn datetime(value: i64) -> Result<chrono::DateTime<chrono::Utc>, String> {
|
||||
chrono::DateTime::from_timestamp_millis(value)
|
||||
.filter(|value| (0..=9999).contains(&value.year()))
|
||||
.ok_or_else(|| "DateTime is outside the four-digit UTC calendar range".into())
|
||||
}
|
||||
|
||||
pub fn parse_datetime(text: &str) -> Result<i64, String> {
|
||||
if let Some(fraction) = text.get(19..).and_then(|tail| tail.strip_prefix('.'))
|
||||
&& fraction.bytes().take_while(u8::is_ascii_digit).count() > 3
|
||||
{
|
||||
return Err("DateTime requires at most three fractional digits".into());
|
||||
}
|
||||
let parsed = chrono::DateTime::parse_from_rfc3339(text)
|
||||
.map_err(|_| "invalid RFC3339 DateTime".to_owned())?;
|
||||
if parsed.nanosecond() >= 1_000_000_000 || parsed.nanosecond() % 1_000_000 != 0 {
|
||||
return Err("DateTime requires millisecond precision and no leap seconds".into());
|
||||
}
|
||||
let value = parsed.timestamp_millis();
|
||||
datetime(value)?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub fn format_datetime(value: i64) -> Result<String, String> {
|
||||
let value = datetime(value)?;
|
||||
let precision = if value.timestamp_subsec_millis() == 0 {
|
||||
chrono::SecondsFormat::Secs
|
||||
} else {
|
||||
chrono::SecondsFormat::Millis
|
||||
};
|
||||
Ok(value.to_rfc3339_opts(precision, true))
|
||||
}
|
||||
|
||||
pub fn parse_date(text: &str) -> Result<i64, String> {
|
||||
if text.len() != 10
|
||||
|| text.as_bytes()[4] != b'-'
|
||||
|| text.as_bytes()[7] != b'-'
|
||||
|| !text
|
||||
.bytes()
|
||||
.enumerate()
|
||||
.all(|(index, byte)| matches!(index, 4 | 7) || byte.is_ascii_digit())
|
||||
{
|
||||
return Err("Date must use YYYY-MM-DD".into());
|
||||
}
|
||||
let date = chrono::NaiveDate::parse_from_str(text, "%Y-%m-%d")
|
||||
.map_err(|_| "invalid calendar Date".to_owned())?;
|
||||
Ok(date
|
||||
.and_hms_opt(0, 0, 0)
|
||||
.expect("midnight")
|
||||
.and_utc()
|
||||
.timestamp_millis())
|
||||
}
|
||||
|
||||
pub fn format_date(value: i64) -> Result<String, String> {
|
||||
if value.rem_euclid(MILLIS_PER_DAY) != 0 {
|
||||
return Err("Date must be UTC midnight milliseconds".into());
|
||||
}
|
||||
Ok(datetime(value)?.format("%Y-%m-%d").to_string())
|
||||
}
|
||||
|
||||
pub fn parse_time(text: &str) -> Result<i64, String> {
|
||||
let bytes = text.as_bytes();
|
||||
if !matches!(bytes.len(), 8 | 12)
|
||||
|| bytes[2] != b':'
|
||||
|| bytes[5] != b':'
|
||||
|| (bytes.len() == 12 && bytes[8] != b'.')
|
||||
|| !bytes
|
||||
.iter()
|
||||
.enumerate()
|
||||
.all(|(index, byte)| matches!(index, 2 | 5 | 8) || byte.is_ascii_digit())
|
||||
{
|
||||
return Err("Time must use HH:MM:SS[.sss]".into());
|
||||
}
|
||||
let time = chrono::NaiveTime::parse_from_str(text, "%H:%M:%S%.f")
|
||||
.map_err(|_| "invalid clock Time".to_owned())?;
|
||||
if time.nanosecond() >= 1_000_000_000 {
|
||||
return Err("Time does not accept leap seconds".into());
|
||||
}
|
||||
Ok(i64::from(time.num_seconds_from_midnight()) * 1000
|
||||
+ i64::from(time.nanosecond() / 1_000_000))
|
||||
}
|
||||
|
||||
pub fn format_time(value: i64) -> Result<String, String> {
|
||||
if !(0..MILLIS_PER_DAY).contains(&value) {
|
||||
return Err("Time must be milliseconds within one day".into());
|
||||
}
|
||||
let time = chrono::NaiveTime::from_num_seconds_from_midnight_opt(
|
||||
(value / 1000) as u32,
|
||||
(value % 1000) as u32 * 1_000_000,
|
||||
)
|
||||
.expect("validated time of day");
|
||||
Ok(time
|
||||
.format(if value % 1000 == 0 {
|
||||
"%H:%M:%S"
|
||||
} else {
|
||||
"%H:%M:%S%.3f"
|
||||
})
|
||||
.to_string())
|
||||
}
|
||||
|
||||
pub fn duration(millis: i64) -> i64 {
|
||||
millis
|
||||
}
|
||||
|
||||
pub fn add(value: i64, duration: i64) -> Result<i64, String> {
|
||||
value
|
||||
.checked_add(duration)
|
||||
.ok_or_else(|| "DateTime addition overflow".into())
|
||||
}
|
||||
|
||||
pub fn subtract(value: i64, duration: i64) -> Result<i64, String> {
|
||||
value
|
||||
.checked_sub(duration)
|
||||
.ok_or_else(|| "DateTime subtraction overflow".into())
|
||||
}
|
||||
|
||||
pub fn difference(later: i64, earlier: i64) -> Result<i64, String> {
|
||||
later
|
||||
.checked_sub(earlier)
|
||||
.ok_or_else(|| "Duration difference overflow".into())
|
||||
}
|
||||
|
||||
pub fn unix_millis() -> Result<i64, String> {
|
||||
let elapsed = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_err(|error| error.to_string())?;
|
||||
i64::try_from(elapsed.as_millis()).map_err(|_| "clock value is out of Int range".into())
|
||||
}
|
||||
|
||||
pub fn monotonic_nanos() -> Result<i64, String> {
|
||||
static ORIGIN: OnceLock<Instant> = OnceLock::new();
|
||||
i64::try_from(ORIGIN.get_or_init(Instant::now).elapsed().as_nanos())
|
||||
.map_err(|_| "clock value is out of Int range".into())
|
||||
}
|
||||
|
||||
pub fn sleep(millis: i64) -> Result<(), String> {
|
||||
let millis = u64::try_from(millis).map_err(|_| "sleep duration must not be negative")?;
|
||||
std::thread::sleep(Duration::from_millis(millis));
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, pem::PemObject};
|
||||
use rustls::{ClientConfig, RootCertStore, ServerConfig};
|
||||
|
||||
use crate::{bytes::Bytes, transport::Transport};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ClientTls(pub(crate) Arc<ClientConfig>);
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ServerTls(pub(crate) Arc<ServerConfig>);
|
||||
|
||||
pub fn system() -> ClientTls {
|
||||
static CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
|
||||
ClientTls(
|
||||
CONFIG
|
||||
.get_or_init(|| {
|
||||
let roots = RootCertStore {
|
||||
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
|
||||
};
|
||||
client_config(roots)
|
||||
})
|
||||
.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Custom trust replaces public roots; it never disables certificate/name verification.
|
||||
pub fn client(ca_pem: &Bytes) -> Result<ClientTls, String> {
|
||||
let certificates = certificates(ca_pem)?;
|
||||
let mut roots = RootCertStore::empty();
|
||||
for certificate in certificates {
|
||||
roots.add(certificate).map_err(|error| error.to_string())?;
|
||||
}
|
||||
Ok(ClientTls(client_config(roots)))
|
||||
}
|
||||
|
||||
fn client_config(roots: RootCertStore) -> Arc<ClientConfig> {
|
||||
let mut config = ClientConfig::builder()
|
||||
.with_root_certificates(roots)
|
||||
.with_no_client_auth();
|
||||
// Share trust material and keep rustls's session-cache capacity small.
|
||||
config.resumption = rustls::client::Resumption::in_memory_sessions(32);
|
||||
Arc::new(config)
|
||||
}
|
||||
|
||||
pub fn server(cert_pem: &Bytes, key_pem: &Bytes) -> Result<ServerTls, String> {
|
||||
let key = PrivateKeyDer::from_pem_slice(key_pem.values()).map_err(|error| error.to_string())?;
|
||||
let mut config = ServerConfig::builder()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(certificates(cert_pem)?, key)
|
||||
.map_err(|error| error.to_string())?;
|
||||
config.session_storage = rustls::server::ServerSessionMemoryCache::new(32);
|
||||
config.send_tls13_tickets = 1;
|
||||
Ok(ServerTls(Arc::new(config)))
|
||||
}
|
||||
|
||||
fn certificates(pem: &Bytes) -> Result<Vec<CertificateDer<'static>>, String> {
|
||||
let certificates = CertificateDer::pem_slice_iter(pem.values())
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|error| error.to_string())?;
|
||||
if certificates.is_empty() {
|
||||
return Err("TLS PEM contains no certificates".into());
|
||||
}
|
||||
Ok(certificates)
|
||||
}
|
||||
|
||||
pub(crate) async fn connect(
|
||||
stream: tokio::net::TcpStream,
|
||||
host: &str,
|
||||
config: &ClientTls,
|
||||
) -> Result<Transport, String> {
|
||||
let name = ServerName::try_from(host.to_owned()).map_err(|error| error.to_string())?;
|
||||
let stream = tokio_rustls::TlsConnector::from(config.0.clone())
|
||||
.connect(name, stream)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(Transport::Tls(Box::new(stream.into())))
|
||||
}
|
||||
|
||||
pub(crate) async fn accept(
|
||||
stream: tokio::net::TcpStream,
|
||||
config: Option<&ServerTls>,
|
||||
millis: i64,
|
||||
) -> Result<Transport, String> {
|
||||
let Some(config) = config else {
|
||||
return Ok(Transport::Tcp(stream));
|
||||
};
|
||||
let stream = tokio::time::timeout(
|
||||
crate::task::positive_duration(millis)?,
|
||||
tokio_rustls::TlsAcceptor::from(config.0.clone()).accept(stream),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| "TLS handshake timed out".to_owned())?
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(Transport::Tls(Box::new(stream.into())))
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ClientTls {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("ClientTls(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ServerTls {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("ServerTls(<resource>)")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
use std::io::{self, IoSlice};
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
use tokio::net::TcpStream;
|
||||
|
||||
// TLS state is large; plain TCP keeps its direct, allocation-free representation.
|
||||
pub(crate) enum Transport {
|
||||
Tcp(TcpStream),
|
||||
Tls(Box<tokio_rustls::TlsStream<TcpStream>>),
|
||||
}
|
||||
|
||||
macro_rules! dispatch {
|
||||
($value:expr, $method:ident $(, $argument:expr)*) => {
|
||||
match $value {
|
||||
Transport::Tcp(stream) => Pin::new(stream).$method($($argument),*),
|
||||
Transport::Tls(stream) => Pin::new(stream.as_mut()).$method($($argument),*),
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
impl AsyncRead for Transport {
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
context: &mut Context<'_>,
|
||||
buffer: &mut ReadBuf<'_>,
|
||||
) -> Poll<io::Result<()>> {
|
||||
dispatch!(self.get_mut(), poll_read, context, buffer)
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for Transport {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
context: &mut Context<'_>,
|
||||
bytes: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
dispatch!(self.get_mut(), poll_write, context, bytes)
|
||||
}
|
||||
fn poll_write_vectored(
|
||||
self: Pin<&mut Self>,
|
||||
context: &mut Context<'_>,
|
||||
buffers: &[IoSlice<'_>],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
dispatch!(self.get_mut(), poll_write_vectored, context, buffers)
|
||||
}
|
||||
fn is_write_vectored(&self) -> bool {
|
||||
match self {
|
||||
Self::Tcp(stream) => stream.is_write_vectored(),
|
||||
Self::Tls(stream) => stream.is_write_vectored(),
|
||||
}
|
||||
}
|
||||
fn poll_flush(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
dispatch!(self.get_mut(), poll_flush, context)
|
||||
}
|
||||
fn poll_shutdown(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
dispatch!(self.get_mut(), poll_shutdown, context)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn connect(
|
||||
host: &str,
|
||||
port: i64,
|
||||
tls: Option<&crate::tls::ClientTls>,
|
||||
) -> Result<Transport, String> {
|
||||
let stream = crate::net::connect_stream(host, port).await?;
|
||||
match tls {
|
||||
Some(config) => crate::tls::connect(stream, host, config).await,
|
||||
None => Ok(Transport::Tcp(stream)),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,328 @@
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::{
|
||||
SinkExt, StreamExt,
|
||||
stream::{self, SplitSink, SplitStream},
|
||||
};
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_tungstenite::{
|
||||
WebSocketStream,
|
||||
tungstenite::{
|
||||
self,
|
||||
protocol::{CloseFrame, Role, WebSocketConfig, frame::coding::CloseCode},
|
||||
},
|
||||
};
|
||||
|
||||
use crate::transport::Transport;
|
||||
use crate::{
|
||||
async_stream::AsyncStream, bytes::Bytes, endpoint::Endpoint, http::HttpReply, net,
|
||||
resource::Pull,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct Limits {
|
||||
pub message_bytes: i64,
|
||||
pub idle_ms: i64,
|
||||
pub write_ms: i64,
|
||||
}
|
||||
|
||||
impl Limits {
|
||||
pub(crate) fn check(self) -> Result<WebSocketConfig, String> {
|
||||
let maximum = usize::try_from(self.message_bytes)
|
||||
.ok()
|
||||
.filter(|value| *value > 0 && value.checked_add(14).is_some())
|
||||
.ok_or("WebSocket message_bytes must be positive and leave space for a frame header")?;
|
||||
if self.idle_ms <= 0 || self.write_ms <= 0 {
|
||||
return Err("WebSocket idle_ms and write_ms must be positive".into());
|
||||
}
|
||||
Ok(WebSocketConfig::default()
|
||||
.read_buffer_size(4096)
|
||||
.write_buffer_size(0)
|
||||
.max_write_buffer_size(maximum.max(125) + 14)
|
||||
.max_message_size(Some(maximum))
|
||||
.max_frame_size(Some(maximum.max(125))))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum Message {
|
||||
Text(String),
|
||||
Binary(Bytes),
|
||||
Ping(Bytes),
|
||||
Pong(Bytes),
|
||||
}
|
||||
|
||||
type Wire = WebSocketStream<Transport>;
|
||||
struct Parts {
|
||||
reader: Mutex<SplitStream<Wire>>,
|
||||
writer: Mutex<SplitSink<Wire, tungstenite::Message>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct WebSocket {
|
||||
endpoint: Arc<Endpoint<Parts>>,
|
||||
limits: Limits,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for WebSocket {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("WebSocket(<resource>)")
|
||||
}
|
||||
}
|
||||
|
||||
impl WebSocket {
|
||||
fn new(stream: Wire, limits: Limits) -> Self {
|
||||
let (writer, reader) = stream.split();
|
||||
Self {
|
||||
endpoint: Arc::new(Endpoint::new(Parts {
|
||||
reader: Mutex::new(reader),
|
||||
writer: Mutex::new(writer),
|
||||
})),
|
||||
limits,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn abort(&self) {
|
||||
let _closed = self.endpoint.close();
|
||||
}
|
||||
|
||||
async fn operation<T>(
|
||||
&self,
|
||||
millis: i64,
|
||||
operation: impl Future<Output = Result<T, String>>,
|
||||
) -> Result<T, String> {
|
||||
tokio::time::timeout(
|
||||
Duration::from_millis(millis as u64),
|
||||
self.endpoint.while_open(operation),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| "WebSocket operation timed out".to_owned())?
|
||||
}
|
||||
}
|
||||
|
||||
// 中途取消或失败的帧写入不能重试；关闭整个传输，避免留下半帧。
|
||||
struct AbortOnDrop<'a>(Option<&'a WebSocket>);
|
||||
impl Drop for AbortOnDrop<'_> {
|
||||
fn drop(&mut self) {
|
||||
if let Some(socket) = self.0 {
|
||||
socket.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn accept(reply: &HttpReply, limits: Limits) -> Result<WebSocket, String> {
|
||||
reply.upgrade(limits).await
|
||||
}
|
||||
|
||||
pub(crate) async fn from_server(
|
||||
stream: Transport,
|
||||
buffered: Vec<u8>,
|
||||
limits: Limits,
|
||||
) -> Result<WebSocket, String> {
|
||||
let stream =
|
||||
WebSocketStream::from_partially_read(stream, buffered, Role::Server, Some(limits.check()?))
|
||||
.await;
|
||||
Ok(WebSocket::new(stream, limits))
|
||||
}
|
||||
|
||||
/// address 是数字 IP；target 是 / 开头的路径，握手、掩码与响应校验交给 tungstenite。
|
||||
pub async fn connect(
|
||||
address: &str,
|
||||
port: i64,
|
||||
target: &str,
|
||||
limits: Limits,
|
||||
) -> Result<WebSocket, String> {
|
||||
let config = limits.check()?;
|
||||
if !target.starts_with('/') || target.contains('#') {
|
||||
return Err("WebSocket target must be an origin-form path without a fragment".into());
|
||||
}
|
||||
tokio::time::timeout(Duration::from_millis(limits.write_ms as u64), async {
|
||||
let socket = net::connect_stream_timeout(address, port, limits.write_ms).await?;
|
||||
let host = if address.contains(':') {
|
||||
format!("[{address}]")
|
||||
} else {
|
||||
address.to_owned()
|
||||
};
|
||||
let uri = format!("ws://{host}:{port}{target}");
|
||||
let (stream, _) =
|
||||
tokio_tungstenite::client_async_with_config(uri, Transport::Tcp(socket), Some(config))
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(WebSocket::new(stream, limits))
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "WebSocket handshake timed out".to_owned())?
|
||||
}
|
||||
|
||||
pub async fn open(
|
||||
url: &str,
|
||||
tls: crate::tls::ClientTls,
|
||||
limits: Limits,
|
||||
) -> Result<WebSocket, String> {
|
||||
let config = limits.check()?;
|
||||
let uri: hyper::Uri = url
|
||||
.parse()
|
||||
.map_err(|error| format!("invalid WebSocket URL: {error}"))?;
|
||||
if url.contains('#')
|
||||
|| uri
|
||||
.authority()
|
||||
.is_some_and(|authority| authority.as_str().contains('@'))
|
||||
{
|
||||
return Err("WebSocket URL cannot contain credentials or a fragment".into());
|
||||
}
|
||||
let secure = match uri.scheme_str() {
|
||||
Some("ws") => false,
|
||||
Some("wss") => true,
|
||||
_ => return Err("WebSocket URL must use ws or wss".into()),
|
||||
};
|
||||
let host = uri
|
||||
.host()
|
||||
.ok_or("WebSocket URL requires a host")?
|
||||
.trim_matches(['[', ']']);
|
||||
let port = i64::from(uri.port_u16().unwrap_or(if secure { 443 } else { 80 }));
|
||||
tokio::time::timeout(crate::task::positive_duration(limits.write_ms)?, async {
|
||||
let transport = crate::transport::connect(host, port, secure.then_some(&tls)).await?;
|
||||
let (stream, _) = tokio_tungstenite::client_async_with_config(url, transport, Some(config))
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(WebSocket::new(stream, limits))
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "WebSocket handshake timed out".to_owned())?
|
||||
}
|
||||
|
||||
pub async fn send(socket: &WebSocket, message: Message) -> Result<(), String> {
|
||||
let message = match message {
|
||||
Message::Text(text) => tungstenite::Message::Text(text.into()),
|
||||
Message::Binary(bytes) => tungstenite::Message::Binary(bytes.into_http()),
|
||||
Message::Ping(bytes) => tungstenite::Message::Ping(bytes.into_http()),
|
||||
Message::Pong(bytes) => tungstenite::Message::Pong(bytes.into_http()),
|
||||
};
|
||||
let limit = if message.is_ping() || message.is_pong() {
|
||||
125
|
||||
} else {
|
||||
socket.limits.message_bytes as usize
|
||||
};
|
||||
if message.len() > limit {
|
||||
return Err("WebSocket message exceeds its size limit".into());
|
||||
}
|
||||
let parts = socket.endpoint.get()?;
|
||||
write_wire(socket, &parts, Some(message)).await
|
||||
}
|
||||
|
||||
async fn write_wire(
|
||||
socket: &WebSocket,
|
||||
parts: &Parts,
|
||||
message: Option<tungstenite::Message>,
|
||||
) -> Result<(), String> {
|
||||
socket
|
||||
.operation(socket.limits.write_ms, async {
|
||||
let mut writer = parts.writer.lock().await;
|
||||
// 等待方向锁时尚未写入；此时取消不应中断另一个任务正在发送的帧。
|
||||
let mut guard = AbortOnDrop(Some(socket));
|
||||
match message {
|
||||
Some(message) => writer.send(message).await,
|
||||
None => writer.flush().await,
|
||||
}
|
||||
.map_err(|error| error.to_string())?;
|
||||
guard.0 = None;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn receive(socket: &WebSocket) -> Result<Option<Message>, String> {
|
||||
let parts = socket.endpoint.get()?;
|
||||
let incoming = socket
|
||||
.operation(socket.limits.idle_ms, async {
|
||||
parts
|
||||
.reader
|
||||
.lock()
|
||||
.await
|
||||
.next()
|
||||
.await
|
||||
.transpose()
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.await;
|
||||
let incoming = match incoming {
|
||||
Ok(incoming) => incoming,
|
||||
Err(error) => {
|
||||
socket.abort();
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
use tungstenite::Message as WireMessage;
|
||||
let message = match incoming {
|
||||
Some(WireMessage::Text(text)) => Message::Text(text.to_string()),
|
||||
Some(WireMessage::Binary(bytes)) => Message::Binary(Bytes::from_http(bytes)),
|
||||
Some(WireMessage::Pong(bytes)) => Message::Pong(Bytes::from_http(bytes)),
|
||||
Some(WireMessage::Ping(bytes)) => {
|
||||
write_wire(socket, &parts, None).await?;
|
||||
Message::Ping(Bytes::from_http(bytes))
|
||||
}
|
||||
Some(WireMessage::Close(_)) => {
|
||||
let _guard = AbortOnDrop(Some(socket));
|
||||
write_wire(socket, &parts, None).await?;
|
||||
return Ok(None);
|
||||
}
|
||||
None => {
|
||||
socket.abort();
|
||||
return Ok(None);
|
||||
}
|
||||
Some(WireMessage::Frame(_)) => unreachable!("tungstenite never returns raw frames"),
|
||||
};
|
||||
Ok(Some(message))
|
||||
}
|
||||
|
||||
pub fn messages(socket: WebSocket) -> AsyncStream<Result<Message, String>> {
|
||||
AsyncStream::new(stream::unfold(socket, |socket| async move {
|
||||
let value = match receive(&socket).await {
|
||||
Ok(Some(message)) => Pull::Item(Ok(message)),
|
||||
Ok(None) => return None,
|
||||
Err(error) => Pull::Last(Err(error)),
|
||||
};
|
||||
Some((value, socket))
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn close(socket: &WebSocket, code: i64, reason: &str) -> Result<(), String> {
|
||||
let code = u16::try_from(code)
|
||||
.map(CloseCode::from)
|
||||
.ok()
|
||||
.filter(|code| code.is_allowed())
|
||||
.ok_or("invalid WebSocket close code")?;
|
||||
if reason.len() > 123 {
|
||||
return Err("WebSocket close reason exceeds 123 UTF-8 bytes".into());
|
||||
}
|
||||
// 先唤醒并禁止普通读写，再由此操作独占最后的传输所有权，完成关闭握手。
|
||||
let parts = socket.endpoint.take()?;
|
||||
tokio::time::timeout(
|
||||
Duration::from_millis(socket.limits.write_ms as u64),
|
||||
async {
|
||||
parts
|
||||
.writer
|
||||
.lock()
|
||||
.await
|
||||
.send(tungstenite::Message::Close(Some(CloseFrame {
|
||||
code,
|
||||
reason: reason.to_owned().into(),
|
||||
})))
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
while let Some(message) = parts.reader.lock().await.next().await {
|
||||
if matches!(
|
||||
message.map_err(|error| error.to_string())?,
|
||||
tungstenite::Message::Close(_)
|
||||
) {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|_| "WebSocket close handshake timed out".to_owned())?
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
//! Strict wire boundary. Parsed nodes never escape into generated business signatures.
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::fmt;
|
||||
|
||||
use base64::Engine as _;
|
||||
use serde::Deserializer;
|
||||
use serde::de::{self, MapAccess, SeqAccess, Visitor};
|
||||
use serde_json::value::RawValue;
|
||||
|
||||
pub const MAX_BYTES: usize = 16 * 1024 * 1024;
|
||||
pub const MAX_DEPTH: usize = 64;
|
||||
pub const MAX_ELEMENTS: usize = 65_536;
|
||||
|
||||
pub struct Node<'a> {
|
||||
raw: &'a str,
|
||||
kind: Kind<'a>,
|
||||
}
|
||||
|
||||
enum Kind<'a> {
|
||||
Null,
|
||||
Bool(bool),
|
||||
Number,
|
||||
Text(String),
|
||||
List(Vec<Node<'a>>),
|
||||
Object(BTreeMap<String, Node<'a>>),
|
||||
}
|
||||
|
||||
pub fn parse(text: &str) -> Result<Node<'_>, String> {
|
||||
if text.len() > MAX_BYTES {
|
||||
return Err("wire JSON exceeds byte limit".into());
|
||||
}
|
||||
let raw: &RawValue = serde_json::from_str(text).map_err(|error| error.to_string())?;
|
||||
let mut remaining = MAX_ELEMENTS;
|
||||
parse_node(raw, 0, &mut remaining)
|
||||
}
|
||||
|
||||
fn parse_node<'a>(
|
||||
raw: &'a RawValue,
|
||||
depth: usize,
|
||||
remaining: &mut usize,
|
||||
) -> Result<Node<'a>, String> {
|
||||
if depth > MAX_DEPTH {
|
||||
return Err("wire JSON exceeds depth limit".into());
|
||||
}
|
||||
*remaining = remaining
|
||||
.checked_sub(1)
|
||||
.ok_or("wire JSON exceeds element limit")?;
|
||||
let text = raw.get();
|
||||
let kind = match text.as_bytes()[0] {
|
||||
b'n' => Kind::Null,
|
||||
b't' => Kind::Bool(true),
|
||||
b'f' => Kind::Bool(false),
|
||||
b'"' => Kind::Text(serde_json::from_str(text).map_err(|error| error.to_string())?),
|
||||
b'[' | b'{' => {
|
||||
if depth == MAX_DEPTH {
|
||||
return Err("wire JSON exceeds depth limit".into());
|
||||
}
|
||||
let mut decoder = serde_json::Deserializer::from_str(text);
|
||||
let kind = decoder
|
||||
.deserialize_any(Container { depth, remaining })
|
||||
.map_err(|error| error.to_string())?;
|
||||
decoder.end().map_err(|error| error.to_string())?;
|
||||
kind
|
||||
}
|
||||
_ => Kind::Number,
|
||||
};
|
||||
Ok(Node { raw: text, kind })
|
||||
}
|
||||
|
||||
struct Container<'b> {
|
||||
depth: usize,
|
||||
remaining: &'b mut usize,
|
||||
}
|
||||
|
||||
impl<'de> Visitor<'de> for Container<'_> {
|
||||
type Value = Kind<'de>;
|
||||
fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("a bounded JSON container")
|
||||
}
|
||||
fn visit_seq<A: SeqAccess<'de>>(self, mut sequence: A) -> Result<Self::Value, A::Error> {
|
||||
let mut values = Vec::new();
|
||||
while let Some(raw) = sequence.next_element::<&'de RawValue>()? {
|
||||
values
|
||||
.push(parse_node(raw, self.depth + 1, self.remaining).map_err(de::Error::custom)?);
|
||||
}
|
||||
Ok(Kind::List(values))
|
||||
}
|
||||
fn visit_map<A: MapAccess<'de>>(self, mut object: A) -> Result<Self::Value, A::Error> {
|
||||
let mut values = BTreeMap::new();
|
||||
while let Some(key) = object.next_key::<String>()? {
|
||||
if values.contains_key(&key) {
|
||||
return Err(de::Error::custom("duplicate JSON field"));
|
||||
}
|
||||
let raw = object.next_value::<&'de RawValue>()?;
|
||||
values.insert(
|
||||
key,
|
||||
parse_node(raw, self.depth + 1, self.remaining).map_err(de::Error::custom)?,
|
||||
);
|
||||
}
|
||||
Ok(Kind::Object(values))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Node<'a> {
|
||||
pub fn bytes(&self) -> Result<crate::bytes::Bytes, String> {
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(self.text()?)
|
||||
.map(crate::bytes::Bytes::new)
|
||||
.map_err(|_| "invalid wire Bytes base64".into())
|
||||
}
|
||||
|
||||
pub fn is_null(&self) -> bool {
|
||||
matches!(self.kind, Kind::Null)
|
||||
}
|
||||
pub fn raw(&self) -> &'a str {
|
||||
self.raw
|
||||
}
|
||||
pub fn text(&self) -> Result<&str, String> {
|
||||
match &self.kind {
|
||||
Kind::Text(value) => Ok(value),
|
||||
_ => Err("expected wire string".into()),
|
||||
}
|
||||
}
|
||||
pub fn boolean(&self) -> Result<bool, String> {
|
||||
match self.kind {
|
||||
Kind::Bool(value) => Ok(value),
|
||||
_ => Err("expected wire Bool".into()),
|
||||
}
|
||||
}
|
||||
pub fn int(&self) -> Result<i64, String> {
|
||||
if !matches!(self.kind, Kind::Number) {
|
||||
return Err("expected wire integer".into());
|
||||
}
|
||||
self.raw
|
||||
.parse()
|
||||
.map_err(|_| "wire integer is not an exact i64".into())
|
||||
}
|
||||
pub fn float(&self) -> Result<f64, String> {
|
||||
if !matches!(self.kind, Kind::Number) {
|
||||
return Err("expected wire Float".into());
|
||||
}
|
||||
let value: f64 = self.raw.parse().map_err(|_| "invalid wire Float")?;
|
||||
if !value.is_finite() {
|
||||
return Err("non-finite wire Float".into());
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
pub fn list(&self) -> Result<&[Node<'a>], String> {
|
||||
match &self.kind {
|
||||
Kind::List(values) => Ok(values),
|
||||
_ => Err("expected wire List".into()),
|
||||
}
|
||||
}
|
||||
pub fn fields(&self, names: &[&str]) -> Result<&BTreeMap<String, Node<'a>>, String> {
|
||||
let fields = self.object()?;
|
||||
if fields.keys().any(|name| !names.contains(&name.as_str())) {
|
||||
return Err("unknown wire record field".into());
|
||||
}
|
||||
Ok(fields)
|
||||
}
|
||||
pub fn object(&self) -> Result<&BTreeMap<String, Node<'a>>, String> {
|
||||
match &self.kind {
|
||||
Kind::Object(fields) => Ok(fields),
|
||||
_ => Err("expected wire record".into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Only validated JSON can be placed in an API envelope without re-encoding its numbers.
|
||||
pub struct Encoded(String);
|
||||
impl Encoded {
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
pub fn null() -> Self {
|
||||
Self("null".into())
|
||||
}
|
||||
}
|
||||
|
||||
enum Frame {
|
||||
Array(usize),
|
||||
Object {
|
||||
names: BTreeSet<String>,
|
||||
awaiting_value: bool,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct Encoder {
|
||||
output: String,
|
||||
stack: Vec<Frame>,
|
||||
elements: usize,
|
||||
failed: bool,
|
||||
}
|
||||
|
||||
impl Encoder {
|
||||
pub fn bytes(&mut self, value: &crate::bytes::Bytes) -> Result<(), String> {
|
||||
let length = base64::encoded_len(value.values().len(), true)
|
||||
.and_then(|length| length.checked_add(2))
|
||||
.ok_or("wire JSON exceeds byte limit")?;
|
||||
if length > MAX_BYTES.saturating_sub(self.output.len()) {
|
||||
return Err("wire JSON exceeds byte limit".into());
|
||||
}
|
||||
self.text(&base64::engine::general_purpose::STANDARD.encode(value.values()))
|
||||
}
|
||||
|
||||
fn append(&mut self, text: &str) -> Result<(), String> {
|
||||
if text.len() > MAX_BYTES.saturating_sub(self.output.len()) {
|
||||
// Frame bookkeeping may already have advanced before an append fails.
|
||||
// Such a partial write must never become a validated Encoded value.
|
||||
self.failed = true;
|
||||
return Err("wire JSON exceeds byte limit".into());
|
||||
}
|
||||
self.output.push_str(text);
|
||||
Ok(())
|
||||
}
|
||||
fn value(&mut self) -> Result<(), String> {
|
||||
if self.elements == MAX_ELEMENTS {
|
||||
return Err("wire JSON exceeds element limit".into());
|
||||
}
|
||||
self.elements += 1;
|
||||
let comma = match self.stack.last_mut() {
|
||||
Some(Frame::Array(count)) => {
|
||||
*count += 1;
|
||||
*count > 1
|
||||
}
|
||||
Some(Frame::Object { awaiting_value, .. }) if *awaiting_value => {
|
||||
*awaiting_value = false;
|
||||
false
|
||||
}
|
||||
Some(Frame::Object { .. }) => return Err("wire object requires a field name".into()),
|
||||
None if self.elements != 1 => return Err("wire JSON has multiple roots".into()),
|
||||
None => false,
|
||||
};
|
||||
if comma {
|
||||
self.append(",")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
pub fn null(&mut self) -> Result<(), String> {
|
||||
self.scalar("null")
|
||||
}
|
||||
pub fn boolean(&mut self, value: bool) -> Result<(), String> {
|
||||
self.scalar(if value { "true" } else { "false" })
|
||||
}
|
||||
pub fn int(&mut self, value: i64) -> Result<(), String> {
|
||||
self.scalar(&value.to_string())
|
||||
}
|
||||
pub fn float(&mut self, value: f64) -> Result<(), String> {
|
||||
if !value.is_finite() {
|
||||
return Err("non-finite wire Float".into());
|
||||
}
|
||||
self.scalar(&serde_json::to_string(&value).map_err(|error| error.to_string())?)
|
||||
}
|
||||
pub fn text(&mut self, value: &str) -> Result<(), String> {
|
||||
if value.len() > MAX_BYTES.saturating_sub(self.output.len()) {
|
||||
return Err("wire JSON exceeds byte limit".into());
|
||||
}
|
||||
self.scalar(&serde_json::to_string(value).map_err(|error| error.to_string())?)
|
||||
}
|
||||
fn scalar(&mut self, text: &str) -> Result<(), String> {
|
||||
self.value()?;
|
||||
self.append(text)
|
||||
}
|
||||
pub fn json(&mut self, text: &str) -> Result<(), String> {
|
||||
if text.len() > MAX_BYTES.saturating_sub(self.output.len()) {
|
||||
return Err("wire JSON exceeds byte limit".into());
|
||||
}
|
||||
let raw: &RawValue = serde_json::from_str(text).map_err(|error| error.to_string())?;
|
||||
let mut remaining = MAX_ELEMENTS - self.elements;
|
||||
parse_node(raw, self.stack.len(), &mut remaining)?;
|
||||
self.value()?;
|
||||
self.elements = MAX_ELEMENTS - remaining;
|
||||
self.append(raw.get())
|
||||
}
|
||||
pub fn begin_array(&mut self) -> Result<(), String> {
|
||||
self.begin(false)
|
||||
}
|
||||
pub fn begin_object(&mut self) -> Result<(), String> {
|
||||
self.begin(true)
|
||||
}
|
||||
fn begin(&mut self, object: bool) -> Result<(), String> {
|
||||
if self.stack.len() >= MAX_DEPTH {
|
||||
return Err("wire JSON exceeds depth limit".into());
|
||||
}
|
||||
self.value()?;
|
||||
self.append(if object { "{" } else { "[" })?;
|
||||
self.stack.push(if object {
|
||||
Frame::Object {
|
||||
names: BTreeSet::new(),
|
||||
awaiting_value: false,
|
||||
}
|
||||
} else {
|
||||
Frame::Array(0)
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
pub fn key(&mut self, name: &str) -> Result<(), String> {
|
||||
if name.len() > MAX_BYTES.saturating_sub(self.output.len()) {
|
||||
return Err("wire JSON exceeds byte limit".into());
|
||||
}
|
||||
let Some(Frame::Object {
|
||||
names,
|
||||
awaiting_value,
|
||||
}) = self.stack.last_mut()
|
||||
else {
|
||||
return Err("wire field outside record".into());
|
||||
};
|
||||
if *awaiting_value {
|
||||
return Err("wire field has no value".into());
|
||||
}
|
||||
if !names.insert(name.to_owned()) {
|
||||
return Err("duplicate JSON field".into());
|
||||
}
|
||||
*awaiting_value = true;
|
||||
let comma = names.len() > 1;
|
||||
if comma {
|
||||
self.append(",")?;
|
||||
}
|
||||
self.append(&serde_json::to_string(name).map_err(|error| error.to_string())?)?;
|
||||
self.append(":")
|
||||
}
|
||||
pub fn end(&mut self) -> Result<(), String> {
|
||||
if matches!(
|
||||
self.stack.last(),
|
||||
Some(Frame::Object {
|
||||
awaiting_value: true,
|
||||
..
|
||||
})
|
||||
) {
|
||||
return Err("incomplete wire container".into());
|
||||
}
|
||||
match self.stack.pop() {
|
||||
Some(Frame::Array(_)) => self.append("]"),
|
||||
Some(Frame::Object {
|
||||
awaiting_value: false,
|
||||
..
|
||||
}) => self.append("}"),
|
||||
_ => Err("incomplete wire container".into()),
|
||||
}
|
||||
}
|
||||
pub fn finish(self) -> Result<Encoded, String> {
|
||||
if self.failed || !self.stack.is_empty() || self.elements == 0 {
|
||||
return Err("incomplete wire JSON".into());
|
||||
}
|
||||
Ok(Encoded(self.output))
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user