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> {
        // Fast-path: no entries or single entry can't conflict with itself
        let entry_count = {
            let file_opens = self.file_opens.read().await;
            file_opens.get(path).map_or(0, |e| e.len())
        };
        if entry_count <= 1 {
            return Vec::new();
        }

        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() {
                if !share_access_compatible(entry.share_access, new_share_access) {
                    let new_level = OplockLevel::Ii as u8;

                    notifications.push(OplockBreakNotification {
                        structure_size: 24,
                        oplock_level: new_level,
                        reserved: 0,
                        reserved2: 0,
                        file_id: entry.file_id,
                    });

                    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;
                }
            }
        }
    }
}

/// Lease state flags (MS-SMB2 §2.2.13.2).
pub const SMB2_LEASE_READ: u32 = 0x01;
pub const SMB2_LEASE_HANDLE: u32 = 0x02;
pub const SMB2_LEASE_WRITE: u32 = 0x04;

/// Lease entry for LeaseManager.
#[derive(Debug, Clone)]
pub struct LeaseEntry {
    pub lease_key: [u8; 16],
    pub lease_state: u32,
    pub lease_flags: u32,
    pub file_id: FileId,
    pub path: SmbPath,
    pub session_id: u64,
    pub tree_id: u32,
}

/// Global lease manager for SMB 3.x (MS-SMB2 §3.3.1.9).
pub struct LeaseManager {
    /// LeaseKey → LeaseEntry.
    leases: RwLock<HashMap<[u8; 16], LeaseEntry>>,
}

impl LeaseManager {
    pub fn new() -> Self {
        Self {
            leases: RwLock::new(HashMap::new()),
        }
    }
    
    /// Register a lease on CREATE (MS-SMB2 §3.3.5.9).
    pub async fn register(&self, entry: LeaseEntry) {
        let mut leases = self.leases.write().await;
        leases.insert(entry.lease_key, entry);
    }
    
    /// Remove a lease on CLOSE.
    pub async fn unregister(&self, lease_key: &[u8; 16]) {
        let mut leases = self.leases.write().await;
        leases.remove(lease_key);
    }
    
    /// Check if lease can be granted (MS-SMB2 §3.3.5.9).
    pub async fn can_grant(&self, requested_state: u32) -> bool {
        // Simple check: allow lease if no conflicting leases exist
        let leases = self.leases.read().await;
        for entry in leases.values() {
            // Check for conflicts
            if (entry.lease_state & SMB2_LEASE_WRITE) != 0 && (requested_state & SMB2_LEASE_READ) != 0 {
                return false; // WRITE lease conflicts with READ request
            }
            if (entry.lease_state & SMB2_LEASE_HANDLE) != 0 && (requested_state & SMB2_LEASE_HANDLE) != 0 {
                return false; // HANDLE lease conflicts with HANDLE request
            }
        }
        true
    }
    
    /// Break lease when conflicting access occurs (MS-SMB2 §3.3.5.10).
    pub async fn break_lease(&self, requested_state: u32) -> Vec<LeaseBreakNotification> {
        // Fast-path: no leases to break
        if self.leases.read().await.is_empty() {
            return Vec::new();
        }

        let mut leases = self.leases.write().await;
        let mut notifications = Vec::new();
        
        for (key, entry) in leases.iter_mut() {
            // Check if lease needs to break
            let needs_break = (entry.lease_state & SMB2_LEASE_WRITE) != 0 && (requested_state & SMB2_LEASE_READ) != 0;
            
            if needs_break {
                // Break to READ lease (or none)
                entry.lease_state = SMB2_LEASE_READ;
                entry.lease_flags |= 0x02; // SMB2_LEASE_FLAG_BREAKING
                
                notifications.push(LeaseBreakNotification {
                    structure_size: 36,
                    lease_key: *key,
                    current_lease_state: entry.lease_state,
                    new_lease_state: SMB2_LEASE_READ,
                    break_reason: 0,
                    lease_flags: entry.lease_flags,
                    access_mask: 0,
                    share_mask: 0,
                });
            }
        }
        
        notifications
    }
}

/// SMB2_LEASE_BREAK_NOTIFICATION (MS-SMB2 §2.2.26).
#[derive(Debug, Clone)]
pub struct LeaseBreakNotification {
    pub structure_size: u16,
    pub lease_key: [u8; 16],
    pub current_lease_state: u32,
    pub new_lease_state: u32,
    pub break_reason: u32,
    pub lease_flags: u32,
    pub access_mask: u32,
    pub share_mask: u32,
}

impl LeaseBreakNotification {
    pub fn write_to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(36);
        buf.extend_from_slice(&self.structure_size.to_le_bytes());
        buf.extend_from_slice(&self.lease_key);
        buf.extend_from_slice(&self.current_lease_state.to_le_bytes());
        buf.extend_from_slice(&self.new_lease_state.to_le_bytes());
        buf.extend_from_slice(&self.break_reason.to_le_bytes());
        buf.extend_from_slice(&self.lease_flags.to_le_bytes());
        buf.extend_from_slice(&self.access_mask.to_le_bytes());
        buf.extend_from_slice(&self.share_mask.to_le_bytes());
        buf
    }
}

/// 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()
    }
}