use std::collections::HashMap;
use axum::extract::State;
use axum::http::{header, StatusCode};
use axum::response::{IntoResponse, Response};
use laterite_auth::AuthenticatedUser;
use laterite_core::t;
use crate::list::{self, ListConfig, ListQuery};
use crate::session::{FlashLevel, SessionHandle};
use crate::AdminState;
pub(crate) const MAX_ROWS: i64 = 20_000;
const PAGE: i64 = 1_000;
#[derive(Clone, Copy, PartialEq, Eq)]
enum Format {
Csv,
Json,
}
impl Format {
fn read(raw: &HashMap<String, String>) -> Self {
match raw.get("format").map(String::as_str) {
Some("json") => Format::Json,
_ => Format::Csv,
}
}
fn extension(self) -> &'static str {
match self {
Format::Csv => "csv",
Format::Json => "json",
}
}
fn content_type(self) -> &'static str {
match self {
Format::Csv => "text/csv; charset=utf-8",
Format::Json => "application/json",
}
}
}
pub(crate) async fn handle(
State(state): State<AdminState>,
config: &ListConfig,
path: &str,
raw: &HashMap<String, String>,
user: &AuthenticatedUser,
shell: &crate::Shell,
session: &SessionHandle,
) -> Response {
let format = Format::read(raw);
let back = format!("{}{}", state.admin_path, path);
let stored = list::stored_columns(&state, path, user).await;
let config = list::narrowed(config, stored.as_deref(), user);
let (order_by, order_dir) = list::resolve_sort(
&config,
raw.get("sort").map(String::as_str),
raw.get("dir").map(String::as_str),
);
let filters = list::resolve_filters(&config, raw);
let q = raw.get("q").map(String::as_str).unwrap_or("").trim();
let mut rows: Vec<Vec<String>> = Vec::new();
let mut offset = 0i64;
loop {
let req = ListQuery {
offset,
order_by: &order_by,
order_dir,
q,
filters: &filters,
per_page: Some(PAGE),
children_of: None,
};
let page = match list::query(&state.db, &config, &req).await {
Ok(page) => page,
Err(e) => {
tracing::error!(entity = %config.entity, error = %e, "export query failed");
session.push_flash(FlashLevel::Error, t!("The export could not be produced."));
return axum::response::Redirect::to(&back).into_response();
}
};
if page.total > MAX_ROWS {
session.push_flash(
FlashLevel::Error,
t!("That is too many records to export at once. Narrow the list and try again."),
);
return axum::response::Redirect::to(&back).into_response();
}
let count = page.rows.len();
rows.extend(page.rows.into_iter().map(|r| r.cells));
if count < PAGE as usize {
break;
}
offset += PAGE;
}
let labels: Vec<String> = config.columns.iter().map(|c| shell.tt(&c.label)).collect();
let names: Vec<String> = config.columns.iter().map(|c| c.field.clone()).collect();
let body = match format {
Format::Csv => match to_csv(&labels, &rows) {
Ok(body) => body,
Err(e) => {
tracing::error!(error = %e, "writing the export failed");
session.push_flash(FlashLevel::Error, t!("The export could not be produced."));
return axum::response::Redirect::to(&back).into_response();
}
},
Format::Json => to_json(&names, &rows),
};
let filename = format!("{}.{}", config.entity, format.extension());
(
StatusCode::OK,
[
(header::CONTENT_TYPE, format.content_type().to_string()),
(
header::CONTENT_DISPOSITION,
format!("attachment; filename=\"{filename}\""),
),
],
body,
)
.into_response()
}
const FORMULA_TRIGGERS: [char; 5] = ['=', '+', '@', '\t', '\r'];
fn defuse(cell: &str) -> String {
let dangerous = cell.starts_with(FORMULA_TRIGGERS)
|| (cell.starts_with('-') && cell.parse::<f64>().is_err());
if dangerous {
format!("'{cell}")
} else {
cell.to_string()
}
}
fn to_csv(headers: &[String], rows: &[Vec<String>]) -> Result<String, csv::Error> {
let mut writer = csv::Writer::from_writer(Vec::new());
writer.write_record(headers.iter().map(|h| defuse(h)))?;
for row in rows {
writer.write_record(row.iter().map(|cell| defuse(cell)))?;
}
let bytes = writer.into_inner().map_err(|e| e.into_error())?;
let mut out = String::from("\u{feff}");
out.push_str(&String::from_utf8_lossy(&bytes));
Ok(out)
}
fn to_json(headers: &[String], rows: &[Vec<String>]) -> String {
let out: Vec<serde_json::Value> = rows
.iter()
.map(|row| {
let pairs: serde_json::Map<String, serde_json::Value> = headers
.iter()
.zip(row)
.map(|(k, v)| (k.clone(), serde_json::Value::String(v.clone())))
.collect();
serde_json::Value::Object(pairs)
})
.collect();
serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".to_string())
}
#[cfg(test)]
mod tests {
use super::*;
fn rows() -> Vec<Vec<String>> {
vec![
vec!["ada".to_string(), "Yes".to_string()],
vec!["b, comma".to_string(), "say \"hi\"".to_string()],
vec!["line\nbreak".to_string(), "No".to_string()],
]
}
fn headers() -> Vec<String> {
vec!["username".to_string(), "superuser".to_string()]
}
use laterite_core::testing::connect_test;
use laterite_core::Db;
async fn seeded() -> (Db, laterite_core::testing::TestGuard) {
let (db, guard) = connect_test(&[laterite_auth::migrations()]).await;
for (name, superuser) in [("ada", true), ("brendan", false), ("clara", false)] {
let hash = laterite_auth::password::hash_password("x").unwrap();
laterite_auth::store::create_user(
&db,
name,
&format!("{name}@example.test"),
name,
None,
&hash,
superuser,
)
.await
.unwrap();
}
(db, guard)
}
fn users_list() -> ListConfig {
ListConfig {
entity: "backend_users".to_string(),
title: "Users".into(),
columns: vec![
crate::list::ListColumn::new("username", "Username"),
crate::list::ListColumn::new("is_superuser", "Superuser").yes_no(),
],
order_by: "username".to_string(),
order_dir: crate::list::SortDir::Asc,
filters: vec![crate::list::ListFilter::boolean(
"is_superuser",
"Superuser",
)],
..Default::default()
}
}
async fn body_of(raw: HashMap<String, String>) -> (String, String) {
let (db, _guard) = seeded().await;
let state = AdminState::new(
laterite_auth::AuthService::new(db.clone(), laterite_auth::AuthConfig::default()),
db,
);
let resp = handle(
State(state),
&users_list(),
"/users",
&raw,
&crate::audit::test_actor(),
&crate::Shell::test(),
&crate::session::SessionHandle::from_blob(None),
)
.await;
let disposition = resp
.headers()
.get(header::CONTENT_DISPOSITION)
.map(|v| v.to_str().unwrap().to_string())
.unwrap_or_default();
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
(String::from_utf8(bytes.to_vec()).unwrap(), disposition)
}
#[tokio::test]
async fn the_file_matches_the_view_that_produced_it() {
let (all, disposition) = body_of(HashMap::new()).await;
assert!(disposition.contains("backend_users.csv"), "{disposition}");
let lines: Vec<&str> = all.trim().lines().collect();
assert_eq!(lines[0], "\u{feff}Username,Superuser");
assert_eq!(lines.len(), 4);
assert!(
lines[1].starts_with("ada"),
"sorted ascending: {}",
lines[1]
);
let mut raw = HashMap::new();
raw.insert("f_is_superuser".to_string(), "1".to_string());
let (filtered, _) = body_of(raw).await;
assert_eq!(filtered.trim().lines().count(), 2, "header plus one row");
assert!(filtered.contains("ada"));
assert!(!filtered.contains("brendan"));
let mut raw = HashMap::new();
raw.insert("q".to_string(), "bren".to_string());
let (searched, _) = body_of(raw).await;
assert_eq!(searched.trim().lines().count(), 2);
assert!(searched.contains("brendan"));
}
#[tokio::test]
async fn json_asks_for_the_same_rows_in_the_other_shape() {
let mut raw = HashMap::new();
raw.insert("format".to_string(), "json".to_string());
let (body, disposition) = body_of(raw).await;
assert!(disposition.contains("backend_users.json"), "{disposition}");
let parsed: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(parsed.as_array().unwrap().len(), 3);
assert_eq!(parsed[0]["username"], "ada");
}
#[test]
fn a_cell_a_spreadsheet_would_execute_is_defused() {
for value in [
"=1+1",
"+1",
"@SUM(A1)",
"=cmd|' /C calc'!A0",
"\tleading tab",
"\rleading return",
] {
assert_eq!(
defuse(value),
format!("'{value}"),
"{value:?} was left live"
);
}
}
#[test]
fn a_negative_number_is_still_a_number() {
assert_eq!(defuse("-5"), "-5");
assert_eq!(defuse("-12.50"), "-12.50");
assert_eq!(defuse("-1e3"), "-1e3");
assert_eq!(defuse("-cmd|' /C calc'!A0"), "'-cmd|' /C calc'!A0");
}
#[test]
fn ordinary_text_is_untouched() {
for value in ["hello", "", "a=b", "5", "user@acme.example"] {
assert_eq!(defuse(value), value);
}
}
#[test]
fn a_header_is_defused_too() {
let csv = to_csv(&["=BAD".to_string()], &[]).unwrap();
assert!(csv.starts_with("\u{feff}'=BAD"), "{csv}");
}
#[test]
fn csv_quotes_and_escapes_every_awkward_value() {
let out = to_csv(&headers(), &rows()).unwrap();
assert!(out.starts_with("\u{feff}username,superuser\n"));
assert!(out.contains("\"b, comma\""), "a comma is quoted: {out}");
assert!(
out.contains("\"say \"\"hi\"\"\""),
"a quote is doubled: {out}"
);
assert!(
out.contains("\"line\nbreak\""),
"a newline is quoted: {out}"
);
}
#[test]
fn json_is_an_array_of_objects_keyed_by_column() {
let out = to_json(&headers(), &rows());
let parsed: serde_json::Value = serde_json::from_str(&out).unwrap();
assert_eq!(parsed.as_array().unwrap().len(), 3);
assert_eq!(parsed[0]["username"], "ada");
assert_eq!(parsed[1]["superuser"], "say \"hi\"");
}
#[test]
fn the_format_defaults_to_csv() {
let mut raw = HashMap::new();
assert!(Format::read(&raw) == Format::Csv);
raw.insert("format".to_string(), "json".to_string());
assert!(Format::read(&raw) == Format::Json);
raw.insert("format".to_string(), "xlsx".to_string());
assert!(Format::read(&raw) == Format::Csv);
}
}