AkurAI Build
Menu

akurai-tasks

public

Latest change 350a37d39889e09d093509a69e7f06106958ffa2 - Harden task coordination and recovery by Ólafur Búi Ólafsson

use super::*;
use sha2::{Digest, Sha256};
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write};

pub struct ExclusiveLock {
    _file: File,
}

impl Store {
    pub fn path(&self) -> &Path {
        self.path.as_path()
    }

    pub fn acquire_exclusive_lock(database: &Path) -> Result<ExclusiveLock> {
        let path = database.with_extension("lock");
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).map_err(storage)?;
        }
        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .truncate(false)
            .open(&path)
            .map_err(storage)?;
        file.try_lock().map_err(|error| {
            Error::Conflict(format!(
                "database is already owned or lock cannot be acquired: {error}"
            ))
        })?;
        file.set_len(0).map_err(storage)?;
        writeln!(file, "{}", std::process::id()).map_err(storage)?;
        file.sync_all().map_err(storage)?;
        Ok(ExclusiveLock { _file: file })
    }

    pub fn doctor(&self) -> Result<Value> {
        let mut db = self
            .db
            .lock()
            .map_err(|_| Error::Storage("lock poisoned".into()))?;
        let entries = db.range(b"", &[u8::MAX]).map_err(storage)?;
        let mut records = 0i64;
        let mut items = 0i64;
        let mut events = 0i64;
        let mut last_event = 0i64;
        for (key, bytes) in entries {
            let value = decode(&bytes)?;
            records += 1;
            if key.starts_with(b"item/") {
                items += 1;
                let project = field(&value, "project")
                    .ok_or_else(|| Error::Storage("item missing project".into()))?;
                if get_value(&mut db, &format!("project/{project}"))?.is_none() {
                    return Err(Error::Storage(format!(
                        "item references missing project {project}"
                    )));
                }
            }
            if key.starts_with(b"event/") {
                events += 1;
                let sequence = value
                    .get("sequence")
                    .and_then(Value::as_i64)
                    .ok_or_else(|| Error::Storage("event missing sequence".into()))?;
                if sequence <= last_event {
                    return Err(Error::Storage(
                        "event sequence is not strictly increasing".into(),
                    ));
                }
                last_event = sequence;
            }
        }
        Ok(object(vec![
            ("ok", Value::Bool(true)),
            ("records", Value::Int(records)),
            ("items", Value::Int(items)),
            ("events", Value::Int(events)),
            ("lastEvent", Value::Int(last_event)),
        ]))
    }

    pub fn backup(&self, output: &Path) -> Result<Value> {
        let parent = output.parent().unwrap_or_else(|| Path::new("."));
        fs::create_dir_all(parent).map_err(storage)?;
        let staged = output.with_extension(format!("backup-{}.tmp", std::process::id()));
        let manifest_path = manifest_path(output);
        let staged_manifest =
            manifest_path.with_extension(format!("manifest-{}.tmp", std::process::id()));
        let mut db = self
            .db
            .lock()
            .map_err(|_| Error::Storage("lock poisoned".into()))?;
        db.commit().map_err(storage)?;
        fs::copy(self.path.as_path(), &staged).map_err(storage)?;
        File::open(&staged)
            .and_then(|file| file.sync_all())
            .map_err(storage)?;
        let (bytes, sha256) = file_digest(&staged)?;
        let manifest = object(vec![
            ("version", Value::Int(1)),
            ("database", string(&self.path.display().to_string())),
            ("backup", string(&output.display().to_string())),
            ("bytes", Value::Int(bytes as i64)),
            ("sha256", string(&sha256)),
            ("createdAt", Value::Int(now())),
        ]);
        fs::write(&staged_manifest, manifest.to_json()).map_err(storage)?;
        File::open(&staged_manifest)
            .and_then(|file| file.sync_all())
            .map_err(storage)?;
        fs::rename(&staged, output).map_err(storage)?;
        fs::rename(&staged_manifest, &manifest_path).map_err(storage)?;
        sync_directory(parent)?;
        Ok(manifest)
    }

    pub fn backup_snapshot(&self, actor: &str) -> Result<Value> {
        if !self.authorize(actor, None, Permission::Admin)? {
            return Err(Error::Forbidden(
                "principal lacks the administrator role".into(),
            ));
        }
        let directory = std::env::var_os("AKURAI_TASKS_BACKUP_STAGING")
            .map(PathBuf::from)
            .unwrap_or_else(|| PathBuf::from("/var/backups/akurai-tasks/staging"));
        let output = directory.join(format!("akurai-tasks-{}-{}.db", now(), std::process::id()));
        self.backup(&output)
    }

    pub fn restore(database: &Path, backup: &Path) -> Result<Value> {
        let _lock = Self::acquire_exclusive_lock(database)?;
        let parent = database.parent().unwrap_or_else(|| Path::new("."));
        fs::create_dir_all(parent).map_err(storage)?;
        let manifest_text = fs::read_to_string(manifest_path(backup))
            .map_err(|error| Error::Invalid(format!("backup manifest is required: {error}")))?;
        let manifest = parse(&manifest_text)
            .map_err(|error| Error::Invalid(format!("invalid manifest: {error}")))?;
        if manifest.get("version").and_then(Value::as_i64) != Some(1) {
            return Err(Error::Invalid("unsupported backup manifest version".into()));
        }
        let (source_bytes, source_sha256) = file_digest(backup)?;
        if manifest.get("bytes").and_then(Value::as_i64) != Some(source_bytes as i64)
            || field(&manifest, "sha256") != Some(&source_sha256)
        {
            return Err(Error::Conflict("backup manifest digest mismatch".into()));
        }

        let staged = database.with_extension(format!("restore-{}.tmp", std::process::id()));
        fs::copy(backup, &staged).map_err(storage)?;
        File::open(&staged)
            .and_then(|file| file.sync_all())
            .map_err(storage)?;
        let staged_store = Store::open(&staged)?;
        staged_store.doctor()?;
        drop(staged_store);

        let replaced = database.with_extension(format!("replaced-{}", now()));
        if database.exists() {
            fs::rename(database, &replaced).map_err(storage)?;
        }
        if let Err(error) = fs::rename(&staged, database) {
            if replaced.exists() {
                let _ = fs::rename(&replaced, database);
            }
            return Err(storage(error));
        }
        if let Err(error) = sync_directory(parent) {
            let _ = fs::remove_file(database);
            if replaced.exists() {
                let _ = fs::rename(&replaced, database);
                let _ = sync_directory(parent);
            }
            return Err(error);
        }
        let restored = Store::open(database)?;
        let doctor = restored.doctor()?;
        drop(restored);
        let (bytes, sha256) = file_digest(database)?;
        Ok(object(vec![
            ("ok", Value::Bool(true)),
            ("database", string(&database.display().to_string())),
            (
                "replaced",
                if replaced.exists() {
                    string(&replaced.display().to_string())
                } else {
                    Value::Null
                },
            ),
            ("bytes", Value::Int(bytes as i64)),
            ("sha256", string(&sha256)),
            ("doctor", doctor),
        ]))
    }
}

fn manifest_path(output: &Path) -> PathBuf {
    output.with_extension(format!(
        "{}manifest.json",
        output
            .extension()
            .and_then(|value| value.to_str())
            .map(|value| format!("{value}."))
            .unwrap_or_default()
    ))
}

fn sync_directory(path: &Path) -> Result<()> {
    File::open(path)
        .and_then(|directory| directory.sync_all())
        .map_err(storage)
}

fn file_digest(path: &Path) -> Result<(u64, String)> {
    let mut file = File::open(path).map_err(storage)?;
    let mut hash = Sha256::new();
    let mut buffer = [0u8; 64 * 1024];
    let mut bytes = 0u64;
    loop {
        let read = file.read(&mut buffer).map_err(storage)?;
        if read == 0 {
            break;
        }
        bytes += read as u64;
        hash.update(&buffer[..read]);
    }
    let mut digest = String::with_capacity(64);
    for byte in hash.finalize() {
        digest.push_str(&format!("{byte:02x}"));
    }
    Ok((bytes, digest))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn backup_doctor_and_restore_preserve_records() {
        let base = std::env::temp_dir().join(format!(
            "akurai-tasks-maint-{}-{}",
            std::process::id(),
            now()
        ));
        fs::create_dir_all(&base).unwrap();
        let lock_target = base.join("locked.db");
        let lock = Store::acquire_exclusive_lock(&lock_target).unwrap();
        assert!(matches!(
            Store::acquire_exclusive_lock(&lock_target),
            Err(Error::Conflict(_))
        ));
        drop(lock);
        drop(Store::acquire_exclusive_lock(&lock_target).unwrap());
        let database = base.join("tasks.db");
        let backup = base.join("backup.db");
        let restored = base.join("restored.db");
        let store = Store::open(&database).unwrap();
        store
            .create_project("OPS", "Operations", vec![], "admin", "project")
            .unwrap();
        assert_eq!(store.doctor().unwrap().get("ok"), Some(&Value::Bool(true)));
        let manifest = store.backup(&backup).unwrap();
        assert!(
            manifest
                .get("sha256")
                .and_then(Value::as_str)
                .unwrap()
                .len()
                == 64
        );
        drop(store);
        Store::restore(&restored, &backup).unwrap();
        let restored_store = Store::open(&restored).unwrap();
        assert_eq!(
            field(&restored_store.get_project("OPS").unwrap(), "name"),
            Some("Operations")
        );
        OpenOptions::new()
            .append(true)
            .open(&backup)
            .unwrap()
            .write_all(b"tampered")
            .unwrap();
        assert!(matches!(
            Store::restore(&base.join("tampered.db"), &backup),
            Err(Error::Conflict(_))
        ));
        drop(restored_store);
        let _ = fs::remove_dir_all(base);
    }
}