markbase/vendor/smb-server/src/oplock.rs

//! Oplock Manager — global state tracking for opportunistic locking.
//!
//! MS-SMB2 §2.2.13 / §2.2.14: Oplocks allow clients to cache file data locally,
//! reducing network round-trips. The server tracks all opens per file and
//! triggers OPLOCK_BREAK_NOTIFICATION when conflicting opens occur.
//!
//! Also includes LockManager for byte-range locking (MS-SMB2 §2.2.26).

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

use crate::builder::Access;
use crate::path::SmbPath;
use crate::proto::messages::{FileId, OplockBreakNotification, OplockLevel};

/// An entry tracking one client's oplock on a file.
#[derive(Debug, Clone)]
pub struct OplockEntry {
    pub file_id: FileId,
    pub tree_id: u32,
    pub session_id: u64,
    pub oplock_level: u8,
    pub share_access: u32,
    pub granted_access: Access,
    pub connection_id: u64, // For notification routing
}

/// Global oplock state manager (MS-SMB2 §3.3.1.6).
pub struct OplockManager {
    /// File path → all opens with oplocks on that file.
    file_opens: RwLock<HashMap<SmbPath, Vec<OplockEntry>>>,
}

impl OplockManager {
    pub fn new() -> Self {
        Self {
            file_opens: RwLock::new(HashMap::new()),
        }
    }

    /// Check if requested oplock can be granted (MS-SMB2 §3.3.5.9).
    /// Returns the granted level (may be lower than requested).
    pub async fn can_grant(
        &self,
        path: &SmbPath,
        requested_level: u8,
        share_access: u32,
        granted_access: Access,
    ) -> Option<u8> {
        let file_opens = self.file_opens.read().await;
        let existing = file_opens.get(path);

        // No existing opens → grant requested level
        if existing.is_none() || existing.unwrap().is_empty() {
            return Some(requested_level);
        }

        let existing_opens = existing.unwrap();

        // Check ShareAccess conflicts (MS-SMB2 §3.3.5.9)
        for entry in existing_opens {
            // If existing open doesn't allow sharing, deny oplock
            if !share_access_compatible(entry.share_access, share_access) {
                return None;
            }

            // If existing has exclusive/batch oplock, can only grant Level II
            if entry.oplock_level == OplockLevel::Exclusive as u8
                || entry.oplock_level == OplockLevel::Batch as u8
            {
                // Can grant Level II if share access compatible
                if requested_level == OplockLevel::Ii as u8
                    && share_access_compatible(entry.share_access, share_access)
                {
                    return Some(OplockLevel::Ii as u8);
                }
                // Otherwise deny
                return None;
            }
        }

        // All existing opens are Level II → grant requested level
        Some(requested_level)
    }

    /// Register a new open with oplock (MS-SMB2 §3.3.5.9).
    pub async fn register(&self, path: &SmbPath, entry: OplockEntry) {
        let mut file_opens = self.file_opens.write().await;
        file_opens
            .entry(path.clone())
            .or_insert_with(Vec::new)
            .push(entry);
    }

    /// Remove an open when closed (MS-SMB2 §3.3.5.7).
    pub async fn unregister(&self, path: &SmbPath, file_id: &FileId) {
        let mut file_opens = self.file_opens.write().await;
        if let Some(entries) = file_opens.get_mut(path) {
            entries.retain(|e| e.file_id != *file_id);
            if entries.is_empty() {
                file_opens.remove(path);
            }
        }
    }

    /// Trigger oplock break when conflicting open occurs (MS-SMB2 §3.3.5.9).
    /// Returns notifications to send to affected clients.
    pub async fn break_oplock(
        &self,
        path: &SmbPath,
        new_share_access: u32,
        new_granted_access: Access,
    ) -> Vec<OplockBreakNotification> {
        let mut notifications = Vec::new();
        let mut file_opens = self.file_opens.write().await;

        if let Some(entries) = file_opens.get_mut(path) {
            for entry in entries.iter_mut() {
                // Check if new open conflicts with existing oplock
                if !share_access_compatible(entry.share_access, new_share_access) {
                    // Need to break the oplock
                    let new_level = OplockLevel::Ii as u8; // Downgrade to Level II

                    // Build notification (MS-SMB2 §2.2.23.1)
                    notifications.push(OplockBreakNotification {
                        structure_size: 24,
                        oplock_level: new_level,
                        reserved: 0,
                        reserved2: 0,
                        file_id: entry.file_id,
                    });

                    // Update entry's oplock level
                    entry.oplock_level = new_level;
                }
            }
        }

        notifications
    }

    /// Get all opens for a file (for diagnostics).
    pub async fn get_opens(&self, path: &SmbPath) -> Vec<OplockEntry> {
        let file_opens = self.file_opens.read().await;
        file_opens.get(path).cloned().unwrap_or_default()
    }
}

impl Default for OplockManager {
    fn default() -> Self {
        Self::new()
    }
}

/// Check ShareAccess compatibility (MS-SMB2 §3.3.5.9).
pub fn share_access_compatible(existing: u32, new: u32) -> bool {
    const FILE_SHARE_READ: u32 = 0x00000001;
    const FILE_SHARE_WRITE: u32 = 0x00000002;
    const FILE_SHARE_DELETE: u32 = 0x00000004;

    // If existing denies read sharing and new wants read → conflict
    if (existing & FILE_SHARE_READ) == 0 && (new & FILE_SHARE_READ) != 0 {
        return false;
    }

    // If existing denies write sharing and new wants write → conflict
    if (existing & FILE_SHARE_WRITE) == 0 && (new & FILE_SHARE_WRITE) != 0 {
        return false;
    }

    // If existing denies delete sharing and new wants delete → conflict
    if (existing & FILE_SHARE_DELETE) == 0 && (new & FILE_SHARE_DELETE) != 0 {
        return false;
    }

    true
}

// ---------------------------------------------------------------------------
// Byte-range Lock Manager (MS-SMB2 §2.2.26)
// ---------------------------------------------------------------------------

/// A byte-range lock entry (MS-SMB2 §2.2.26.1).
#[derive(Debug, Clone)]
pub struct LockRange {
    pub offset: u64,
    pub length: u64,
    pub exclusive: bool, // FLAG_EXCLUSIVE_LOCK vs FLAG_SHARED_LOCK
    pub session_id: u64,
    pub tree_id: u32,
}

impl OplockManager {
    /// Update oplock level after client acknowledges a break (MS-SMB2 §2.2.24).
    pub async fn update_oplock_level(&self, path: &SmbPath, file_id: FileId, new_level: u8) {
        let mut file_opens = self.file_opens.write().await;
        if let Some(entries) = file_opens.get_mut(path) {
            for entry in entries.iter_mut() {
                if entry.file_id == file_id {
                    entry.oplock_level = new_level;
                    break;
                }
            }
        }
    }
}

/// Global byte-range lock manager (MS-SMB2 §3.3.1.9).
pub struct LockManager {
    /// FileId → active locks on that file.
    file_locks: RwLock<HashMap<FileId, Vec<LockRange>>>,
}

impl LockManager {
    pub fn new() -> Self {
        Self {
            file_locks: RwLock::new(HashMap::new()),
        }
    }

    /// Acquire a lock (MS-SMB2 §3.3.5.14).
    /// Returns Ok(()) if lock acquired, Err if conflict.
    pub async fn acquire(
        &self,
        file_id: &FileId,
        offset: u64,
        length: u64,
        exclusive: bool,
        session_id: u64,
        tree_id: u32,
    ) -> Result<(), String> {
        let mut file_locks = self.file_locks.write().await;
        
        // Check for conflicts with existing locks
        if let Some(locks) = file_locks.get(file_id) {
            for lock in locks {
                // Check if ranges overlap
                if Self::ranges_overlap(offset, length, lock.offset, lock.length) {
                    // If either is exclusive, conflict
                    if exclusive || lock.exclusive {
                        // Same session can upgrade lock
                        if lock.session_id == session_id && lock.tree_id == tree_id {
                            continue; // Allow same session to overlap
                        }
                        return Err("Lock conflict".to_string());
                    }
                }
            }
        }
        
        // No conflict → add lock
        file_locks
            .entry(*file_id)
            .or_insert_with(Vec::new)
            .push(LockRange {
                offset,
                length,
                exclusive,
                session_id,
                tree_id,
            });
        
        Ok(())
    }

    /// Release a lock (MS-SMB2 §3.3.5.14).
    pub async fn release(
        &self,
        file_id: &FileId,
        offset: u64,
        length: u64,
        session_id: u64,
        tree_id: u32,
    ) {
        let mut file_locks = self.file_locks.write().await;
        
        if let Some(locks) = file_locks.get_mut(file_id) {
            locks.retain(|lock| {
                // Keep locks that don't match this release
                !(lock.offset == offset 
                    && lock.length == length 
                    && lock.session_id == session_id 
                    && lock.tree_id == tree_id)
            });
            
            if locks.is_empty() {
                file_locks.remove(file_id);
            }
        }
    }

    /// Check if two byte ranges overlap.
    fn ranges_overlap(offset1: u64, length1: u64, offset2: u64, length2: u64) -> bool {
        let end1 = offset1 + length1;
        let end2 = offset2 + length2;
        
        // Overlap if one range starts before the other ends
        offset1 < end2 && offset2 < end1
    }

    /// Get all locks for a file (for diagnostics).
    pub async fn get_locks(&self, file_id: &FileId) -> Vec<LockRange> {
        let file_locks = self.file_locks.read().await;
        file_locks.get(file_id).cloned().unwrap_or_default()
    }

    /// Clear all locks for a file (when file is closed).
    pub async fn clear(&self, file_id: &FileId) {
        let mut file_locks = self.file_locks.write().await;
        file_locks.remove(file_id);
    }
}

impl Default for LockManager {
    fn default() -> Self {
        Self::new()
    }
}