markbase/markbase-webdav/src/webdav/lock_manager.rs

use dav_server::davpath::DavPath;
use dav_server::ls::{DavLock, DavLockSystem, LsFuture};
use rusqlite::{params, Connection};
use std::fmt;
use std::path::PathBuf;
use std::time::{Duration, SystemTime};
use uuid::Uuid;
use xmltree::Element;

#[derive(Debug, Clone)]
pub struct LockManager {
    db_path: PathBuf,
    user_id: String,
}

impl LockManager {
    pub fn new(user_id: String, db_path: PathBuf) -> Self {
        LockManager { db_path, user_id }
    }

    pub fn init_db(&self) -> Result<(), rusqlite::Error> {
        let conn = Connection::open(&self.db_path)?;
        conn.execute_batch(
            "CREATE TABLE IF NOT EXISTS file_locks (
                lock_id INTEGER PRIMARY KEY AUTOINCREMENT,
                token TEXT UNIQUE NOT NULL,
                path TEXT NOT NULL,
                user_id TEXT NOT NULL,
                principal TEXT,
                owner_xml TEXT,
                timeout_at INTEGER,
                timeout_secs INTEGER,
                shared INTEGER NOT NULL DEFAULT 0,
                deep INTEGER NOT NULL DEFAULT 0,
                created_at INTEGER NOT NULL,
                refreshed_at INTEGER
            );
            CREATE INDEX IF NOT EXISTS idx_locks_path ON file_locks(path);
            CREATE INDEX IF NOT EXISTS idx_locks_token ON file_locks(token);
            CREATE INDEX IF NOT EXISTS idx_locks_user ON file_locks(user_id);
            
            CREATE TABLE IF NOT EXISTS lock_history (
                history_id INTEGER PRIMARY KEY AUTOINCREMENT,
                token TEXT NOT NULL,
                path TEXT NOT NULL,
                user_id TEXT NOT NULL,
                action TEXT NOT NULL,
                timestamp INTEGER NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_history_token ON lock_history(token);",
        )?;
        Ok(())
    }

    fn get_conn(&self) -> Result<Connection, rusqlite::Error> {
        Connection::open(&self.db_path)
    }

    fn lock_to_dav_lock(&self, row: &rusqlite::Row) -> Result<DavLock, rusqlite::Error> {
        let path_str: String = row.get(2)?;
        let principal: Option<String> = row.get(4)?;
        let owner_xml: Option<String> = row.get(5)?;
        let timeout_at_ts: Option<i64> = row.get(6)?;
        let timeout_secs: Option<i64> = row.get(7)?;
        let shared: i32 = row.get(8)?;
        let deep: i32 = row.get(9)?;

        let timeout_at =
            timeout_at_ts.map(|ts| SystemTime::UNIX_EPOCH + Duration::from_secs(ts as u64));

        let timeout = timeout_secs.map(|s| Duration::from_secs(s as u64));

        let owner = owner_xml.and_then(|xml| Element::parse(xml.as_bytes()).ok());

        let token: String = row.get(1)?;

        Ok(DavLock {
            token,
            path: Box::new(DavPath::new(&path_str).unwrap_or_else(|_| DavPath::new("/").unwrap())),
            principal,
            owner: owner.map(Box::new),
            timeout_at,
            timeout,
            shared: shared != 0,
            deep: deep != 0,
        })
    }

    fn lock_to_dav_lock_from_select(
        &self,
        row: &rusqlite::Row,
    ) -> Result<DavLock, rusqlite::Error> {
        let token: String = row.get(0)?;
        let path_str: String = row.get(1)?;
        let principal: Option<String> = row.get(2)?;
        let owner_xml: Option<String> = row.get(3)?;
        let timeout_at_ts: Option<i64> = row.get(4)?;
        let timeout_secs: Option<i64> = row.get(5)?;
        let shared: i32 = row.get(6)?;
        let deep: i32 = row.get(7)?;

        let timeout_at =
            timeout_at_ts.map(|ts| SystemTime::UNIX_EPOCH + Duration::from_secs(ts as u64));

        let timeout = timeout_secs.map(|s| Duration::from_secs(s as u64));

        let owner = owner_xml.and_then(|xml| Element::parse(xml.as_bytes()).ok());

        Ok(DavLock {
            token,
            path: Box::new(DavPath::new(&path_str).unwrap_or_else(|_| DavPath::new("/").unwrap())),
            principal,
            owner: owner.map(Box::new),
            timeout_at,
            timeout,
            shared: shared != 0,
            deep: deep != 0,
        })
    }

    fn cleanup_expired_locks(&self, conn: &Connection) -> Result<(), rusqlite::Error> {
        let now = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;

        conn.execute(
            "DELETE FROM file_locks WHERE timeout_at IS NOT NULL AND timeout_at < ?1",
            params![now],
        )?;
        Ok(())
    }
}

impl DavLockSystem for LockManager {
    fn lock(
        &'_ self,
        path: &DavPath,
        principal: Option<&str>,
        owner: Option<&Element>,
        timeout: Option<Duration>,
        shared: bool,
        deep: bool,
    ) -> LsFuture<'_, Result<DavLock, DavLock>> {
        let path_str = path.to_string();
        let path_owned = path.clone();
        let token = format!("urn:uuid:{}", Uuid::new_v4());
        let principal_str = principal.map(|s| s.to_string());
        let owner_clone = owner.cloned();
        let owner_xml = owner.and_then(|e| {
            let mut buf = Vec::new();
            e.write(&mut buf).ok()?;
            String::from_utf8(buf).ok()
        });

        let timeout_secs = timeout.map(|d| d.as_secs() as i64);
        let timeout_at = timeout.map(|d| {
            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;
            now + d.as_secs() as i64
        });

        let now = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;

        Box::pin(async move {
            let conn = match self.get_conn() {
                Ok(c) => c,
                Err(_) => {
                    return Err(DavLock {
                        token: String::new(),
                        path: Box::new(path_owned.clone()),
                        principal: principal_str.clone(),
                        owner: owner_clone.map(Box::new),
                        timeout_at: None,
                        timeout,
                        shared,
                        deep,
                    });
                }
            };

            self.cleanup_expired_locks(&conn).ok();

            let existing_lock = conn.query_row(
                "SELECT token, path, principal, owner_xml, timeout_at, timeout_secs, shared, deep 
                 FROM file_locks 
                 WHERE path = ?1 AND user_id = ?2",
                params![path_str, &self.user_id],
                |row| self.lock_to_dav_lock_from_select(row),
            );

            if let Ok(conflict) = existing_lock {
                if !(shared && conflict.shared) {
                    return Err(conflict);
                }
            }

            conn.execute(
                "INSERT INTO file_locks 
                 (token, path, user_id, principal, owner_xml, timeout_at, timeout_secs, shared, deep, created_at)
                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
                params![
                    &token,
                    &path_str,
                    &self.user_id,
                    &principal_str,
                    &owner_xml,
                    timeout_at,
                    timeout_secs,
                    if shared { 1 } else { 0 },
                    if deep { 1 } else { 0 },
                    now,
                ],
            ).ok();

            conn.execute(
                "INSERT INTO lock_history (token, path, user_id, action, timestamp)
                 VALUES (?1, ?2, ?3, 'lock', ?4)",
                params![&token, &path_str, &self.user_id, now],
            )
            .ok();

            Ok(DavLock {
                token,
                path: Box::new(path_owned.clone()),
                principal: principal_str,
                owner: owner_clone.map(Box::new),
                timeout_at: timeout_at
                    .map(|t| SystemTime::UNIX_EPOCH + Duration::from_secs(t as u64)),
                timeout,
                shared,
                deep,
            })
        })
    }

    fn unlock(&'_ self, path: &DavPath, token: &str) -> LsFuture<'_, Result<(), ()>> {
        let path_str = path.to_string();
        let token_str = token.to_string();

        Box::pin(async move {
            let conn = match self.get_conn() {
                Ok(c) => c,
                Err(_) => return Err(()),
            };

            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;

            let rows = conn.execute(
                "DELETE FROM file_locks WHERE token = ?1 AND path = ?2 AND user_id = ?3",
                params![&token_str, &path_str, &self.user_id],
            );

            if let Ok(deleted) = rows {
                if deleted > 0 {
                    conn.execute(
                        "INSERT INTO lock_history (token, path, user_id, action, timestamp)
                         VALUES (?1, ?2, ?3, 'unlock', ?4)",
                        params![&token_str, &path_str, &self.user_id, now],
                    )
                    .ok();
                    return Ok(());
                }
            }

            Err(())
        })
    }

    fn refresh(
        &'_ self,
        path: &DavPath,
        token: &str,
        timeout: Option<Duration>,
    ) -> LsFuture<'_, Result<DavLock, ()>> {
        let path_str = path.to_string();
        let token_str = token.to_string();
        let timeout_secs = timeout.map(|d| d.as_secs() as i64);
        let timeout_at = timeout.map(|d| {
            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;
            now + d.as_secs() as i64
        });

        Box::pin(async move {
            let conn = match self.get_conn() {
                Ok(c) => c,
                Err(_) => return Err(()),
            };

            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;

            let updated = conn.execute(
                "UPDATE file_locks 
                 SET timeout_at = ?1, timeout_secs = ?2, refreshed_at = ?3 
                 WHERE token = ?4 AND path = ?5 AND user_id = ?6",
                params![
                    timeout_at,
                    timeout_secs,
                    now,
                    &token_str,
                    &path_str,
                    &self.user_id
                ],
            );

            if let Ok(rows) = updated {
                if rows > 0 {
                    conn.execute(
                        "INSERT INTO lock_history (token, path, user_id, action, timestamp)
                         VALUES (?1, ?2, ?3, 'refresh', ?4)",
                        params![&token_str, &path_str, &self.user_id, now],
                    )
                    .ok();

                    return conn
                        .query_row(
                            "SELECT * FROM file_locks WHERE token = ?1",
                            params![&token_str],
                            |row| self.lock_to_dav_lock(row),
                        )
                        .map(|lock| {
                            if let Some(t) = timeout {
                                DavLock {
                                    timeout: Some(t),
                                    ..lock
                                }
                            } else {
                                lock
                            }
                        })
                        .map_err(|_| ());
                }
            }

            Err(())
        })
    }

    fn check(
        &'_ self,
        path: &DavPath,
        principal: Option<&str>,
        ignore_principal: bool,
        deep: bool,
        submitted_tokens: &[String],
    ) -> LsFuture<'_, Result<(), DavLock>> {
        let path_str = path.to_string();
        let path_owned = path.clone();
        let principal_str = principal.map(|s| s.to_string());
        let tokens = submitted_tokens.to_vec();
        let user_id = self.user_id.clone();

        Box::pin(async move {
            let conn = match self.get_conn() {
                Ok(c) => c,
                Err(_) => return Ok(()),
            };

            self.cleanup_expired_locks(&conn).ok();

            let mut stmt = conn
                .prepare("SELECT * FROM file_locks WHERE path = ?1 AND user_id = ?2")
                .map_err(|_| DavLock {
                    token: String::new(),
                    path: Box::new(path_owned.clone()),
                    principal: None,
                    owner: None,
                    timeout_at: None,
                    timeout: None,
                    shared: false,
                    deep: false,
                })?;

            let locks = stmt
                .query_map(params![&path_str, &user_id], |row| {
                    self.lock_to_dav_lock(row)
                })
                .map_err(|_| DavLock {
                    token: String::new(),
                    path: Box::new(path_owned.clone()),
                    principal: None,
                    owner: None,
                    timeout_at: None,
                    timeout: None,
                    shared: false,
                    deep: false,
                })?;

            for lock in locks.flatten() {
                if tokens.contains(&lock.token) {
                    continue;
                }

                if ignore_principal {
                    continue;
                }

                if let Some(ref lock_principal) = lock.principal {
                    if let Some(ref check_principal) = principal_str {
                        if lock_principal == check_principal {
                            continue;
                        }
                    }
                }

                if deep && lock.deep {
                    return Err(lock);
                }

                if !deep {
                    return Err(lock);
                }
            }

            Ok(())
        })
    }

    fn discover(&'_ self, path: &DavPath) -> LsFuture<'_, Vec<DavLock>> {
        let path_str = path.to_string();
        let user_id = self.user_id.clone();

        Box::pin(async move {
            let conn = match self.get_conn() {
                Ok(c) => c,
                Err(_) => return Vec::new(),
            };

            self.cleanup_expired_locks(&conn).ok();

            let mut stmt =
                match conn.prepare("SELECT * FROM file_locks WHERE path = ?1 AND user_id = ?2") {
                    Ok(s) => s,
                    Err(_) => return Vec::new(),
                };

            let locks = stmt.query_map(params![&path_str, &user_id], |row| {
                self.lock_to_dav_lock(row)
            });

            match locks {
                Ok(l) => l.filter_map(|r| r.ok()).collect(),
                Err(_) => Vec::new(),
            }
        })
    }

    fn delete(&'_ self, path: &DavPath) -> LsFuture<'_, Result<(), ()>> {
        let path_str = path.to_string();
        let user_id = self.user_id.clone();

        Box::pin(async move {
            let conn = match self.get_conn() {
                Ok(c) => c,
                Err(_) => return Err(()),
            };

            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;

            conn.execute(
                "INSERT INTO lock_history (token, path, user_id, action, timestamp)
                 SELECT token, path, user_id, 'delete', ?1 
                 FROM file_locks 
                 WHERE path LIKE ?2 AND user_id = ?3",
                params![now, format!("{}%", path_str), &user_id],
            )
            .ok();

            conn.execute(
                "DELETE FROM file_locks WHERE path LIKE ?1 AND user_id = ?2",
                params![format!("{}%", path_str), &user_id],
            )
            .map(|_| ())
            .map_err(|_| ())
        })
    }
}

impl fmt::Display for LockManager {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "LockManager(user={}, db={:?})",
            self.user_id, self.db_path
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use dav_server::davpath::DavPath;
    use std::time::Duration;
    use tempfile::tempdir;

    #[test]
    fn test_lock_manager_creation() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_locks.sqlite");
        let manager = LockManager::new("test_user".to_string(), db_path.clone());

        assert_eq!(manager.user_id, "test_user");
        assert_eq!(manager.db_path, db_path);
    }

    #[test]
    fn test_init_db() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_locks.sqlite");
        let manager = LockManager::new("test_user".to_string(), db_path);

        manager.init_db().expect("Failed to initialize database");

        let conn = Connection::open(&manager.db_path).unwrap();
        conn.execute("SELECT * FROM file_locks LIMIT 1", [])
            .unwrap();
        conn.execute("SELECT * FROM lock_history LIMIT 1", [])
            .unwrap();
    }

    #[tokio::test]
    async fn test_lock_and_unlock() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_locks.sqlite");
        let manager = LockManager::new("test_user".to_string(), db_path);
        manager.init_db().unwrap();

        let path = DavPath::new("/test/file.txt").unwrap();

        let lock_result = manager.lock(&path, None, None, None, false, false).await;

        match lock_result {
            Ok(lock) => {
                assert!(lock.token.starts_with("urn:uuid:"));
                assert_eq!(lock.path.as_ref(), &path);

                let unlock_result = manager.unlock(&path, &lock.token).await;
                assert!(unlock_result.is_ok());
            }
            Err(_) => {
                panic!("Lock should succeed on first attempt");
            }
        }
    }

    #[tokio::test]
    async fn test_lock_conflict() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_locks.sqlite");
        let manager = LockManager::new("test_user".to_string(), db_path);
        manager.init_db().unwrap();

        let path = DavPath::new("/test/file.txt").unwrap();

        let lock1 = manager
            .lock(&path, Some("user1"), None, None, false, false)
            .await;
        assert!(lock1.is_ok());

        let lock2 = manager
            .lock(&path, Some("user2"), None, None, false, false)
            .await;
        assert!(lock2.is_err());
    }

    #[tokio::test]
    async fn test_lock_discover() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_locks.sqlite");
        let manager = LockManager::new("test_user".to_string(), db_path);
        manager.init_db().unwrap();

        let path = DavPath::new("/test/file.txt").unwrap();

        let lock = manager
            .lock(&path, None, None, None, false, false)
            .await
            .unwrap();

        let discovered = manager.discover(&path).await;
        assert_eq!(discovered.len(), 1);
        assert_eq!(discovered[0].token, lock.token);
    }

    #[tokio::test]
    async fn test_lock_refresh() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_locks.sqlite");
        let manager = LockManager::new("test_user".to_string(), db_path);
        manager.init_db().unwrap();

        let path = DavPath::new("/test/file.txt").unwrap();
        let timeout = Duration::from_secs(60);

        let lock = manager
            .lock(&path, None, None, Some(timeout), false, false)
            .await
            .unwrap();

        let refreshed = manager
            .refresh(&path, &lock.token, Some(Duration::from_secs(120)))
            .await;
        assert!(refreshed.is_ok());

        let refreshed_lock = refreshed.unwrap();
        assert_eq!(refreshed_lock.timeout, Some(Duration::from_secs(120)));
    }
}