485 lines
16 KiB
Rust
485 lines
16 KiB
Rust
//! Lossless, bounded Prosody internal key/value persistence.
|
|
use anyhow::{Result, bail, ensure};
|
|
use std::{
|
|
fs::{self, OpenOptions},
|
|
io::{Read, Write},
|
|
path::{Path, PathBuf},
|
|
sync::Mutex,
|
|
};
|
|
static TRANSACTION: Mutex<()> = Mutex::new(());
|
|
const LIMIT: usize = 16 * 1024 * 1024;
|
|
#[derive(Clone, Debug)]
|
|
pub enum LuaValue {
|
|
Nil,
|
|
Bool(bool),
|
|
Integer(i64),
|
|
Number(f64),
|
|
String(Vec<u8>),
|
|
Table(Vec<(LuaValue, LuaValue)>),
|
|
}
|
|
impl PartialEq for LuaValue {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
match (self, other) {
|
|
(Self::Nil, Self::Nil) => true,
|
|
(Self::Bool(a), Self::Bool(b)) => a == b,
|
|
(Self::Integer(a), Self::Integer(b)) => a == b,
|
|
(Self::Number(a), Self::Number(b)) => a == b,
|
|
(Self::Integer(a), Self::Number(b)) | (Self::Number(b), Self::Integer(a)) => {
|
|
b.is_finite()
|
|
&& b.fract() == 0.0
|
|
&& *b >= i64::MIN as f64
|
|
&& *b < (i64::MAX as f64)
|
|
&& *a == *b as i64
|
|
}
|
|
(Self::String(a), Self::String(b)) => a == b,
|
|
(Self::Table(a), Self::Table(b)) => {
|
|
a.len() == b.len()
|
|
&& a.iter()
|
|
.all(|(k, v)| b.iter().any(|(k2, v2)| k == k2 && v == v2))
|
|
}
|
|
_ => false,
|
|
}
|
|
}
|
|
}
|
|
impl Default for LuaValue {
|
|
fn default() -> Self {
|
|
Self::Table(vec![])
|
|
}
|
|
}
|
|
impl From<&str> for LuaValue {
|
|
fn from(v: &str) -> Self {
|
|
Self::String(v.as_bytes().to_vec())
|
|
}
|
|
}
|
|
impl From<String> for LuaValue {
|
|
fn from(v: String) -> Self {
|
|
Self::String(v.into_bytes())
|
|
}
|
|
}
|
|
impl LuaValue {
|
|
pub fn table() -> Self {
|
|
Self::default()
|
|
}
|
|
pub fn get_key(&self, k: &Self) -> Option<&Self> {
|
|
if let Self::Table(t) = self {
|
|
t.iter().find(|(a, _)| a == k).map(|(_, v)| v)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
pub fn get(&self, k: &str) -> Option<&Self> {
|
|
self.get_key(&k.into())
|
|
}
|
|
pub fn set_key(&mut self, k: Self, v: Self) {
|
|
if let Self::Table(t) = self {
|
|
t.retain(|(a, _)| a != &k);
|
|
if v != Self::Nil {
|
|
t.push((k, v));
|
|
}
|
|
}
|
|
}
|
|
pub fn set(&mut self, k: &str, v: Self) {
|
|
self.set_key(k.into(), v)
|
|
}
|
|
pub fn as_str(&self) -> Option<&str> {
|
|
if let Self::String(s) = self {
|
|
std::str::from_utf8(s).ok()
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
pub fn entries(&self) -> &[(Self, Self)] {
|
|
if let Self::Table(t) = self { t } else { &[] }
|
|
}
|
|
pub fn as_i64(&self) -> Option<i64> {
|
|
if let Self::Integer(v) = self {
|
|
Some(*v)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|
|
pub fn encode_name(name: &str) -> String {
|
|
encode(name, false)
|
|
}
|
|
fn encode(name: &str, store: bool) -> String {
|
|
let mut out = String::new();
|
|
for b in name.bytes() {
|
|
if b.is_ascii_alphanumeric() || (store && b == b'_') {
|
|
out.push(b as char)
|
|
} else {
|
|
out.push_str(&format!("%{b:02x}"))
|
|
}
|
|
}
|
|
out
|
|
}
|
|
pub fn path(root: &Path, host: &str, user: Option<&str>, store: &str) -> PathBuf {
|
|
let d = root.join(encode_name(host));
|
|
match user {
|
|
Some(u) => d
|
|
.join(encode(store, true))
|
|
.join(format!("{}.dat", encode_name(u))),
|
|
None => d.join(format!("{}.dat", encode(store, true))),
|
|
}
|
|
}
|
|
pub fn parse(input: &[u8]) -> Result<LuaValue> {
|
|
ensure!(input.len() <= LIMIT, "storage limit exceeded");
|
|
let lua = crate::lua::bounded_vm().map_err(|e| anyhow::anyhow!(e.to_string()))?;
|
|
let env = lua
|
|
.create_table()
|
|
.map_err(|e| anyhow::anyhow!(e.to_string()))?;
|
|
let v = lua
|
|
.load(input)
|
|
.set_mode(mlua::chunk::ChunkMode::Text)
|
|
.set_environment(env)
|
|
.eval::<mlua::Value>()
|
|
.map_err(|e| anyhow::anyhow!(e.to_string()))?;
|
|
from_lua(v)
|
|
}
|
|
pub fn from_lua(v: mlua::Value) -> Result<LuaValue> {
|
|
fn convert(v: mlua::Value, depth: usize, count: &mut usize) -> Result<LuaValue> {
|
|
*count += 1;
|
|
ensure!(depth < 64 && *count <= 200_000, "storage complexity limit");
|
|
Ok(match v {
|
|
mlua::Value::Nil => LuaValue::Nil,
|
|
mlua::Value::Boolean(v) => LuaValue::Bool(v),
|
|
mlua::Value::Integer(v) => LuaValue::Integer(v),
|
|
mlua::Value::Number(v) => {
|
|
ensure!(v.is_finite(), "non-finite storage number");
|
|
LuaValue::Number(v)
|
|
}
|
|
mlua::Value::String(v) => LuaValue::String(v.as_bytes().to_vec()),
|
|
mlua::Value::Table(t) => {
|
|
let mut o = vec![];
|
|
for p in t.pairs::<mlua::Value, mlua::Value>() {
|
|
let (k, v) = p.map_err(|e| anyhow::anyhow!(e.to_string()))?;
|
|
ensure!(
|
|
matches!(
|
|
k,
|
|
mlua::Value::Boolean(_)
|
|
| mlua::Value::String(_)
|
|
| mlua::Value::Integer(_)
|
|
| mlua::Value::Number(_)
|
|
),
|
|
"unsupported table key"
|
|
);
|
|
o.push((convert(k, depth + 1, count)?, convert(v, depth + 1, count)?));
|
|
}
|
|
LuaValue::Table(o)
|
|
}
|
|
_ => bail!("unsupported storage value"),
|
|
})
|
|
}
|
|
convert(v, 0, &mut 0)
|
|
}
|
|
pub fn serialize(value: &LuaValue) -> Result<String> {
|
|
fn put(v: &LuaValue, s: &mut String, d: usize) -> Result<()> {
|
|
ensure!(d < 64 && s.len() <= LIMIT, "storage complexity limit");
|
|
match v {
|
|
LuaValue::Nil => s.push_str("nil"),
|
|
LuaValue::Bool(v) => s.push_str(if *v { "true" } else { "false" }),
|
|
LuaValue::Integer(v) => s.push_str(&v.to_string()),
|
|
LuaValue::Number(v) => {
|
|
ensure!(v.is_finite(), "non-finite storage number");
|
|
s.push_str(&format!("{v:?}"))
|
|
}
|
|
LuaValue::String(v) => {
|
|
s.push('"');
|
|
for b in v {
|
|
match b {
|
|
b'"' => s.push_str("\\\""),
|
|
b'\\' => s.push_str("\\\\"),
|
|
32..=126 => s.push(*b as char),
|
|
_ => s.push_str(&format!("\\{b:03}")),
|
|
}
|
|
}
|
|
s.push('"')
|
|
}
|
|
LuaValue::Table(t) => {
|
|
s.push('{');
|
|
for (k, v) in t {
|
|
ensure!(
|
|
matches!(
|
|
k,
|
|
LuaValue::Bool(_)
|
|
| LuaValue::Integer(_)
|
|
| LuaValue::Number(_)
|
|
| LuaValue::String(_)
|
|
),
|
|
"invalid key"
|
|
);
|
|
s.push('[');
|
|
put(k, s, d + 1)?;
|
|
s.push_str("]=");
|
|
put(v, s, d + 1)?;
|
|
s.push(';')
|
|
}
|
|
s.push('}')
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
let mut s = "return ".to_string();
|
|
put(value, &mut s, 0)?;
|
|
s.push_str(";\n");
|
|
ensure!(s.len() <= LIMIT, "storage limit exceeded");
|
|
Ok(s)
|
|
}
|
|
pub fn load(root: &Path, host: &str, user: Option<&str>, store: &str) -> Result<Option<LuaValue>> {
|
|
let p = path(root, host, user, store);
|
|
let f = match fs::File::open(&p) {
|
|
Ok(f) => f,
|
|
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
let mut b = vec![];
|
|
f.take((LIMIT + 1) as u64).read_to_end(&mut b)?;
|
|
Ok(Some(parse(&b)?))
|
|
}
|
|
pub fn atomic_write(p: &Path, data: &[u8]) -> Result<()> {
|
|
let parent = p
|
|
.parent()
|
|
.ok_or_else(|| anyhow::anyhow!("missing parent"))?;
|
|
fs::create_dir_all(parent)?;
|
|
let tmp = parent.join(format!(".{}.tmp", uuid::Uuid::new_v4()));
|
|
let result = (|| -> Result<()> {
|
|
let mut opts = OpenOptions::new();
|
|
opts.write(true).create_new(true);
|
|
#[cfg(unix)]
|
|
{
|
|
use std::os::unix::fs::OpenOptionsExt;
|
|
opts.mode(0o600);
|
|
}
|
|
let mut f = opts.open(&tmp)?;
|
|
f.write_all(data)?;
|
|
f.sync_all()?;
|
|
fs::rename(&tmp, p)?;
|
|
fs::File::open(parent)?.sync_all()?;
|
|
Ok(())
|
|
})();
|
|
if result.is_err() {
|
|
let _ = fs::remove_file(tmp);
|
|
}
|
|
result
|
|
}
|
|
pub fn save(
|
|
root: &Path,
|
|
host: &str,
|
|
user: Option<&str>,
|
|
store: &str,
|
|
value: &LuaValue,
|
|
) -> Result<()> {
|
|
if matches!(value, LuaValue::Nil) || matches!(value,LuaValue::Table(v) if v.is_empty()) {
|
|
let p = path(root, host, user, store);
|
|
match fs::remove_file(&p) {
|
|
Ok(()) => {
|
|
if let Some(parent) = p.parent() {
|
|
fs::File::open(parent)?.sync_all()?;
|
|
}
|
|
}
|
|
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
return Ok(());
|
|
}
|
|
let s = serialize(value)?;
|
|
atomic_write(&path(root, host, user, store), s.as_bytes())
|
|
}
|
|
/// Process-local serialized read/modify/write. Never run two servers on one data directory.
|
|
pub fn update<F: FnOnce(&mut LuaValue) -> Result<()>>(
|
|
root: &Path,
|
|
host: &str,
|
|
user: Option<&str>,
|
|
store: &str,
|
|
f: F,
|
|
) -> Result<LuaValue> {
|
|
let _g = TRANSACTION
|
|
.lock()
|
|
.map_err(|_| anyhow::anyhow!("storage lock poisoned"))?;
|
|
let mut v = load(root, host, user, store)?.unwrap_or_default();
|
|
ensure!(matches!(v, LuaValue::Table(_)), "store is not a table");
|
|
f(&mut v)?;
|
|
save(root, host, user, store, &v)?;
|
|
Ok(v)
|
|
}
|
|
/// util.stanza.preserialize-compatible representation (mixed numeric/string keys).
|
|
pub fn stanza_value(n: roxmltree::Node<'_, '_>) -> LuaValue {
|
|
if n.is_text() {
|
|
return n.text().unwrap_or("").into();
|
|
}
|
|
let mut t = LuaValue::table();
|
|
t.set("name", n.tag_name().name().into());
|
|
let mut attr = LuaValue::table();
|
|
attr.set("xmlns", n.tag_name().namespace().unwrap_or("").into());
|
|
for a in n.attributes() {
|
|
let key = match a.namespace() {
|
|
Some("http://www.w3.org/XML/1998/namespace") => format!("xml:{}", a.name()),
|
|
Some(ns) => format!("{ns}\u{1}{}", a.name()),
|
|
None => a.name().into(),
|
|
};
|
|
attr.set(&key, a.value().into());
|
|
}
|
|
t.set("attr", attr);
|
|
for (i, c) in n
|
|
.children()
|
|
.filter(|x| x.is_element() || x.is_text())
|
|
.enumerate()
|
|
{
|
|
t.set_key(LuaValue::Integer(i as i64 + 1), stanza_value(c));
|
|
}
|
|
t
|
|
}
|
|
pub fn stanza_xml(v: &LuaValue) -> Result<String> {
|
|
use crate::wire::escape;
|
|
if let LuaValue::String(s) = v {
|
|
return Ok(escape(std::str::from_utf8(s)?));
|
|
}
|
|
let name = v
|
|
.get("name")
|
|
.and_then(LuaValue::as_str)
|
|
.ok_or_else(|| anyhow::anyhow!("missing stanza name"))?;
|
|
ensure!(
|
|
!name.is_empty()
|
|
&& name
|
|
.chars()
|
|
.all(|c| c.is_alphanumeric() || "_-.".contains(c)),
|
|
"invalid stanza name"
|
|
);
|
|
let mut s = format!("<{name}");
|
|
let mut namespace_index = 0;
|
|
if let Some(a) = v.get("attr") {
|
|
for (k, val) in a.entries() {
|
|
let key = k
|
|
.as_str()
|
|
.ok_or_else(|| anyhow::anyhow!("invalid attribute key"))?;
|
|
let val = val
|
|
.as_str()
|
|
.ok_or_else(|| anyhow::anyhow!("invalid attribute value"))?;
|
|
let split = key.split_once('\u{1}').or_else(|| key.split_once('|'));
|
|
if let Some((ns, local)) = split {
|
|
ensure!(
|
|
!local.is_empty()
|
|
&& local
|
|
.chars()
|
|
.all(|c| c.is_alphanumeric() || "_-.".contains(c)),
|
|
"invalid attribute"
|
|
);
|
|
if ns == "http://www.w3.org/XML/1998/namespace" {
|
|
s.push_str(&format!(" xml:{local}='{}'", escape(val)))
|
|
} else {
|
|
namespace_index += 1;
|
|
s.push_str(&format!(
|
|
" xmlns:a{namespace_index}='{}' a{namespace_index}:{local}='{}'",
|
|
escape(ns),
|
|
escape(val)
|
|
))
|
|
}
|
|
} else {
|
|
ensure!(
|
|
!key.is_empty()
|
|
&& key
|
|
.chars()
|
|
.all(|c| c.is_alphanumeric() || "_-.:".contains(c)),
|
|
"invalid attribute"
|
|
);
|
|
s.push_str(&format!(" {key}='{}'", escape(val)))
|
|
}
|
|
}
|
|
}
|
|
s.push('>');
|
|
let mut i = 1;
|
|
while let Some(child) = v.get_key(&LuaValue::Integer(i)) {
|
|
s.push_str(&stanza_xml(child)?);
|
|
i += 1;
|
|
}
|
|
s.push_str(&format!("</{name}>"));
|
|
Ok(s)
|
|
}
|
|
/// Prosody datamanager.purge semantics. Caller must stop this user's sessions first.
|
|
/// Unknown stores are included, but only exact per-user storage files are removed.
|
|
pub fn purge_user(root: &Path, host: &str, user: &str) -> Result<()> {
|
|
let _guard = TRANSACTION
|
|
.lock()
|
|
.map_err(|_| anyhow::anyhow!("storage lock poisoned"))?;
|
|
ensure!(
|
|
!host.is_empty() && !user.is_empty(),
|
|
"host and user required"
|
|
);
|
|
let dir = root.join(encode_name(host));
|
|
let meta = match fs::symlink_metadata(&dir) {
|
|
Ok(m) => m,
|
|
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
ensure!(
|
|
meta.is_dir() && !meta.file_type().is_symlink(),
|
|
"refusing symlinked data host"
|
|
);
|
|
let user = encode_name(user);
|
|
let mut files = vec![];
|
|
for entry in fs::read_dir(&dir)? {
|
|
let entry = entry?;
|
|
let ty = entry.file_type()?;
|
|
ensure!(
|
|
!ty.is_symlink(),
|
|
"refusing symlinked host store during account purge"
|
|
);
|
|
if !ty.is_dir() {
|
|
continue;
|
|
}
|
|
for extension in ["dat", "list", "lidx", "rappend"] {
|
|
let p = entry.path().join(format!("{user}.{extension}"));
|
|
match fs::symlink_metadata(&p) {
|
|
Ok(m) => {
|
|
ensure!(
|
|
m.is_file() || m.file_type().is_symlink(),
|
|
"account store path is not a file"
|
|
);
|
|
files.push(p)
|
|
}
|
|
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
}
|
|
// Revoke login before deleting the remainder; failures are returned, never hidden.
|
|
files.sort_by_key(|p| {
|
|
if p.parent()
|
|
.and_then(|p| p.file_name())
|
|
.is_some_and(|s| s == "accounts")
|
|
{
|
|
0
|
|
} else {
|
|
1
|
|
}
|
|
});
|
|
let mut errors = vec![];
|
|
for p in files {
|
|
match fs::remove_file(&p) {
|
|
Ok(()) => {
|
|
if let Some(parent) = p.parent()
|
|
&& let Err(e) = fs::File::open(parent).and_then(|f| f.sync_all())
|
|
{
|
|
errors.push(e.to_string())
|
|
}
|
|
}
|
|
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
|
|
Err(e) => errors.push(e.to_string()),
|
|
}
|
|
}
|
|
ensure!(
|
|
errors.is_empty(),
|
|
"account purge failed: {}",
|
|
errors.join("; ")
|
|
);
|
|
Ok(())
|
|
}
|
|
/// Coordinate account creation with key/value updates and user purges.
|
|
/// The callback must not call `update`, `purge_user`, or this function recursively.
|
|
pub fn with_exclusive_access<T>(f: impl FnOnce() -> Result<T>) -> Result<T> {
|
|
let _guard = TRANSACTION
|
|
.lock()
|
|
.map_err(|_| anyhow::anyhow!("storage lock poisoned"))?;
|
|
f()
|
|
}
|