initial commit

src/storage.rs

@@ -0,0 +1,484 @@
+//! 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()
+}