admin/trailofbits-buttercup
1
0
Fork 0
mirror of https://github.com/trailofbits/buttercup synced 2026-06-21 14:11:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Modify fuzzy search (#484)

Browse Source
This commit is contained in:
Evan Downing
2025-04-29 15:37:25 -04:00
committed by GitHub
parent 558c0aeb97
commit b1a1c19d99
3 changed files with 64 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
program-model/src/buttercup/program_model/__cli__.py
+2 -2
View File
@@ -31,8 +31,8 @@ def main() -> None:
command = get_subcommand(settings)
setup_package_logger(__name__, settings.log_level)
if isinstance(command, ServeCommand):
init_telemetry("program-model")
redis = Redis.from_url(command.redis_url, decode_responses=False) # type: ignore[unreachable]
init_telemetry("program-model") # type: ignore[unreachable]
redis = Redis.from_url(command.redis_url, decode_responses=False)
with ProgramModel(
sleep_time=command.sleep_time,
redis=redis,
program-model/src/buttercup/program_model/codequery.py
+57 -31
View File
@@ -10,7 +10,6 @@ from pathlib import Path
from itertools import groupby
from typing import ClassVar
import rapidfuzz
from functools import lru_cache
from buttercup.common.challenge_task import ChallengeTask, ChallengeTaskError
from buttercup.program_model.api.tree_sitter import CodeTS
@@ -101,10 +100,6 @@ class CodeQuery:
self._create_codequery_db()
logger.info("CodeQuery DB created successfully.")
# TODO(Evan): Is this too much to keep in memory?
self._get_all_functions = lru_cache(maxsize=1000)(self._get_all_functions) # type: ignore [method-assign]
self._get_all_types = lru_cache(maxsize=1000)(self._get_all_types) # type: ignore [method-assign]
def _verify_requirements(self) -> None:
"""Verify that the required commands are installed."""
required_commands = ["cscope", "ctags", "cqmakedb", "cqsearch"]
@@ -343,9 +338,20 @@ class CodeQuery:
) -> list[Function]:
"""Get the definition(s) of a function in the codebase or in a specific
file. File paths are based on the challenge task container structure
(e.g. /src)."""
(e.g. /src).
The order of the results is (1) exact matches and (2) fuzzy matches sorted in descending order of similarity.
NOTE: Fuzzy search will be disabled if a file path is provided.
"""
if fuzzy and file_path:
logger.warning(
"Fuzzy search will be disabled because file path %s was provided.",
file_path,
)
# FIXME(Evan): Sometimes cscope doesn't identify a function. They can be found by looking for symbols.
results_all: list[CQSearchResult] = []
results: list[CQSearchResult] = []
for search_type in ["1", "2"]: # 1 for symbols, 2 for functions
cqsearch_args = [
"-s",
@@ -354,18 +360,16 @@ class CodeQuery:
search_type,
"-t",
function_name,
"-f" if fuzzy else "-e",
"-e",
"-u",
]
if file_path:
cqsearch_args += ["-b", file_path.as_posix()]
results = self._run_cqsearch(*cqsearch_args)
results_all.extend(results)
results.extend(self._run_cqsearch(*cqsearch_args))
# Extended fuzzy matching
if fuzzy:
if fuzzy and file_path is None:
# Fuzzy match the function name against all functions in the codebase
fuzzy_matches: list[tuple[CQSearchResult, float]] = sorted(
[
@@ -383,15 +387,15 @@ class CodeQuery:
len(fuzzy_matches),
function_name,
)
results_all.extend(fuzzy_matches)
results.extend(fuzzy_matches)
logger.info(
"Found %d instances (fuzzy or exact) of function %s",
len(results_all),
len(results),
function_name,
)
res: set[Function] = set()
results_by_file = groupby(results_all, key=lambda x: x.file)
results_by_file = groupby(results, key=lambda x: x.file)
for file, file_results in results_by_file:
file_results_list = list(file_results)
functions_found = list(set(result.value for result in file_results_list))
@@ -440,7 +444,13 @@ class CodeQuery:
"Found %d functions for %s after filtering", len(res), function_name
)
return self._rebase_functions_file_paths(list(res))
# Sort in same order as results
results_value = [r.value for r in results]
res_sorted: list[Function] = sorted(
res, key=lambda x: results_value.index(x.name)
)
return self._rebase_functions_file_paths(res_sorted)
def get_callers(self, function: Function | str) -> list[Function]:
"""Get the callers of a function. File paths are based on the challenge
@@ -451,13 +461,13 @@ class CodeQuery:
function_name = function.name
cqsearch_args = [
"-s",
self.CODEQUERY_DB, # Specify the database file path
self.CODEQUERY_DB,
"-p",
"6",
"-t",
function_name,
"-e",
"-u", # use full paths
"-u",
]
results = self._run_cqsearch(*cqsearch_args)
@@ -480,13 +490,13 @@ class CodeQuery:
function_name = function.name
cqsearch_args = [
"-s",
self.CODEQUERY_DB, # Specify the database file path
self.CODEQUERY_DB,
"-p",
"7",
"-t",
function_name,
"-e",
"-u", # use full paths
"-u",
]
results = self._run_cqsearch(*cqsearch_args)
@@ -513,16 +523,27 @@ class CodeQuery:
fuzzy_threshold: int = 80,
) -> list[TypeDefinition]:
"""Finds and return the definition of type named `typename`. File paths
are based on the challenge task container structure (e.g. /src)."""
are based on the challenge task container structure (e.g. /src).
The order of the results is (1) exact matches and (2) fuzzy matches sorted in descending order of similarity.
NOTE: Fuzzy search will be disabled if a file path is provided.
"""
if fuzzy and file_path:
logger.warning(
"Fuzzy search will be disabled because file path %s was provided.",
file_path,
)
cqsearch_args = [
"-s",
self.CODEQUERY_DB, # Specify the database file path
self.CODEQUERY_DB,
"-p",
"1", # '1' for symbol
"1",
"-t",
str(type_name), # Convert to string to ensure type safety
"-f" if fuzzy else "-e",
"-u", # use full paths
str(type_name),
"-e",
"-u",
]
if file_path:
cqsearch_args += ["-b", file_path.as_posix()]
@@ -530,7 +551,7 @@ class CodeQuery:
results: list[CQSearchResult] = list(self._run_cqsearch(*cqsearch_args))
# Extended fuzzy matching
if fuzzy:
if fuzzy and file_path is None:
# Fuzzy match the function name against all functions in the codebase
fuzzy_matches: list[tuple[CQSearchResult, float]] = sorted(
[
@@ -601,22 +622,27 @@ class CodeQuery:
logger.info("Found %d types for %s after filtering", len(res), type_name)
# Sort in same order as results
results_value = [r.value for r in results]
res_sorted: list[TypeDefinition] = sorted(
res, key=lambda x: results_value.index(x.name)
)
# Rebase the file paths
res = self._rebase_types_file_paths(res)
return res
return self._rebase_types_file_paths(res_sorted)
def get_type_calls(self, type_definition: TypeDefinition) -> list[TypeUsageInfo]:
"""Get the calls to a type definition. File paths are based on the challenge
task container structure (e.g. /src)."""
cqsearch_args = [
"-s",
self.CODEQUERY_DB, # Specify the database file path
self.CODEQUERY_DB,
"-p",
"8",
"-t",
type_definition.name,
"-e",
"-u", # use full paths
"-u",
]
results = self._run_cqsearch(*cqsearch_args)
program-model/tests/test_codequery.py
+5 -3
View File
@@ -203,9 +203,9 @@ def test_get_functions_fuzzy(mock_c_challenge_task: ChallengeTask):
functions = codequery.get_functions("function", fuzzy=True)
assert len(functions) == 4
functions = codequery.get_functions("function", Path("test3.c"), fuzzy=True)
assert len(functions) == 4
assert len(functions) == 0
functions = codequery.get_functions("function3", Path("test3.c"), fuzzy=True)
assert len(functions) == 4
assert len(functions) == 1
def test_keep_status(
@@ -281,12 +281,14 @@ def test_get_types_fuzzy(mock_c_challenge_task: ChallengeTask):
with patch("subprocess.run", side_effect=mock_docker_run(mock_c_challenge_task)):
codequery = CodeQuery(mock_c_challenge_task)
types = codequery.get_types("my", Path("test4.c"), fuzzy=True)
assert len(types) == 2
assert len(types) == 0
types = codequery.get_types("myInt", Path("test4.c"), fuzzy=True)
assert len(types) == 1
types = codequery.get_types("myOtherInt", Path("test4.c"), fuzzy=True)
assert len(types) == 1
types = codequery.get_types("my", fuzzy=True)
assert len(types) == 0
types = codequery.get_types("my", fuzzy=True, fuzzy_threshold=10)
assert len(types) == 2
types = codequery.get_types("myOtherInt", Path("test4.c"), "function5", fuzzy=True)
assert len(types) == 1