Sessions.

This commit is contained in:
Mauro D
2023-04-18 10:43:45 +00:00
parent 3751133c7d
commit c807883374
36 changed files with 2037 additions and 425 deletions

View File

@@ -7,7 +7,7 @@ use std::slice::Iter;
use store::{
write::{IntoBitmap, Operation, ToBitmaps},
BlobHash, Deserialize, Serialize,
Deserialize, Serialize,
};
use utils::{
codec::leb128::{Leb128Iterator, Leb128Vec},
@@ -15,14 +15,8 @@ use utils::{
};
use crate::types::{
acl::Acl,
blob::{BlobId, BlobSection},
date::UTCDate,
id::Id,
keyword::Keyword,
property::Property,
type_state::TypeState,
value::Value,
acl::Acl, blob::BlobId, date::UTCDate, id::Id, keyword::Keyword, property::Property,
type_state::TypeState, value::Value,
};
#[derive(Debug, Clone, Default, serde::Serialize, PartialEq, Eq)]
@@ -209,8 +203,7 @@ impl SerializeValue for Value {
}
Value::BlobId(v) => {
buf.push(BLOB_ID);
buf.extend_from_slice(&v.hash.hash);
buf.push_leb128(v.section.as_ref().map_or(0, |s| s.offset_start));
v.serialize_value(buf);
}
Value::Keyword(v) => {
buf.push(KEYWORD);
@@ -253,21 +246,7 @@ impl DeserializeValue for Value {
DATE => Some(Value::Date(UTCDate::from_timestamp(
bytes.next_leb128::<u64>()? as i64,
))),
BLOB_ID => {
let mut hash = BlobHash::default();
for byte in hash.hash.iter_mut() {
*byte = *bytes.next()?;
}
let offset_start = bytes.next_leb128::<usize>()?;
Some(Value::BlobId(BlobId {
hash,
section: Some(BlobSection {
offset_start,
size: 0,
encoding: 0,
}),
}))
}
BLOB_ID => Some(Value::BlobId(BlobId::deserialize_value(bytes)?)),
KEYWORD => Some(Value::Keyword(Keyword::deserialize_value(bytes)?)),
TYPE_STATE => Some(Value::TypeState(TypeState::deserialize_value(bytes)?)),
ACL => Some(Value::Acl(Acl::deserialize_value(bytes)?)),

View File

@@ -21,19 +21,32 @@
* for more details.
*/
use std::io::Write;
use std::{borrow::Borrow, io::Write};
use store::{BlobHash, BLOB_HASH_LEN};
use store::{
rand::{self, Rng},
write::now,
BlobKind,
};
use utils::codec::{
base32_custom::Base32Writer,
leb128::{Leb128Iterator, Leb128Writer},
};
use crate::parser::{base32::JsonBase32Reader, json::Parser, JsonObjectParser};
use crate::{
object::{DeserializeValue, SerializeValue},
parser::{base32::JsonBase32Reader, json::Parser, JsonObjectParser},
};
use super::date::UTCDate;
const B_LINKED: u8 = 0x10;
const B_LINKED_MAILDIR: u8 = 0x20;
const B_TEMPORARY: u8 = 0x40;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct BlobId {
pub hash: BlobHash,
pub kind: BlobKind,
pub section: Option<BlobSection>,
}
@@ -44,61 +57,158 @@ pub struct BlobSection {
pub encoding: u8,
}
impl BlobId {
pub fn maildir(account_id: u32, document_id: u32) -> Self {
Self {
kind: BlobKind::LinkedMaildir {
account_id,
document_id,
},
section: None,
}
}
pub fn temporary(account_id: u32) -> Self {
let now_secs = now();
let now = UTCDate::from_timestamp(now_secs as i64);
Self {
kind: BlobKind::Temporary {
account_id,
creation_year: now.year,
creation_month: now.month,
creation_day: now.day,
seq: ((now_secs % 86400) as u32) << 15
| rand::thread_rng().gen_range(0u32..=32767u32),
},
section: None,
}
}
pub fn has_access(&self, account_id: u32) -> bool {
match &self.kind {
BlobKind::Linked { account_id: a, .. } => *a == account_id,
BlobKind::LinkedMaildir { account_id: a, .. } => *a == account_id,
BlobKind::Temporary { account_id: a, .. } => *a == account_id,
}
}
}
impl JsonObjectParser for BlobId {
fn parse(parser: &mut Parser<'_>) -> crate::parser::Result<Self>
where
Self: Sized,
{
let encoding = match parser
.next_unescaped()?
.ok_or_else(|| parser.error_value())?
{
b'a' => None,
b @ b'b'..=b'g' => Some(b - b'b'),
_ => {
return Err(parser.error_value());
}
};
let mut it = JsonBase32Reader::new(parser);
let mut hash = [0; BLOB_HASH_LEN];
for byte in hash.iter_mut().take(BLOB_HASH_LEN) {
*byte = it.next().ok_or_else(|| it.error())?;
BlobId::from_iter(&mut it).ok_or_else(|| it.error())
}
}
Ok(BlobId {
hash: BlobHash { hash },
section: if let Some(encoding) = encoding {
impl BlobId {
pub fn new(kind: BlobKind) -> Self {
BlobId {
kind,
section: None,
}
}
#[allow(clippy::should_implement_trait)]
pub fn from_iter<T, U>(it: &mut T) -> Option<Self>
where
T: Iterator<Item = U> + Leb128Iterator<U>,
U: Borrow<u8>,
{
let kind = *it.next()?.borrow();
let encoding = kind & 0x0F;
BlobId {
kind: match kind & 0xF0 {
B_LINKED => BlobKind::Linked {
account_id: it.next_leb128()?,
collection: *it.next()?.borrow(),
document_id: it.next_leb128()?,
},
B_LINKED_MAILDIR => BlobKind::LinkedMaildir {
account_id: it.next_leb128()?,
document_id: it.next_leb128()?,
},
B_TEMPORARY => BlobKind::Temporary {
account_id: it.next_leb128()?,
creation_year: u16::from_be_bytes([*it.next()?.borrow(), *it.next()?.borrow()]),
creation_month: *it.next()?.borrow(),
creation_day: *it.next()?.borrow(),
seq: it.next_leb128()?,
},
_ => return None,
},
section: if encoding != 0 {
BlobSection {
offset_start: it.next_leb128().ok_or_else(|| it.error())?,
size: it.next_leb128().ok_or_else(|| it.error())?,
encoding,
offset_start: it.next_leb128()?,
size: it.next_leb128()?,
encoding: encoding - 1,
}
.into()
} else {
None
},
})
}
.into()
}
fn serialize_into(&self, writer: &mut (impl Write + Leb128Writer)) {
let kind = self
.section
.as_ref()
.map_or(0, |section| section.encoding + 1);
match &self.kind {
BlobKind::Linked {
account_id,
collection,
document_id,
} => {
let _ = writer.write(&[kind | B_LINKED]);
let _ = writer.write_leb128(*account_id);
let _ = writer.write(&[*collection]);
let _ = writer.write_leb128(*document_id);
}
BlobKind::LinkedMaildir {
account_id,
document_id,
} => {
let _ = writer.write(&[kind | B_LINKED_MAILDIR]);
let _ = writer.write_leb128(*account_id);
let _ = writer.write_leb128(*document_id);
}
BlobKind::Temporary {
account_id,
creation_year,
creation_month,
creation_day,
seq,
} => {
let _ = writer.write(&[kind | B_TEMPORARY]);
let _ = writer.write_leb128(*account_id);
let _ = writer.write(&creation_year.to_be_bytes()[..]);
let _ = writer.write(&[*creation_month]);
let _ = writer.write(&[*creation_day]);
let _ = writer.write_leb128(*seq);
}
}
impl BlobId {
pub fn new(hash: BlobHash) -> Self {
BlobId {
hash,
section: None,
if let Some(section) = &self.section {
let _ = writer.write_leb128(section.offset_start);
let _ = writer.write_leb128(section.size);
}
}
pub fn new_section(
hash: BlobHash,
kind: BlobKind,
offset_start: usize,
offset_end: usize,
encoding: impl Into<u8>,
) -> Self {
BlobId {
hash,
kind,
section: BlobSection {
offset_start,
size: offset_end - offset_start,
@@ -117,29 +227,30 @@ impl BlobId {
}
}
impl From<&BlobHash> for BlobId {
fn from(id: &BlobHash) -> Self {
BlobId::new(*id)
}
}
impl From<BlobHash> for BlobId {
fn from(id: BlobHash) -> Self {
BlobId::new(id)
}
}
impl Default for BlobId {
fn default() -> Self {
Self {
hash: BlobHash {
hash: [0; BLOB_HASH_LEN],
BlobId {
kind: store::BlobKind::LinkedMaildir {
account_id: u32::MAX,
document_id: u32::MAX,
},
section: None,
}
}
}
impl From<&BlobKind> for BlobId {
fn from(kind: &BlobKind) -> Self {
BlobId::new(*kind)
}
}
impl From<BlobKind> for BlobId {
fn from(id: BlobKind) -> Self {
BlobId::new(id)
}
}
impl serde::Serialize for BlobId {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
@@ -152,20 +263,20 @@ impl serde::Serialize for BlobId {
impl std::fmt::Display for BlobId {
#[allow(clippy::unused_io_amount)]
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut writer;
if let Some(section) = &self.section {
writer =
Base32Writer::with_capacity(BLOB_HASH_LEN + (std::mem::size_of::<u32>() * 2) + 1);
writer.push_char(char::from(b'b' + section.encoding));
writer.write(&self.hash.hash).unwrap();
writer.write_leb128(section.offset_start).unwrap();
writer.write_leb128(section.size).unwrap();
} else {
writer = Base32Writer::with_capacity(BLOB_HASH_LEN + 1);
writer.push_char('a');
writer.write(&self.hash.hash).unwrap();
}
let mut writer = Base32Writer::with_capacity(std::mem::size_of::<BlobId>() * 2);
self.serialize_into(&mut writer);
f.write_str(&writer.finalize())
}
}
impl SerializeValue for BlobId {
fn serialize_value(self, buf: &mut Vec<u8>) {
self.serialize_into(buf)
}
}
impl DeserializeValue for BlobId {
fn deserialize_value(bytes: &mut std::slice::Iter<'_, u8>) -> Option<Self> {
BlobId::from_iter(bytes)
}
}

View File

@@ -101,6 +101,21 @@ impl Id {
Self { id }
}
pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
let mut id = 0;
for &ch in bytes {
let i = BASE32_INVERSE[ch as usize];
if i != u8::MAX {
id = (id << 5) | i as u64;
} else {
return None;
}
}
Id { id }.into()
}
pub fn singleton() -> Self {
Self::new(20080258862541)
}

View File

@@ -2,7 +2,7 @@ use std::{borrow::Cow, fmt::Display};
use mail_parser::{Addr, DateTime, Group};
use serde::Serialize;
use store::BlobHash;
use store::BlobKind;
use crate::{
error::method::MethodError,
@@ -275,8 +275,8 @@ impl From<BlobId> for Value {
}
}
impl From<BlobHash> for Value {
fn from(value: BlobHash) -> Self {
impl From<BlobKind> for Value {
fn from(value: BlobKind) -> Self {
Value::BlobId(BlobId::new(value))
}
}

View File

@@ -12,3 +12,8 @@ mail-builder = { git = "https://github.com/stalwartlabs/mail-builder", features
mail-send = { git = "https://github.com/stalwartlabs/mail-send" }
serde = { version = "1.0", features = ["derive"]}
serde_json = "1.0"
hyper = { version = "1.0.0-rc.3", features = ["server", "http1", "http2"] }
http-body-util = "0.1.0-rc.2"
form_urlencoded = "1.1.0"
tracing = "0.1"
tokio = { version = "1.23", features = ["rt"] }

303
crates/jmap/src/api/http.rs Normal file
View File

@@ -0,0 +1,303 @@
use std::sync::Arc;
use http_body_util::{combinators::BoxBody, BodyExt, Full};
use hyper::{
body::{self, Bytes},
header::{self, CONTENT_TYPE},
server::conn::http1,
service::service_fn,
Method, StatusCode,
};
use jmap_proto::{
error::request::{RequestError, RequestLimitError},
response::Response,
types::{blob::BlobId, id::Id},
};
use tokio::{
io::{AsyncRead, AsyncWrite},
net::TcpStream,
};
use utils::listener::{SessionData, SessionManager};
use crate::{
blob::{DownloadResponse, UploadResponse},
JMAP,
};
impl JMAP {
pub async fn parse_request(
&self,
req: &mut hyper::Request<hyper::body::Incoming>,
) -> hyper::Response<BoxBody<Bytes, hyper::Error>> {
let mut path = req.uri().path().split('/');
path.next();
match path.next().unwrap_or("") {
"jmap" => match (path.next().unwrap_or(""), req.method()) {
("", &Method::POST) => {
return match fetch_body(req, self.config.request_max_size).await {
Ok(bytes) => match self.handle_request(&bytes).await {
Ok(response) => response.into_http_response(),
Err(err) => err.into_http_response(),
},
Err(err) => err.into_http_response(),
}
}
("download", &Method::GET) => {
if let (Some(account_id), Some(blob_id), Some(name)) = (
path.next().and_then(|p| Id::from_bytes(p.as_bytes())),
path.next()
.and_then(|p| BlobId::from_iter(&mut p.as_bytes().iter())),
path.next(),
) {
return match self.blob_download(&blob_id, account_id.document_id()).await {
Ok(Some(blob)) => DownloadResponse {
filename: name.to_string(),
content_type: req
.uri()
.query()
.and_then(|q| {
form_urlencoded::parse(q.as_bytes())
.find(|(k, _)| k == "accept")
.map(|(_, v)| v.into_owned())
})
.unwrap_or("application/octet-stream".to_string()),
blob,
}
.into_http_response(),
Ok(None) => RequestError::not_found().into_http_response(),
Err(err) => RequestError::internal_server_error().into_http_response(),
};
}
}
("upload", &Method::POST) => {
if let Some(account_id) = path.next().and_then(|p| Id::from_bytes(p.as_bytes()))
{
return match fetch_body(req, self.config.upload_max_size).await {
Ok(bytes) => {
match self
.blob_upload(
account_id,
req.headers()
.get(CONTENT_TYPE)
.and_then(|h| h.to_str().ok())
.unwrap_or("application/octet-stream"),
&bytes,
)
.await
{
Ok(response) => response.into_http_response(),
Err(err) => err.into_http_response(),
}
}
Err(err) => err.into_http_response(),
};
}
}
("eventsource", &Method::GET) => {
todo!()
}
("ws", &Method::GET) => {
todo!()
}
_ => (),
},
".well-known" => match (path.next().unwrap_or(""), req.method()) {
("jmap", &Method::GET) => {
todo!()
}
("oauth-authorization-server", &Method::GET) => {
todo!()
}
_ => (),
},
"auth" => match (path.next().unwrap_or(""), req.method()) {
("", &Method::GET) => {
todo!()
}
("", &Method::POST) => {
todo!()
}
("code", &Method::GET) => {
todo!()
}
("code", &Method::POST) => {
todo!()
}
("device", &Method::POST) => {
todo!()
}
("token", &Method::POST) => {
todo!()
}
_ => (),
},
_ => (),
}
RequestError::not_found().into_http_response()
}
}
impl SessionManager for super::SessionManager {
fn spawn(&self, session: SessionData<TcpStream>) {
let jmap = self.inner.clone();
tokio::spawn(async move {
if let Some(tls_acceptor) = &session.instance.tls_acceptor {
let span = session.span;
match tls_acceptor.accept(session.stream).await {
Ok(stream) => {
handle_request(
jmap,
SessionData {
stream,
local_ip: session.local_ip,
remote_ip: session.remote_ip,
span,
in_flight: session.in_flight,
instance: session.instance,
shutdown_rx: session.shutdown_rx,
},
)
.await;
}
Err(err) => {
tracing::debug!(
parent: &span,
context = "tls",
event = "error",
"Failed to accept TLS connection: {}",
err
);
}
}
} else {
handle_request(jmap, session).await;
}
});
}
}
async fn handle_request<T: AsyncRead + AsyncWrite + Unpin + 'static>(
jmap: Arc<JMAP>,
session: SessionData<T>,
) {
let span = session.span;
if let Err(http_err) = http1::Builder::new()
.keep_alive(true)
.serve_connection(
session.stream,
service_fn(|mut req: hyper::Request<body::Incoming>| {
let jmap = jmap.clone();
let span = span.clone();
async move {
let response = jmap.parse_request(&mut req).await;
tracing::debug!(
parent: &span,
event = "request",
uri = req.uri().to_string(),
status = response.status().to_string(),
);
Ok::<_, hyper::Error>(response)
}
}),
)
.await
{
tracing::debug!(
parent: &span,
event = "http-error",
reason = %http_err,
);
}
}
async fn fetch_body(
req: &mut hyper::Request<hyper::body::Incoming>,
max_size: usize,
) -> Result<Vec<u8>, RequestError> {
let mut bytes = Vec::with_capacity(1024);
while let Some(Ok(frame)) = req.frame().await {
if let Some(data) = frame.data_ref() {
if bytes.len() + data.len() < max_size {
bytes.extend_from_slice(data);
} else {
return Err(RequestError::limit(RequestLimitError::Size));
}
}
}
Ok(bytes)
}
trait ToHttpResponse {
fn into_http_response(self) -> hyper::Response<BoxBody<Bytes, hyper::Error>>;
}
impl ToHttpResponse for Response {
fn into_http_response(self) -> hyper::Response<BoxBody<Bytes, hyper::Error>> {
hyper::Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
.body(
Full::new(Bytes::from(serde_json::to_string(&self).unwrap()))
.map_err(|never| match never {})
.boxed(),
)
.unwrap()
}
}
impl ToHttpResponse for DownloadResponse {
fn into_http_response(self) -> hyper::Response<BoxBody<Bytes, hyper::Error>> {
hyper::Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, self.content_type)
.header(
header::CONTENT_DISPOSITION,
format!(
"attachment; filename=\"{}\"",
self.filename.replace('\"', "\\\"")
),
)
.header(
header::CACHE_CONTROL,
"private, immutable, max-age=31536000",
)
.body(
Full::new(Bytes::from(self.blob))
.map_err(|never| match never {})
.boxed(),
)
.unwrap()
}
}
impl ToHttpResponse for UploadResponse {
fn into_http_response(self) -> hyper::Response<BoxBody<Bytes, hyper::Error>> {
hyper::Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
.body(
Full::new(Bytes::from(serde_json::to_string(&self).unwrap()))
.map_err(|never| match never {})
.boxed(),
)
.unwrap()
}
}
impl ToHttpResponse for RequestError {
fn into_http_response(self) -> hyper::Response<BoxBody<Bytes, hyper::Error>> {
hyper::Response::builder()
.status(self.status)
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
.body(
Full::new(Bytes::from(serde_json::to_string(&self).unwrap()))
.map_err(|never| match never {})
.boxed(),
)
.unwrap()
}
}

View File

@@ -1 +1,12 @@
use std::sync::Arc;
use crate::JMAP;
pub mod http;
pub mod request;
pub mod session;
#[derive(Clone)]
pub struct SessionManager {
pub inner: Arc<JMAP>,
}

View File

@@ -0,0 +1,360 @@
use jmap_proto::{request::capability::Capability, response::serialize_hex, types::id::Id};
use store::ahash::AHashSet;
use utils::map::vec_map::VecMap;
use crate::Config;
#[derive(Debug, Clone, serde::Serialize)]
pub struct Session {
#[serde(rename(serialize = "capabilities"))]
capabilities: VecMap<Capability, Capabilities>,
#[serde(rename(serialize = "accounts"))]
accounts: VecMap<Id, Account>,
#[serde(rename(serialize = "primaryAccounts"))]
primary_accounts: VecMap<Capability, Id>,
#[serde(rename(serialize = "username"))]
username: String,
#[serde(rename(serialize = "apiUrl"))]
api_url: String,
#[serde(rename(serialize = "downloadUrl"))]
download_url: String,
#[serde(rename(serialize = "uploadUrl"))]
upload_url: String,
#[serde(rename(serialize = "eventSourceUrl"))]
event_source_url: String,
#[serde(rename(serialize = "state"))]
#[serde(serialize_with = "serialize_hex")]
state: u32,
#[serde(skip)]
base_url: String,
}
#[derive(Debug, Clone, serde::Serialize)]
struct Account {
#[serde(rename(serialize = "name"))]
name: String,
#[serde(rename(serialize = "isPersonal"))]
is_personal: bool,
#[serde(rename(serialize = "isReadOnly"))]
is_read_only: bool,
#[serde(rename(serialize = "accountCapabilities"))]
account_capabilities: VecMap<Capability, Capabilities>,
}
#[derive(Debug, Clone, serde::Serialize)]
#[serde(untagged)]
#[allow(dead_code)]
enum Capabilities {
Core(CoreCapabilities),
Mail(MailCapabilities),
Submission(SubmissionCapabilities),
VacationResponse(VacationResponseCapabilities),
WebSocket(WebSocketCapabilities),
Sieve(SieveCapabilities),
}
#[derive(Debug, Clone, serde::Serialize)]
struct CoreCapabilities {
#[serde(rename(serialize = "maxSizeUpload"))]
max_size_upload: usize,
#[serde(rename(serialize = "maxConcurrentUpload"))]
max_concurrent_upload: usize,
#[serde(rename(serialize = "maxSizeRequest"))]
max_size_request: usize,
#[serde(rename(serialize = "maxConcurrentRequests"))]
max_concurrent_requests: usize,
#[serde(rename(serialize = "maxCallsInRequest"))]
max_calls_in_request: usize,
#[serde(rename(serialize = "maxObjectsInGet"))]
max_objects_in_get: usize,
#[serde(rename(serialize = "maxObjectsInSet"))]
max_objects_in_set: usize,
#[serde(rename(serialize = "collationAlgorithms"))]
collation_algorithms: Vec<String>,
}
#[derive(Debug, Clone, serde::Serialize)]
struct WebSocketCapabilities {
#[serde(rename(serialize = "url"))]
url: String,
#[serde(rename(serialize = "supportsPush"))]
supports_push: bool,
}
#[derive(Debug, Clone, serde::Serialize)]
struct SieveCapabilities {
#[serde(rename(serialize = "maxSizeScriptName"))]
max_script_name: usize,
#[serde(rename(serialize = "maxSizeScript"))]
max_script_size: usize,
#[serde(rename(serialize = "maxNumberScripts"))]
max_scripts: usize,
#[serde(rename(serialize = "maxNumberRedirects"))]
max_redirects: usize,
#[serde(rename(serialize = "sieveExtensions"))]
extensions: Vec<String>,
#[serde(rename(serialize = "notificationMethods"))]
notification_methods: Option<Vec<String>>,
#[serde(rename(serialize = "externalLists"))]
ext_lists: Option<Vec<String>>,
}
#[derive(Debug, Clone, serde::Serialize)]
struct MailCapabilities {
#[serde(rename(serialize = "maxMailboxesPerEmail"))]
max_mailboxes_per_email: Option<usize>,
#[serde(rename(serialize = "maxMailboxDepth"))]
max_mailbox_depth: usize,
#[serde(rename(serialize = "maxSizeMailboxName"))]
max_size_mailbox_name: usize,
#[serde(rename(serialize = "maxSizeAttachmentsPerEmail"))]
max_size_attachments_per_email: usize,
#[serde(rename(serialize = "emailQuerySortOptions"))]
email_query_sort_options: Vec<String>,
#[serde(rename(serialize = "mayCreateTopLevelMailbox"))]
may_create_top_level_mailbox: bool,
}
#[derive(Debug, Clone, serde::Serialize)]
struct SubmissionCapabilities {
#[serde(rename(serialize = "maxDelayedSend"))]
max_delayed_send: usize,
#[serde(rename(serialize = "submissionExtensions"))]
submission_extensions: Vec<String>,
}
#[derive(Debug, Clone, serde::Serialize)]
struct VacationResponseCapabilities {}
struct BaseCapabilities {
capabilities: VecMap<Capability, Capabilities>,
}
impl BaseCapabilities {
pub fn new(config: &crate::Config, raw_config: &Config) -> Self {
Self {
capabilities: VecMap::from_iter([
(
Capability::Core,
Capabilities::Core(CoreCapabilities::new(config)),
),
(
Capability::Mail,
Capabilities::Mail(MailCapabilities::new(config)),
),
(
Capability::Sieve,
Capabilities::Sieve(SieveCapabilities::new(config, raw_config)),
),
]),
}
}
}
impl Session {
pub fn new(base_url: &str, base_capabilities: &BaseCapabilities) -> Session {
let mut capabilities = base_capabilities.capabilities.clone();
capabilities.append(
Capability::WebSocket,
Capabilities::WebSocket(WebSocketCapabilities::new(&base_url)),
);
Session {
capabilities,
accounts: VecMap::new(),
primary_accounts: VecMap::new(),
username: "".to_string(),
api_url: format!("{}/jmap/", base_url),
download_url: format!(
"{}/jmap/download/{{accountId}}/{{blobId}}/{{name}}?accept={{type}}",
base_url
),
upload_url: format!("{}/jmap/upload/{{accountId}}/", base_url),
event_source_url: format!(
"{}/jmap/eventsource/?types={{types}}&closeafter={{closeafter}}&ping={{ping}}",
base_url
),
base_url: base_url.to_string(),
state: 0,
}
}
pub fn set_primary_account(
&mut self,
account_id: Id,
username: String,
name: String,
capabilities: Option<&[Capability]>,
) {
self.username = username;
if let Some(capabilities) = capabilities {
for capability in capabilities {
self.primary_accounts.append(capability.clone(), account_id);
}
} else {
for capability in self.capabilities.keys() {
self.primary_accounts.append(capability.clone(), account_id);
}
}
self.accounts.set(
account_id,
Account::new(name, true, false).add_capabilities(capabilities, &self.capabilities),
);
}
pub fn add_account(
&mut self,
account_id: Id,
name: String,
is_personal: bool,
is_read_only: bool,
capabilities: Option<&[Capability]>,
) {
self.accounts.set(
account_id,
Account::new(name, is_personal, is_read_only)
.add_capabilities(capabilities, &self.capabilities),
);
}
pub fn set_state(&mut self, state: u32) {
self.state = state;
}
pub fn api_url(&self) -> &str {
&self.api_url
}
pub fn base_url(&self) -> &str {
&self.base_url
}
}
impl Account {
pub fn new(name: String, is_personal: bool, is_read_only: bool) -> Account {
Account {
name,
is_personal,
is_read_only,
account_capabilities: VecMap::new(),
}
}
pub fn add_capabilities(
mut self,
capabilities: Option<&[Capability]>,
core_capabilities: &VecMap<Capability, Capabilities>,
) -> Account {
if let Some(capabilities) = capabilities {
for capability in capabilities {
self.account_capabilities.append(
capability.clone(),
core_capabilities.get(capability).unwrap().clone(),
);
}
} else {
self.account_capabilities = core_capabilities.clone();
}
self
}
}
impl CoreCapabilities {
pub fn new(config: &crate::Config) -> Self {
CoreCapabilities {
max_size_upload: config.max_size_upload,
max_concurrent_upload: config.max_concurrent_uploads,
max_size_request: config.max_size_request,
max_concurrent_requests: config.max_concurrent_requests,
max_calls_in_request: config.max_calls_in_request,
max_objects_in_get: config.max_objects_in_get,
max_objects_in_set: config.max_objects_in_set,
collation_algorithms: vec![
"i;ascii-numeric".to_string(),
"i;ascii-casemap".to_string(),
"i;unicode-casemap".to_string(),
],
}
}
}
impl WebSocketCapabilities {
pub fn new(base_url: &str) -> Self {
WebSocketCapabilities {
url: format!("ws{}/jmap/ws", base_url.strip_prefix("http").unwrap()),
supports_push: true,
}
}
}
impl SieveCapabilities {
pub fn new(config: &crate::Config, raw_config: &Config) -> Self {
let mut notification_methods = Vec::new();
for part in settings
.get("sieve-notification-uris")
.unwrap_or_else(|| "mailto".to_string())
.split_ascii_whitespace()
{
if !part.is_empty() {
notification_methods.push(part.to_string());
}
}
let mut capabilities: AHashSet<Capability> =
AHashSet::from_iter(Capability::all().iter().cloned());
if let Some(disable) = settings.get("sieve-disable-capabilities") {
for item in disable.split_ascii_whitespace() {
capabilities.remove(&Capability::parse(item));
}
}
let mut extensions = capabilities
.into_iter()
.map(|c| c.to_string())
.collect::<Vec<String>>();
extensions.sort_unstable();
SieveCapabilities {
max_script_name: config.sieve_max_script_name,
max_script_size: settings
.parse("sieve-max-script-size")
.unwrap_or(1024 * 1024),
max_scripts: config.sieve_max_scripts,
max_redirects: settings.parse("sieve-max-redirects").unwrap_or(1),
extensions,
notification_methods: if !notification_methods.is_empty() {
notification_methods.into()
} else {
None
},
ext_lists: None,
}
}
}
impl MailCapabilities {
pub fn new(config: &crate::Config) -> Self {
MailCapabilities {
max_mailboxes_per_email: None,
max_mailbox_depth: config.mailbox_max_depth,
max_size_mailbox_name: config.mailbox_name_max_len,
max_size_attachments_per_email: config.mail_attachments_max_size,
email_query_sort_options: [
"receivedAt",
"size",
"from",
"to",
"subject",
"sentAt",
"hasKeyword",
"allInThreadHaveKeyword",
"someInThreadHaveKeyword",
]
.iter()
.map(|s| s.to_string())
.collect(),
may_create_top_level_mailbox: true,
}
}
}

View File

@@ -7,16 +7,12 @@ use mail_parser::{
use crate::JMAP;
impl JMAP {
pub async fn blob_retrieve(
pub async fn blob_download(
&self,
blob_id: &BlobId,
account_id: u32,
) -> store::Result<Option<Vec<u8>>> {
if !self
.store
.has_blob_access(&blob_id.hash, vec![account_id])
.await?
{
if !blob_id.has_access(account_id) {
// TODO: validate ACL
let acl = "true";
return Ok(None);
@@ -26,7 +22,7 @@ impl JMAP {
Ok(self
.store
.get_blob(
&blob_id.hash,
&blob_id.kind,
(section.offset_start as u32)
..(section.offset_start.saturating_add(section.size) as u32),
)
@@ -37,7 +33,7 @@ impl JMAP {
Encoding::QuotedPrintable => quoted_printable_decode(&bytes),
}))
} else {
self.store.get_blob(&blob_id.hash, 0..u32::MAX).await
self.store.get_blob(&blob_id.kind, 0..u32::MAX).await
}
}
}

View File

@@ -1 +1,21 @@
pub mod get;
use jmap_proto::types::{blob::BlobId, id::Id};
pub mod download;
pub mod upload;
#[derive(Debug, serde::Serialize)]
pub struct UploadResponse {
#[serde(rename(serialize = "accountId"))]
account_id: Id,
#[serde(rename(serialize = "blobId"))]
blob_id: BlobId,
#[serde(rename(serialize = "type"))]
c_type: String,
size: usize,
}
pub struct DownloadResponse {
pub filename: String,
pub content_type: String,
pub blob: Vec<u8>,
}

View File

@@ -0,0 +1,31 @@
use jmap_proto::{
error::request::RequestError,
types::{blob::BlobId, id::Id},
};
use crate::JMAP;
use super::UploadResponse;
impl JMAP {
pub async fn blob_upload(
&self,
account_id: Id,
content_type: &str,
data: &[u8],
) -> Result<UploadResponse, RequestError> {
let blob_id = BlobId::temporary(account_id.document_id());
self.store
.put_blob(&blob_id.kind, data)
.await
.map_err(|err| RequestError::internal_server_error())?;
Ok(UploadResponse {
account_id,
blob_id,
c_type: content_type.to_string(),
size: data.len(),
})
}
}

View File

@@ -44,7 +44,7 @@ impl ToBodyPart for Vec<MessagePart<'_>> {
Property::BlobId if multipart.is_none() => {
let base_offset = blob_id.start_offset();
BlobId::new_section(
blob_id.hash,
blob_id.kind,
part.offset_body + base_offset,
part.offset_end + base_offset,
part.encoding as u8,

View File

@@ -123,7 +123,7 @@ impl JMAP {
};
// Retrieve raw message if needed
let blob_id = values.get(&Property::BlobId).as_blob()?;
let blob_id = BlobId::maildir(account_id, id.document_id());
let raw_message = if needs_body || needs_headers {
let offset = if !needs_body {
blob_id
@@ -135,7 +135,7 @@ impl JMAP {
u32::MAX
};
if let Some(raw_message) = self.store.get_blob(&blob_id.hash, 0..offset).await? {
if let Some(raw_message) = self.store.get_blob(&blob_id.kind, 0..offset).await? {
raw_message
} else {
let log = "true";
@@ -154,7 +154,6 @@ impl JMAP {
} else {
None
};
let blob_id = BlobId::new(blob_id.hash);
// Prepare response
let mut email = Object::with_capacity(properties.len());

View File

@@ -70,7 +70,7 @@ impl JMAP {
// Fetch raw message to import
let raw_message =
if let Some(raw_message) = self.blob_retrieve(&email.blob_id, account_id).await? {
if let Some(raw_message) = self.blob_download(&email.blob_id, account_id).await? {
raw_message
} else {
not_created.append(

View File

@@ -3,7 +3,6 @@ use std::borrow::Cow;
use jmap_proto::{
object::Object,
types::{
blob::{BlobId, BlobSection},
date::UTCDate,
keyword::Keyword,
property::{HeaderForm, Property},
@@ -18,7 +17,6 @@ use mail_parser::{
use store::{
fts::{builder::FtsIndexBuilder, Language},
write::{BatchBuilder, F_BITMAP, F_INDEX, F_VALUE},
BlobHash,
};
use crate::email::headers::IntoForm;
@@ -33,7 +31,6 @@ pub(super) trait IndexMessage {
fn index_message(
&mut self,
message: Message,
blob_hash: BlobHash,
keywords: Vec<Keyword>,
mailbox_ids: Vec<u32>,
received_at: u64,
@@ -41,22 +38,10 @@ pub(super) trait IndexMessage {
) -> store::Result<()>;
}
/*
o id
o blobId
o threadId
o mailboxIds
o keywords
o receivedAt
*/
impl IndexMessage for BatchBuilder {
fn index_message(
&mut self,
message: Message,
blob_hash: BlobHash,
keywords: Vec<Keyword>,
mailbox_ids: Vec<u32>,
received_at: u64,
@@ -64,20 +49,6 @@ impl IndexMessage for BatchBuilder {
) -> store::Result<()> {
let mut object = Object::with_capacity(15);
// Add blobHash and body offset
object.append(
Property::BlobId,
Value::BlobId(BlobId {
hash: blob_hash,
section: BlobSection {
offset_start: message.root_part().offset_body,
size: 0,
encoding: 0,
}
.into(),
}),
);
// Index keywords
self.value(
Property::Keywords,

View File

@@ -11,7 +11,7 @@ use mail_parser::{
use store::{
query::Filter,
write::{log::ChangeLogBuilder, now, BatchBuilder, F_BITMAP, F_CLEAR, F_VALUE},
BlobHash, ValueKey,
ValueKey,
};
use utils::map::vec_map::VecMap;
@@ -25,7 +25,7 @@ use super::index::{TrimTextValue, MAX_SORT_FIELD_LENGTH};
pub struct IngestedEmail {
pub id: Id,
pub change_id: u64,
pub blob_hash: BlobHash,
pub blob_id: BlobId,
pub size: usize,
}
@@ -88,9 +88,6 @@ impl JMAP {
None
};
// Store blob
let blob_hash = self.store.write_blob(account_id, raw_message).await?;
// Obtain a documentId and changeId
let document_id = self
.store
@@ -101,6 +98,10 @@ impl JMAP {
.assign_change_id(account_id, Collection::Email)
.await?;
// Store blob
let blob_id = BlobId::maildir(account_id, document_id);
self.store.put_blob(&blob_id.kind, raw_message).await?;
// Build change log
let mut changes = ChangeLogBuilder::with_change_id(change_id);
let thread_id = if let Some(thread_id) = thread_id {
@@ -124,7 +125,6 @@ impl JMAP {
let mut batch = BatchBuilder::new();
batch.index_message(
message,
blob_hash,
keywords,
mailbox_ids,
received_at.unwrap_or_else(now),
@@ -137,7 +137,7 @@ impl JMAP {
Ok(IngestedEmail {
id,
change_id,
blob_hash,
blob_id,
size: raw_message.len(),
})
}
@@ -259,6 +259,6 @@ impl From<IngestedEmail> for Object<Value> {
Object::with_capacity(3)
.with_property(Property::Id, email.id)
.with_property(Property::ThreadId, email.id.prefix_id())
.with_property(Property::BlobId, BlobId::new(email.blob_hash))
.with_property(Property::BlobId, email.blob_id)
}
}

View File

@@ -15,6 +15,7 @@ pub struct Config {
pub query_max_results: usize,
pub request_max_size: usize,
pub request_max_calls: usize,
pub upload_max_size: usize,
}
pub enum MaybeError {

View File

@@ -1,8 +1,8 @@
use rusqlite::{params, OptionalExtension};
use crate::{
write::{now, Batch, Operation},
AclKey, BitmapKey, BlobKey, IndexKey, LogKey, Serialize, Store, ValueKey,
write::{Batch, Operation},
AclKey, BitmapKey, IndexKey, LogKey, Serialize, Store, ValueKey,
};
use super::{BITS_MASK, BITS_PER_BLOCK};
@@ -163,30 +163,6 @@ impl Store {
.execute(params![bitmap_value_clear, &key])?;
};
}
Operation::Blob { key, set } => {
let key = BlobKey {
account_id,
collection,
document_id,
hash: key,
}
.serialize();
if *set {
let now_;
let value = if document_id != u32::MAX {
&[]
} else {
now_ = now().to_be_bytes();
&now_[..]
};
trx.prepare_cached("INSERT OR REPLACE INTO o (k, v) VALUES (?, ?)")?
.execute([&key[..], value])?;
} else {
trx.prepare_cached("DELETE FROM o WHERE k = ?")?
.execute([&key])?;
}
}
Operation::Acl {
grant_account_id,
set,

View File

@@ -1,36 +1,23 @@
pub mod purge;
//pub mod purge;
pub mod read;
pub mod write;
use std::{
io::Write,
path::{Path, PathBuf},
};
use std::path::{Path, PathBuf};
use utils::{codec::base32_custom::Base32Writer, config::Config};
use utils::config::Config;
use crate::{BlobHash, Serialize};
use crate::BlobKind;
pub enum BlobStore {
Local {
base_path: PathBuf,
hash_levels: usize,
},
Local(PathBuf),
Remote(String),
}
impl BlobStore {
pub async fn new(config: &Config) -> crate::Result<Self> {
Ok(BlobStore::Local {
base_path: config.value_require("blob.store.path")?.into(),
hash_levels: config.property("blob.store.hash")?.unwrap_or(1),
})
}
}
impl Serialize for &BlobHash {
fn serialize(self) -> Vec<u8> {
self.hash.to_vec()
Ok(BlobStore::Local(
config.value_require("blob.store.path")?.into(),
))
}
}
@@ -40,17 +27,41 @@ impl From<std::io::Error> for crate::Error {
}
}
fn get_path(base_path: &Path, hash_levels: usize, blob_id: &BlobHash) -> crate::Result<PathBuf> {
fn get_path(base_path: &Path, kind: &BlobKind) -> crate::Result<PathBuf> {
let mut path = base_path.to_path_buf();
let hash = &blob_id.hash;
for byte in hash.iter().take(hash_levels) {
path.push(format!("{:x}", byte));
match kind {
BlobKind::Linked {
account_id,
collection,
document_id,
} => {
path.push(format!("{:x}", account_id));
path.push(format!("{:x}", collection));
path.push(format!("{:x}", document_id));
}
BlobKind::LinkedMaildir {
account_id,
document_id,
} => {
path.push(format!("{:x}", account_id));
path.push("Maildir");
path.push("cur");
path.push(format!("{:x}", document_id));
}
BlobKind::Temporary {
account_id,
creation_year,
creation_month,
creation_day,
seq,
} => {
path.push("tmp");
path.push(creation_year.to_string());
path.push(creation_month.to_string());
path.push(creation_day.to_string());
path.push(format!("{:x}_{:x}", account_id, seq));
}
}
// Base32 encode the hash
let mut writer = Base32Writer::with_capacity(hash.len());
writer.write_all(hash).unwrap();
path.push(&writer.finalize());
Ok(path)
}

View File

@@ -82,7 +82,7 @@ impl Store {
}
for hash in results.delete {
self.blob.delete(&crate::BlobHash { hash }).await?;
self.blob.delete(&crate::BlobKind { hash }).await?;
}
Ok(())

View File

@@ -1,27 +1,23 @@
use std::{io::SeekFrom, ops::Range};
use roaring::RoaringBitmap;
use tokio::{
fs::{self, File},
io::{AsyncReadExt, AsyncSeekExt},
};
use crate::{write::key::DeserializeBigEndian, BlobHash, BlobKey, Store, BLOB_HASH_LEN};
use crate::{BlobKind, Store};
use super::{get_path, BlobStore};
impl Store {
pub async fn get_blob(
&self,
id: &BlobHash,
kind: &BlobKind,
range: Range<u32>,
) -> crate::Result<Option<Vec<u8>>> {
match &self.blob {
BlobStore::Local {
base_path,
hash_levels,
} => {
let blob_path = get_path(base_path, *hash_levels, id)?;
BlobStore::Local(base_path) => {
let blob_path = get_path(base_path, kind)?;
let blob_size = match fs::metadata(&blob_path).await {
Ok(m) => m.len(),
Err(_) => return Ok(None),
@@ -54,70 +50,4 @@ impl Store {
BlobStore::Remote(_) => todo!(),
}
}
pub async fn has_blob_access(
&self,
blob_hash: &BlobHash,
account_ids: Vec<u32>,
) -> crate::Result<bool> {
// Check if the blob already exists
let from_key = BlobKey {
account_id: 0,
collection: 0,
document_id: 0,
hash: blob_hash.hash,
};
let to_key = BlobKey {
account_id: u32::MAX,
collection: u8::MAX,
document_id: u32::MAX,
hash: blob_hash.hash,
};
self.iterate(false, from_key, to_key, true, false, move |acc, key, _| {
let account_id = key.deserialize_be_u32(BLOB_HASH_LEN)?;
if account_ids.contains(&account_id) {
*acc = true;
Ok(false)
} else {
Ok(true)
}
})
.await
}
pub async fn has_blob_access_doc(
&self,
blob_hash: &BlobHash,
account_id: u32,
collection: impl Into<u8>,
document_ids: RoaringBitmap,
) -> crate::Result<bool> {
// Check if the blob already exists
let collection = collection.into();
let from_key = BlobKey {
account_id,
collection,
document_id: 0,
hash: blob_hash.hash,
};
let to_key = BlobKey {
account_id,
collection,
document_id: u32::MAX,
hash: blob_hash.hash,
};
self.iterate(false, from_key, to_key, true, false, move |acc, key, _| {
let document_id =
key.deserialize_be_u32(BLOB_HASH_LEN + std::mem::size_of::<u32>() + 1)?;
if document_ids.contains(document_id) {
*acc = true;
Ok(false)
} else {
Ok(true)
}
})
.await
}
}

View File

@@ -3,66 +3,18 @@ use tokio::{
io::AsyncWriteExt,
};
use crate::{write::BatchBuilder, BlobHash, BlobKey, Store, BLOB_HASH_LEN};
use crate::{BlobKind, Store};
use super::{get_path, BlobStore};
impl Store {
pub async fn write_blob(&self, account_id: u32, data: &[u8]) -> crate::Result<BlobHash> {
let id = BlobHash::from(data);
pub async fn put_blob(&self, kind: &BlobKind, data: &[u8]) -> crate::Result<bool> {
match &self.blob {
BlobStore::Local(base_path) => {
let blob_path = get_path(base_path, kind)?;
// Check if the blob already exists
let from_key = BlobKey {
account_id: 0,
collection: 0,
document_id: 0,
hash: [0; BLOB_HASH_LEN],
};
let to_key = BlobKey {
account_id: u32::MAX,
collection: u8::MAX,
document_id: u32::MAX,
hash: id.hash,
};
let found = self
.iterate(false, from_key, to_key, true, false, |acc, _, _| {
*acc = true;
Ok(false)
})
.await?;
if !found {
// Write the blob
self.blob.put(&id, data).await?;
// Write a temporary link to the blob
self.write(
BatchBuilder::new()
.with_account_id(account_id)
.with_collection(u8::MAX)
.update_document(u32::MAX)
.blob(&id, 0)
.build_batch(),
)
.await?;
}
Ok(id)
}
}
impl BlobStore {
pub async fn put(&self, id: &BlobHash, data: &[u8]) -> crate::Result<bool> {
match self {
BlobStore::Local {
base_path,
hash_levels,
} => {
let blob_path = get_path(base_path, *hash_levels, id)?;
if blob_path.exists() {
let metadata = fs::metadata(&blob_path).await?;
let metadata = fs::metadata(&blob_path).await;
if let Ok(metadata) = metadata {
if metadata.len() as usize == data.len() {
return Ok(false);
}
@@ -79,13 +31,29 @@ impl BlobStore {
}
}
pub async fn delete(&self, id: &BlobHash) -> crate::Result<bool> {
match self {
BlobStore::Local {
base_path,
hash_levels,
} => {
let blob_path = get_path(base_path, *hash_levels, id)?;
pub async fn copy_blob(&self, src: &BlobKind, dest: &BlobKind) -> crate::Result<bool> {
match &self.blob {
BlobStore::Local(base_path) => {
let src_path = get_path(base_path, src)?;
let dest_path = get_path(base_path, dest)?;
if fs::metadata(&src_path).await.is_err() {
return Ok(false);
}
fs::create_dir_all(dest_path.parent().unwrap()).await?;
fs::copy(src_path, dest_path).await?;
Ok(true)
}
BlobStore::Remote(_) => todo!(),
}
}
pub async fn delete_blob(&self, kind: &BlobKind) -> crate::Result<bool> {
match &self.blob {
BlobStore::Local(base_path) => {
let blob_path = get_path(base_path, kind)?;
if blob_path.exists() {
fs::remove_file(&blob_path).await?;
@@ -98,13 +66,3 @@ impl BlobStore {
}
}
}
impl From<&[u8]> for BlobHash {
fn from(data: &[u8]) -> Self {
let mut hasher = blake3::Hasher::new();
hasher.update(data);
Self {
hash: hasher.finalize().into(),
}
}
}

View File

@@ -9,6 +9,7 @@ pub mod query;
pub mod write;
pub use ahash;
pub use rand;
pub use roaring;
#[cfg(feature = "rocks")]
@@ -99,14 +100,6 @@ pub struct ValueKey {
pub field: u8,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct BlobKey<T: AsRef<[u8]>> {
pub account_id: u32,
pub collection: u8,
pub document_id: u32,
pub hash: T,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct AclKey {
pub grant_account_id: u32,
@@ -123,16 +116,23 @@ pub struct LogKey {
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct BlobHash {
pub hash: [u8; BLOB_HASH_LEN],
}
impl Default for BlobHash {
fn default() -> Self {
Self {
hash: [0; BLOB_HASH_LEN],
}
}
pub enum BlobKind {
Linked {
account_id: u32,
collection: u8,
document_id: u32,
},
LinkedMaildir {
account_id: u32,
document_id: u32,
},
Temporary {
account_id: u32,
creation_year: u16,
creation_month: u8,
creation_day: u8,
seq: u32,
},
}
pub type Result<T> = std::result::Result<T, Error>;
@@ -175,8 +175,6 @@ pub const TAG_ID: u8 = 0;
pub const TAG_TEXT: u8 = 1 << 0;
pub const TAG_STATIC: u8 = 1 << 1;
pub const BLOB_HASH_LEN: usize = 32;
pub const SUBSPACE_BITMAPS: u8 = b'b';
pub const SUBSPACE_VALUES: u8 = b'v';
pub const SUBSPACE_LOGS: u8 = b'l';

View File

@@ -125,14 +125,6 @@ impl BatchBuilder {
self
}
pub fn blob(&mut self, blob_id: impl Serialize, options: u32) -> &mut Self {
self.ops.push(Operation::Blob {
key: blob_id.serialize(),
set: !options.has_flag(F_CLEAR),
});
self
}
pub fn custom(&mut self, value: impl IntoOperations) -> crate::Result<()> {
value.build(self)
}

View File

@@ -1,10 +1,7 @@
use std::convert::TryInto;
use utils::codec::leb128::Leb128_;
use crate::{
AclKey, BitmapKey, BlobKey, IndexKey, IndexKeyPrefix, Key, LogKey, Serialize, ValueKey,
BLOB_HASH_LEN,
};
use crate::{AclKey, BitmapKey, IndexKey, IndexKeyPrefix, LogKey, Serialize, ValueKey};
pub struct KeySerializer {
buf: Vec<u8>,
@@ -119,12 +116,6 @@ impl DeserializeBigEndian for &[u8] {
}
}
impl<T: AsRef<[u8]> + Sync + Send + 'static> Key for BlobKey<T> {
fn subspace(&self) -> u8 {
crate::SUBSPACE_BLOBS
}
}
impl ValueKey {
pub fn new(
account_id: u32,
@@ -232,30 +223,6 @@ impl<T: AsRef<[u8]>> Serialize for &BitmapKey<T> {
}
}
impl<T: AsRef<[u8]>> Serialize for &BlobKey<T> {
fn serialize(self) -> Vec<u8> {
let hash = self.hash.as_ref();
#[cfg(feature = "key_subspace")]
{
KeySerializer::new(std::mem::size_of::<BlobKey<T>>() + BLOB_HASH_LEN + 1)
.write(crate::SUBSPACE_BLOBS)
}
#[cfg(not(feature = "key_subspace"))]
{ KeySerializer::new(std::mem::size_of::<BlobKey<T>>() + BLOB_HASH_LEN) }
.write(hash)
.write(self.account_id)
.write(self.collection)
.write(self.document_id)
.finalize()
}
}
impl<T: AsRef<[u8]>> Serialize for BlobKey<T> {
fn serialize(self) -> Vec<u8> {
(&self).serialize()
}
}
impl Serialize for &AclKey {
fn serialize(self) -> Vec<u8> {
#[cfg(feature = "key_subspace")]

View File

@@ -55,10 +55,6 @@ pub enum Operation {
key: Vec<u8>,
set: bool,
},
Blob {
key: Vec<u8>,
set: bool,
},
Acl {
grant_account_id: u32,
set: Option<Vec<u8>>,

View File

@@ -4,4 +4,9 @@ version = "0.1.0"
edition = "2021"
[dependencies]
rustls = "0.21.0"
rustls-pemfile = "1.0"
tokio = { version = "1.23", features = ["net"] }
tokio-rustls = { version = "0.24.0"}
serde = { version = "1.0", features = ["derive"]}
tracing = "0.1"

View File

@@ -178,3 +178,9 @@ impl_unsigned_leb128!(u16, [0, 7, 14]);
impl_unsigned_leb128!(u32, [0, 7, 14, 21, 28]);
impl_unsigned_leb128!(u64, [0, 7, 14, 21, 28, 35, 42, 49, 56, 63]);
impl_unsigned_leb128!(usize, [0, 7, 14, 21, 28, 35, 42, 49, 56, 63]);
impl Leb128Writer for Vec<u8> {
fn write_leb128<T: Leb128_>(&mut self, value: T) -> std::io::Result<usize> {
T::to_leb128_writer(value, self)
}
}

View File

@@ -0,0 +1,97 @@
/*
* Copyright (c) 2023 Stalwart Labs Ltd.
*
* This file is part of the Stalwart SMTP Server.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as
* published by the Free Software Foundation, either version 3 of
* the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
* in the LICENSE file at the top-level directory of this distribution.
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*
* You can be released from the requirements of the AGPLv3 license by
* purchasing a commercial license. Please contact licensing@stalw.art
* for more details.
*/
use std::{io::Cursor, sync::Arc};
use rustls::{
server::{ClientHello, ResolvesServerCert, ResolvesServerCertUsingSni},
sign::CertifiedKey,
version::{TLS12, TLS13},
Certificate, PrivateKey, SupportedProtocolVersion,
};
use rustls_pemfile::{certs, read_one, Item};
use super::Config;
pub static TLS13_VERSION: &[&SupportedProtocolVersion] = &[&TLS13];
pub static TLS12_VERSION: &[&SupportedProtocolVersion] = &[&TLS12];
pub struct CertificateResolver {
pub resolver: Option<ResolvesServerCertUsingSni>,
pub default_cert: Option<Arc<CertifiedKey>>,
}
impl ResolvesServerCert for CertificateResolver {
fn resolve(&self, hello: ClientHello<'_>) -> Option<Arc<CertifiedKey>> {
self.resolver
.as_ref()
.and_then(|r| r.resolve(hello))
.or_else(|| self.default_cert.clone())
}
}
impl Config {
pub fn rustls_certificate(&self, cert_id: &str) -> super::Result<Vec<Certificate>> {
let certs = certs(&mut Cursor::new(self.file_contents((
"certificate",
cert_id,
"cert",
))?))
.map_err(|err| {
format!("Failed to read certificates in \"certificate.{cert_id}.cert\": {err}")
})?
.into_iter()
.map(Certificate)
.collect::<Vec<_>>();
if !certs.is_empty() {
Ok(certs)
} else {
Err(format!(
"No certificates found in \"certificate.{cert_id}.cert\"."
))
}
}
pub fn rustls_private_key(&self, cert_id: &str) -> super::Result<PrivateKey> {
match read_one(&mut Cursor::new(self.file_contents((
"certificate",
cert_id,
"private-key",
))?))
.map_err(|err| {
format!("Failed to read private keys in \"certificate.{cert_id}.private-key\": {err}",)
})?
.into_iter()
.next()
{
Some(Item::PKCS8Key(key) | Item::RSAKey(key) | Item::ECKey(key)) => Ok(PrivateKey(key)),
Some(_) => Err(format!(
"Unsupported private keys found in \"certificate.{cert_id}.private-key\".",
)),
None => Err(format!(
"No private keys found in \"certificate.{cert_id}.private-key\".",
)),
}
}
}

View File

@@ -0,0 +1,504 @@
/*
* Copyright (c) 2023 Stalwart Labs Ltd.
*
* This file is part of the Stalwart SMTP Server.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as
* published by the Free Software Foundation, either version 3 of
* the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
* in the LICENSE file at the top-level directory of this distribution.
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*
* You can be released from the requirements of the AGPLv3 license by
* purchasing a commercial license. Please contact licensing@stalw.art
* for more details.
*/
use std::{net::SocketAddr, sync::Arc, time::Duration};
use rustls::{
cipher_suite::{
TLS13_AES_128_GCM_SHA256, TLS13_AES_256_GCM_SHA384, TLS13_CHACHA20_POLY1305_SHA256,
TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256, TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256, TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384, TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
},
server::{NoClientAuth, ResolvesServerCertUsingSni},
sign::{any_supported_type, CertifiedKey},
ServerConfig, SupportedCipherSuite, ALL_CIPHER_SUITES, ALL_KX_GROUPS, ALL_VERSIONS,
};
use tokio::net::TcpSocket;
use crate::UnwrapFailure;
use super::{
certificate::{CertificateResolver, TLS12_VERSION, TLS13_VERSION},
utils::{AsKey, ParseKey, ParseValue},
Config, Listener, Server, ServerProtocol, Servers,
};
impl Config {
pub fn parse_servers(&self) -> super::Result<Servers> {
let mut servers: Vec<Server> = Vec::new();
for (internal_id, id) in self.sub_keys("server.listener").enumerate() {
let mut server = self.parse_server(id)?;
if !servers.iter().any(|s| s.id == server.id) {
server.internal_id = internal_id as u16;
servers.push(server);
} else {
return Err(format!("Duplicate listener id {:?}.", server.id));
}
}
if !servers.is_empty() {
Ok(Servers { inner: servers })
} else {
Err("No server directives found in config file.".to_string())
}
}
fn parse_server(&self, id: &str) -> super::Result<Server> {
// Build TLS config
let (tls, tls_implicit) = if self
.property_or_default(("server.listener", id, "tls.enable"), "server.tls.enable")?
.unwrap_or(false)
{
// Parse protocol versions
let mut tls_v2 = false;
let mut tls_v3 = false;
for (key, protocol) in self.values_or_default(
("server.listener", id, "tls.protocols"),
"server.tls.protocols",
) {
match protocol {
"TLSv1.2" | "0x0303" => tls_v2 = true,
"TLSv1.3" | "0x0304" => tls_v3 = true,
protocol => {
return Err(format!(
"Unsupported TLS protocol {protocol:?} found in key {key:?}",
))
}
}
}
// Parse cipher suites
let mut ciphers = Vec::new();
for (key, protocol) in
self.values_or_default(("server.listener", id, "tls.ciphers"), "server.tls.ciphers")
{
ciphers.push(protocol.parse_key(key)?);
}
// Obtain default certificate
let cert_id = self
.value_or_default(
("server.listener", id, "tls.certificate"),
"server.tls.certificate",
)
.ok_or_else(|| format!("Undefined certificate id for listener {id:?}."))?;
let cert = self.rustls_certificate(cert_id)?;
let pki = self.rustls_private_key(cert_id)?;
// Add SNI certificates
let mut resolver = ResolvesServerCertUsingSni::new();
let mut has_sni = false;
for (key, value) in
self.values_or_default(("server.listener", id, "tls.sni"), "server.tls.sni")
{
if let Some(prefix) = key.strip_suffix(".subject") {
has_sni = true;
resolver
.add(
value,
match self.value((prefix, "certificate")) {
Some(sni_cert_id) if sni_cert_id != cert_id => CertifiedKey {
cert: self.rustls_certificate(sni_cert_id)?,
key: any_supported_type(&self.rustls_private_key(sni_cert_id)?)
.map_err(|err| {
format!(
"Failed to sign SNI certificate for {key:?}: {err}",
)
})?,
ocsp: None,
sct_list: None,
},
_ => CertifiedKey {
cert: cert.clone(),
key:
any_supported_type(&pki).map_err(|err| {
format!(
"Failed to sign SNI certificate for {key:?}: {err}",
)
})?,
ocsp: None,
sct_list: None,
},
},
)
.map_err(|err| {
format!("Failed to add SNI certificate for {key:?}: {err}")
})?;
}
}
// Add default certificate
let default_cert = Some(Arc::new(CertifiedKey {
cert,
key: any_supported_type(&pki)
.map_err(|err| format!("Failed to sign certificate id {cert_id:?}: {err}"))?,
ocsp: None,
sct_list: None,
}));
// Build server config
let mut config = ServerConfig::builder()
.with_cipher_suites(if !ciphers.is_empty() {
&ciphers
} else {
ALL_CIPHER_SUITES
})
.with_kx_groups(&ALL_KX_GROUPS)
.with_protocol_versions(if tls_v3 == tls_v2 {
ALL_VERSIONS
} else if tls_v3 {
TLS13_VERSION
} else {
TLS12_VERSION
})
.map_err(|err| format!("Failed to build TLS config: {err}"))?
.with_client_cert_verifier(NoClientAuth::boxed())
.with_cert_resolver(Arc::new(CertificateResolver {
resolver: if has_sni { resolver.into() } else { None },
default_cert,
}));
//config.key_log = Arc::new(KeyLogger::default());
config.ignore_client_order = self
.property_or_default(
("server.listener", id, "tls.ignore-client-order"),
"server.tls.ignore-client-order",
)?
.unwrap_or(true);
(
config.into(),
self.property_or_default(
("server.listener", id, "tls.implicit"),
"server.tls.implicit",
)?
.unwrap_or(true),
)
} else {
(None, false)
};
// Build listeners
let mut listeners = Vec::new();
for result in self.properties::<SocketAddr>(("server.listener", id, "bind")) {
// Parse bind address and build socket
let (_, addr) = result?;
let socket = if addr.is_ipv4() {
TcpSocket::new_v4()
} else {
TcpSocket::new_v6()
}
.map_err(|err| format!("Failed to create socket: {err}"))?;
let mut backlog = None;
let mut ttl = None;
// Set socket options
for option in [
"reuse-addr",
"reuse-port",
"send-buffer-size",
"recv-buffer-size",
"linger",
"tos",
"backlog",
"ttl",
] {
if let Some(value) = self.value_or_default(
("server.listener", id, "socket", option),
("server.socket", option),
) {
let key = ("server.listener", id, "socket", option);
match option {
"reuse-addr" => socket.set_reuseaddr(value.parse_key(key)?),
#[cfg(not(target_env = "msvc"))]
"reuse-port" => socket.set_reuseport(value.parse_key(key)?),
"send-buffer-size" => socket.set_send_buffer_size(value.parse_key(key)?),
"recv-buffer-size" => socket.set_recv_buffer_size(value.parse_key(key)?),
"linger" => {
socket.set_linger(Duration::from_millis(value.parse_key(key)?).into())
}
"tos" => socket.set_tos(value.parse_key(key)?),
"backlog" => {
backlog = Some(value.parse_key(key)?);
continue;
}
"ttl" => {
ttl = Some(value.parse_key(key)?);
continue;
}
_ => unreachable!(),
}
.map_err(|err| {
format!("Failed to set socket option '{option}' for listener '{id}': {err}")
})?;
}
}
listeners.push(Listener {
socket,
addr,
ttl,
backlog,
});
}
if listeners.is_empty() {
return Err(format!("No 'bind' directive found for listener id {id:?}"));
}
let protocol = self
.property_or_default(("server.listener", id, "protocol"), "server.protocol")?
.unwrap_or(ServerProtocol::Smtp);
Ok(Server {
id: id.to_string(),
internal_id: 0,
hostname: self
.value_or_default(("server.listener", id, "hostname"), "server.hostname")
.ok_or("Hostname directive not found.")?
.to_string(),
data: if matches!(protocol, ServerProtocol::Smtp | ServerProtocol::Lmtp) {
self.value_or_default(("server.listener", id, "data"), "server.data")
.unwrap_or("Stalwart SMTP at your service")
.to_string()
} else {
self.value_or_default(("server.listener", id, "url"), "server.url")
.failed(&format!("No 'url' directive found for listener {id:?}"))
.to_string()
},
protocol,
listeners,
tls,
tls_implicit,
})
}
}
impl ParseValue for ServerProtocol {
fn parse_value(key: impl AsKey, value: &str) -> super::Result<Self> {
if value.eq_ignore_ascii_case("smtp") {
Ok(Self::Smtp)
} else if value.eq_ignore_ascii_case("lmtp") {
Ok(Self::Lmtp)
} else if value.eq_ignore_ascii_case("jmap") {
Ok(Self::Jmap)
} else if value.eq_ignore_ascii_case("imap") {
Ok(Self::Imap)
} else if value.eq_ignore_ascii_case("http") {
Ok(Self::Http)
} else {
Err(format!(
"Invalid server protocol type {:?} for property {:?}.",
value,
key.as_key()
))
}
}
}
impl ParseValue for SocketAddr {
fn parse_value(key: impl AsKey, value: &str) -> super::Result<Self> {
value.parse().map_err(|_| {
format!(
"Invalid socket address {:?} for property {:?}.",
value,
key.as_key()
)
})
}
}
impl ParseValue for SupportedCipherSuite {
fn parse_value(key: impl AsKey, value: &str) -> super::Result<Self> {
Ok(match value {
// TLS1.3 suites
"TLS13_AES_256_GCM_SHA384" => TLS13_AES_256_GCM_SHA384,
"TLS13_AES_128_GCM_SHA256" => TLS13_AES_128_GCM_SHA256,
"TLS13_CHACHA20_POLY1305_SHA256" => TLS13_CHACHA20_POLY1305_SHA256,
// TLS1.2 suites
"TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384" => TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
"TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256" => TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
"TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256" => {
TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256
}
"TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384" => TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
"TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256" => TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
"TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256" => {
TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256
}
cipher => {
return Err(format!(
"Unsupported TLS cipher suite {:?} found in key {:?}",
cipher,
key.as_key()
))
}
})
}
}
#[cfg(test)]
mod tests {
use std::{fs, path::PathBuf};
use tokio::net::TcpSocket;
use crate::config::{Config, Listener, Server, ServerProtocol};
fn add_test_certs(config: &str) -> String {
let mut cert_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
cert_path.push("resources");
cert_path.push("tests");
cert_path.push("certs");
let mut cert = cert_path.clone();
cert.push("tls_cert.pem");
let mut pk = cert_path.clone();
pk.push("tls_privatekey.pem");
config
.replace("{CERT}", cert.as_path().to_str().unwrap())
.replace("{PK}", pk.as_path().to_str().unwrap())
}
#[test]
fn parse_servers() {
let mut file = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
file.push("resources");
file.push("tests");
file.push("config");
file.push("servers.toml");
let toml = add_test_certs(&fs::read_to_string(file).unwrap());
// Parse servers
let config = Config::parse(&toml).unwrap();
let servers = config.parse_servers().unwrap();
let expected_servers = vec![
Server {
id: "smtp".to_string(),
internal_id: 0,
hostname: "mx.example.org".to_string(),
data: "Stalwart SMTP - hi there!".to_string(),
protocol: ServerProtocol::Smtp,
listeners: vec![Listener {
socket: TcpSocket::new_v4().unwrap(),
addr: "127.0.0.1:9925".parse().unwrap(),
ttl: 3600.into(),
backlog: 1024.into(),
}],
tls: None,
tls_implicit: false,
},
Server {
id: "smtps".to_string(),
internal_id: 1,
hostname: "mx.example.org".to_string(),
data: "Stalwart SMTP - hi there!".to_string(),
protocol: ServerProtocol::Smtp,
listeners: vec![
Listener {
socket: TcpSocket::new_v4().unwrap(),
addr: "127.0.0.1:9465".parse().unwrap(),
ttl: 4096.into(),
backlog: 1024.into(),
},
Listener {
socket: TcpSocket::new_v4().unwrap(),
addr: "127.0.0.1:9466".parse().unwrap(),
ttl: 4096.into(),
backlog: 1024.into(),
},
],
tls: None,
tls_implicit: true,
},
Server {
id: "submission".to_string(),
internal_id: 2,
hostname: "submit.example.org".to_string(),
data: "Stalwart SMTP submission at your service".to_string(),
protocol: ServerProtocol::Smtp,
listeners: vec![Listener {
socket: TcpSocket::new_v4().unwrap(),
addr: "127.0.0.1:9991".parse().unwrap(),
ttl: 3600.into(),
backlog: 2048.into(),
}],
tls: None,
tls_implicit: true,
},
];
for (server, expected_server) in servers.inner.into_iter().zip(expected_servers) {
assert_eq!(
server.id, expected_server.id,
"failed for {}",
expected_server.id
);
assert_eq!(
server.internal_id, expected_server.internal_id,
"failed for {}",
expected_server.id
);
assert_eq!(
server.hostname, expected_server.hostname,
"failed for {}",
expected_server.id
);
assert_eq!(
server.data, expected_server.data,
"failed for {}",
expected_server.id
);
assert_eq!(
server.protocol, expected_server.protocol,
"failed for {}",
expected_server.id
);
assert_eq!(
server.tls_implicit, expected_server.tls_implicit,
"failed for {}",
expected_server.id
);
for (listener, expected_listener) in
server.listeners.into_iter().zip(expected_server.listeners)
{
assert_eq!(
listener.addr, expected_listener.addr,
"failed for {}",
expected_server.id
);
assert_eq!(
listener.ttl, expected_listener.ttl,
"failed for {}",
expected_server.id
);
assert_eq!(
listener.backlog, expected_listener.backlog,
"failed for {}",
expected_server.id
);
}
}
}
}

View File

@@ -21,14 +21,65 @@
* for more details.
*/
pub mod certificate;
pub mod listener;
pub mod parser;
pub mod utils;
use std::collections::BTreeMap;
use std::{collections::BTreeMap, fmt::Display, net::SocketAddr};
use rustls::ServerConfig;
use tokio::net::TcpSocket;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Config {
keys: BTreeMap<String, String>,
}
#[derive(Debug, Default)]
pub struct Server {
pub id: String,
pub internal_id: u16,
pub hostname: String,
pub data: String,
pub protocol: ServerProtocol,
pub listeners: Vec<Listener>,
pub tls: Option<ServerConfig>,
pub tls_implicit: bool,
}
pub struct Servers {
pub inner: Vec<Server>,
}
#[derive(Debug)]
pub struct Listener {
pub socket: TcpSocket,
pub addr: SocketAddr,
pub ttl: Option<u32>,
pub backlog: Option<u32>,
}
#[derive(Debug, PartialEq, Eq, Clone, Copy, Default)]
pub enum ServerProtocol {
#[default]
Smtp,
Lmtp,
Jmap,
Imap,
Http,
}
impl Display for ServerProtocol {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ServerProtocol::Smtp => write!(f, "smtp"),
ServerProtocol::Lmtp => write!(f, "lmtp"),
ServerProtocol::Jmap => write!(f, "jmap"),
ServerProtocol::Imap => write!(f, "imap"),
ServerProtocol::Http => write!(f, "http"),
}
}
}
pub type Result<T> = std::result::Result<T, String>;

View File

@@ -23,4 +23,38 @@
pub mod codec;
pub mod config;
pub mod listener;
pub mod map;
pub trait UnwrapFailure<T> {
fn failed(self, action: &str) -> T;
}
impl<T> UnwrapFailure<T> for Option<T> {
fn failed(self, message: &str) -> T {
match self {
Some(result) => result,
None => {
eprintln!("{message}");
std::process::exit(1);
}
}
}
}
impl<T, E: std::fmt::Display> UnwrapFailure<T> for Result<T, E> {
fn failed(self, message: &str) -> T {
match self {
Ok(result) => result,
Err(err) => {
eprintln!("{message}: {err}");
std::process::exit(1);
}
}
}
}
pub fn failed(message: &str) -> ! {
eprintln!("{message}");
std::process::exit(1);
}

View File

@@ -0,0 +1,96 @@
use std::{
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
time::{Duration, Instant},
};
#[derive(Debug)]
pub struct RateLimiter {
pub max_requests: f64,
pub max_interval: f64,
limiter: (Instant, f64),
}
#[derive(Debug, Clone)]
pub struct ConcurrencyLimiter {
pub max_concurrent: u64,
pub concurrent: Arc<AtomicU64>,
}
pub struct InFlight {
concurrent: Arc<AtomicU64>,
}
impl Drop for InFlight {
fn drop(&mut self) {
self.concurrent.fetch_sub(1, Ordering::Relaxed);
}
}
impl RateLimiter {
pub fn new(max_requests: u64, max_interval: u64) -> Self {
RateLimiter {
max_requests: max_requests as f64,
max_interval: max_interval as f64,
limiter: (Instant::now(), max_requests as f64),
}
}
pub fn is_allowed(&mut self) -> bool {
// Check rate limit
let elapsed = self.limiter.0.elapsed().as_secs_f64();
self.limiter.1 += elapsed * (self.max_requests / self.max_interval);
if self.limiter.1 > self.max_requests {
self.limiter.1 = self.max_requests;
}
if self.limiter.1 >= 1.0 {
self.limiter.0 = Instant::now();
self.limiter.1 -= 1.0;
true
} else {
false
}
}
pub fn retry_at(&self) -> Instant {
Instant::now()
+ Duration::from_secs(
(self.max_interval as u64).saturating_sub(self.limiter.0.elapsed().as_secs()),
)
}
pub fn elapsed(&self) -> Duration {
self.limiter.0.elapsed()
}
pub fn reset(&mut self) {
self.limiter = (Instant::now(), self.max_requests);
}
}
impl ConcurrencyLimiter {
pub fn new(max_concurrent: u64) -> Self {
ConcurrencyLimiter {
max_concurrent,
concurrent: Arc::new(0.into()),
}
}
pub fn is_allowed(&self) -> Option<InFlight> {
if self.concurrent.load(Ordering::Relaxed) < self.max_concurrent {
// Return in-flight request
self.concurrent.fetch_add(1, Ordering::Relaxed);
Some(InFlight {
concurrent: self.concurrent.clone(),
})
} else {
None
}
}
pub fn check_is_allowed(&self) -> bool {
self.concurrent.load(Ordering::Relaxed) < self.max_concurrent
}
}

View File

@@ -0,0 +1,151 @@
use std::sync::Arc;
use tokio::{net::TcpListener, sync::watch};
use tokio_rustls::TlsAcceptor;
use crate::{
config::{Listener, Server, ServerProtocol, Servers},
failed,
listener::SessionData,
UnwrapFailure,
};
use super::{limiter::ConcurrencyLimiter, ServerInstance, SessionManager};
impl Server {
pub fn spawn(
self,
manager: impl SessionManager,
max_concurrent: u64,
shutdown_rx: watch::Receiver<bool>,
) -> Result<(), String> {
// Prepare instance
let instance = Arc::new(ServerInstance {
data: if matches!(self.protocol, ServerProtocol::Smtp | ServerProtocol::Lmtp) {
format!("220 {} {}\r\n", self.hostname, self.data)
} else {
self.data
},
id: self.id,
listener_id: self.internal_id,
protocol: self.protocol,
hostname: self.hostname,
tls_acceptor: self.tls.map(|config| TlsAcceptor::from(Arc::new(config))),
is_tls_implicit: self.tls_implicit,
});
// Start concurrency limiter
let limiter = Arc::new(ConcurrencyLimiter::new(max_concurrent));
// Spawn listeners
for listener in self.listeners {
tracing::info!(
id = instance.id,
protocol = ?instance.protocol,
bind.ip = listener.addr.ip().to_string(),
bind.port = listener.addr.port(),
tls = instance.is_tls_implicit,
"Starting listener"
);
let local_ip = listener.addr.ip();
// Bind socket
let listener = listener.listen();
// Spawn listener
let mut shutdown_rx = shutdown_rx.clone();
let manager = manager.clone();
let instance = instance.clone();
let limiter = limiter.clone();
tokio::spawn(async move {
loop {
tokio::select! {
stream = listener.accept() => {
match stream {
Ok((stream, remote_addr)) => {
// Enforce concurrency
if let Some(in_flight) = limiter.is_allowed() {
let span = tracing::info_span!(
"session",
instance = instance.id,
protocol = ?instance.protocol,
remote.ip = remote_addr.ip().to_string(),
remote.port = remote_addr.port(),
);
// Spawn connection
manager.spawn(SessionData {
stream,
local_ip,
remote_ip: remote_addr.ip(),
span,
in_flight,
instance: instance.clone(),
shutdown_rx: shutdown_rx.clone(),
});
} else {
tracing::info!(
context = "throttle",
event = "too-many-requests",
instance = instance.id,
protocol = ?instance.protocol,
remote.ip = remote_addr.ip().to_string(),
remote.port = remote_addr.port(),
max_concurrent = max_concurrent,
"Too many concurrent connections."
);
};
}
Err(err) => {
tracing::debug!(context = "io",
event = "error",
instance = instance.id,
protocol = ?instance.protocol,
"Failed to accept TCP connection: {}", err);
}
}
},
_ = shutdown_rx.changed() => {
tracing::debug!(
event = "shutdown",
instance = instance.id,
protocol = ?instance.protocol,
"Listener shutting down.");
break;
}
};
}
});
}
Ok(())
}
}
impl Servers {
pub fn bind(&self) {
for server in &self.inner {
for listener in &server.listeners {
listener
.socket
.bind(listener.addr)
.failed(&format!("Failed to bind to {}", listener.addr));
}
}
}
}
impl Listener {
pub fn listen(self) -> TcpListener {
let listener = self
.socket
.listen(self.backlog.unwrap_or(1024))
.unwrap_or_else(|err| failed(&format!("Failed to listen on {}: {}", self.addr, err)));
if let Some(ttl) = self.ttl {
listener.set_ttl(ttl).unwrap_or_else(|err| {
failed(&format!("Failed to set TTL on {}: {}", self.addr, err))
});
}
listener
}
}

View File

@@ -0,0 +1,39 @@
use std::{net::IpAddr, sync::Arc};
use tokio::{
io::{AsyncRead, AsyncWrite},
net::TcpStream,
sync::watch,
};
use tokio_rustls::TlsAcceptor;
use crate::config::ServerProtocol;
use self::limiter::InFlight;
pub mod limiter;
pub mod listen;
pub struct ServerInstance {
pub id: String,
pub listener_id: u16,
pub protocol: ServerProtocol,
pub hostname: String,
pub data: String,
pub tls_acceptor: Option<TlsAcceptor>,
pub is_tls_implicit: bool,
}
pub struct SessionData<T: AsyncRead + AsyncWrite + Unpin + 'static> {
pub stream: T,
pub local_ip: IpAddr,
pub remote_ip: IpAddr,
pub span: tracing::Span,
pub in_flight: InFlight,
pub instance: Arc<ServerInstance>,
pub shutdown_rx: watch::Receiver<bool>,
}
pub trait SessionManager: Sync + Send + 'static + Clone {
fn spawn(&self, session: SessionData<TcpStream>);
}

View File

@@ -1,20 +1,18 @@
use std::{sync::Arc, time::Duration};
use std::sync::Arc;
use store::ahash::AHashMap;
use store::{
write::{BatchBuilder, F_CLEAR},
BlobHash, BlobKey, Store, BLOB_HASH_LEN,
};
use store::Store;
pub async fn test(db: Arc<Store>) {
unimplemented!()
}
/*
let ttl = 1_u64;
let blob_1 = vec![b'a'; 1024];
let blob_2 = vec![b'b'; 1024];
let blob_id_1 = BlobHash::from(&blob_1[..]);
let blob_id_2 = BlobHash::from(&blob_2[..]);
let blob_id_1 = BlobKind::from(&blob_1[..]);
let blob_id_2 = BlobKind::from(&blob_2[..]);
// Insert the same blobs concurrently
let handles = (1..=100)
@@ -91,13 +89,13 @@ pub async fn test(db: Arc<Store>) {
}
struct BlobPurge {
result: AHashMap<BlobHash, (u32, u32)>,
result: AHashMap<BlobKind, (u32, u32)>,
link_count: u32,
ephemeral_count: u32,
id: [u8; BLOB_HASH_LEN],
}
async fn get_all_blobs(store: &Store) -> AHashMap<BlobHash, (u32, u32)> {
async fn get_all_blobs(store: &Store) -> AHashMap<BlobKind, (u32, u32)> {
let results = BlobPurge {
result: AHashMap::new(),
id: [0u8; BLOB_HASH_LEN],
@@ -122,7 +120,7 @@ async fn get_all_blobs(store: &Store) -> AHashMap<BlobHash, (u32, u32)> {
.iterate(results, from_key, to_key, false, true, move |b, k, v| {
if !k.starts_with(&b.id) {
if b.link_count != u32::MAX {
let id = BlobHash { hash: b.id };
let id = BlobKind { hash: b.id };
b.result.insert(id, (b.link_count, b.ephemeral_count));
}
b.link_count = 0;
@@ -142,9 +140,10 @@ async fn get_all_blobs(store: &Store) -> AHashMap<BlobHash, (u32, u32)> {
.unwrap();
if b.link_count != u32::MAX {
let id = BlobHash { hash: b.id };
let id = BlobKind { hash: b.id };
b.result.insert(id, (b.link_count, b.ephemeral_count));
}
b.result
}
*/