use base64::Engine;
use sqlx::PgPool;
use uuid::Uuid;
use crate::{
auth::AuthCtx,
error::{BusError, BusResult},
model::{AttachmentContent, AttachmentMeta, ts},
};
pub const MAX_ATTACHMENT_BYTES: usize = 256 * 1024;
pub const MAX_ATTACHMENTS_PER_TARGET: i64 = 8;
pub struct NewAttachment {
pub filename: String,
pub content_type: String,
pub data: Vec<u8>,
}
pub fn decode_input(
filename: &str,
content_type: Option<String>,
data_base64: &str,
) -> BusResult<NewAttachment> {
let filename = filename.trim();
if filename.is_empty() {
return Err(BusError::invalid("attachment filename cannot be empty"));
}
if filename.len() > 255 {
return Err(BusError::invalid(
"attachment filename is limited to 255 characters",
));
}
let data = base64::engine::general_purpose::STANDARD
.decode(data_base64.trim())
.map_err(|_| {
BusError::invalid(format!(
"attachment '{filename}': data_base64 is not valid base64"
))
})?;
if data.is_empty() {
return Err(BusError::invalid(format!(
"attachment '{filename}' is empty"
)));
}
if data.len() > MAX_ATTACHMENT_BYTES {
return Err(BusError::invalid(format!(
"attachment '{filename}' is {} bytes; the limit is {MAX_ATTACHMENT_BYTES} \
(share a link or trim the file instead)",
data.len()
)));
}
let content_type = content_type
.map(|c| c.trim().to_owned())
.filter(|c| !c.is_empty())
.unwrap_or_else(|| "application/octet-stream".to_owned());
Ok(NewAttachment {
filename: filename.to_owned(),
content_type,
data,
})
}
pub async fn insert_for_message(
conn: &mut sqlx::PgConnection,
team_id: Uuid,
message_id: i64,
uploader: Uuid,
att: &NewAttachment,
) -> BusResult<AttachmentMeta> {
let (id,): (i64,) = sqlx::query_as(
r#"
INSERT INTO attachments
(team_id, message_id, uploader_agent_id, filename, content_type, size_bytes, data)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id
"#,
)
.bind(team_id)
.bind(message_id)
.bind(uploader)
.bind(&att.filename)
.bind(&att.content_type)
.bind(att.data.len() as i64)
.bind(&att.data)
.fetch_one(conn)
.await?;
Ok(AttachmentMeta {
id,
filename: att.filename.clone(),
content_type: att.content_type.clone(),
size_bytes: att.data.len() as i64,
})
}
pub async fn attach_to_task(
pool: &PgPool,
auth: &AuthCtx,
task_key: &str,
att: &NewAttachment,
) -> BusResult<AttachmentMeta> {
let task_key = task_key.trim();
let task: Option<(Uuid,)> =
sqlx::query_as("SELECT id FROM tasks WHERE team_id = $1 AND key = $2")
.bind(auth.team_id)
.bind(task_key)
.fetch_optional(pool)
.await?;
let (task_id,) = task.ok_or_else(|| BusError::not_found(format!("task '{task_key}'")))?;
let (count,): (i64,) = sqlx::query_as("SELECT count(*) FROM attachments WHERE task_id = $1")
.bind(task_id)
.fetch_one(pool)
.await?;
if count >= MAX_ATTACHMENTS_PER_TARGET {
return Err(BusError::invalid(format!(
"task '{task_key}' already has {MAX_ATTACHMENTS_PER_TARGET} attachments; \
delete-and-recreate is not supported, start a new task or link externally"
)));
}
let (id,): (i64,) = sqlx::query_as(
r#"
INSERT INTO attachments
(team_id, task_id, uploader_agent_id, filename, content_type, size_bytes, data)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id
"#,
)
.bind(auth.team_id)
.bind(task_id)
.bind(auth.agent_id)
.bind(&att.filename)
.bind(&att.content_type)
.bind(att.data.len() as i64)
.bind(&att.data)
.fetch_one(pool)
.await?;
Ok(AttachmentMeta {
id,
filename: att.filename.clone(),
content_type: att.content_type.clone(),
size_bytes: att.data.len() as i64,
})
}
pub async fn attach_to_message(
pool: &PgPool,
auth: &AuthCtx,
message_id: i64,
att: &NewAttachment,
) -> BusResult<AttachmentMeta> {
let row: Option<(Uuid,)> =
sqlx::query_as("SELECT sender_agent_id FROM messages WHERE id = $1 AND team_id = $2")
.bind(message_id)
.bind(auth.team_id)
.fetch_optional(pool)
.await?;
let (sender,) = row.ok_or_else(|| BusError::not_found(format!("message {message_id}")))?;
if sender != auth.agent_id {
return Err(BusError::invalid(
"you can only attach files to your own messages",
));
}
let (count,): (i64,) = sqlx::query_as("SELECT count(*) FROM attachments WHERE message_id = $1")
.bind(message_id)
.fetch_one(pool)
.await?;
if count >= MAX_ATTACHMENTS_PER_TARGET {
return Err(BusError::invalid(format!(
"message {message_id} already has {MAX_ATTACHMENTS_PER_TARGET} attachments"
)));
}
let mut conn = pool.acquire().await?;
insert_for_message(&mut conn, auth.team_id, message_id, auth.agent_id, att).await
}
pub async fn get_attachment(
pool: &PgPool,
auth: &AuthCtx,
id: i64,
) -> BusResult<AttachmentContent> {
let row: Option<(
String,
String,
i64,
Vec<u8>,
String,
chrono::DateTime<chrono::Utc>,
Option<Uuid>,
Option<Uuid>,
)> = sqlx::query_as(
r#"
SELECT a.filename, a.content_type, a.size_bytes, a.data, u.name, a.created_at,
m.recipient_agent_id, m.sender_agent_id
FROM attachments a
JOIN agents u ON u.id = a.uploader_agent_id
LEFT JOIN messages m ON m.id = a.message_id
WHERE a.id = $1 AND a.team_id = $2
"#,
)
.bind(id)
.bind(auth.team_id)
.fetch_optional(pool)
.await?;
let (filename, content_type, size_bytes, data, uploaded_by, created_at, dm_recipient, sender) =
row.ok_or_else(|| BusError::not_found(format!("attachment {id}")))?;
if let Some(recipient) = dm_recipient
&& recipient != auth.agent_id
&& sender != Some(auth.agent_id)
{
return Err(BusError::not_found(format!("attachment {id}")));
}
Ok(AttachmentContent {
id,
filename,
content_type,
size_bytes,
uploaded_by,
created_at: ts(created_at),
data_base64: base64::engine::general_purpose::STANDARD.encode(data),
})
}