angr/tests/knowledge_plugins/functions/test_function_manager_lmdb.py

#!/usr/bin/env python3
# pylint:disable=missing-class-docstring,protected-access
"""Test cases for FunctionManager LMDB save/load and LRU cache functionality."""

from __future__ import annotations

__package__ = __package__ or "tests.knowledge_plugins.functions"  # pylint:disable=redefined-builtin

import os
import unittest

import angr
from tests.common import bin_location

test_location = os.path.join(bin_location, "tests")


class TestFunctionManagerLMDB(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.bin_path = os.path.join(test_location, "x86_64", "fauxware")

    def test_default_unlimited_cache(self):
        """Test that default cache is unlimited (no eviction)."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions

        assert fm.cache_limit is None, "Default should be None (unlimited)"
        assert fm.spilled_function_count == 0, "No functions should be spilled by default"
        assert fm.cached_function_count == fm.total_function_count, "All functions should be in memory"

    def test_set_cache_limit(self):
        """Test setting cache limit triggers eviction."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total_count = len(fm)

        # Set cache limit
        cache_limit = 5
        fm.cache_limit = cache_limit

        assert fm.cached_function_count <= cache_limit, (
            f"Cache limit not respected: {fm.cached_function_count} > {cache_limit}"
        )
        assert fm.total_function_count == total_count, "Total function count should be preserved"
        assert fm.spilled_function_count == total_count - fm.cached_function_count, "Spilled count incorrect"

    def test_cache_properties(self):
        """Test cache monitoring properties."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total = len(fm)

        # Set small cache limit
        fm.cache_limit = 3

        # Verify properties
        assert fm.cached_function_count <= 3
        assert fm.spilled_function_count >= 0
        assert fm.total_function_count == total
        assert fm.cached_function_count + fm.spilled_function_count == fm.total_function_count

    def test_access_spilled_function(self):
        """Test that accessing a spilled function loads it from LMDB."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions

        # Set small cache limit
        fm.cache_limit = 3

        # Find a spilled function
        if fm.spilled_function_count == 0:
            self.skipTest("No spilled functions to test")

        spilled_addr = next(iter(fm._spilled_addrs))

        # Access the spilled function
        func = fm[spilled_addr]

        # Verify it was loaded
        assert func is not None, "Failed to load spilled function"
        assert func.addr == spilled_addr, "Loaded function has wrong address"

    def test_dynamic_cache_limit_decrease(self):
        """Test decreasing cache limit dynamically."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total = len(fm)

        # Start with larger limit
        fm.cache_limit = 10
        assert fm.cached_function_count <= 10

        # Decrease limit
        fm.cache_limit = 5
        assert fm.cached_function_count <= 5

        # Further decrease
        fm.cache_limit = 2
        assert fm.cached_function_count <= 2

        # Total should be preserved
        assert fm.total_function_count == total

    def test_dynamic_cache_limit_increase(self):
        """Test increasing cache limit dynamically."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total = len(fm)

        # Start with small limit
        fm.cache_limit = 3
        cached_before = fm.cached_function_count

        # Increase limit (shouldn't auto-load more)
        fm.cache_limit = 10
        assert fm.cached_function_count >= cached_before  # May increase due to access

        # Total preserved
        assert fm.total_function_count == total

    def test_contains_with_spilled(self):
        """Test __contains__ checks both in-memory and spilled functions."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        all_addrs = set(fm)

        # Set small cache limit
        fm.cache_limit = 3

        # All addresses should still be "in" the function manager
        for addr in all_addrs:
            assert addr in fm, f"Address {hex(addr)} should be in function manager"

    def test_iter_with_spilled(self):
        """Test __iter__ includes both in-memory and spilled functions."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        all_addrs_before = set(fm)

        # Set small cache limit
        fm.cache_limit = 3

        # Iteration should still include all addresses
        all_addrs_after = set(fm)
        assert all_addrs_before == all_addrs_after, "Iteration should include all functions"

    def test_len_with_spilled(self):
        """Test __len__ returns total count including spilled."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total_before = len(fm)

        # Set small cache limit
        fm.cache_limit = 3

        # Length should still return total count
        assert len(fm) == total_before, "len() should return total count"

    def test_spilled_function_call_sites(self):
        """Test that call sites survive spilling to LMDB and reloading."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        main_func = fm.function(name="main")
        assert main_func is not None
        main_addr = main_func.addr
        fm[main_addr]._add_call_site(0x4007D3, None, None)
        call_sites_before = dict(fm[main_addr]._call_sites)
        assert len(call_sites_before) > 1

        # Evict main from the cache
        fm.cache_limit = 1
        other_addr = next(addr for addr in fm if addr != main_addr)
        _ = fm[other_addr]
        assert main_addr in fm._spilled_addrs

        loaded = fm[main_addr]
        assert dict(loaded._call_sites) == call_sites_before

    def test_delete_spilled_function(self):
        """Test deleting a spilled function."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total_before = len(fm)

        # Set small cache limit
        fm.cache_limit = 3

        if fm.spilled_function_count == 0:
            self.skipTest("No spilled functions to test")

        # Get a spilled address
        spilled_addr = next(iter(fm._spilled_addrs))

        # Delete it
        del fm[spilled_addr]

        # Verify it's removed
        assert spilled_addr not in fm
        assert len(fm) == total_before - 1

    def test_delete_cached_function(self):
        """Test deleting a cached (in-memory) function."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total_before = len(fm)

        # Set cache limit
        fm.cache_limit = 5

        # Get a cached address
        cached_addr = next(iter(fm._function_map.keys()))

        # Delete it
        del fm[cached_addr]

        # Verify it's removed
        assert cached_addr not in fm
        assert len(fm) == total_before - 1

    def test_clear_with_spilled(self):
        """Test clear() removes both in-memory and spilled functions."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions

        # Set cache limit
        fm.cache_limit = 3

        # Clear
        fm.clear()

        assert len(fm) == 0
        assert fm.cached_function_count == 0
        assert fm.spilled_function_count == 0

    def test_copy_with_spilled(self):
        """Test copy() works with spilled functions."""
        proj = angr.Project(self.bin_path, auto_load_libs=False)
        proj.analyses.CFGFast()

        fm = proj.kb.functions
        total = len(fm)

        # Set cache limit
        fm.cache_limit = 5

        # Copy
        fm_copy = fm.copy()

        # Verify copy has all functions
        assert fm_copy.total_function_count == total


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