Files
trailofbits-buttercup/program-model/tests/test_tree_sitter.py
T
Riccardo Schirone 47f38b8863 fix: apply ruff import sorting fixes to all source files
Run `ruff check --fix` across all components to fix I001 import sorting
violations in src/ and test/ directories.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-26 15:17:12 +01:00

824 lines
25 KiB
Python

from dataclasses import dataclass
from pathlib import Path
import pytest
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.task_meta import TaskMeta
from buttercup.program_model.api.tree_sitter import CodeTS, TypeDefinitionType
@dataclass(frozen=True)
class FunctionInfo:
num_bodies: int
body_excerpts: list[str]
@pytest.fixture
def task_dir(tmp_path: Path) -> Path:
"""Create a mock challenge task directory structure."""
# Create the main directories
base_path = tmp_path / "task_rw"
oss_fuzz = base_path / "fuzz-tooling" / "fuzz-tooling"
source = base_path / "src" / "example_project"
diffs = base_path / "diff" / "my-diff"
oss_fuzz.mkdir(parents=True, exist_ok=True)
source.mkdir(parents=True, exist_ok=True)
diffs.mkdir(parents=True, exist_ok=True)
# Create mock project.yaml file
project_yaml_path = oss_fuzz / "projects" / "example_project" / "project.yaml"
project_yaml_path.parent.mkdir(parents=True, exist_ok=True)
project_yaml_path.write_text(
"language: c\n"
"sanitizers:\n"
" - address\n"
" - memory\n"
" - undefined\n"
"architectures:\n"
" - x86_64\n"
"fuzzing_engines:\n"
" - afl\n"
" - honggfuzz\n"
" - libfuzzer\n",
)
# Create some mock patch files
(diffs / "patch1.diff").write_text("mock patch 1")
(diffs / "patch2.diff").write_text("mock patch 2")
# Create a mock helper.py file
helper_path = oss_fuzz / "infra" / "helper.py"
helper_path.parent.mkdir(parents=True, exist_ok=True)
helper_path.write_text("import sys;\nsys.exit(0)\n")
# Create a mock test.txt file
(source / "test.txt").write_text("mock test content")
# Create a test C file with two functions
test_c_content = """#include <stdio.h>
// Forward declarations - these should not be matched
struct forward_struct;
union forward_union;
enum forward_enum;
// Preprocessor type definitions
#define MY_TYPE my_struct_t
#define ANOTHER_TYPE struct my_struct
struct struct_name {
int a;
int b;
};
int add(int a, int b) {
return a + b;
}
void print_hello(void) {
printf("Hello, World!\\n");
}
"""
(source / "test.c").write_text(test_c_content)
test2_c_content = """#include <stdio.h>
#ifdef TEST
int add(int a, int b) {
return a + b;
}
#else
double add(double a, double b) {
return a + b;
}
#endif
"""
(source / "test2.c").write_text(test2_c_content)
# Create task metadata
TaskMeta(
project_name="example_project",
focus="example_project",
task_id="task-id-tree-sitter",
metadata={
"task_id": "task-id-tree-sitter",
"round_id": "testing",
"team_id": "tob",
},
).save(base_path)
return base_path
@pytest.fixture
def task_dir_cpp(tmp_path: Path) -> Path:
"""Create a mock challenge task directory structure."""
base_path = tmp_path / "task_rw"
oss_fuzz = base_path / "fuzz-tooling" / "fuzz-tooling"
source = base_path / "src" / "example_project"
oss_fuzz.mkdir(parents=True, exist_ok=True)
source.mkdir(parents=True, exist_ok=True)
# Create mock project.yaml file
project_yaml_path = oss_fuzz / "projects" / "example_project" / "project.yaml"
project_yaml_path.parent.mkdir(parents=True, exist_ok=True)
project_yaml_path.write_text(
"language: cpp\n"
"sanitizers:\n"
" - address\n"
" - memory\n"
" - undefined\n"
"architectures:\n"
" - x86_64\n"
"fuzzing_engines:\n"
" - afl\n"
" - honggfuzz\n"
" - libfuzzer\n",
)
# Create a mock helper.py file
helper_path = oss_fuzz / "infra" / "helper.py"
helper_path.parent.mkdir(parents=True, exist_ok=True)
helper_path.write_text("import sys;\nsys.exit(0)\n")
# Create a test C file (as C++ project) with two functions
test_c_content = """#include <stdio.h>
// Forward declarations - these should not be matched
struct forward_struct;
union forward_union;
enum forward_enum;
// Preprocessor type definitions
#define MY_TYPE my_struct_t
#define ANOTHER_TYPE struct my_struct
struct struct_name {
int a;
int b;
};
int add(int a, int b) {
return a + b;
}
void print_hello(void) {
printf("Hello, World!\\n");
}
"""
(source / "test.cpp").write_text(test_c_content)
test2_c_content = """#include <stdio.h>
#ifdef TEST
int add(int a, int b) {
return a + b;
}
#else
double add(double a, double b) {
return a + b;
}
#endif
"""
(source / "test2.cpp").write_text(test2_c_content)
test3_cpp_content = """#include <cstdio>
class Calculator {
private:
int value;
public:
Calculator(int initial_value = 0) : value(initial_value) {}
int add(int a, int b) {
return a + b;
}
void setValue(int new_value) {
value = new_value;
}
int getValue() const {
return value;
}
};
int main() {
Calculator calc;
int result = calc.add(5, 3);
calc.setValue(result);
printf("Result: %d\n", calc.getValue());
return 0;
}
"""
(source / "test3.cpp").write_text(test3_cpp_content)
# Create task metadata
TaskMeta(
project_name="example_project",
focus="example_project",
task_id="task-id-tree-sitter-cpp",
metadata={
"task_id": "task-id-tree-sitter-cpp",
"round_id": "testing",
"team_id": "tob",
},
).save(base_path)
return base_path
@pytest.fixture
def challenge_task_readonly(task_dir: Path) -> ChallengeTask:
"""Create a mock challenge task for testing."""
return ChallengeTask(
read_only_task_dir=task_dir,
)
@pytest.fixture
def cpp_challenge_task_readonly(task_dir_cpp: Path) -> ChallengeTask:
"""Create a mock challenge task for testing."""
return ChallengeTask(
read_only_task_dir=task_dir_cpp,
)
@pytest.fixture
def java_task_dir(tmp_path: Path) -> Path:
"""Create a mock challenge task directory structure."""
# Create the main directories
base_path = tmp_path / "task_rw"
oss_fuzz = base_path / "fuzz-tooling" / "fuzz-tooling"
source = base_path / "src" / "example_project"
diffs = base_path / "diff" / "my-diff"
oss_fuzz.mkdir(parents=True, exist_ok=True)
source.mkdir(parents=True, exist_ok=True)
diffs.mkdir(parents=True, exist_ok=True)
# Create mock project.yaml file
project_yaml_path = oss_fuzz / "projects" / "example_project" / "project.yaml"
project_yaml_path.parent.mkdir(parents=True, exist_ok=True)
project_yaml_path.write_text("language: java\n")
# Create a mock helper.py file
helper_path = oss_fuzz / "infra" / "helper.py"
helper_path.parent.mkdir(parents=True, exist_ok=True)
helper_path.write_text("import sys;\nsys.exit(0)\n")
# Create a mock test.txt file
(source / "test.txt").write_text("mock test content")
# Create task metadata
TaskMeta(
project_name="example_project",
focus="example_project",
task_id="task-id-tree-sitter",
metadata={
"task_id": "task-id-tree-sitter",
"round_id": "testing",
"team_id": "tob",
},
).save(base_path)
return base_path
@pytest.fixture
def java_challenge_task_readonly(java_task_dir: Path) -> ChallengeTask:
"""Create a mock challenge task for testing."""
return ChallengeTask(
read_only_task_dir=java_task_dir,
)
def test_get_functions_code_c(challenge_task_readonly: ChallengeTask):
"""Test getting function code from a C file."""
code_ts = CodeTS(challenge_task_readonly)
functions = code_ts.get_functions(Path("src/example_project/test.c"))
assert "add" in functions
assert "print_hello" in functions
add_function = functions["add"]
assert len(add_function.bodies) == 1
assert "int add(int a, int b)" in add_function.bodies[0].body
assert "return a + b;" in add_function.bodies[0].body
print_hello_function = functions["print_hello"]
assert len(print_hello_function.bodies) == 1
assert "void print_hello(void)" in print_hello_function.bodies[0].body
assert 'printf("Hello, World!\\n");' in print_hello_function.bodies[0].body
def test_get_function_c(challenge_task_readonly: ChallengeTask):
"""Test getting a function from a C file."""
code_ts = CodeTS(challenge_task_readonly)
function = code_ts.get_function("add", Path("src/example_project/test.c"))
assert function is not None
assert function.name == "add"
assert function.file_path == Path("src/example_project/test.c")
assert len(function.bodies) == 1
assert "int add(int a, int b)" in function.bodies[0].body
assert "return a + b;" in function.bodies[0].body
assert function.bodies[0].start_line == 17
assert function.bodies[0].end_line == 19
def test_get_function_multiple_definitions_c(challenge_task_readonly: ChallengeTask):
"""Test getting a function from a C file with multiple definitions."""
code_ts = CodeTS(challenge_task_readonly)
function = code_ts.get_function("add", Path("src/example_project/test2.c"))
assert function is not None
assert function.name == "add"
assert function.file_path == Path("src/example_project/test2.c")
assert len(function.bodies) == 2
assert "#ifdef TEST" in function.bodies[0].body
assert "int add(int a, int b)" in function.bodies[0].body
assert "double add(double a, double b)" in function.bodies[1].body
assert "#else" in function.bodies[1].body
assert function.bodies[0].start_line == 3
assert function.bodies[0].end_line == 6
assert function.bodies[1].start_line == 7
assert function.bodies[1].end_line == 10
def test_get_type_definition_types(challenge_task_readonly: ChallengeTask):
"""Test getting different types of definitions."""
code_ts = CodeTS(challenge_task_readonly)
types = code_ts.parse_types_in_code(Path("src/example_project/test.c"))
# Test preprocessor type definitions
type_def = types["MY_TYPE"]
assert type_def is not None
assert type_def.type == TypeDefinitionType.PREPROC_TYPE
assert "#define MY_TYPE my_struct_t" in type_def.definition
type_def = types["ANOTHER_TYPE"]
assert type_def is not None
assert type_def.type == TypeDefinitionType.PREPROC_TYPE
assert "#define ANOTHER_TYPE struct my_struct" in type_def.definition
@pytest.mark.parametrize(
"function_name,file_path,function_info",
[
(
"png_icc_check_length",
"src/example-libpng/png.c",
FunctionInfo(
num_bodies=1,
body_excerpts=[
"""if (!icc_check_length(png_ptr, colorspace, name, profile_length))
return 0;
/* This needs to be here because the 'normal' check is in
* png_decompress_chunk, yet this happens after the attempt to
* png_malloc_base the required data. We only need this on read; on write
* the caller supplies the profile buffer so libpng doesn't allocate it. See
* the call to icc_check_length below (the write case).
*/
# ifdef PNG_SET_USER_LIMITS_SUPPORTED
else if (png_ptr->user_chunk_malloc_max > 0 &&
png_ptr->user_chunk_malloc_max < profile_length)
return png_icc_profile_error(png_ptr, colorspace, name, profile_length,
"exceeds application limits");""",
],
),
),
(
"png_pow10",
"src/example-libpng/png.c",
FunctionInfo(
num_bodies=1,
body_excerpts=[
"""/* Utility used below - a simple accurate power of ten from an integral
* exponent.
*/
static double
png_pow10(int power)
{
int recip = 0;
double d = 1;
/* Handle negative exponent with a reciprocal at the end because
* 10 is exact whereas .1 is inexact in base 2
*/
if (power < 0)
{
if (power < DBL_MIN_10_EXP) return 0;
recip = 1; power = -power;
}""",
],
),
),
(
"png_check_IHDR",
"src/example-libpng/png.c",
FunctionInfo(
num_bodies=1,
body_excerpts=[
"""int error = 0;
/* Check for width and height valid values */
if (width == 0)
{
png_warning(png_ptr, "Image width is zero in IHDR");
error = 1;
}
if (width > PNG_UINT_31_MAX)
{
png_warning(png_ptr, "Invalid image width in IHDR");
error = 1;
}""",
],
),
),
],
)
@pytest.mark.integration
def test_libpng_indexing(
libpng_oss_fuzz_task: ChallengeTask,
function_name: str,
file_path: str,
function_info: FunctionInfo,
):
"""Test that we can parse libpng code using tree-sitter."""
code_ts = CodeTS(libpng_oss_fuzz_task)
function = code_ts.get_function(function_name, Path(file_path))
assert function is not None
assert len(function.bodies) == function_info.num_bodies
for body in function_info.body_excerpts:
assert any([body in x.body for x in function.bodies])
def test_get_field_type(java_challenge_task_readonly: ChallengeTask):
"""Test getting the type of a field of a type definition."""
code_ts = CodeTS(java_challenge_task_readonly)
typedef = b"""class Person {
age = 30;
String something;
public String child() {
return this.child.toString();
}
Person2 child;
}
"""
type_name = code_ts.get_field_type_name(typedef, "child")
assert type_name == "Person2"
def test_get_method_return_type(java_challenge_task_readonly: ChallengeTask):
"""Test getting the return type of a method of a type definition."""
code_ts = CodeTS(java_challenge_task_readonly)
typedef = b"""class Person {
int getName = 40;
public SuperClass getname() {
int getName = 1;
return new SuperClass(getName);
}
public String getName() {
return "John";
}
}
"""
type_name = code_ts.get_method_return_type_name(typedef, "getName")
assert type_name == "String"
typedef = b"""public interface LoggerRepository {
/**
* Add a {@link HierarchyEventListener} event to the repository.
*
* @param listener The listener
*/
void addHierarchyEventListener(HierarchyEventListener listener);
/**
* Returns whether this repository is disabled for a given
* level. The answer depends on the repository threshold and the
* <code>level</code> parameter. See also {@link #setThreshold}
* method.
*
* @param level The level
* @return whether this repository is disabled.
*/
boolean isDisabled(int level);
/**
* Set the repository-wide threshold. All logging requests below the
* threshold are immediately dropped. By default, the threshold is
* set to <code>Level.ALL</code> which has the lowest possible rank.
*
* @param level The level
*/
void setThreshold(Level level);
/**
* Another form of {@link #setThreshold(Level)} accepting a string
* parameter instead of a <code>Level</code>.
*
* @param val The threshold value
*/
void setThreshold(String val);
void emitNoAppenderWarning(Category cat);
/**
* Get the repository-wide threshold. See {@link #setThreshold(Level)} for an explanation.
*
* @return the level.
*/
Level getThreshold();
Logger getLogger(String name);
Logger getLogger(String name, LoggerFactory factory);
Logger getRootLogger();
Logger exists(String name);
void shutdown();
@SuppressWarnings("rawtypes")
Enumeration getCurrentLoggers();
/**
* Deprecated. Please use {@link #getCurrentLoggers} instead.
*
* @return an enumeration of loggers.
*/
@SuppressWarnings("rawtypes")
Enumeration getCurrentCategories();
void fireAddAppenderEvent(Category logger, Appender appender);
void resetConfiguration();
}"""
type_name = code_ts.get_method_return_type_name(typedef, "getLogger")
assert type_name == "Logger"
def test_get_functions_code_cpp_c_content(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting function C code from a C++ file."""
code_ts = CodeTS(cpp_challenge_task_readonly)
functions = code_ts.get_functions(Path("src/example_project/test.cpp"))
assert "add" in functions
assert "print_hello" in functions
add_function = functions["add"]
assert len(add_function.bodies) == 1
assert "int add(int a, int b)" in add_function.bodies[0].body
assert "return a + b;" in add_function.bodies[0].body
print_hello_function = functions["print_hello"]
assert len(print_hello_function.bodies) == 1
assert "void print_hello(void)" in print_hello_function.bodies[0].body
assert 'printf("Hello, World!\\n");' in print_hello_function.bodies[0].body
def test_get_function_multiple_definitions_cpp(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting a function from a C file with multiple definitions."""
code_ts = CodeTS(cpp_challenge_task_readonly)
function = code_ts.get_function("add", Path("src/example_project/test2.cpp"))
assert function is not None
assert function.name == "add"
assert function.file_path == Path("src/example_project/test2.cpp")
assert len(function.bodies) == 2
assert "#ifdef TEST" in function.bodies[0].body
assert "int add(int a, int b)" in function.bodies[0].body
assert "double add(double a, double b)" in function.bodies[1].body
assert "#else" in function.bodies[1].body
assert function.bodies[0].start_line == 3
assert function.bodies[0].end_line == 6
assert function.bodies[1].start_line == 7
assert function.bodies[1].end_line == 10
def test_get_functions_code_cpp_cpp_content(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting C++ functions and methods from a C++ file."""
code_ts = CodeTS(cpp_challenge_task_readonly)
functions = code_ts.get_functions(Path("src/example_project/test3.cpp"))
# Test that we found the expected functions
assert "add" in functions
assert "setValue" in functions
assert "getValue" in functions
assert "main" in functions
# Test the add method
add_function = functions["add"]
assert len(add_function.bodies) == 1
assert "int add(int a, int b)" in add_function.bodies[0].body
assert "return a + b;" in add_function.bodies[0].body
# Test the setValue method
set_value_function = functions["setValue"]
assert len(set_value_function.bodies) == 1
assert "void setValue(int new_value)" in set_value_function.bodies[0].body
assert "value = new_value;" in set_value_function.bodies[0].body
# Test the getValue method
get_value_function = functions["getValue"]
assert len(get_value_function.bodies) == 1
assert "int getValue() const" in get_value_function.bodies[0].body
assert "return value;" in get_value_function.bodies[0].body
# Test the main function
main_function = functions["main"]
assert len(main_function.bodies) == 1
assert "int main()" in main_function.bodies[0].body
assert "Calculator calc;" in main_function.bodies[0].body
def test_get_types_cpp(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting C++ type definitions from a C++ file."""
code_ts = CodeTS(cpp_challenge_task_readonly)
types = code_ts.parse_types_in_code(Path("src/example_project/test3.cpp"))
# Test that we found the Calculator class
assert "Calculator" in types
calculator_type = types["Calculator"]
assert calculator_type.type == TypeDefinitionType.CLASS
assert "class Calculator" in calculator_type.definition
assert "int value;" in calculator_type.definition
def test_get_function_cpp(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting a specific C++ function from a file."""
code_ts = CodeTS(cpp_challenge_task_readonly)
function = code_ts.get_function("add", Path("src/example_project/test3.cpp"))
assert function is not None
assert function.name == "add"
assert function.file_path == Path("src/example_project/test3.cpp")
assert len(function.bodies) == 1
assert "int add(int a, int b)" in function.bodies[0].body
assert "return a + b;" in function.bodies[0].body
def test_get_functions_cpp_with_struct(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting functions from a C++ file with struct definitions."""
# Create a test file with struct
test_struct_content = """#include <iostream>
struct Point {
int x;
int y;
Point(int x, int y) : x(x), y(y) {}
int getX() const {
return x;
}
void setX(int new_x) {
x = new_x;
}
};
int distance(Point p1, Point p2) {
return (p1.getX() - p2.getX()) * (p1.getX() - p2.getX());
}
"""
# Write the test file
test_file_path = cpp_challenge_task_readonly.task_dir / "src" / "example_project" / "test_struct.cpp"
test_file_path.write_text(test_struct_content)
code_ts = CodeTS(cpp_challenge_task_readonly)
functions = code_ts.get_functions(Path("src/example_project/test_struct.cpp"))
# Test that we found the expected functions
assert "getX" in functions
assert "setX" in functions
assert "distance" in functions
# Test the getX method
get_x_function = functions["getX"]
assert len(get_x_function.bodies) == 1
assert "int getX() const" in get_x_function.bodies[0].body
assert "return x;" in get_x_function.bodies[0].body
# Test the setX method
set_x_function = functions["setX"]
assert len(set_x_function.bodies) == 1
assert "void setX(int new_x)" in set_x_function.bodies[0].body
assert "x = new_x;" in set_x_function.bodies[0].body
# Test the distance function
distance_function = functions["distance"]
assert len(distance_function.bodies) == 1
assert "int distance(Point p1, Point p2)" in distance_function.bodies[0].body
assert "return (p1.getX() - p2.getX())" in distance_function.bodies[0].body
def test_get_types_cpp_with_struct(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting type definitions from a C++ file with struct."""
# Create a test file with struct
test_struct_content = """#include <iostream>
struct Point {
int x;
int y;
};
class Rectangle {
Point topLeft;
Point bottomRight;
};
"""
# Write the test file
test_file_path = cpp_challenge_task_readonly.task_dir / "src" / "example_project" / "test_types.cpp"
test_file_path.write_text(test_struct_content)
code_ts = CodeTS(cpp_challenge_task_readonly)
types = code_ts.parse_types_in_code(Path("src/example_project/test_types.cpp"))
# Test that we found the expected types
assert "Point" in types
assert "Rectangle" in types
# Test the Point struct
point_type = types["Point"]
assert point_type.type == TypeDefinitionType.STRUCT
assert "struct Point" in point_type.definition
assert "int x;" in point_type.definition
assert "int y;" in point_type.definition
# Test the Rectangle class
rectangle_type = types["Rectangle"]
assert rectangle_type.type == TypeDefinitionType.CLASS
assert "class Rectangle" in rectangle_type.definition
assert "Point topLeft;" in rectangle_type.definition
assert "Point bottomRight;" in rectangle_type.definition
def test_get_functions_cpp_qualified_name(cpp_challenge_task_readonly: ChallengeTask):
"""Test getting C++ functions with qualified names (class::method syntax)."""
# Create a test file with qualified function names
test_qualified_content = """#include <iostream>
class LibRaw {
public:
void parse_exif(INT64 base);
};
void LibRaw::parse_exif(INT64 base)
{
unsigned entries, tag, type, len, c;
double expo, ape;
INT64 save;
// Function implementation
entries = get2();
for (c = 0; c < entries; c++) {
tag = get2();
type = get2();
len = get4();
save = ftell(ifp);
if (type == 5 && len == 1) {
expo = getreal(type);
ape = pow(2, expo);
}
}
}
int main() {
LibRaw raw;
raw.parse_exif(0);
return 0;
}
"""
# Write the test file
test_file_path = cpp_challenge_task_readonly.task_dir / "src" / "example_project" / "test_qualified.cpp"
test_file_path.write_text(test_qualified_content)
code_ts = CodeTS(cpp_challenge_task_readonly)
functions = code_ts.get_functions(Path("src/example_project/test_qualified.cpp"))
# Test that we found the expected functions
assert "parse_exif" in functions
assert "main" in functions
# Test the parse_exif function (should capture the implementation, not the declaration)
parse_exif_function = functions["parse_exif"]
assert len(parse_exif_function.bodies) == 1
assert "void LibRaw::parse_exif(INT64 base)" in parse_exif_function.bodies[0].body
assert "unsigned entries, tag, type, len, c;" in parse_exif_function.bodies[0].body
assert "entries = get2();" in parse_exif_function.bodies[0].body
# Test the main function
main_function = functions["main"]
assert len(main_function.bodies) == 1
assert "int main()" in main_function.bodies[0].body
assert "LibRaw raw;" in main_function.bodies[0].body