trailofbits-dropkit/dropkit/lock.py

"""File-based locking to prevent concurrent dropkit write operations."""

import fcntl
import json
import os
import time
from collections.abc import Generator
from contextlib import contextmanager, suppress
from functools import wraps
from pathlib import Path

import typer
from rich.console import Console

# Lock file location
LOCK_FILE = Path("/tmp/dropkit.lock")

# Default timeout for acquiring lock (seconds)
DEFAULT_TIMEOUT = 30.0

# Polling interval when waiting for lock (seconds)
POLL_INTERVAL = 0.5

console = Console()


class LockError(Exception):
    """Raised when lock acquisition fails."""


class LockInfo:
    """Information about the current lock holder."""

    def __init__(self, pid: int, command: str):
        self.pid = pid
        self.command = command

    @classmethod
    def from_file(cls, lock_file: Path) -> "LockInfo | None":
        """Read lock info from file, returns None if unreadable."""
        try:
            if lock_file.exists():
                content = lock_file.read_text().strip()
                if content:
                    data = json.loads(content)
                    return cls(pid=data.get("pid", 0), command=data.get("command", "unknown"))
        except (json.JSONDecodeError, OSError, KeyError):
            pass
        return None

    def to_json(self) -> str:
        """Serialize lock info to JSON."""
        return json.dumps({"pid": self.pid, "command": self.command})

    def is_process_alive(self) -> bool:
        """Check if the lock holder process is still running."""
        try:
            os.kill(self.pid, 0)  # Signal 0 checks existence without killing
            return True
        except OSError:
            return False


def _write_lock_info(lock_fd: int, command: str) -> None:
    """Write lock holder information to the lock file."""
    info = LockInfo(pid=os.getpid(), command=command)
    content = info.to_json().encode()
    os.ftruncate(lock_fd, 0)
    os.lseek(lock_fd, 0, os.SEEK_SET)
    os.write(lock_fd, content)


def _clear_lock_info(lock_fd: int) -> None:
    """Clear lock holder information (on release)."""
    with suppress(OSError):
        os.ftruncate(lock_fd, 0)


@contextmanager
def operation_lock(
    command: str,
    timeout: float | None = None,
) -> Generator[None, None, None]:
    """
    Context manager that acquires an exclusive lock for dropkit operations.

    Uses fcntl.flock() which automatically releases the lock when:
    - The context manager exits normally
    - An exception is raised
    - The process crashes or is killed

    Args:
        command: Name of the command acquiring the lock (for debugging)
        timeout: Maximum seconds to wait for lock (default: 30)

    Raises:
        LockError: If lock cannot be acquired within timeout

    Example:
        with operation_lock("create"):
            # Exclusive operation here
            pass
    """
    # Resolve default timeout at runtime (allows patching in tests)
    if timeout is None:
        timeout = DEFAULT_TIMEOUT

    # Ensure parent directory exists (should always exist for /tmp)
    LOCK_FILE.parent.mkdir(parents=True, exist_ok=True)

    # Open file for reading and writing (creates if doesn't exist)
    lock_fd = os.open(str(LOCK_FILE), os.O_RDWR | os.O_CREAT, 0o600)

    try:
        start_time = time.monotonic()
        acquired = False

        while time.monotonic() - start_time < timeout:
            try:
                # Try non-blocking exclusive lock
                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
                acquired = True
                break
            except OSError:
                # Lock is held by another process, wait and retry
                time.sleep(POLL_INTERVAL)

        if not acquired:
            # Timeout - read lock info for error message
            lock_info = LockInfo.from_file(LOCK_FILE)
            if lock_info and lock_info.is_process_alive():
                raise LockError(
                    f"Another dropkit operation is in progress.\n"
                    f"  Command: {lock_info.command}\n"
                    f"  PID: {lock_info.pid}\n"
                    f"Wait for it to complete or check if it's stuck."
                )
            else:
                # Stale lock - try one more time (blocking briefly)
                try:
                    fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
                    acquired = True
                except OSError:
                    raise LockError(
                        f"Could not acquire lock after {timeout} seconds.\nLock file: {LOCK_FILE}"
                    )

        # Write lock holder info
        _write_lock_info(lock_fd, command)

        yield  # Execute the protected code

    finally:
        # Clear lock info and release lock
        _clear_lock_info(lock_fd)
        fcntl.flock(lock_fd, fcntl.LOCK_UN)
        os.close(lock_fd)


def requires_lock(command_name: str):
    """
    Decorator that wraps a function with operation_lock and handles errors.

    Args:
        command_name: Name of the command (for lock info)

    Example:
        @app.command()
        @requires_lock("create")
        def create(...):
            pass
    """

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            try:
                with operation_lock(command_name):
                    return func(*args, **kwargs)
            except LockError as e:
                console.print(f"[red]Error: {e}[/red]")
                raise typer.Exit(1)

        return wrapper

    return decorator