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

Add retrieving types in codequery (#260)

Browse Source 
* Add retrieving types in codequery

* Add fuzzy search option (#271)
This commit is contained in:
Evan Downing
2025-03-27 08:23:01 -04:00
committed by GitHub
parent b44a7dc3ed
commit 5d60b606b7
5 changed files with 207 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
program-model/src/buttercup/program_model/api/tree_sitter.py
+14 -24
View File
@@ -4,10 +4,13 @@ import logging
from dataclasses import dataclass
from pathlib import Path
from functools import lru_cache
from enum import Enum
from buttercup.common.challenge_task import ChallengeTask
from buttercup.program_model.utils.common import Function, FunctionBody
from buttercup.program_model.utils.common import (
Function,
FunctionBody,
TypeDefinition,
TypeDefinitionType,
)
from tree_sitter_language_pack import get_language, get_parser
from buttercup.common.project_yaml import ProjectYaml
@@ -81,26 +84,6 @@ QUERY_STR_TYPES_JAVA = """
"""
@dataclass
class TypeDefinitionType(str, Enum):
"""Enum to store type definition type."""
STRUCT = "struct"
UNION = "union"
ENUM = "enum"
TYPEDEF = "typedef"
PREPROC_TYPE = "preproc_type"
@dataclass
class TypeDefinition:
"""Class to store type definition information."""
name: str
type: TypeDefinitionType
definition: str
@dataclass
class CodeTS:
"""Class to extract information about functions in a challenge project using TreeSitter."""
@@ -199,7 +182,9 @@ class CodeTS:
functions = self.get_functions(file_path)
return functions.get(function_name)
def parse_types_in_code(self, file_path: Path) -> dict[str, TypeDefinition]:
def parse_types_in_code(
self, file_path: Path, typename: str | None = None, fuzzy: bool | None = False
) -> dict[str, TypeDefinition]:
"""Parse the definition of a type in a piece of code."""
logger.debug("Parsing types in code")
code = self.challenge_task.task_dir.joinpath(file_path).read_bytes()
@@ -233,6 +218,10 @@ class CodeTS:
type_definition = code[start_byte : definition_node.end_byte].decode()
name = name_node.text.decode()
if typename and not fuzzy and name != typename:
continue
if typename and fuzzy and typename not in name:
continue
logger.debug("Type name: %s", name)
logger.debug("Type definition: %s", type_definition)
@@ -255,6 +244,7 @@ class CodeTS:
name=name,
type=type_def_type,
definition=type_definition,
definition_line=definition_node.start_point[0],
)
return res
program-model/src/buttercup/program_model/codequery.py
+105 -12
View File
@@ -8,11 +8,15 @@ import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from itertools import groupby
from typing import ClassVar
from typing import ClassVar, Union
from buttercup.common.challenge_task import ChallengeTask
from buttercup.program_model.api.tree_sitter import CodeTS
from buttercup.program_model.utils.common import Function
from buttercup.program_model.utils.common import (
Function,
TypeDefinition,
)
from buttercup.common.project_yaml import ProjectYaml
logger = logging.getLogger(__name__)
@@ -215,7 +219,10 @@ class CodeQuery:
return [result for result in results if result is not None]
def get_functions(
self, function_name: str, file_path: Path | None = None
self,
function_name: str,
file_path: Path | None = None,
fuzzy: bool | None = False,
) -> list[Function]:
"""Get the definition(s) of a function in the codebase or in a specific file."""
cqsearch_args = [
@@ -225,31 +232,117 @@ class CodeQuery:
"2",
"-t",
function_name,
"-e",
"-f" if fuzzy else "-e",
"-u",
]
if file_path:
cqsearch_args += ["-b", file_path.as_posix()]
results = self._run_cqsearch(*cqsearch_args)
logger.debug("cqsearch output: %s", results)
res = []
res: list[Function] = []
results_by_file = groupby(results, key=lambda x: x.file)
for file, results in results_by_file:
if not all(result.value == function_name for result in results):
functions_found = [result.value for result in results]
if not fuzzy and not all(function_name == f for f in functions_found):
logger.warning(
"Function name mismatch, this should not happen: %s != %s",
results[0].value,
"Function name mismatch, this should not happen: %s",
function_name,
)
continue
if fuzzy and not all(function_name in f for f in functions_found):
logger.warning(
"Function name mismatch, this should not happen: %s",
function_name,
)
continue
function = self.ts.get_function(function_name, file)
if function is None:
logger.warning("Function not found in tree-sitter: %s", function_name)
for function in functions_found:
f = self.ts.get_function(function, file)
if f is None:
logger.warning("Function not found in tree-sitter: %s", function)
continue
res.append(f)
return res
def get_types(
self,
type_name: Union[bytes, str],
file_path: Path | None = None,
function_name: str | None = None,
fuzzy: bool | None = False,
) -> list[TypeDefinition]:
"""Finds and return the definition of type named `typename`."""
# Build the cqsearch command to find occurences of the typename in the code
cqsearch_args = [
"-s",
self.CODEQUERY_DB, # Specify the database file path
"-p",
"1", # '1' for symbol
"-t",
type_name, # The name of the type
"-f" if fuzzy else "-e",
"-u", # use full paths
]
if file_path:
cqsearch_args += ["-b", file_path.as_posix()]
results = self._run_cqsearch(*cqsearch_args)
logger.debug("cqsearch output: %s", results)
res: list[TypeDefinition] = []
results_by_file = groupby(results, key=lambda x: x.file)
for file, results in results_by_file:
types_found = [result.value for result in results]
if not fuzzy and not all(type_name == t for t in types_found):
logger.warning(
"Type name mismatch, this should not happen: %s",
type_name,
)
continue
if fuzzy and not all(type_name in t for t in types_found):
logger.warning(
"Type name mismatch, this should not happen: %s",
type_name,
)
continue
res.append(function)
typedefs: dict[str, TypeDefinition] = {}
for typename in types_found:
t = self.ts.parse_types_in_code(file, typename, fuzzy)
if not t:
logger.warning(
"Type definition not found in tree-sitter: %s", typename
)
continue
typedefs.update(t)
if function_name:
# Get the function definition to find its scope
function = self.ts.get_function(function_name, file)
if function:
# Filter type definitions to only include those within the function's scope
filtered_typedefs = {}
for name, typedef in typedefs.items():
# Check if the type definition is within the function's scope
for body in function.bodies:
if (
body.start_line
<= typedef.definition_line
<= body.end_line
):
filtered_typedefs[name] = typedef
break
typedefs = filtered_typedefs
else:
typedefs = {}
res.extend(typedefs.values())
return res
program-model/src/buttercup/program_model/utils/common.py
+29
View File
@@ -1,5 +1,6 @@
import logging
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
@@ -35,3 +36,31 @@ class Function:
bodies: list[FunctionBody] = field(default_factory=list)
"""List of function bodies."""
@dataclass
class TypeDefinitionType(str, Enum):
"""Enum to store type definition type."""
STRUCT = "struct"
UNION = "union"
ENUM = "enum"
TYPEDEF = "typedef"
PREPROC_TYPE = "preproc_type"
@dataclass
class TypeDefinition:
"""Class to store type definition information."""
name: str
"""Name of the type."""
type: TypeDefinitionType
"""Type of the type."""
definition: str
"""Definition of the type."""
definition_line: int
"""Line number of the definition of the type."""
program-model/tests/test_codequery.py
+59
View File
@@ -7,6 +7,7 @@ from pathlib import Path
from buttercup.common.challenge_task import ChallengeTask
from buttercup.program_model.codequery import CodeQuery, CodeQueryPersistent
from buttercup.common.task_meta import TaskMeta
from buttercup.program_model.utils.common import TypeDefinitionType
def setup_dirs(tmp_path: Path) -> Path:
@@ -49,6 +50,13 @@ def setup_dirs(tmp_path: Path) -> Path:
int function4(char *s) {
return strlen(s);
}
""")
(source / "test4.c").write_text("""typedef int myInt;
myInt function5(myInt a, myInt b) {
typedef int myOtherInt;
myOtherInt c = a + b;
return a + b + c;
}
""")
# Create task metadata
@@ -123,6 +131,17 @@ def test_get_functions_multiple(mock_challenge_task: ChallengeTask):
)
def test_get_functions_fuzzy(mock_challenge_task: ChallengeTask):
"""Test that we can get functions (fuzzy search) in codebase"""
codequery = CodeQuery(mock_challenge_task)
functions = codequery.get_functions("function", fuzzy=True)
assert len(functions) == 4
functions = codequery.get_functions("function", Path("test3.c"), fuzzy=True)
assert len(functions) == 2
functions = codequery.get_functions("function3", Path("test3.c"), fuzzy=True)
assert len(functions) == 1
def test_keep_status(
mock_challenge_task: ChallengeTask,
mock_challenge_task_ro: ChallengeTask,
@@ -161,6 +180,46 @@ def test_keep_status(
assert mock_challenge_task_ro.task_dir.exists()
def test_get_types(mock_challenge_task: ChallengeTask):
"""Test that we can get types in codebase"""
codequery = CodeQuery(mock_challenge_task)
types = codequery.get_types("myInt", Path("test3.c"))
assert len(types) == 0
types = codequery.get_types("myInt")
assert len(types) == 1
types = codequery.get_types("myInt", Path("test4.c"))
assert len(types) == 1
assert types[0].name == "myInt"
assert types[0].type == TypeDefinitionType.TYPEDEF
assert types[0].definition == "typedef int myInt;"
assert types[0].definition_line == 0
types = codequery.get_types("myInt", Path("test4.c"), function_name="function5")
assert len(types) == 0
types = codequery.get_types(
"myOtherInt", Path("test4.c"), function_name="function5"
)
assert len(types) == 1
assert types[0].name == "myOtherInt"
assert types[0].type == TypeDefinitionType.TYPEDEF
assert types[0].definition == " typedef int myOtherInt;"
assert types[0].definition_line == 2
def test_get_types_fuzzy(mock_challenge_task: ChallengeTask):
"""Test that we can get types (fuzzy search) in codebase"""
codequery = CodeQuery(mock_challenge_task)
types = codequery.get_types("my", Path("test4.c"), fuzzy=True)
assert len(types) == 2
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) == 2
types = codequery.get_types("myOtherInt", Path("test4.c"), "function5", fuzzy=True)
assert len(types) == 1
@pytest.fixture
def libjpeg_oss_fuzz_task(tmp_path: Path) -> ChallengeTask:
"""Create a challenge task using a real OSS-Fuzz repository."""
program-model/uv.lock
Generated
-1
View File
@@ -1,5 +1,4 @@
version = 1
revision = 1
requires-python = ">=3.10, <3.13"
resolution-markers = [
"python_full_version >= '3.12.4'",