use crate::{OP_ID_PREFIX, PROTOCOL_VERSION};
use base64::Engine as _;
use chrono::{DateTime, Utc};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
use std::fmt;
use uuid::Uuid;
pub const BODY_TEXT_MAX_BYTES: usize = 1024 * 1024;
pub const IMAGE_DATA_MAX_BYTES: usize = 2 * 1024 * 1024;
pub const IMAGE_DATA_MAX_ENCODED_BYTES: usize = IMAGE_DATA_MAX_BYTES * 4 / 3 + 1024;
pub const IMAGE_MIMES: [&str; 4] = ["image/png", "image/jpeg", "image/gif", "image/webp"];
pub const CAUSALITY_LABEL_MAX_ENTRIES: usize = 8;
pub const CAUSALITY_LABEL_KEY_MAX_BYTES: usize = 32;
pub const CAUSALITY_LABEL_VALUE_MAX_BYTES: usize = 256;
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum Principal {
Role {
role: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
session: Option<String>,
},
Gateway {
gateway: String,
channel: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
conversation: Option<String>,
},
Cluster { cluster: String },
}
impl Principal {
pub fn role(role: impl Into<String>) -> Self {
Principal::Role {
role: role.into(),
session: None,
}
}
pub fn role_session(role: impl Into<String>, session: impl Into<String>) -> Self {
Principal::Role {
role: role.into(),
session: Some(session.into()),
}
}
pub fn gateway(
gateway: impl Into<String>,
channel: impl Into<String>,
conversation: Option<String>,
) -> Self {
Principal::Gateway {
gateway: gateway.into(),
channel: channel.into(),
conversation,
}
}
pub fn role_name(&self) -> Option<&str> {
match self {
Principal::Role { role, .. } => Some(role),
Principal::Gateway { .. } | Principal::Cluster { .. } => None,
}
}
pub fn kind(&self) -> &'static str {
match self {
Principal::Role { .. } => "role",
Principal::Gateway { .. } => "gateway",
Principal::Cluster { .. } => "cluster",
}
}
}
impl fmt::Display for Principal {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Principal::Role { role, session } => match session {
Some(s) => write!(f, "{role}#{s}"),
None => write!(f, "{role}"),
},
Principal::Gateway {
gateway,
channel,
conversation,
} => match conversation {
Some(c) => write!(f, "gw:{gateway}:{channel}:{c}"),
None => write!(f, "gw:{gateway}:{channel}"),
},
Principal::Cluster { cluster } => write!(f, "cluster:{cluster}"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum MsgKind {
Task,
Completion,
Note,
Control,
}
impl MsgKind {
pub fn is_control_plane(self) -> bool {
matches!(self, MsgKind::Task | MsgKind::Completion | MsgKind::Control)
}
pub fn as_str(self) -> &'static str {
match self {
MsgKind::Task => "task",
MsgKind::Completion => "completion",
MsgKind::Note => "note",
MsgKind::Control => "control",
}
}
}
impl fmt::Display for MsgKind {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case", tag = "op")]
pub enum ControlOp {
Recycle { task_id: String, reason: String },
Probe { task_id: String },
Snapshot { task_id: String },
Cancel { task_id: String, reason: String },
Focus { task_id: String },
}
impl ControlOp {
pub fn task_id(&self) -> &str {
match self {
ControlOp::Recycle { task_id, .. }
| ControlOp::Probe { task_id }
| ControlOp::Snapshot { task_id }
| ControlOp::Cancel { task_id, .. }
| ControlOp::Focus { task_id } => task_id,
}
}
pub fn as_str(&self) -> &'static str {
match self {
ControlOp::Recycle { .. } => "recycle",
ControlOp::Probe { .. } => "probe",
ControlOp::Snapshot { .. } => "snapshot",
ControlOp::Cancel { .. } => "cancel",
ControlOp::Focus { .. } => "focus",
}
}
pub fn name(&self) -> &'static str {
self.as_str()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum Outcome {
Done,
Failed,
Cancelled,
Blocked,
}
impl Outcome {
pub fn as_str(self) -> &'static str {
match self {
Outcome::Done => "done",
Outcome::Failed => "failed",
Outcome::Cancelled => "cancelled",
Outcome::Blocked => "blocked",
}
}
}
impl fmt::Display for Outcome {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
#[serde(rename_all = "snake_case", default)]
pub struct ImagePart {
pub data_base64: String,
pub mime: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
}
impl ImagePart {
pub fn decode(&self) -> Result<Vec<u8>> {
if self.data_base64.len() > IMAGE_DATA_MAX_ENCODED_BYTES {
return Err(image_over_budget());
}
let bytes = base64::engine::general_purpose::STANDARD
.decode(&self.data_base64)
.map_err(|e| Error::invalid("body.image.data_base64", format!("bad base64: {e}")))?;
Ok(bytes)
}
}
fn image_over_budget() -> Error {
Error::invalid(
"body.image.data_base64",
format!("image exceeds {IMAGE_DATA_MAX_BYTES} bytes"),
)
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case", default)]
pub struct Body {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub head: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub image: Option<ImagePart>,
}
impl Body {
pub fn text(text: impl Into<String>) -> Self {
Body {
text: Some(text.into()),
head: None,
image: None,
}
}
pub fn image(data_base64: impl Into<String>, mime: impl Into<String>) -> Self {
Body {
text: None,
head: None,
image: Some(ImagePart {
data_base64: data_base64.into(),
mime: mime.into(),
name: None,
}),
}
}
pub fn is_empty(&self) -> bool {
self.text.is_none() && self.image.is_none()
}
fn validate(&self) -> Result<()> {
if self.is_empty() {
return Err(Error::invalid(
"body",
"body requires text or image".to_string(),
));
}
if let Some(text) = &self.text {
let len = text.len();
if len > BODY_TEXT_MAX_BYTES {
return Err(Error::invalid(
"body.text",
format!("text exceeds {BODY_TEXT_MAX_BYTES} bytes"),
));
}
}
if let Some(image) = &self.image {
if !IMAGE_MIMES.contains(&image.mime.as_str()) {
return Err(Error::invalid(
"body.image.mime",
format!("mime must be one of: {}", IMAGE_MIMES.join(", ")),
));
}
let decoded = image.decode()?;
if decoded.len() > IMAGE_DATA_MAX_BYTES {
return Err(image_over_budget());
}
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
#[serde(rename_all = "snake_case", default)]
pub struct Causality {
pub task: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_task: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reply_to: Option<String>,
pub hop: u32,
pub attempt: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub family: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hop_budget: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub deadline: Option<DateTime<Utc>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub labels: Option<BTreeMap<String, String>>,
}
impl Causality {
pub fn root(task: impl Into<String>) -> Self {
let task = task.into();
Causality {
family: Some(task.clone()),
task,
parent_task: None,
reply_to: None,
hop: 0,
attempt: 0,
hop_budget: None,
origin: None,
deadline: None,
labels: None,
}
}
pub fn child_of(&self) -> Causality {
Causality {
task: new_task_id(),
parent_task: Some(self.task.clone()),
reply_to: None,
hop: self.hop.saturating_add(1),
attempt: 0,
family: Some(self.family.clone().unwrap_or_else(|| self.task.clone())),
hop_budget: self.hop_budget,
origin: self.origin.clone(),
deadline: self.deadline,
labels: self.labels.clone(),
}
}
pub fn validate(&self) -> Result<()> {
let Some(labels) = &self.labels else {
return Ok(());
};
if labels.len() > CAUSALITY_LABEL_MAX_ENTRIES {
return Err(Error::invalid(
"causality.labels",
format!(
"{} labels carried, at most {CAUSALITY_LABEL_MAX_ENTRIES}",
labels.len()
),
));
}
for (key, value) in labels {
if key.is_empty() || key.len() > CAUSALITY_LABEL_KEY_MAX_BYTES {
return Err(Error::invalid(
"causality.labels",
format!("label key {key:?} must be 1..={CAUSALITY_LABEL_KEY_MAX_BYTES} bytes"),
));
}
if value.len() > CAUSALITY_LABEL_VALUE_MAX_BYTES {
return Err(Error::invalid(
"causality.labels",
format!(
"label {key:?} carries {} bytes, at most {CAUSALITY_LABEL_VALUE_MAX_BYTES}",
value.len()
),
));
}
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub struct Envelope {
pub protocol: u16,
pub id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub op_id: Option<String>,
pub kind: MsgKind,
pub from: Principal,
pub to: Principal,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub control: Option<ControlOp>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub causality: Option<Causality>,
pub body: Body,
pub ts: DateTime<Utc>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ttl_ms: Option<u64>,
pub admin: bool,
}
pub fn new_envelope(
kind: MsgKind,
from: Principal,
to: Principal,
body: Body,
causality: Option<Causality>,
) -> Result<Envelope> {
let env = Envelope {
protocol: PROTOCOL_VERSION,
id: new_id(),
op_id: match kind {
MsgKind::Note => None,
_ => Some(new_op_id()),
},
kind,
from,
to,
control: None,
causality,
body,
ts: Utc::now(),
ttl_ms: None,
admin: false,
};
env.validate()?;
Ok(env)
}
impl Envelope {
pub fn redelivered(&self) -> Envelope {
let mut next = self.clone();
if let Some(causality) = &mut next.causality {
causality.attempt = causality.attempt.saturating_add(1);
}
next
}
pub fn task_id(&self) -> Option<&str> {
self.causality.as_ref().map(|c| c.task.as_str())
}
pub fn validate(&self) -> Result<()> {
if self.protocol != PROTOCOL_VERSION {
return Err(Error::invalid(
"protocol",
format!(
"protocol {} unsupported, expected {PROTOCOL_VERSION}",
self.protocol
),
));
}
if Uuid::parse_str(&self.id).is_err() {
return Err(Error::invalid("id", "id must be a uuid".to_string()));
}
match self.kind {
MsgKind::Note => {}
other => {
let Some(op_id) = &self.op_id else {
return Err(Error::invalid(
"op_id",
format!("op_id is required for kind {other}"),
));
};
validate_op_id(op_id)?;
}
}
if self.from.role_name() == Some("") {
return Err(Error::invalid(
"from.role",
"role name must not be empty".to_string(),
));
}
if self.to.role_name() == Some("") {
return Err(Error::invalid(
"to.role",
"role name must not be empty".to_string(),
));
}
self.body.validate()?;
if let Some(causality) = &self.causality {
causality.validate()?;
}
match self.kind {
MsgKind::Control => {
if self.control.is_none() {
return Err(Error::invalid(
"control",
"control kind requires a control op".to_string(),
));
}
}
MsgKind::Note | MsgKind::Task | MsgKind::Completion => {
if self.control.is_some() {
return Err(Error::invalid(
"control",
format!("control is reserved for kind {}", MsgKind::Control),
));
}
}
}
if self.kind.is_control_plane() && self.causality.is_none() {
return Err(Error::invalid(
"causality",
crate::text::causality_required(&self.kind.to_string()),
));
}
if let Some(causality) = &self.causality {
if Uuid::parse_str(&causality.task).is_err() {
return Err(Error::invalid(
"causality.task",
"task must be a uuid".to_string(),
));
}
}
Ok(())
}
pub fn fingerprint(&self) -> String {
let mut value = serde_json::to_value(self).unwrap_or(serde_json::Value::Null);
if let Some(map) = value.as_object_mut() {
map.remove("id");
map.remove("ts");
map.remove("op_id");
}
if let Some(causality) = value.get_mut("causality").and_then(|c| c.as_object_mut()) {
causality.insert("attempt".to_string(), serde_json::Value::from(0));
}
let bytes = serde_json::to_vec(&value).unwrap_or_default();
sha256_hex(&bytes)
}
}
pub fn sha256_hex(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
let digest = hasher.finalize();
let mut out = String::with_capacity(digest.len() * 2);
for byte in digest {
out.push(char::from_digit((byte >> 4) as u32, 16).expect("hex digit"));
out.push(char::from_digit((byte & 0xf) as u32, 16).expect("hex digit"));
}
out
}
pub fn new_op_id() -> String {
format!("{OP_ID_PREFIX}{}", Uuid::new_v4())
}
pub fn new_id() -> String {
Uuid::new_v4().to_string()
}
pub fn new_task_id() -> String {
Uuid::new_v4().to_string()
}
fn validate_op_id(op_id: &str) -> Result<()> {
let Some(rest) = op_id.strip_prefix(OP_ID_PREFIX) else {
return Err(Error::invalid(
"op_id",
format!("op_id must start with {OP_ID_PREFIX}"),
));
};
if Uuid::parse_str(rest).is_err() {
return Err(Error::invalid(
"op_id",
"op_id must carry a uuid".to_string(),
));
}
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Error {
Invalid { field: String, message: String },
}
impl Error {
pub fn invalid(field: impl Into<String>, message: impl Into<String>) -> Self {
Error::Invalid {
field: field.into(),
message: message.into(),
}
}
pub fn field(&self) -> &str {
match self {
Error::Invalid { field, .. } => field,
}
}
pub fn message(&self) -> &str {
match self {
Error::Invalid { message, .. } => message,
}
}
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Error::Invalid { field, message } => write!(f, "{field}: {message}"),
}
}
}
impl std::error::Error for Error {}
pub type Result<T, E = Error> = std::result::Result<T, E>;