trailofbits-buttercup/fuzzer_runner/tests/test_temp_dir.py

import os
import tempfile
import threading
import time
import unittest
from unittest.mock import MagicMock, patch

from buttercup.fuzzer_runner.temp_dir import _scratch_path_var, get_temp_dir, patched_temp_dir


class TestGetTempDir(unittest.TestCase):
    """Test cases for the get_temp_dir function."""

    def setUp(self):
        """Set up test fixtures."""
        # Clear any existing context variable value
        _scratch_path_var.set(None)

    def test_get_temp_dir_with_scratch_path(self):
        """Test get_temp_dir uses scratch path when available."""
        test_scratch_path = "/test/scratch/path"

        with patch("buttercup.fuzzer_runner.temp_dir.shell.create_directory") as mock_create_dir:
            # Set the context variable
            token = _scratch_path_var.set(test_scratch_path)

            try:
                result = get_temp_dir()

                # Should use scratch path as prefix
                expected_path = os.path.join(test_scratch_path, f"temp-{os.getpid()}")
                self.assertEqual(result, expected_path)
                mock_create_dir.assert_called_once_with(expected_path)
            finally:
                _scratch_path_var.reset(token)

    def test_get_temp_dir_fallback_with_fuzz_inputs_disk(self):
        """Test get_temp_dir falls back to FUZZ_INPUTS_DISK when no scratch path."""
        test_fuzz_inputs_path = "/test/fuzz/inputs"

        with (
            patch("buttercup.fuzzer_runner.temp_dir.environment.get_value") as mock_get_value,
            patch("buttercup.fuzzer_runner.temp_dir.shell.create_directory") as mock_create_dir,
        ):
            mock_get_value.return_value = test_fuzz_inputs_path

            result = get_temp_dir(use_fuzz_inputs_disk=True)

            expected_path = os.path.join(test_fuzz_inputs_path, f"temp-{os.getpid()}")
            self.assertEqual(result, expected_path)
            mock_get_value.assert_called_once_with("FUZZ_INPUTS_DISK", tempfile.gettempdir())
            mock_create_dir.assert_called_once_with(expected_path)

    def test_get_temp_dir_fallback_without_fuzz_inputs_disk(self):
        """Test get_temp_dir falls back to system temp when use_fuzz_inputs_disk=False."""
        with patch("buttercup.fuzzer_runner.temp_dir.shell.create_directory") as mock_create_dir:
            result = get_temp_dir(use_fuzz_inputs_disk=False)

            expected_path = os.path.join(tempfile.gettempdir(), f"temp-{os.getpid()}")
            self.assertEqual(result, expected_path)
            mock_create_dir.assert_called_once_with(expected_path)

    def test_get_temp_dir_fallback_to_system_temp(self):
        """Test get_temp_dir falls back to system temp when FUZZ_INPUTS_DISK not set."""
        with (
            patch("buttercup.fuzzer_runner.temp_dir.environment.get_value") as mock_get_value,
            patch("buttercup.fuzzer_runner.temp_dir.shell.create_directory") as mock_create_dir,
        ):
            # Return system temp dir when FUZZ_INPUTS_DISK is not set
            mock_get_value.return_value = tempfile.gettempdir()

            result = get_temp_dir(use_fuzz_inputs_disk=True)

            expected_path = os.path.join(tempfile.gettempdir(), f"temp-{os.getpid()}")
            self.assertEqual(result, expected_path)
            mock_create_dir.assert_called_once_with(expected_path)


class TestPatchedTempDir(unittest.TestCase):
    """Test cases for the patched_temp_dir context manager."""

    def test_patched_temp_dir_basic_functionality(self):
        """Test basic functionality of patched_temp_dir context manager."""
        with patch("buttercup.fuzzer_runner.temp_dir.scratch_dir") as mock_scratch_dir:
            # Mock scratch_dir context manager
            mock_scratch = MagicMock()
            mock_scratch.path = "/mock/scratch/path"
            mock_scratch_dir.return_value.__enter__.return_value = mock_scratch
            mock_scratch_dir.return_value.__exit__.return_value = None

            with patched_temp_dir() as scratch:
                # Verify we get the scratch directory
                self.assertEqual(scratch, mock_scratch)

                # Verify the context variable is set correctly
                self.assertEqual(_scratch_path_var.get(), "/mock/scratch/path")

    def test_patched_temp_dir_patches_clusterfuzz_function(self):
        """Test that patched_temp_dir properly patches the clusterfuzz function."""
        with patch("buttercup.fuzzer_runner.temp_dir.scratch_dir") as mock_scratch_dir:
            # Mock scratch_dir context manager
            mock_scratch = MagicMock()
            mock_scratch.path = "/mock/scratch/path"
            mock_scratch_dir.return_value.__enter__.return_value = mock_scratch
            mock_scratch_dir.return_value.__exit__.return_value = None

            # Test that the patching works by calling the function directly
            with patched_temp_dir():
                with patch("buttercup.fuzzer_runner.temp_dir.shell.create_directory") as mock_create_dir:
                    # Import and call the clusterfuzz function - it should be patched to use our implementation
                    from clusterfuzz._internal.bot.fuzzers import utils

                    result = utils.get_temp_dir()

                    # Should return path using our scratch directory
                    expected_path = os.path.join("/mock/scratch/path", f"temp-{os.getpid()}")
                    self.assertEqual(result, expected_path)
                    mock_create_dir.assert_called_once_with(expected_path)

    def test_patched_temp_dir_context_variable_cleanup(self):
        """Test that context variable is properly cleaned up after context exit."""
        with patch("buttercup.fuzzer_runner.temp_dir.scratch_dir") as mock_scratch_dir:
            # Mock scratch_dir context manager
            mock_scratch = MagicMock()
            mock_scratch.path = "/mock/scratch/path"
            mock_scratch_dir.return_value.__enter__.return_value = mock_scratch
            mock_scratch_dir.return_value.__exit__.return_value = None

            # Verify context variable is None before
            self.assertIsNone(_scratch_path_var.get())

            with patched_temp_dir():
                # Verify context variable is set inside context
                self.assertEqual(_scratch_path_var.get(), "/mock/scratch/path")

            # Verify context variable is cleaned up after
            self.assertIsNone(_scratch_path_var.get())

    def test_patched_temp_dir_nested_contexts(self):
        """Test that nested patched_temp_dir contexts work correctly."""
        with patch("buttercup.fuzzer_runner.temp_dir.scratch_dir") as mock_scratch_dir:
            # Mock different scratch directories for nested contexts
            mock_scratch1 = MagicMock()
            mock_scratch1.path = "/mock/scratch/path1"
            mock_scratch2 = MagicMock()
            mock_scratch2.path = "/mock/scratch/path2"

            mock_scratch_dir.return_value.__enter__.side_effect = [mock_scratch1, mock_scratch2]
            mock_scratch_dir.return_value.__exit__.return_value = None

            with patched_temp_dir():
                self.assertEqual(_scratch_path_var.get(), "/mock/scratch/path1")

                with patched_temp_dir():
                    self.assertEqual(_scratch_path_var.get(), "/mock/scratch/path2")

                # Should restore to outer context
                self.assertEqual(_scratch_path_var.get(), "/mock/scratch/path1")

            # Should be cleaned up completely
            self.assertIsNone(_scratch_path_var.get())

    def test_patched_temp_dir_exception_handling(self):
        """Test that context variable is cleaned up even when exceptions occur."""
        with patch("buttercup.fuzzer_runner.temp_dir.scratch_dir") as mock_scratch_dir:
            mock_scratch = MagicMock()
            mock_scratch.path = "/mock/scratch/path"
            mock_scratch_dir.return_value.__enter__.return_value = mock_scratch
            mock_scratch_dir.return_value.__exit__.return_value = None

            try:
                with patched_temp_dir():
                    self.assertEqual(_scratch_path_var.get(), "/mock/scratch/path")
                    raise ValueError("Test exception")
            except ValueError:
                pass

            # Should still be cleaned up after exception
            self.assertIsNone(_scratch_path_var.get())


class TestThreadSafety(unittest.TestCase):
    """Test cases for thread safety of the temp_dir module."""

    def test_thread_safety_multiple_contexts(self):
        """Test that multiple threads can use patched_temp_dir simultaneously."""
        results = {}
        errors = []

        def thread_worker(thread_id: int):
            try:
                with patch("buttercup.fuzzer_runner.temp_dir.scratch_dir") as mock_scratch_dir:
                    mock_scratch = MagicMock()
                    mock_scratch.path = f"/mock/scratch/thread{thread_id}"
                    mock_scratch_dir.return_value.__enter__.return_value = mock_scratch
                    mock_scratch_dir.return_value.__exit__.return_value = None

                    with patched_temp_dir():
                        # Store the context variable value for this thread
                        results[thread_id] = _scratch_path_var.get()

                        # Sleep to allow other threads to run
                        time.sleep(0.1)

                        # Verify the value hasn't changed
                        if _scratch_path_var.get() != f"/mock/scratch/thread{thread_id}":
                            errors.append(f"Thread {thread_id} context variable changed!")

            except Exception as e:
                errors.append(f"Thread {thread_id} error: {e}")

        # Start multiple threads
        threads = []
        for i in range(5):
            thread = threading.Thread(target=thread_worker, args=(i,))
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Verify no errors occurred
        self.assertEqual(errors, [], f"Thread safety errors: {errors}")

        # Verify each thread had its own context value
        for i in range(5):
            self.assertEqual(results[i], f"/mock/scratch/thread{i}")

        # Verify context is clean after all threads
        self.assertIsNone(_scratch_path_var.get())


if __name__ == "__main__":
    unittest.main()