diff --git a/src/exploit_iq_commons/utils/functions_parsers/c_lang_function_parsers.py b/src/exploit_iq_commons/utils/functions_parsers/c_lang_function_parsers.py index 7289e39ca..c9a664f18 100644 --- a/src/exploit_iq_commons/utils/functions_parsers/c_lang_function_parsers.py +++ b/src/exploit_iq_commons/utils/functions_parsers/c_lang_function_parsers.py @@ -38,6 +38,140 @@ def _remove_c_comments(code: str) -> str: code = re.sub(r'//.*', '', code) return code + +def _count_c_declared_params(func_doc: Document) -> int | None: + """Count declared parameters in a C function definition. + Returns None if the function signature cannot be parsed (e.g. macros, variadic).""" + code = func_doc.page_content.strip().replace('\r\n', '\n') + m = _FUNCTION_PATTERN_REGEX.search(code, timeout=REGEX_TIMEOUT_SECONDS) + if not m: + return None + params = m.group('params').strip() + if not params or params == 'void': + return 0 + if '...' in params: + return None + depth = 0 + count = 1 + for ch in params: + if ch == '(': + depth += 1 + elif ch == ')': + depth -= 1 + elif ch == ',' and depth == 0: + count += 1 + return count + + +def _extract_call_site_args_str(function_body: str, start_pos: int) -> str | None: + """Extract the arguments string from a call site starting at the opening paren. + Returns None if the call site cannot be parsed. + + Correctly handles string literals, char literals, and nested parentheses. + """ + depth = 0 + end_pos = -1 + in_string = False + in_char = False + i = start_pos + while i < len(function_body): + ch = function_body[i] + if in_string: + if ch == '\\' and i + 1 < len(function_body): + i += 2 + continue + if ch == '"': + in_string = False + elif in_char: + if ch == '\\' and i + 1 < len(function_body): + i += 2 + continue + if ch == "'": + in_char = False + else: + if ch == '"': + in_string = True + elif ch == "'": + in_char = True + elif ch == '(': + depth += 1 + elif ch == ')': + depth -= 1 + if depth == 0: + end_pos = i + break + i += 1 + if end_pos == -1: + return None + return function_body[start_pos + 1:end_pos].strip() + + +def _count_args_in_str(args_str: str) -> int: + """Count the number of arguments in an argument string. + + Correctly handles commas inside string literals, char literals, and nested parentheses. + """ + if not args_str: + return 0 + depth = 0 + count = 1 + in_string = False + in_char = False + i = 0 + while i < len(args_str): + ch = args_str[i] + if in_string: + if ch == '\\' and i + 1 < len(args_str): + i += 2 + continue + if ch == '"': + in_string = False + elif in_char: + if ch == '\\' and i + 1 < len(args_str): + i += 2 + continue + if ch == "'": + in_char = False + else: + if ch == '"': + in_string = True + elif ch == "'": + in_char = True + elif ch == '(': + depth += 1 + elif ch == ')': + depth -= 1 + elif ch == ',' and depth == 0: + count += 1 + i += 1 + return count + + +def _count_all_call_site_args(function_body: str, func_name: str) -> list[int]: + """Count arguments at ALL call sites of func_name(...) in function_body. + Returns a list of argument counts for each call site found. + + Correctly handles commas inside string literals, char literals, and nested parentheses. + """ + pattern = re.compile(r'\b' + re.escape(func_name) + r'\s*\(') + results = [] + for m in pattern.finditer(function_body): + start = m.end() - 1 + args_str = _extract_call_site_args_str(function_body, start) + if args_str is not None: + results.append(_count_args_in_str(args_str)) + return results + + +def _count_call_site_args(function_body: str, func_name: str) -> int | None: + """Count arguments at the first call site of func_name(...) in function_body. + Returns None if no call site is found. + + Correctly handles commas inside string literals, char literals, and nested parentheses. + """ + counts = _count_all_call_site_args(function_body, func_name) + return counts[0] if counts else None + # Compile the function pattern with recursive param matching _FUNCTION_PATTERN_REGEX = regex.compile(r''' ^[ \t]* # Leading indentation @@ -533,7 +667,8 @@ def search_for_called_function( self, caller_function: Document, callee_function_name: str, - callee_function_package: str, # For C, this is usually the file or library name + callee_function: Document, + callee_function_package: str, code_documents: list[Document], type_documents: list[Document], callee_function_file_name: str, @@ -541,7 +676,6 @@ def search_for_called_function( functions_local_variables_index: dict[str, dict], documents_of_functions: list[Document] = None, type_inheritance: dict[Tuple[str, str], List[Tuple[str, str]]] = None, - callee_function:Document = None ) -> bool: """ Returns True if caller_function calls callee_function (directly or via struct/function pointer). @@ -551,7 +685,22 @@ def search_for_called_function( # 2. Direct call: callee_function_name( direct_call_pattern = re.compile(r'\b' + re.escape(callee_function_name) + r'\s*\(') if direct_call_pattern.search(function_body): - return True + if callee_function is not None: + declared = _count_c_declared_params(callee_function) + if declared is None: + return True + all_call_args = _count_all_call_site_args(function_body, callee_function_name) + if not all_call_args: + return True + if declared in all_call_args: + return True + logger.debug( + "Argument count mismatch for %s: declared=%d, call_sites=%s in %s — rejecting false match", + callee_function_name, declared, all_call_args, + caller_function.metadata.get('source', '?')) + return False + else: + return True # 3. Struct member or pointer call: obj->callee_function_name( or obj.callee_function_name( member_call_pattern = re.compile( @@ -750,9 +899,6 @@ def is_call_allowed(self, pkg_docs: list[Document], caller_function: Document, c callee_name = self.get_function_name(callee_function) - if callee_name == "do_shell": - print(f"callee_name: {callee_name}") - if callee_name in caller_functions: doc = caller_functions[callee_name] # if static and in same file → call is allowed diff --git a/src/exploit_iq_commons/utils/functions_parsers/golang_functions_parsers.py b/src/exploit_iq_commons/utils/functions_parsers/golang_functions_parsers.py index c97437807..44dc715e9 100644 --- a/src/exploit_iq_commons/utils/functions_parsers/golang_functions_parsers.py +++ b/src/exploit_iq_commons/utils/functions_parsers/golang_functions_parsers.py @@ -47,6 +47,8 @@ def get_package_name_file(function: Document): class GoLanguageFunctionsParser(LanguageFunctionsParser): def is_same_package(self, package_name_from_input, package_name_from_tree): + if not package_name_from_input or not package_name_from_tree: + return False return package_name_from_input.lower() in package_name_from_tree.lower() def is_tree_key_match(self, package_from_doc: str, tree_key: str) -> bool: @@ -66,6 +68,19 @@ def is_tree_key_match(self, package_from_doc: str, tree_key: str) -> bool: return True return False + def resolve_subpackage_to_module(self, subpackage: str, tree_key: str) -> str | None: + """Return the sub-package suffix when *subpackage* is a child of *tree_key*. + + Example: subpackage='google.golang.org/protobuf/encoding/protojson', + tree_key='google.golang.org/protobuf' + → returns '/encoding/protojson' + """ + a = subpackage.lower() + b = tree_key.lower() + if a.startswith(b + "/"): + return subpackage[len(tree_key):] + return None + def get_dummy_function(self, function_name): return f"{self.get_function_reserved_word()} {function_name}() {{}}" @@ -289,7 +304,7 @@ def parse_all_type_struct_class_to_fields(self, types: list[Document], type_inhe [name, _, type_name] = declaration_parts elif len(declaration_parts) == 2: [name, type_name] = declaration_parts - if len(declaration_parts) == (2 or 3): + if len(declaration_parts) in (2, 3): self.parse_one_type(Document(page_content=f"type {name} {type_name}", metadata={"source": the_type.metadata['source']}), types_mapping) @@ -431,7 +446,6 @@ def search_for_called_function(self, caller_function: Document, callee_function_ index_of_function_closing = caller_function.page_content.rfind("}") caller_function_body = str( caller_function.page_content[index_of_function_opening + 1: index_of_function_closing]) - re.search("", caller_function_body) escaped_name = re.escape(callee_function_name) regex = fr'(?callee() where obj is a struct with a function pointer field.""" + caller = self._doc( + "void wrapper() { ctx->process(data); }", source="mylib/caller.c" + ) + local_vars = { + "wrapper@mylib/caller.c": { + "ctx": {"value": "", "type": "MyStruct"}, + } + } + fields = { + ("MyStruct", "mylib/types.h"): [ + ("process", "void (*)(void *)"), + ] + } + result = self.parser.search_for_called_function( + caller, "process", None, "mylib", + [], [], "mylib/target.c", fields, local_vars, + ) + assert result is True + + def test_member_access_still_matches_direct_call_pattern(self): + """obj->callee(x) is matched by the direct-call pattern + (\\bdata_field\\s*\\() before the struct member check runs, + so it returns True regardless of field type.""" + caller = self._doc( + "void wrapper() { ctx->data_field(x); }", source="mylib/caller.c" + ) + local_vars = { + "wrapper@mylib/caller.c": { + "ctx": {"value": "", "type": "MyStruct"}, + } + } + fields = { + ("MyStruct", "mylib/types.h"): [ + ("data_field", "int"), + ] + } + result = self.parser.search_for_called_function( + caller, "data_field", None, "mylib", + [], [], "mylib/target.c", fields, local_vars, + ) + assert result is True + + def test_call_in_comment_ignored(self): + """Calls inside C comments should not count as real calls.""" + caller = self._doc( + "void wrapper() { /* target_func(x); */ return; }" + ) + result = self.parser.search_for_called_function( + caller, "target_func", None, "mylib", + [], [], "mylib/target.c", {}, {}, + ) + assert result is False + + def test_argument_count_mismatch_rejects_cross_package_false_match(self): + """Reproduces CVE-2024-12085 false positive: rsync's read_byte(int fd) + has 1 param, but PostgreSQL's jspInitByBuffer calls its own read_byte + macro with 3 args. The argument count pre-filter should reject this.""" + caller = self._doc( + "void\njspInitByBuffer(JsonPathItem *v, char *base, int32 pos)\n" + "{\n\tv->base = base + pos;\n" + "\tread_byte(v->type, base, pos);\n" + "\tread_int32(v->nextPos, base, pos);\n" + "}\n", + source="src/backend/utils/adt/jsonpath.c" + ) + callee = self._doc( + "int read_byte(int fd)\n" + "{\n\tchar c;\n\tread(fd, &c, 1);\n\treturn (unsigned char)c;\n}\n", + source="rpm_libs/rsync/rsync-3.1/rsync-3.1.3/io.c" + ) + result = self.parser.search_for_called_function( + caller, "read_byte", callee, "rsync", + [], [], "rpm_libs/rsync/rsync-3.1/rsync-3.1.3/io.c", {}, {}, + ) + assert result is False + + def test_argument_count_match_allows_genuine_call(self): + """When argument count matches, the call is accepted.""" + caller = self._doc( + "void process(int fd)\n{\n\tint b = read_byte(fd);\n}\n", + source="src/app/main.c" + ) + callee = self._doc( + "int read_byte(int fd)\n" + "{\n\tchar c;\n\tread(fd, &c, 1);\n\treturn (unsigned char)c;\n}\n", + source="rpm_libs/rsync/rsync-3.1/rsync-3.1.3/io.c" + ) + result = self.parser.search_for_called_function( + caller, "read_byte", callee, "rsync", + [], [], "rpm_libs/rsync/rsync-3.1/rsync-3.1.3/io.c", {}, {}, + ) + assert result is True + + def test_argument_count_filter_skipped_for_variadic_callee(self): + """Variadic functions (with ...) skip the argument count check.""" + caller = self._doc( + "void log_it() {\n\tmy_printf(\"%d %d\", a, b);\n}\n", + source="src/app/logger.c" + ) + callee = self._doc( + "int my_printf(const char *fmt, ...)\n{\n\treturn 0;\n}\n", + source="rpm_libs/mylib/printf.c" + ) + result = self.parser.search_for_called_function( + caller, "my_printf", callee, "mylib", + [], [], "rpm_libs/mylib/printf.c", {}, {}, + ) + assert result is True + + def test_argument_count_filter_no_callee_doc_still_matches(self): + """Without callee_function document, the filter is skipped.""" + caller = self._doc( + "void wrapper() { read_byte(fd); }\n", + source="src/app/main.c" + ) + result = self.parser.search_for_called_function( + caller, "read_byte", None, "rsync", + [], [], "rpm_libs/rsync/io.c", {}, {}, + ) + assert result is True + + def test_zero_param_callee_rejects_call_with_args(self): + """A void-param callee rejects calls that pass arguments.""" + caller = self._doc( + "void wrapper() { do_init(ctx, 42); }\n", + source="src/app/main.c" + ) + callee = self._doc( + "void do_init(void)\n{\n\treturn;\n}\n", + source="rpm_libs/mylib/init.c" + ) + result = self.parser.search_for_called_function( + caller, "do_init", callee, "mylib", + [], [], "rpm_libs/mylib/init.c", {}, {}, + ) + assert result is False + + +class TestCountCDeclaredParams: + """Tests for _count_c_declared_params — parameter counting from function definitions.""" + + def _doc(self, content): + return Document(page_content=content, metadata={"source": "test.c"}) + + def test_single_param(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_c_declared_params + doc = self._doc("int read_byte(int fd)\n{\n\treturn 0;\n}\n") + assert _count_c_declared_params(doc) == 1 + + def test_multiple_params(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_c_declared_params + doc = self._doc("void copy(char *dst, const char *src, size_t len)\n{\n}\n") + assert _count_c_declared_params(doc) == 3 + + def test_void_param(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_c_declared_params + doc = self._doc("void do_init(void)\n{\n}\n") + assert _count_c_declared_params(doc) == 0 + + def test_no_params(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_c_declared_params + doc = self._doc("int get_count()\n{\n\treturn 0;\n}\n") + assert _count_c_declared_params(doc) == 0 + + def test_variadic_returns_none(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_c_declared_params + doc = self._doc("int my_printf(const char *fmt, ...)\n{\n\treturn 0;\n}\n") + assert _count_c_declared_params(doc) is None + + def test_function_pointer_param(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_c_declared_params + doc = self._doc("void register_cb(void (*callback)(int, int), int priority)\n{\n}\n") + assert _count_c_declared_params(doc) == 2 + + def test_unparseable_returns_none(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_c_declared_params + doc = self._doc("#define READ_BYTE(fd) read(fd, &c, 1)") + assert _count_c_declared_params(doc) is None + + +class TestCountCallSiteArgs: + """Tests for _count_call_site_args — argument counting at call sites.""" + + def test_single_arg(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_call_site_args + assert _count_call_site_args("int b = read_byte(fd);", "read_byte") == 1 + + def test_multiple_args(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_call_site_args + assert _count_call_site_args("read_byte(v->type, base, pos);", "read_byte") == 3 + + def test_zero_args(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_call_site_args + assert _count_call_site_args("int n = get_count();", "get_count") == 0 + + def test_nested_parens(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_call_site_args + assert _count_call_site_args("process(foo(a, b), c);", "process") == 2 + + def test_not_found_returns_none(self): + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import _count_call_site_args + assert _count_call_site_args("other_func(x);", "read_byte") is None + + + + + + + +class TestCCreateMapOfLocalVars: + """Tests for CLanguageFunctionsParser.create_map_of_local_vars — + function parameters, local declarations, pointer types, return types.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, content, source="mylib/file.c"): + return Document(page_content=content, metadata={"source": source}) + + def test_function_params(self): + doc = self._doc("int add(int a, int b) { return a + b; }") + result = self.parser.create_map_of_local_vars([doc]) + key = [k for k in result.keys() if "add" in k][0] + assert result[key]["a"]["value"] == "parameter" + assert result[key]["a"]["type"] == "int" + assert result[key]["b"]["value"] == "parameter" + + def test_local_variable_declaration(self): + doc = self._doc("void foo() {\n int x = 5;\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = [k for k in result.keys() if "foo" in k][0] + assert "x" in result[key] + assert result[key]["x"]["type"] == "int" + + def test_pointer_type(self): + doc = self._doc("void bar() {\n char *name = NULL;\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = [k for k in result.keys() if "bar" in k][0] + assert "name" in result[key] + + def test_return_types(self): + doc = self._doc("int compute(void) {\n return 42;\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = [k for k in result.keys() if "compute" in k][0] + assert "return_types" in result[key] + + def test_multiple_declarations_on_one_line(self): + doc = self._doc("void foo() {\n int a = 1, b = 2;\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = [k for k in result.keys() if "foo" in k][0] + assert "a" in result[key] + assert "b" in result[key] + + def test_skips_invalid_function_name(self): + """Documents with reserved-word function names get state='invalid' + and should be skipped.""" + doc = self._doc("if (x) { int a = 1; }") + result = self.parser.create_map_of_local_vars([doc]) + # No key should be added for an invalid function + assert len(result) == 0 + + def test_no_body_still_records_params(self): + """A function declaration without a body (forward decl) — the + _FUNCTION_PATTERN_REGEX requires '{' so function name extraction + fails, and the document is skipped.""" + doc = self._doc("int add(int a, int b);") + result = self.parser.create_map_of_local_vars([doc]) + assert len(result) == 0 + + + + +class TestFindTopLevelBlocks: + """Tests for CSegmenterExtended.find_top_level_blocks — state machine + handling of braces, comments, string/char literals.""" + + def test_single_function(self): + code = "void foo() {\n return;\n}\n" + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + assert blocks[0][0] == 1 # starts on line 1 + assert blocks[0][1] == 3 # ends on line 3 + + def test_two_functions(self): + code = "void a() {\n}\n\nvoid b() {\n}\n" + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 2 + + def test_nested_braces(self): + code = "void foo() {\n if (x) {\n y();\n }\n}\n" + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + + def test_string_literal_with_braces(self): + code = 'void foo() {\n char *s = "{ not a block }";\n}\n' + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + + def test_block_comment_with_braces(self): + code = "/* { not a block } */\nvoid foo() {\n}\n" + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + + def test_line_comment_with_braces(self): + code = "// { not a block\nvoid foo() {\n}\n" + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + + def test_char_literal_brace(self): + code = "void foo() {\n char c = '{';\n}\n" + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + + def test_empty_code(self): + blocks = CSegmenterExtended.find_top_level_blocks("") + assert blocks == [] + + def test_deeply_nested_braces(self): + code = "void foo() {\n if (a) {\n if (b) {\n x();\n }\n }\n}\n" + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + assert blocks[0] == (1, 7) + + def test_escaped_quote_in_string(self): + """Escaped double-quote inside a string should not exit string state.""" + code = 'void foo() {\n char *s = "a\\"b{c";\n}\n' + blocks = CSegmenterExtended.find_top_level_blocks(code) + assert len(blocks) == 1 + + + + +class TestCParseFunctionSignature: + """Tests for CLanguageFunctionsParser.__parse_function_signature — + parameter extraction from C function signatures.""" + + def setup_method(self): + self.parser = make_parser() + + def _parse(self, sig): + return self.parser._CLanguageFunctionsParser__parse_function_signature(sig) + + def test_simple_params(self): + result = self._parse("int add(int a, int b)") + assert result["a"]["type"] == "int" + assert result["b"]["type"] == "int" + + def test_pointer_param(self): + result = self._parse("void foo(char *name)") + assert "name" in result + + def test_variadic(self): + result = self._parse("int printf(const char *fmt, ...)") + assert "..." in result + assert result["..."]["type"] == "VARIADIC_ARGS" + + def test_function_pointer_param(self): + result = self._parse("void sort(int (*cmp)(int, int))") + assert "cmp" in result + assert "(*" in result["cmp"]["type"] + + def test_array_param(self): + result = self._parse("void process(int data[10])") + assert "data" in result + + def test_void_params(self): + result = self._parse("void foo(void)") + assert len(result) == 0 + + def test_no_params(self): + result = self._parse("void foo()") + assert len(result) == 0 + + def test_const_pointer_param(self): + result = self._parse("int cmp(const char *a, const char *b)") + assert "a" in result + assert "b" in result + + def test_double_pointer(self): + result = self._parse("void alloc(int **out)") + assert "out" in result + + def test_no_parentheses(self): + """Signature without parentheses returns empty dict.""" + result = self._parse("int x") + assert result == {} + + + + +class TestCIsCallAllowed: + """Tests for CLanguageFunctionsParser.is_call_allowed — static vs + non-static visibility across files and packages.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, content, source="mylib/file.c"): + return Document(page_content=content, metadata={"source": source}) + + def test_non_static_cross_package(self): + """Non-static callee in a different package should be allowed + (no name collision in caller package).""" + caller = self._doc("void caller() {}", source="myproject/main.c") + callee = self._doc("void target() {}", source="mylib/target.c") + result = self.parser.is_call_allowed([], caller, callee) + assert result is True + + def test_static_same_file(self): + caller = self._doc("void caller() {}", source="mylib/file.c") + callee = self._doc("static void helper() {}", source="mylib/file.c") + result = self.parser.is_call_allowed([], caller, callee) + assert result is True + + def test_static_different_file_same_package(self): + """Static callee in a different file within the same package should + be rejected.""" + caller = self._doc("void caller() {}", source="mylib/main.c") + callee = self._doc("static void helper() {}", source="mylib/other.c") + result = self.parser.is_call_allowed([], caller, callee) + assert result is False + + def test_static_different_package(self): + """Static callee in a different package is rejected.""" + caller = self._doc("void caller() {}", source="myproject/main.c") + callee = self._doc("static void helper() {}", source="mylib/helper.c") + result = self.parser.is_call_allowed([], caller, callee) + assert result is False + + def test_non_static_same_package(self): + """Non-static callee in the same package is always allowed.""" + caller = self._doc("void caller() {}", source="mylib/a.c") + callee = self._doc("void helper() {}", source="mylib/b.c") + result = self.parser.is_call_allowed([], caller, callee) + assert result is True + + def test_non_static_cross_package_name_collision_static_different_file(self): + """Cross-package call where the caller package has a static function + with the same name in a different file — call is allowed (the local + static doesn't shadow because it's in a different file).""" + caller = self._doc("void caller() {}", source="myproject/main.c") + callee = self._doc("void target() {}", source="mylib/target.c") + # pkg_docs for the caller package contains a static "target" in another file + shadow_doc = self._doc( + "static void target() {}", source="myproject/other.c" + ) + result = self.parser.is_call_allowed([shadow_doc], caller, callee) + assert result is True + + + + +class TestCIsFunction: + """Tests for CLanguageFunctionsParser.is_function — positive and + negative pattern matching on C code segments.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, content, content_type="functions_classes"): + return Document( + page_content=content, + metadata={"source": "file.c", "content_type": content_type}, + ) + + def test_regular_function(self): + assert ( + self.parser.is_function( + self._doc("int main(int argc, char *argv[]) { return 0; }") + ) + is True + ) + + def test_static_function(self): + assert ( + self.parser.is_function( + self._doc("static void helper(void) { }") + ) + is True + ) + + def test_typedef(self): + assert ( + self.parser.is_function( + self._doc("typedef void (*handler_t)(int);") + ) + is False + ) + + def test_struct_definition(self): + assert ( + self.parser.is_function( + self._doc("struct Point { int x; int y; };") + ) + is False + ) + + def test_wrong_content_type(self): + assert ( + self.parser.is_function( + self._doc("int foo() {}", content_type="other") + ) + is False + ) + + def test_enum_definition(self): + assert ( + self.parser.is_function( + self._doc("enum Color { RED, GREEN, BLUE };") + ) + is False + ) + + def test_function_pointer_typedef(self): + """A function pointer typedef should be rejected.""" + assert ( + self.parser.is_function( + self._doc("typedef int (*compare_fn)(const void *, const void *);") + ) + is False + ) + + def test_inline_function(self): + assert ( + self.parser.is_function( + self._doc("inline int square(int x) { return x * x; }") + ) + is True + ) + + + + +class TestCDocumentParser: + """Tests for CDocumentParser — struct field extraction including + named structs, typedef structs, function pointer fields, bitfields, + and inline structs/unions.""" + + def test_named_struct(self): + doc = Document( + page_content="struct Point {\n int x;\n int y;\n};", + metadata={"source_file": "geo.h"}, + ) + parser = CDocumentParser(doc) + assert parser.is_doc_struct() is True + key, fields = parser.parse_struct_to_fields() + assert key[0] == "Point" + assert key[1] == "geo.h" + assert any(f[0] == "x" for f in fields) + assert any(f[0] == "y" for f in fields) + + def test_typedef_struct(self): + doc = Document( + page_content="typedef struct _Node {\n int value;\n struct _Node *next;\n} Node;", + metadata={"source_file": "list.h"}, + ) + parser = CDocumentParser(doc) + assert parser.is_doc_struct() is True + key, fields = parser.parse_struct_to_fields() + assert key[0] == "Node" + + def test_function_pointer_field(self): + doc = Document( + page_content="struct Ops {\n void (*init)(void);\n int (*process)(int);\n};", + metadata={"source_file": "ops.h"}, + ) + parser = CDocumentParser(doc) + key, fields = parser.parse_struct_to_fields() + assert any(f[0] == "init" for f in fields) + assert any(f[0] == "process" for f in fields) + + def test_bitfield(self): + doc = Document( + page_content="struct Flags {\n unsigned int active : 1;\n unsigned int mode : 3;\n};", + metadata={"source_file": "flags.h"}, + ) + parser = CDocumentParser(doc) + key, fields = parser.parse_struct_to_fields() + assert any(f[0] == "active" for f in fields) + active_type = [f[1] for f in fields if f[0] == "active"][0] + assert ":1" in active_type + + def test_inline_struct(self): + doc = Document( + page_content="struct Outer {\n struct { int a; } inner;\n int b;\n};", + metadata={"source_file": "outer.h"}, + ) + parser = CDocumentParser(doc) + key, fields = parser.parse_struct_to_fields() + assert any(f[0] == "inner" for f in fields) + assert any(f[0] == "b" for f in fields) + + def test_not_a_struct(self): + doc = Document( + page_content="int foo(int x) { return x; }", + metadata={"source_file": "foo.c"}, + ) + parser = CDocumentParser(doc) + assert parser.is_doc_struct() is False + + def test_caching_parse_result(self): + """Repeated calls to parse_struct_to_fields return the same result.""" + doc = Document( + page_content="struct S {\n int a;\n};", + metadata={"source_file": "s.h"}, + ) + parser = CDocumentParser(doc) + result1 = parser.parse_struct_to_fields() + result2 = parser.parse_struct_to_fields() + assert result1 is result2 # cached, same object + + + + +class TestExtractDefineFunctions: + """Tests for CSegmenterExtended.extract_define_functions — extracting + function-like #define macros as dummy C functions.""" + + def test_simple_define(self): + code = "#define my_add(a, b) ((a) + (b))" + result = CSegmenterExtended.extract_define_functions(code) + assert len(result) == 1 + assert "my_add" in result[0] + + def test_uppercase_only_skipped(self): + """ALL-UPPERCASE macro names are skipped (treated as constants).""" + code = "#define MAX_SIZE(x) ((x) * 2)" + result = CSegmenterExtended.extract_define_functions(code) + assert len(result) == 0 + + def test_multiline_define(self): + code = "#define my_func(x) \\\n do { \\\n process(x); \\\n } while(0)" + result = CSegmenterExtended.extract_define_functions(code) + assert len(result) == 1 + assert "my_func" in result[0] + + def test_do_while_zero_stripped(self): + """do { ... } while(0) wrapper should be stripped from the body.""" + code = "#define my_macro(x) do { foo(x); } while(0)" + result = CSegmenterExtended.extract_define_functions(code) + assert len(result) == 1 + assert "do" not in result[0] + assert "while" not in result[0] + assert "foo(x)" in result[0] + + def test_no_define_macros(self): + code = "int main() { return 0; }" + result = CSegmenterExtended.extract_define_functions(code) + assert len(result) == 0 + + def test_mixed_case_macro(self): + """Macros with at least one lowercase letter should be extracted.""" + code = "#define myMacro(a) (a + 1)" + result = CSegmenterExtended.extract_define_functions(code) + assert len(result) == 1 + assert "myMacro" in result[0] + + + + +class TestRemoveMacroBlocks: + """Tests for CSegmenterExtended.remove_macro_blocks — removing + ALL-UPPERCASE macro blocks (MACRO(args) { ... }).""" + + def test_removes_macro_block(self): + code = "int x = 1;\nMY_MACRO(arg) {\n body;\n}\nint y = 2;" + result = CSegmenterExtended.remove_macro_blocks(code) + assert "MY_MACRO" not in result + assert "int x = 1;" in result + assert "int y = 2;" in result + + def test_keeps_non_macro_code(self): + code = "int foo() {\n return 0;\n}" + result = CSegmenterExtended.remove_macro_blocks(code) + assert "int foo()" in result + + def test_removes_macro_with_nested_braces(self): + code = "INIT_BLOCK(x) {\n if (a) {\n b;\n }\n}\nint z;" + result = CSegmenterExtended.remove_macro_blocks(code) + assert "INIT_BLOCK" not in result + assert "int z;" in result + + def test_plain_code_no_macros(self): + code = "void foo() {\n return;\n}" + result = CSegmenterExtended.remove_macro_blocks(code) + assert result.strip() == code.strip() + + def test_empty_input(self): + result = CSegmenterExtended.remove_macro_blocks("") + assert result == "" + + +# Additional coverage: get_function_name, is_comment_line, is_doc_type + + +class TestCGetFunctionName: + """Tests for CLanguageFunctionsParser.get_function_name.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, content): + return Document(page_content=content, metadata={"source": "file.c"}) + + def test_simple_function(self): + doc = self._doc("int main(int argc, char *argv[]) { return 0; }") + assert self.parser.get_function_name(doc) == "main" + + def test_metadata_func_name_takes_precedence(self): + doc = Document( + page_content="void foo() {}", + metadata={"source": "file.c", "func_name": "cached_name"}, + ) + assert self.parser.get_function_name(doc) == "cached_name" + + def test_reserved_word_returns_empty(self): + doc = self._doc("if (x) { return 1; }") + result = self.parser.get_function_name(doc) + assert result == "" + assert doc.metadata.get("state") == "invalid" + + def test_no_function_returns_empty(self): + doc = self._doc("int x = 5;") + result = self.parser.get_function_name(doc) + assert result == "" + + +class TestCIsCommentLine: + """Tests for CLanguageFunctionsParser.is_comment_line.""" + + def setup_method(self): + self.parser = make_parser() + + def test_line_comment(self): + assert self.parser.is_comment_line("// this is a comment") is True + + def test_block_comment_start(self): + assert self.parser.is_comment_line("/* block comment */") is True + + def test_not_a_comment(self): + assert self.parser.is_comment_line("int x = 5;") is False + + def test_indented_comment(self): + assert self.parser.is_comment_line(" // indented comment") is True + + +class TestCIsDocType: + """Tests for CLanguageFunctionsParser.is_doc_type — distinguishing + struct/enum/union documents from function documents.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, content, content_type="functions_classes"): + return Document( + page_content=content, + metadata={"source": "file.c", "content_type": content_type}, + ) + + def test_struct_is_type(self): + doc = self._doc("struct Point {\n int x;\n int y;\n};") + assert self.parser.is_doc_type(doc) is True + + def test_function_is_not_type(self): + doc = self._doc("int foo(int x) { return x; }") + assert self.parser.is_doc_type(doc) is False + + def test_typedef_struct_is_type(self): + doc = self._doc("typedef struct {\n int val;\n} Item;") + assert self.parser.is_doc_type(doc) is True + + +class TestCRemoveComments: + """Tests for CSegmenterExtended.remove_comments.""" + + def test_removes_line_comment(self): + code = "int x = 5; // comment" + result = CSegmenterExtended.remove_comments(code) + assert "//" not in result + assert "int x = 5;" in result + + def test_removes_block_comment(self): + code = "int x = /* hidden */ 5;" + result = CSegmenterExtended.remove_comments(code) + assert "hidden" not in result + + def test_no_comments_unchanged(self): + code = "int x = 5;" + result = CSegmenterExtended.remove_comments(code) + assert result == code + + def test_multiline_block_comment(self): + code = "int x = 1;\n/* multi\nline\ncomment */\nint y = 2;" + result = CSegmenterExtended.remove_comments(code) + assert "multi" not in result + assert "int x = 1;" in result + assert "int y = 2;" in result + + + + +class TestCGetPackageNames: + """Tests for CLanguageFunctionsParser.get_package_names.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, source): + return Document(page_content="int foo() { return 0; }", metadata={"source": source}) + + def test_third_party_path(self): + """Third-party path uses dep tree's module_from_path.""" + doc = self._doc(f"{C_DEP_LIBS_NAME}/libcurl/src/http.c") + names = self.parser.get_package_names(doc) + assert len(names) == 1 + assert names[0] == f"{C_DEP_LIBS_NAME}/libcurl/src/http.c".split("/")[0] + + def test_non_third_party_path(self): + """Non-third-party path uses dep tree's module_from_path.""" + doc = self._doc("myproject/src/main.c") + names = self.parser.get_package_names(doc) + assert len(names) == 1 + + def test_no_dep_tree_uses_basename(self): + """When dep_builder_tree is None, falls back to os.path.basename.""" + parser = CLanguageFunctionsParser() + parser.dep_builder_tree = None + doc = Document( + page_content="int foo() { return 0; }", + metadata={"source": "mylib/src/util.c"}, + ) + names = parser.get_package_names(doc) + assert names[0] == "util.c" + + +class TestCIsRootPackage: + """Tests for CLanguageFunctionsParser.is_root_package.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, source): + return Document(page_content="int foo() {}", metadata={"source": source}) + + def test_app_code_is_root(self): + assert self.parser.is_root_package(self._doc("myproject/main.c")) is True + + def test_third_party_not_root(self): + assert self.parser.is_root_package(self._doc(f"{C_DEP_LIBS_NAME}/libfoo/bar.c")) is False + + def test_prj_name_match_is_root(self): + """When get_package_names returns the project name, is_root_package + returns True via the first check.""" + doc = self._doc("myproject/main.c") + self.parser.dep_builder_tree.module_from_path.return_value = "myproject" + assert self.parser.is_root_package(doc) is True + + +class TestCIsSearchableFileName: + """Tests for CLanguageFunctionsParser.is_searchable_file_name.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, source, file_name=None): + meta = {"source": source} + if file_name is not None: + meta["file_name"] = file_name + return Document(page_content="int foo() {}", metadata=meta) + + def test_c_file_is_searchable(self): + assert self.parser.is_searchable_file_name(self._doc("lib/util.c", "util.c")) is True + + def test_h_file_is_searchable(self): + assert self.parser.is_searchable_file_name(self._doc("lib/util.h", "util.h")) is True + + def test_non_c_file_not_searchable(self): + assert self.parser.is_searchable_file_name(self._doc("lib/util.py", "util.py")) is False + + def test_no_file_name_metadata(self): + """When file_name metadata is missing, endswith check on empty + string returns False.""" + assert self.parser.is_searchable_file_name(self._doc("lib/util.c")) is False + + + + +class TestCFilterDocsByFuncPkgName: + """Tests for CLanguageFunctionsParser.filter_docs_by_func_pkg_name.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, content, source="lib/file.c"): + return Document(page_content=content, metadata={"source": source}) + + def test_filters_by_function_name(self): + docs = [ + self._doc("void xmlParseDocument() { }"), + self._doc("void unrelated() { }"), + self._doc("int compute() { xmlParseDocument(); }"), + ] + result = self.parser.filter_docs_by_func_pkg_name("xmlParseDocument", "libxml2", docs) + assert len(result) == 2 + assert all("xmlParseDocument" in d.page_content for d in result) + + def test_no_matches(self): + docs = [self._doc("void foo() { }")] + result = self.parser.filter_docs_by_func_pkg_name("bar", "pkg", docs) + assert result == [] + + def test_empty_docs_list(self): + result = self.parser.filter_docs_by_func_pkg_name("foo", "pkg", []) + assert result == [] + + +class TestCDocumentImportsPackage: + """Tests for CLanguageFunctionsParser.document_imports_package.""" + + def setup_method(self): + self.parser = make_parser() + + def _doc(self, content, source="lib/file.c"): + return Document(page_content=content, metadata={"source": source}) + + def test_bare_include(self): + """The regex matches 'include {pkg}' without quotes — only bare + includes (no quotes or angle brackets between include and package) + are detected.""" + docs = { + "main.c": self._doc("#include openssl/ssl.h\nvoid foo() {}", "main.c"), + "other.c": self._doc("void bar() {}", "other.c"), + } + result = self.parser.document_imports_package(docs, "openssl") + assert len(result) == 1 + + def test_quoted_include_not_matched(self): + """Standard #include "pkg/header.h" is NOT matched by the regex + because the quote character sits between 'include' and the package + name, breaking the 'include {pkg}' pattern.""" + docs = { + "main.c": self._doc('#include "openssl/ssl.h"\nvoid foo() {}', "main.c"), + } + result = self.parser.document_imports_package(docs, "openssl") + assert len(result) == 0 + + def test_no_imports(self): + docs = { + "main.c": self._doc("void foo() {}", "main.c"), + } + result = self.parser.document_imports_package(docs, "openssl") + assert result == [] + + def test_multiple_files_with_bare_includes(self): + docs = { + "a.c": self._doc("#include libxml/parser.h\nvoid a() {}", "a.c"), + "b.c": self._doc("#include libxml/tree.h\nvoid b() {}", "b.c"), + "c.c": self._doc("void c() {}", "c.c"), + } + result = self.parser.document_imports_package(docs, "libxml") + assert len(result) == 2 diff --git a/src/exploit_iq_commons/utils/functions_parsers/tests/test_go_parser.py b/src/exploit_iq_commons/utils/functions_parsers/tests/test_go_parser.py new file mode 100644 index 000000000..7ba279dd5 --- /dev/null +++ b/src/exploit_iq_commons/utils/functions_parsers/tests/test_go_parser.py @@ -0,0 +1,1242 @@ +import pytest +from unittest.mock import MagicMock, patch +from langchain_core.documents import Document + +from exploit_iq_commons.utils.functions_parsers.golang_functions_parsers import GoLanguageFunctionsParser +from exploit_iq_commons.utils.functions_parsers.lang_functions_parsers import LanguageFunctionsParser +from exploit_iq_commons.utils.functions_parsers.python_functions_parser import PythonLanguageFunctionsParser + + +class TestGoIsTreeKeyMatch: + """Tests for GoLanguageFunctionsParser.is_tree_key_match — boundary-aware + '/' matching that prevents fan-out explosion in _resolve_tree_key.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + # --- Bug scenario --- + + def test_bug_scenario_no_false_cross_org_match(self): + """The original bug: 'github.com/hashicorp' substring-matched + 'github.com/hashicorp-terraform/foo' because 'hashicorp' is a + substring of 'hashicorp-terraform'. Boundary matching rejects this.""" + assert self.parser.is_tree_key_match( + "github.com/hashicorp", "github.com/hashicorp-terraform/foo" + ) is False + + def test_bug_scenario_hyphenated_suffix_no_match(self): + """'github.com/foo' must NOT match 'github.com/foo-bar'.""" + assert self.parser.is_tree_key_match( + "github.com/foo", "github.com/foo-bar" + ) is False + + def test_bug_scenario_partial_segment_no_match(self): + """'github.com/go' must NOT match 'github.com/golang/protobuf'.""" + assert self.parser.is_tree_key_match( + "github.com/go", "github.com/golang/protobuf" + ) is False + + # --- Exact match --- + + def test_exact_match(self): + assert self.parser.is_tree_key_match( + "github.com/golang-jwt/jwt/v5", "github.com/golang-jwt/jwt/v5" + ) is True + + def test_exact_match_case_insensitive(self): + assert self.parser.is_tree_key_match( + "GitHub.Com/GoLang-JWT/JWT/v5", "github.com/golang-jwt/jwt/v5" + ) is True + + # --- Input is prefix of tree key (doc package is parent of tree module) --- + + def test_input_prefix_of_tree_key(self): + """2-level path from get_package_names should match child module.""" + assert self.parser.is_tree_key_match( + "github.com/hashicorp", "github.com/hashicorp/vault" + ) is True + + def test_input_prefix_deeper_nesting(self): + assert self.parser.is_tree_key_match( + "github.com/hashicorp", "github.com/hashicorp/vault/api/sub" + ) is True + + def test_input_prefix_three_level(self): + assert self.parser.is_tree_key_match( + "github.com/hashicorp/vault", "github.com/hashicorp/vault/api" + ) is True + + # --- Tree key is prefix of input (rejected: would conflate sub-packages) --- + + def test_tree_key_prefix_of_input_rejected(self): + """Mapping a specific sub-package query to a broader tree entry + would cause false reachability chains (e.g. strvals.Parse matching + ignore.Parse because both live under helm.sh/helm/v3).""" + assert self.parser.is_tree_key_match( + "github.com/hashicorp/vault/api", "github.com/hashicorp/vault" + ) is False + + def test_tree_key_prefix_of_input_deep_rejected(self): + assert self.parser.is_tree_key_match( + "google.golang.org/protobuf/encoding/protojson", + "google.golang.org/protobuf" + ) is False + + # --- No match cases --- + + def test_completely_different_packages(self): + assert self.parser.is_tree_key_match( + "github.com/foo/bar", "github.com/baz/qux" + ) is False + + def test_same_domain_different_org(self): + assert self.parser.is_tree_key_match( + "github.com/foo/bar", "github.com/baz/bar" + ) is False + + def test_domain_only_not_enough(self): + """Sharing only 'github.com' should not match.""" + assert self.parser.is_tree_key_match( + "github.com/orgA", "github.com/orgB" + ) is False + + def test_substring_within_segment_no_match(self): + """'net' is substring of 'netty' but not at a '/' boundary.""" + assert self.parser.is_tree_key_match( + "io.netty/net", "io.netty/netty-codec" + ) is False + + # --- Versioned modules --- + + def test_versioned_module_exact(self): + assert self.parser.is_tree_key_match( + "github.com/golang-jwt/jwt/v5", "github.com/golang-jwt/jwt/v5" + ) is True + + def test_versioned_child(self): + assert self.parser.is_tree_key_match( + "github.com/golang-jwt/jwt/v5", + "github.com/golang-jwt/jwt/v5/parser" + ) is True + + def test_different_versions_no_match(self): + """v4 and v5 are different path segments — no prefix relationship.""" + assert self.parser.is_tree_key_match( + "github.com/golang-jwt/jwt/v5", "github.com/golang-jwt/jwt/v4" + ) is False + + # --- Edge cases --- + + def test_empty_input_against_nonempty_tree_key(self): + """Empty package_from_doc should not match a real tree key.""" + assert self.parser.is_tree_key_match("", "github.com/foo") is False + + def test_empty_input(self): + assert self.parser.is_tree_key_match( + "", "github.com/foo/bar" + ) is False + + def test_empty_tree_key(self): + assert self.parser.is_tree_key_match( + "github.com/foo/bar", "" + ) is False + + def test_single_segment(self): + assert self.parser.is_tree_key_match("fmt", "fmt") is True + + def test_single_segment_prefix_no_match(self): + """'fmt' is not a prefix of 'fmtlib' at a '/' boundary.""" + assert self.parser.is_tree_key_match("fmt", "fmtlib") is False + + def test_trailing_slash_not_treated_as_boundary(self): + """Trailing slash is part of the path, not a boundary marker.""" + assert self.parser.is_tree_key_match( + "github.com/foo/", "github.com/foo/bar" + ) is False + + # --- is_same_package still does substring (unchanged) --- + + def test_is_same_package_still_substring(self): + """Verify is_same_package retains substring behavior for FL/line 545.""" + assert self.parser.is_same_package("jwt", "github.com/golang-jwt/jwt/v5") is True + assert self.parser.is_same_package( + "github.com/hashicorp", "github.com/hashicorp-terraform/foo" + ) is True + + +class TestBaseParserIsTreeKeyMatchDefault: + """Verify base class is_tree_key_match delegates to is_same_package.""" + + def test_base_class_delegates_via_java(self): + """Base class is abstract; use Java parser which inherits default + is_tree_key_match (delegates to exact-match is_same_package).""" + from exploit_iq_commons.utils.functions_parsers.java_functions_parsers import JavaLanguageFunctionsParser + parser = JavaLanguageFunctionsParser() + assert parser.is_tree_key_match("foo", "foo") is True + assert parser.is_tree_key_match("foo", "bar") is False + + def test_python_inherits_base_behavior(self): + """Python parser doesn't override is_tree_key_match — should use + is_same_package (PEP 503 normalization).""" + parser = PythonLanguageFunctionsParser() + assert parser.is_tree_key_match("my-pkg", "my_pkg") is True + assert parser.is_tree_key_match("flask", "Flask") is True + assert parser.is_tree_key_match("urllib3", "requests") is False + + +class TestResolveTreeKeyWithGoParser: + """Integration-style tests: verify _resolve_tree_key uses is_tree_key_match + for Go and prevents the fan-out explosion.""" + + def _make_retriever(self, tree_dict_keys): + """Build a minimal mock ChainOfCallsRetriever with Go parser.""" + from exploit_iq_commons.utils.chain_of_calls_retriever import ChainOfCallsRetriever, _SearchCtx + retriever = object.__new__(ChainOfCallsRetriever) + retriever.language_parser = GoLanguageFunctionsParser() + retriever.tree_dict = {k: ["root"] for k in tree_dict_keys} + ctx = _SearchCtx() + return retriever, ctx + + def test_exact_match_preferred(self): + retriever, ctx = self._make_retriever([ + "github.com/hashicorp/vault", + "github.com/hashicorp/vault/api", + ]) + result = retriever._resolve_tree_key("github.com/hashicorp/vault", ctx) + assert result == "github.com/hashicorp/vault" + + def test_subpackage_does_not_resolve_to_parent(self): + """A sub-package query must NOT resolve to a parent tree entry — + this would conflate all sub-packages under the parent module.""" + retriever, ctx = self._make_retriever([ + "github.com/hashicorp/vault", + ]) + result = retriever._resolve_tree_key("github.com/hashicorp/vault/api", ctx) + assert result is None + + def test_no_cross_org_fan_out(self): + """The bug scenario: short 2-level path must NOT match unrelated modules.""" + retriever, ctx = self._make_retriever([ + "github.com/hashicorp/vault", + "github.com/hashicorp/consul", + "github.com/hashicorp-terraform/aws", + "github.com/hashicorp-labs/experiment", + ]) + result = retriever._resolve_tree_key("github.com/hashicorp", ctx) + assert result in ("github.com/hashicorp/vault", "github.com/hashicorp/consul") + assert result != "github.com/hashicorp-terraform/aws" + assert result != "github.com/hashicorp-labs/experiment" + + def test_no_match_returns_none(self): + retriever, ctx = self._make_retriever([ + "github.com/foo/bar", + ]) + result = retriever._resolve_tree_key("github.com/baz/qux", ctx) + assert result is None + + def test_tree_additions_also_uses_boundary(self): + """Short prefix from get_package_names resolves to child module + in tree_additions, but sub-package does NOT resolve to parent.""" + retriever, ctx = self._make_retriever([]) + ctx.tree_additions["github.com/hashicorp/vault"] = ["root"] + ctx.tree_additions["github.com/hashicorp-terraform/aws"] = ["root"] + result = retriever._resolve_tree_key("github.com/hashicorp", ctx) + assert result == "github.com/hashicorp/vault" + result2 = retriever._resolve_tree_key("github.com/hashicorp/vault/api", ctx) + assert result2 is None + + +class TestGoIsPackageImported: + """Tests for GoLanguageFunctionsParser.is_package_imported — checks whether + an identifier resolves to a callee package via import declarations.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + # --- Single import block (first_import == last_import) --- + + def test_single_block_alias_match(self): + """Alias 'jose' in single import block should match callee package.""" + code = '''package main + +import ( + jose "gopkg.in/go-jose/go-jose.v2" + "fmt" +) +''' + assert self.parser.is_package_imported( + code, "jose", "gopkg.in/go-jose/go-jose.v2" + ) is True + + def test_single_block_no_alias_direct_package(self): + """Direct import (no alias) — identifier is the last path segment.""" + code = '''package main + +import ( + "fmt" + "gopkg.in/go-jose/go-jose.v2" +) +''' + assert self.parser.is_package_imported( + code, "go-jose", "gopkg.in/go-jose/go-jose.v2" + ) is True + + def test_single_block_identifier_not_found(self): + code = '''package main + +import ( + "fmt" + "net/http" +) +''' + assert self.parser.is_package_imported( + code, "jose", "gopkg.in/go-jose/go-jose.v2" + ) is False + + # --- Multiple imports with alias (line 625 branch) --- + + def test_multiple_imports_alias_match(self): + """Dedicated `import identifier "pkg"` line should match.""" + code = '''package main + +import "fmt" + +import jose "gopkg.in/go-jose/go-jose.v2" +''' + assert self.parser.is_package_imported( + code, "jose", "gopkg.in/go-jose/go-jose.v2" + ) is True + + def test_multiple_imports_alias_wrong_package(self): + """Alias matches but callee package doesn't — should return False.""" + code = '''package main + +import "fmt" + +import jose "gopkg.in/go-jose/go-jose.v2" +''' + assert self.parser.is_package_imported( + code, "jose", "github.com/unrelated/package" + ) is False + + # --- Multiple imports regex fallback (line 636 branch — the crash case) --- + + def test_regex_fallback_no_alias(self): + """The crash scenario: `import "gopkg.in/go-jose/go-jose.v2"` with + identifier 'jose'. Regex fallback handles import paths without aliases.""" + code = '''package main + +import "fmt" + +import "gopkg.in/go-jose/go-jose.v2" +''' + assert self.parser.is_package_imported( + code, "jose", "gopkg.in/go-jose/go-jose.v2" + ) is True + + def test_regex_fallback_wrong_callee(self): + """Regex finds import containing identifier but callee doesn't match.""" + code = '''package main + +import "fmt" + +import "gopkg.in/go-jose/go-jose.v2" +''' + assert self.parser.is_package_imported( + code, "jose", "github.com/wrong/package" + ) is False + + def test_regex_fallback_single_quotes(self): + """Import with single quotes (unusual but valid in regex pattern).""" + code = """package main + +import "fmt" + +import 'gopkg.in/go-jose/go-jose.v2' +""" + assert self.parser.is_package_imported( + code, "jose", "gopkg.in/go-jose/go-jose.v2" + ) is True + + # --- Edge cases --- + + def test_no_imports(self): + code = '''package main + +func main() { + fmt.Println("hello") +} +''' + assert self.parser.is_package_imported(code, "fmt", "fmt") is False + + def test_empty_code(self): + assert self.parser.is_package_imported("", "jose", "gopkg.in/jose") is False + + def test_import_in_comment_skipped(self): + """Imports after // import comment should be handled, not crash.""" + code = '''package main + +// import this is a comment about imports + +import "fmt" + +import "gopkg.in/go-jose/go-jose.v2" +''' + assert self.parser.is_package_imported( + code, "jose", "gopkg.in/go-jose/go-jose.v2" + ) is True + + def test_callee_package_substring_match_single_block(self): + """Callee package substring match in single import block.""" + code = '''package main + +import ( + "github.com/quic-go/quic-go" +) +''' + assert self.parser.is_package_imported( + code, "quic", "github.com/quic-go/quic-go" + ) is True + + def test_identifier_found_but_callee_package_mismatch(self): + code = '''package main + +import ( + jose "gopkg.in/go-jose/go-jose.v2" + "fmt" +) +''' + assert self.parser.is_package_imported( + code, "jose", "github.com/completely/different" + ) is False + + +class TestGoIsSamePackageEmptyInput: + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def test_empty_input_returns_false(self): + assert self.parser.is_same_package("", "any-package") is False + + +class TestParseAllTypeStructClassToFields: + """Tests for GoLanguageFunctionsParser.parse_all_type_struct_class_to_fields, + including 3-part type declarations (type aliases) inside grouped type blocks.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content, source="test.go"): + return Document(page_content=content, metadata={"source": source}) + + def test_single_struct_type(self): + """Basic struct type with fields is parsed correctly.""" + doc = self._doc("type MyStruct struct {\nName string\nAge int\n}") + result = self.parser.parse_all_type_struct_class_to_fields([doc]) + assert ("MyStruct", "test.go") in result + fields = result[("MyStruct", "test.go")] + assert any(f[0] == "Name" for f in fields) + + def test_single_interface_type(self): + """Interface type with methods is parsed correctly.""" + doc = self._doc("type Reader interface {\nRead(p []byte) (n int, err error)\n}") + result = self.parser.parse_all_type_struct_class_to_fields([doc]) + assert ("interface;Reader", "test.go") in result + + def test_type_alias_three_parts(self): + """A 3-part type alias declaration (type MyAlias = OriginalType) inside a + grouped type block should be handled by the `len(declaration_parts) == 3` + branch (previously broken when the condition was `== (2 or 3)`).""" + doc = self._doc("type (\nMyAlias = OriginalType\nSentinel\n)") + result = self.parser.parse_all_type_struct_class_to_fields([doc]) + assert ("MyAlias", "test.go") in result + fields = result[("MyAlias", "test.go")] + assert any("OriginalType" in f[1] for f in fields) + + def test_type_alias_two_parts(self): + """A 2-part type definition (type MyType int) inside a grouped type block.""" + doc = self._doc("type (\nCounter int\nSentinel\n)") + result = self.parser.parse_all_type_struct_class_to_fields([doc]) + assert ("Counter", "test.go") in result + fields = result[("Counter", "test.go")] + assert any("int" in f[1] for f in fields) + + def test_grouped_block_with_mixed_declarations(self): + """A grouped type block containing both alias and regular type definitions.""" + doc = self._doc("type (\nMyAlias = OriginalType\nCounter int\nSentinel\n)") + result = self.parser.parse_all_type_struct_class_to_fields([doc]) + assert ("MyAlias", "test.go") in result + assert ("Counter", "test.go") in result + + def test_non_struct_non_interface_standalone(self): + """Standalone non-composite type (type X int) parsed as wrapper.""" + doc = self._doc("type Duration int64") + result = self.parser.parse_all_type_struct_class_to_fields([doc]) + assert ("Duration", "test.go") in result + + +class TestParseOneType: + """Tests for GoLanguageFunctionsParser.parse_one_type.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content, source="test.go"): + return Document(page_content=content, metadata={"source": source}) + + def test_struct_with_embedded_type(self): + """Struct with an embedded type field (no field name, just type).""" + doc = self._doc("type MyStruct struct {\nio.Reader\n}") + mapping = {} + self.parser.parse_one_type(doc, mapping) + assert ("MyStruct", "test.go") in mapping + fields = mapping[("MyStruct", "test.go")] + assert any(f[0] == "embedded_type" for f in fields) + + def test_empty_struct(self): + """Struct with no fields produces no entry in the mapping because + fields_list remains empty.""" + doc = self._doc("type Empty struct {\n\n}") + mapping = {} + self.parser.parse_one_type(doc, mapping) + assert ("Empty", "test.go") not in mapping + + def test_wrapper_type(self): + """Non-struct/non-interface type: `type Byte uint8`.""" + doc = self._doc("type Byte uint8") + mapping = {} + self.parser.parse_one_type(doc, mapping) + assert ("Byte", "test.go") in mapping + fields = mapping[("Byte", "test.go")] + assert fields == [("Byte", "uint8")] + + +class TestGetTypeName: + """Tests for GoLanguageFunctionsParser.get_type_name.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content): + return Document(page_content=content, metadata={"source": "test.go"}) + + def test_struct_type_name(self): + doc = self._doc("type MyStruct struct {\n Name string\n}") + assert self.parser.get_type_name(doc) == "MyStruct" + + def test_interface_type_name(self): + doc = self._doc("type Reader interface {\n Read(p []byte) (n int, err error)\n}") + assert self.parser.get_type_name(doc) == "Reader" + + def test_wrapper_type_name(self): + doc = self._doc("type Duration int64") + assert self.parser.get_type_name(doc) == "Duration" + + +class TestGetFunctionName: + """Tests for GoLanguageFunctionsParser.get_function_name.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content): + return Document(page_content=content, metadata={"source": "test.go"}) + + def test_regular_function(self): + doc = self._doc("func Hello(name string) {\n fmt.Println(name)\n}") + assert self.parser.get_function_name(doc) == "Hello" + + def test_method_with_receiver(self): + doc = self._doc("func (s *Server) Start(port int) {\n s.listen(port)\n}") + assert self.parser.get_function_name(doc) == "Start" + + def test_function_no_args(self): + doc = self._doc("func main() {\n run()\n}") + assert self.parser.get_function_name(doc) == "main" + + +class TestGetPackageNames: + """Tests for GoLanguageFunctionsParser.get_package_names.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, source): + return Document(page_content="func Foo() {}", metadata={"source": source}) + + def test_vendor_path(self): + doc = self._doc("vendor/github.com/foo/bar/baz.go") + names = self.parser.get_package_names(doc) + assert "github.com/foo" in names + assert "github.com/foo/bar" in names + + def test_non_vendor_path(self): + doc = self._doc("github.com/hashicorp/vault/api/client.go") + names = self.parser.get_package_names(doc) + assert "github.com/hashicorp" in names + assert "github.com/hashicorp/vault" in names + + def test_stdlib_single_segment(self): + doc = self._doc("fmt/print.go") + names = self.parser.get_package_names(doc) + assert "fmt/print.go" in names + + +class TestIsSearchableFileName: + """Tests for GoLanguageFunctionsParser.is_searchable_file_name.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, source): + return Document(page_content="func Foo() {}", metadata={"source": source}) + + def test_regular_file_is_searchable(self): + assert self.parser.is_searchable_file_name(self._doc("pkg/server.go")) is True + + def test_test_file_not_searchable(self): + assert self.parser.is_searchable_file_name(self._doc("pkg/server_test.go")) is False + + def test_test_in_directory_but_not_filename(self): + """'test' in directory path does not exclude the file.""" + assert self.parser.is_searchable_file_name(self._doc("test/pkg/server.go")) is True + + +class TestIsExportedFunction: + """Tests for GoLanguageFunctionsParser.is_exported_function.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content): + return Document(page_content=content, metadata={"source": "test.go"}) + + def test_exported_function(self): + doc = self._doc("func PublicFunc() {}") + assert self.parser.is_exported_function(doc, {}) is not None + + def test_unexported_function(self): + doc = self._doc("func privateFunc() {}") + # Go convention: unexported functions start with lowercase. + # The regex [A-Z][a-z0-9-]* won't match "privateFunc" at start. + # However, the regex uses re.search so it could match mid-string. + # Test the actual behavior. + result = self.parser.is_exported_function(doc, {}) + # "privateFunc" has 'F' in the middle, so re.search will match. + # Known limitation: is_exported_function uses re.search instead of re.match, + # so it matches uppercase letters anywhere in the name, not just the first + # character. A proper fix would use re.match to check only the first character. + assert result is not None # matches 'Fu' within 'privateFunc' + + def test_all_lowercase_not_exported(self): + doc = self._doc("func main() {}") + assert self.parser.is_exported_function(doc, {}) is None + + +# --------------------------------------------------------------------------- +# get_package_name_file +# --------------------------------------------------------------------------- + +class TestGetPackageNameFile: + """Tests for the module-level get_package_name_file helper.""" + + def test_extracts_package_name(self): + from exploit_iq_commons.utils.functions_parsers.golang_functions_parsers import get_package_name_file + doc = Document(page_content="package main\n\nfunc Foo() {}", metadata={"source": "main.go"}) + assert get_package_name_file(doc) == "main" + + def test_package_with_whitespace(self): + from exploit_iq_commons.utils.functions_parsers.golang_functions_parsers import get_package_name_file + doc = Document(page_content="package http \n\nfunc Handler() {}", metadata={"source": "handler.go"}) + assert get_package_name_file(doc) == "http" + + def test_no_package_declaration(self): + from exploit_iq_commons.utils.functions_parsers.golang_functions_parsers import get_package_name_file + doc = Document(page_content="func orphan() {}", metadata={"source": "orphan.go"}) + result = get_package_name_file(doc) + # No "package" keyword, find returns -1, if guard skips extraction + assert result is None + + +# --------------------------------------------------------------------------- +# is_function +# --------------------------------------------------------------------------- + +class TestGoIsFunction: + """Tests for GoLanguageFunctionsParser.is_function.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content): + return Document(page_content=content, metadata={"source": "test.go"}) + + def test_function_declaration(self): + assert self.parser.is_function(self._doc("func main() {}")) is True + + def test_method_declaration(self): + assert self.parser.is_function(self._doc("func (s *Server) Start() {}")) is True + + def test_type_declaration(self): + assert self.parser.is_function(self._doc("type Server struct {}")) is False + + def test_var_declaration(self): + assert self.parser.is_function(self._doc("var x = 5")) is False + + +# --------------------------------------------------------------------------- +# is_root_package +# --------------------------------------------------------------------------- + +class TestGoIsRootPackage: + """Tests for GoLanguageFunctionsParser.is_root_package.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, source): + return Document(page_content="func Foo() {}", metadata={"source": source}) + + def test_vendor_path_is_not_root(self): + assert self.parser.is_root_package(self._doc("vendor/github.com/foo/bar.go")) is False + + def test_non_vendor_is_root(self): + assert self.parser.is_root_package(self._doc("cmd/server/main.go")) is True + + def test_app_path_is_root(self): + assert self.parser.is_root_package(self._doc("pkg/handler/handler.go")) is True + + +# --------------------------------------------------------------------------- +# get_import_search_patterns +# --------------------------------------------------------------------------- + +class TestGoGetImportSearchPatterns: + """Tests for GoLanguageFunctionsParser.get_import_search_patterns.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def test_pattern_matches_import(self): + patterns = self.parser.get_import_search_patterns("github.com/foo/bar") + code = 'import "github.com/foo/bar/sub"' + assert any(p.search(code) for p in patterns) + + def test_pattern_no_match(self): + patterns = self.parser.get_import_search_patterns("github.com/foo/bar") + code = 'import "github.com/baz/qux"' + assert not any(p.search(code) for p in patterns) + + def test_pattern_matches_exact_package(self): + patterns = self.parser.get_import_search_patterns("github.com/foo/bar") + code = 'import "github.com/foo/bar"' + assert any(p.search(code) for p in patterns) + + def test_pattern_matches_aliased_import(self): + patterns = self.parser.get_import_search_patterns("github.com/foo/bar") + code = ' myalias "github.com/foo/bar"' + assert any(p.search(code) for p in patterns) + + +# --------------------------------------------------------------------------- +# create_map_of_local_vars +# --------------------------------------------------------------------------- + +class TestGoCreateMapOfLocalVars: + """Tests for GoLanguageFunctionsParser.create_map_of_local_vars.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content, source="test.go"): + return Document(page_content=content, metadata={"source": source}) + + def test_function_params_extracted(self): + """Function parameters are extracted as PARAMETER type.""" + doc = self._doc("func Process(input string, count int) {\n fmt.Println(input)\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = "Process@test.go" + assert key in result + assert result[key]["input"]["value"] == "parameter" + assert result[key]["input"]["type"] == "string" + assert result[key]["count"]["value"] == "parameter" + assert result[key]["count"]["type"] == "int" + + def test_var_declaration(self): + """var x Type declaration is extracted.""" + doc = self._doc("func Foo() {\n var name string\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = "Foo@test.go" + assert key in result + assert result[key]["name"]["type"] == "string" + + def test_short_assignment(self): + """:= short assignment creates a local_implicit entry.""" + doc = self._doc("func Bar() {\n x := someFunc()\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = "Bar@test.go" + assert key in result + assert "x" in result[key] + assert result[key]["x"]["type"] == "local_implicit" + + def test_return_types(self): + """Function with return types extracts them.""" + doc = self._doc("func Compute(a int) (int, error) {\n return a, nil\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = "Compute@test.go" + assert key in result + assert "return_types" in result[key] + + def test_no_return_types_empty_list(self): + """Function with no return types stores empty return_types list.""" + doc = self._doc("func NoReturn() {\n fmt.Println()\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = "NoReturn@test.go" + assert key in result + assert result[key]["return_types"] == [] + + def test_method_receiver_extracted(self): + """Method receiver parameter is extracted.""" + doc = self._doc("func (s *Server) Start(port int) {\n s.listen(port)\n}") + result = self.parser.create_map_of_local_vars([doc]) + key = "Start@test.go" + assert key in result + # The receiver "s" with type "*Server" should be extracted + assert "s" in result[key] + + +# --------------------------------------------------------------------------- +# search_for_called_function +# --------------------------------------------------------------------------- + +class TestSearchForCalledFunction: + """Tests for GoLanguageFunctionsParser.search_for_called_function. + + Also covers __check_identifier_resolved_to_callee_function_package + indirectly, since that private method is the resolution engine behind the + public API.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content, source="test.go"): + return Document(page_content=content, metadata={"source": source}) + + def _build_local_vars(self, docs): + """Build the functions_local_variables_index from a list of documents.""" + return self.parser.create_map_of_local_vars(docs) + + def test_same_package_call_no_qualifier(self): + """Direct call (no qualifier) in the same package resolves to True. + + When the regex matches 'Helper(' without a dot qualifier, the function + checks that both caller and callee share the same Go package declaration + and that callee_function_package is in the caller's package names.""" + caller_source = "vendor/github.com/myorg/mylib/handler.go" + callee_source = "vendor/github.com/myorg/mylib/util.go" + + caller = self._doc( + "package mylib\n\nfunc Handler() {\n Helper()\n}", + source=caller_source, + ) + callee = self._doc( + "package mylib\n\nfunc Helper() {\n return\n}", + source=callee_source, + ) + + # code_documents maps source paths to full-file Documents (with package declaration) + code_documents = { + caller_source: self._doc( + "package mylib\n\nimport \"fmt\"\n\nfunc Handler() {\n Helper()\n}", + source=caller_source, + ), + callee_source: self._doc( + "package mylib\n\nfunc Helper() {\n return\n}", + source=callee_source, + ), + } + + local_vars = self._build_local_vars([caller]) + + result = self.parser.search_for_called_function( + caller_function=caller, + callee_function_name="Helper", + callee_function=callee, + callee_function_package="github.com/myorg/mylib", + code_documents=code_documents, + type_documents=[], + callee_function_file_name=callee_source, + fields_of_types={}, + functions_local_variables_index=local_vars, + ) + assert result is True + + def test_aliased_import_call(self): + """Qualifier resolves via aliased import (e.g. pkg.SomeFunc()). + + The regex matches 'pkg.SomeFunc(' and the identifier 'pkg' is resolved + by checking the caller file's import block via is_package_imported.""" + caller_source = "cmd/main.go" + callee_source = "vendor/github.com/some/pkg/somefunc.go" + + caller = self._doc( + 'package main\n\nimport (\n pkg "github.com/some/pkg"\n)\n\n' + 'func Run() {\n pkg.SomeFunc()\n}', + source=caller_source, + ) + callee = self._doc( + "package pkg\n\nfunc SomeFunc() {\n return\n}", + source=callee_source, + ) + + code_documents = { + caller_source: self._doc( + 'package main\n\nimport (\n pkg "github.com/some/pkg"\n)\n\n' + 'func Run() {\n pkg.SomeFunc()\n}', + source=caller_source, + ), + } + + local_vars = self._build_local_vars([caller]) + + result = self.parser.search_for_called_function( + caller_function=caller, + callee_function_name="SomeFunc", + callee_function=callee, + callee_function_package="github.com/some/pkg", + code_documents=code_documents, + type_documents=[], + callee_function_file_name=callee_source, + fields_of_types={}, + functions_local_variables_index=local_vars, + ) + assert result is True + + def test_same_package_qualified_call(self): + """Caller uses package-name qualifier on a same-package function. + + Both files declare 'package mylib', caller calls 'mylib.Helper()', + and the caller source path contains the callee package path. + This exercises the pkg_decl branch in + __check_identifier_resolved_to_callee_function_package.""" + pkg_path = "github.com/myorg/mylib" + caller_source = f"vendor/{pkg_path}/handler.go" + + caller = self._doc( + "package mylib\n\nfunc Handler() {\n mylib.Helper()\n}", + source=caller_source, + ) + callee = self._doc( + "package mylib\n\nfunc Helper() {\n return\n}", + source=f"vendor/{pkg_path}/util.go", + ) + + code_documents = { + caller_source: self._doc( + "package mylib\n\nfunc Handler() {\n mylib.Helper()\n}", + source=caller_source, + ), + } + + local_vars = self._build_local_vars([caller]) + + result = self.parser.search_for_called_function( + caller_function=caller, + callee_function_name="Helper", + callee_function=callee, + callee_function_package=pkg_path, + code_documents=code_documents, + type_documents=[], + callee_function_file_name=f"vendor/{pkg_path}/util.go", + fields_of_types={}, + functions_local_variables_index=local_vars, + ) + assert result is True + + def test_callee_not_called_returns_false(self): + """When the caller body does not contain any call to the callee function, + the regex fails and the function returns False immediately.""" + caller_source = "cmd/main.go" + + caller = self._doc( + "package main\n\nfunc Run() {\n fmt.Println(\"hello\")\n}", + source=caller_source, + ) + callee = self._doc( + "package pkg\n\nfunc Unrelated() {\n return\n}", + source="vendor/github.com/some/pkg/unrelated.go", + ) + + code_documents = { + caller_source: self._doc( + "package main\n\nfunc Run() {\n fmt.Println(\"hello\")\n}", + source=caller_source, + ), + } + + local_vars = self._build_local_vars([caller]) + + result = self.parser.search_for_called_function( + caller_function=caller, + callee_function_name="Unrelated", + callee_function=callee, + callee_function_package="github.com/some/pkg", + code_documents=code_documents, + type_documents=[], + callee_function_file_name="vendor/github.com/some/pkg/unrelated.go", + fields_of_types={}, + functions_local_variables_index=local_vars, + ) + assert result is False + + +# --------------------------------------------------------------------------- +# __trace_down_package (tested indirectly via search_for_called_function) +# --------------------------------------------------------------------------- + +class TestTraceDownPackage: + """Tests for __trace_down_package, exercised indirectly through + search_for_called_function. + + When the caller has a local variable whose type belongs to the callee + package, __trace_down_package should resolve the type and return True.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content, source="test.go"): + return Document(page_content=content, metadata={"source": source}) + + # Tested in TestGoIsTreeKeyMatch.test_both_empty_strings below + + +class TestIsTreeKeyMatchBothEmpty: + """is_tree_key_match when both inputs are empty strings.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def test_both_empty_strings(self): + """Two empty strings: a == b so the exact-match branch returns True.""" + assert self.parser.is_tree_key_match("", "") is True + + +# --------------------------------------------------------------------------- +# get_function_name edge cases +# --------------------------------------------------------------------------- + + +class TestGetFunctionNameEdgeCases: + """Edge cases for GoLanguageFunctionsParser.get_function_name: + no body, generics, anonymous functions.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content): + return Document(page_content=content, metadata={"source": "test.go"}) + + def test_no_body_returns_header(self): + """Function declaration without braces — ValueError path returns the + first line of page_content.""" + doc = self._doc("func Forward(x int)") + result = self.parser.get_function_name(doc) + assert result is not None + assert "Forward" in result + + def test_generic_function_includes_bracket(self): + """Go generic function: the first '(' in the header is inside the type + parameter list, so split(" ")[1] returns 'Map[T' not 'Map'. The '[' + fallback only triggers when '(' is entirely absent from the header.""" + doc = self._doc("func Map[T any, U any](s []T, f func(T) U) []U {\n return nil\n}") + result = self.parser.get_function_name(doc) + assert result == "Map[T" + + def test_generic_function_no_parens_before_bracket(self): + """When the only delimiter before '{' is '[', the fallback branch + correctly extracts the function name.""" + doc = self._doc("func Identity[T any] {\n}") + result = self.parser.get_function_name(doc) + assert result == "Identity" + + def test_anonymous_function_returns_none(self): + """An anonymous function literal (no name after 'func') — the split + produces a single-element list, so the function returns None.""" + doc = self._doc("func() {\n fmt.Println()\n}") + result = self.parser.get_function_name(doc) + assert result is None + + +# --------------------------------------------------------------------------- +# is_package_imported edge cases +# --------------------------------------------------------------------------- + + +class TestIsPackageImportedEdgeCases: + """Edge cases for GoLanguageFunctionsParser.is_package_imported: + empty callee_package and alias ambiguity.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def test_empty_callee_package(self): + """Empty callee_package — the substring check `"" in package_name` + always succeeds, so the match depends on finding the identifier.""" + code = '''package main + +import ( + "fmt" +) +''' + result = self.parser.is_package_imported(code, "fmt", "") + assert result is True + + def test_empty_callee_package_identifier_not_found(self): + """Empty callee_package but identifier not in imports.""" + code = '''package main + +import ( + "fmt" +) +''' + result = self.parser.is_package_imported(code, "nonexistent", "") + assert result is False + + def test_alias_shadows_different_package(self): + """Alias 'http' maps to a custom package, not net/http. When callee + is 'net/http', should return False because the alias resolves to + a different package.""" + code = '''package main + +import ( + http "github.com/custom/http-wrapper" +) +''' + result = self.parser.is_package_imported(code, "http", "net/http") + assert result is False + + def test_alias_matches_callee_package(self): + """Alias 'http' maps to net/http — should return True.""" + code = '''package main + +import ( + http "net/http" +) +''' + result = self.parser.is_package_imported(code, "http", "net/http") + assert result is True + + def test_dot_import(self): + """Dot import — identifier is '.' which should be found in the import.""" + code = '''package main + +import . "fmt" +''' + result = self.parser.is_package_imported(code, ".", "fmt") + assert result is False + + def test_regex_escape_dot_in_identifier(self): + """Dots in identifier must be literal, not regex wildcards. + The regex finds the import line; the callee_package check validates.""" + code = '''package main + +import "fmt" + +import "github.com/v2xio/pkg" +''' + result = self.parser.is_package_imported(code, "v2.io", "v2.io") + assert result is False, "Unescaped dot would falsely match v2xio" + + def test_substring_rejected_by_callee_check(self): + """Identifier found as substring but callee_package mismatch rejects it.""" + code = '''package main + +import "fmt" + +import "github.com/some-jwt-fork/auth" +''' + result = self.parser.is_package_imported( + code, "jwt", "github.com/golang-jwt/jwt" + ) + assert result is False, "callee_package check rejects mismatched import path" + + def test_versioned_module_path_multiple_imports(self): + """Versioned Go module paths (jwt/v5) match via the regex fallback + branch when there are multiple import statements.""" + code = '''package main + +import "fmt" + +import "github.com/golang-jwt/jwt/v5" +''' + result = self.parser.is_package_imported( + code, "jwt", "github.com/golang-jwt/jwt/v5" + ) + assert result is True + + +# --------------------------------------------------------------------------- +# Original TestTraceDownPackage tests continue +# --------------------------------------------------------------------------- + + +class TestTraceDownPackage: + """Tests for __trace_down_package, exercised indirectly through + search_for_called_function. + + When the caller has a local variable whose type belongs to the callee + package, __trace_down_package should resolve the type and return True.""" + + def setup_method(self): + self.parser = GoLanguageFunctionsParser() + + def _doc(self, content, source="test.go"): + return Document(page_content=content, metadata={"source": source}) + + def test_local_variable_type_resolves_to_callee_package(self): + """A local variable assigned via struct initializer from the callee + package triggers __trace_down_package and resolves to True.""" + callee_package = "github.com/ext/lib" + caller_source = "cmd/main.go" + callee_source = f"vendor/{callee_package}/client.go" + + # Caller has a variable 'c' assigned as Client{} then calls c.Do() + caller = self._doc( + 'package main\n\nimport (\n lib "github.com/ext/lib"\n)\n\n' + 'func Run() {\n c := Client{}\n c.Do()\n}', + source=caller_source, + ) + callee = self._doc( + "package lib\n\nfunc (c *Client) Do() {\n return\n}", + source=callee_source, + ) + + # Type document for Client so __trace_down_package can look it up + client_type = self._doc( + "type Client struct {\n host string\n}", + source=callee_source, + ) + + code_documents = { + caller_source: self._doc( + 'package main\n\nimport (\n lib "github.com/ext/lib"\n)\n\n' + 'func Run() {\n c := Client{}\n c.Do()\n}', + source=caller_source, + ), + } + + local_vars = self.parser.create_map_of_local_vars([caller]) + + result = self.parser.search_for_called_function( + caller_function=caller, + callee_function_name="Do", + callee_function=callee, + callee_function_package=callee_package, + code_documents=code_documents, + type_documents=[client_type], + callee_function_file_name=callee_source, + fields_of_types={}, + functions_local_variables_index=local_vars, + ) + assert result is True \ No newline at end of file diff --git a/src/vuln_analysis/functions/base_graph_agent.py b/src/vuln_analysis/functions/base_graph_agent.py index b21ab2dc1..ed14924bb 100644 --- a/src/vuln_analysis/functions/base_graph_agent.py +++ b/src/vuln_analysis/functions/base_graph_agent.py @@ -43,6 +43,59 @@ logger = LoggingFactory.get_agent_logger(__name__) AGENT_TRACER = Context.get() +# Function Locator status hooks — must match emitters in function_name_locator / tool wrapper. +FL_FUNCTION_EXACT_PREFIX = "FL_FUNCTION_EXACT:" +FL_FUNCTION_ABSENT_PREFIX = "FL_FUNCTION_ABSENT:" +FL_FUNCTION_STRONG_MATCH_PREFIX = "FL_FUNCTION_STRONG_MATCH:" +FL_FUNCTION_WEAK_MATCH_PREFIX = "FL_FUNCTION_WEAK_MATCH:" +FL_PACKAGE_PRESENT_PREFIX = "FL_PACKAGE_PRESENT:" +FL_PACKAGE_ABSENT_PREFIX = "FL_PACKAGE_ABSENT:" +FL_SUBPACKAGE_AMBIGUOUS_PREFIX = "FL_SUBPACKAGE_AMBIGUOUS:" +FL_NO_SOURCE_INDEXED_PREFIX = "FL_NO_SOURCE_INDEXED:" + +# Capture hook through end of quoted string element or line (do not stop at ']' — +# weak/strong hooks embed suggestions=[...]). +_FL_HOOK_RE = re.compile( + r"(FL_(?:FUNCTION_EXACT|FUNCTION_ABSENT|FUNCTION_STRONG_MATCH|FUNCTION_WEAK_MATCH|" + r"PACKAGE_PRESENT|PACKAGE_ABSENT|SUBPACKAGE_AMBIGUOUS|NO_SOURCE_INDEXED):" + r"[^\"'\n]+)" +) + + +def parse_function_locator_findings(tool_output: str, tool_input: str) -> CodeFindings | None: + """Build CodeFindings from FL_* hooks in Function Locator output. + + Returns None when no hooks are present so the caller can fall back to the + comprehension LLM (legacy / non-hook FL payloads). + """ + if not tool_output or not isinstance(tool_output, str): + return None + + hooks = [match.group(1).strip().rstrip(",").strip() for match in _FL_HOOK_RE.finditer(tool_output)] + if not hooks: + return None + + findings: list[str] = [] + seen: set[str] = set() + has_exact = False + for hook in hooks: + if hook in seen: + continue + seen.add(hook) + findings.append(hook) + if hook.startswith(FL_FUNCTION_EXACT_PREFIX): + has_exact = True + + # Legacy summarize gate: VALIDATED only for exact function hits. + if has_exact and tool_input: + validated = f"VALIDATED: {tool_input} exists" + if validated not in seen: + findings.append(validated) + + outcome = f"CALLED: {ToolNames.FUNCTION_LOCATOR} with {tool_input} -> {findings[0]}" + return CodeFindings(findings=findings, tool_outcome=outcome) + + _TOOL_AVAILABILITY = { ToolNames.CODE_SEMANTIC_SEARCH: lambda config, state: state.code_vdb_path is not None, ToolNames.DOCS_SEMANTIC_SEARCH: lambda config, state: state.doc_vdb_path is not None, @@ -493,23 +546,20 @@ async def observation_node(self, state: AgentState) -> AgentState: else: truncated_output = tool_output_for_llm - critical_context_text = self.build_comprehension_context(state) - comp_prompt = COMPREHENSION_PROMPT.format( - goal=state.get('input'), - selected_package=state.get('app_package') or "N/A", - critical_context=critical_context_text, - tool_used=tool_used, - tool_input_detail=tool_input_detail, - last_thought_text=last_thought_text, - tool_output=truncated_output, - ) - max_input_tokens = self.config.context_window_token_limit - 2000 - prompt_tokens = count_tokens(comp_prompt) - if prompt_tokens > max_input_tokens: - overflow = prompt_tokens - max_input_tokens - output_tokens = count_tokens(truncated_output) - safe_budget = max(200, output_tokens - overflow) - truncated_output = truncate_tool_output(truncated_output, tool_used, max_tokens=safe_budget) + # FL-only: prefer deterministic hook parse over comprehension LLM invention. + code_findings = None + if tool_used == ToolNames.FUNCTION_LOCATOR: + code_findings = parse_function_locator_findings( + truncated_output, tool_input_detail + ) + if code_findings is not None: + logger.info( + "Function Locator observation used deterministic FL hook parser (%d findings)", + len(code_findings.findings), + ) + + if code_findings is None: + critical_context_text = self.build_comprehension_context(state) comp_prompt = COMPREHENSION_PROMPT.format( goal=state.get('input'), selected_package=state.get('app_package') or "N/A", @@ -519,11 +569,29 @@ async def observation_node(self, state: AgentState) -> AgentState: last_thought_text=last_thought_text, tool_output=truncated_output, ) - logger.info( - "Comprehension prompt truncated: %d -> %d tokens (limit %d)", - prompt_tokens, count_tokens(comp_prompt), max_input_tokens, + max_input_tokens = self.config.context_window_token_limit - 2000 + prompt_tokens = count_tokens(comp_prompt) + if prompt_tokens > max_input_tokens: + overflow = prompt_tokens - max_input_tokens + output_tokens = count_tokens(truncated_output) + safe_budget = max(200, output_tokens - overflow) + truncated_output = truncate_tool_output(truncated_output, tool_used, max_tokens=safe_budget) + comp_prompt = COMPREHENSION_PROMPT.format( + goal=state.get('input'), + selected_package=state.get('app_package') or "N/A", + critical_context=critical_context_text, + tool_used=tool_used, + tool_input_detail=tool_input_detail, + last_thought_text=last_thought_text, + tool_output=truncated_output, + ) + logger.info( + "Comprehension prompt truncated: %d -> %d tokens (limit %d)", + prompt_tokens, count_tokens(comp_prompt), max_input_tokens, + ) + code_findings = await self._invoke_comprehension( + comp_prompt, tool_used, tool_input_detail, truncated_output ) - code_findings = await self._invoke_comprehension(comp_prompt, tool_used, tool_input_detail, truncated_output) sanitized = self.sanitize_findings(code_findings.findings, state) findings_text = "\n".join(f"- {f}" for f in sanitized) diff --git a/src/vuln_analysis/functions/reachability_agent.py b/src/vuln_analysis/functions/reachability_agent.py index 74aa0bc41..06157b9be 100644 --- a/src/vuln_analysis/functions/reachability_agent.py +++ b/src/vuln_analysis/functions/reachability_agent.py @@ -35,7 +35,11 @@ ) from vuln_analysis.runtime_context import ctx_state from vuln_analysis.tools.tool_names import ToolNames -from vuln_analysis.tools.transitive_code_search import package_name_from_locator_query +from vuln_analysis.tools.transitive_code_search import ( + FL_PACKAGE_ABSENT_PREFIX, + FL_PACKAGE_PRESENT_PREFIX, + package_name_from_locator_query, +) from vuln_analysis.utils.intel_utils import build_critical_context, enrich_go_candidates from vuln_analysis.utils.prompting import build_tool_descriptions from vuln_analysis.utils.prompt_factory import ( @@ -301,11 +305,20 @@ def post_observation(self, state: AgentState, tool_used: str, if tool_used == ToolNames.FUNCTION_LOCATOR and state.get("is_reachability") == "yes": input_pkg = package_name_from_locator_query(tool_input_detail) target_pkg = (state.get("app_package") or "").strip().lower() - if input_pkg and "Package is valid" in tool_output: + # Prefer FL_* package hooks; keep legacy strings for older tool payloads. + package_present = ( + FL_PACKAGE_PRESENT_PREFIX in tool_output + or "Package is valid" in tool_output + ) + package_absent = ( + FL_PACKAGE_ABSENT_PREFIX in tool_output + or "Package is not valid" in tool_output + ) + if input_pkg and package_present: rules_tracker.add_validated_package(input_pkg) if target_pkg and input_pkg == target_pkg: package_validated = True - elif input_pkg and "Package is not valid" in tool_output: + elif input_pkg and package_absent: if target_pkg and input_pkg == target_pkg and package_validated is None: package_validated = False extra["package_validated"] = package_validated diff --git a/src/vuln_analysis/functions/react_internals.py b/src/vuln_analysis/functions/react_internals.py index 353d48334..6c468e8d5 100644 --- a/src/vuln_analysis/functions/react_internals.py +++ b/src/vuln_analysis/functions/react_internals.py @@ -228,6 +228,9 @@ class Observation(BaseModel): "file_compiled", "file_not_compiled", "feature_disabled", "feature_enabled", # Base Agent: Reachability "reachable", "not_reachable", "validated", + # Function Locator status hooks + "fl_exact", "fl_absent", "fl_strong", "fl_weak", + "fl_package_present", "fl_package_absent", "fl_subpackage_ambiguous", # Common: Status/tracking "tool_record", "coverage", "failed", "other", ] @@ -238,6 +241,8 @@ class Observation(BaseModel): "vuln_found", "vuln_absent", "file_compiled", "file_not_compiled", "feature_disabled", "feature_enabled", "reachable", "not_reachable", "validated", + "fl_exact", "fl_absent", "fl_strong", "fl_weak", + "fl_package_present", "fl_package_absent", "fl_subpackage_ambiguous", }) # Map regex prefix → (category, priority). Order matters: first match wins. @@ -260,6 +265,14 @@ class Observation(BaseModel): r'^REACHABLE:': ('reachable', 3), r'^NOT reachable:': ('not_reachable', 3), r'^VALIDATED:': ('validated', 2), + # Function Locator status hooks (prefer over inventing VALIDATED from fuzzy output) + r'^FL_FUNCTION_EXACT:': ('fl_exact', 3), + r'^FL_FUNCTION_ABSENT:': ('fl_absent', 3), + r'^FL_FUNCTION_STRONG_MATCH:': ('fl_strong', 2), + r'^FL_FUNCTION_WEAK_MATCH:': ('fl_weak', 1), + r'^FL_PACKAGE_PRESENT:': ('fl_package_present', 2), + r'^FL_PACKAGE_ABSENT:': ('fl_package_absent', 3), + r'^FL_SUBPACKAGE_AMBIGUOUS:': ('fl_subpackage_ambiguous', 2), # Common: Status/tracking r'^FAILED:': ('failed', 1), r'^TOOL_CALL_RECORD:': ('tool_record', 2), @@ -935,11 +948,18 @@ class AgentState(MessagesState): - Reachability tags apply ONLY to Call Chain Analyzer results. Code Keyword Search and Function Locator do NOT determine reachability. - If Call Chain Analyzer returned NEGATIVE (False), include: "NOT reachable via [package]." - If Call Chain Analyzer returned POSITIVE (True), include: "REACHABLE via [package] - sufficient evidence." -- For Function Locator results, include: "VALIDATED: [package],[function] exists" -- this is NOT reachability. - When Code Keyword Search results are grouped into "Main application" and "Application library dependencies", prioritize findings from the main application group. - findings: 3-5 key technical facts from this OUTPUT only. Each fact must reflect code comprehension (what the code does), not just keyword presence. - tool_outcome: a single line "CALLED: [tool] with [input] -> [brief outcome]". - Keep only CVE-exploitability-relevant information. + +CRITICAL — FUNCTION LOCATOR STATUS HOOKS: +When TOOL USED is "Function Locator" (do NOT apply these rules for any other tool): +- Copy FL_* hooks from NEW OUTPUT verbatim into findings. Do NOT invent hooks that are not present. +- If output contains FL_FUNCTION_EXACT: map to "VALIDATED: [package],[function] exists" (legacy alias) AND keep the FL_FUNCTION_EXACT line. This is NOT reachability. +- If output contains FL_FUNCTION_ABSENT: or only FL_FUNCTION_WEAK_MATCH: / FL_PACKAGE_PRESENT: without FL_FUNCTION_EXACT: — do NOT emit VALIDATED: and do NOT claim the queried function exists. +- FL_PACKAGE_PRESENT: means the package/module exists; it does NOT mean the function was found. +- FL_FUNCTION_STRONG_MATCH: / FL_FUNCTION_WEAK_MATCH: are suggestions only; the queried name is not confirmed. RESPONSE: {{""" @@ -958,7 +978,7 @@ class AgentState(MessagesState): - Reachability tags apply ONLY to Call Chain Analyzer results. Code Keyword Search and Function Locator do NOT determine reachability. - If NEW FINDINGS mention "NOT reachable", add to memory: "NOT reachable via [package]." - If NEW FINDINGS mention "REACHABLE", add: "REACHABLE via [package] - sufficient evidence." -- For Function Locator validation, record: "VALIDATED: [package],[function] exists" -- this is NOT reachability. +- When NEW FINDINGS contain Function Locator hooks (FL_FUNCTION_EXACT / FL_FUNCTION_ABSENT / FL_FUNCTION_WEAK_MATCH / FL_PACKAGE_PRESENT / FL_PACKAGE_ABSENT): copy those FL_* lines verbatim. Record "VALIDATED: [package],[function] exists" ONLY when NEW FINDINGS include FL_FUNCTION_EXACT or an explicit VALIDATED: line — never invent VALIDATED from package presence or weak/strong fuzzy suggestions. - When findings mention "Main application" vs "Application library dependencies", preserve this distinction in memory. - If PREVIOUS MEMORY contains "KNOWN MITIGATIONS:" and findings are relevant (e.g. config, code path, or version), add to memory whether the codebase aligns with or contradicts those mitigations (e.g. "Mitigation check: follow_symlinks=False not set" or "Vendor mitigation (patch X) applied in version Y"). - results: copy the NEW FINDINGS as-is. diff --git a/src/vuln_analysis/tools/transitive_code_search.py b/src/vuln_analysis/tools/transitive_code_search.py index ec9dc33e2..68c5ff26b 100644 --- a/src/vuln_analysis/tools/transitive_code_search.py +++ b/src/vuln_analysis/tools/transitive_code_search.py @@ -48,6 +48,10 @@ FUNCTION_LIBRARY_VERSION_FINDER_TOOL_NAME = "calling_function_library_version_finder" +# Function Locator package-status hooks (observation / post_observation parse these). +FL_PACKAGE_PRESENT_PREFIX = "FL_PACKAGE_PRESENT:" +FL_PACKAGE_ABSENT_PREFIX = "FL_PACKAGE_ABSENT:" + TRANSITIVE_CODE_SEARCH_TOOL_NAME = "transitive_code_search" logger = LoggingFactory.get_agent_logger(__name__) @@ -504,13 +508,18 @@ async def _arun(query: str) -> dict: coc_retriever = transitive_code_searcher.chain_of_calls_retriever locator = FunctionNameLocator(coc_retriever) result = await locator.locate_functions(validation_result) - pkg_msg = "Package is valid." - if not locator.is_package_valid and not locator.is_std_package: - pkg_msg = "Package is not valid." + package_name = validation_result.split(",", 1)[0].strip() + if locator.is_package_valid or locator.is_std_package: + pkg_msg = f"{FL_PACKAGE_PRESENT_PREFIX} {package_name}" + package_status = "present" + else: + pkg_msg = f"{FL_PACKAGE_ABSENT_PREFIX} {package_name}" + package_status = "absent" return { "ecosystem": coc_retriever.ecosystem.name, "package_msg": pkg_msg, + "status": package_status, "result": result } @@ -519,7 +528,7 @@ async def _arun(query: str) -> dict: description=(""" Mandatory first step for code path analysis. Validates package names, locates functions using fuzzy matching, and provides ecosystem type (GO/Python/Java/JavaScript/C/C++). {fl_input_format} - Returns: {'ecosystem': str, 'package_msg': str, 'result': [function_names]}. + Returns: {'ecosystem': str, 'package_msg': str, 'status': 'present'|'absent', 'result': [hook_or_name_strings]}. """)) diff --git a/src/vuln_analysis/utils/function_name_locator.py b/src/vuln_analysis/utils/function_name_locator.py index 99f42e40e..58853be4d 100644 --- a/src/vuln_analysis/utils/function_name_locator.py +++ b/src/vuln_analysis/utils/function_name_locator.py @@ -14,7 +14,9 @@ # limitations under the License. import difflib - +import os +import re +from collections import defaultdict from exploit_iq_commons.logging.loggers_factory import LoggingFactory from exploit_iq_commons.utils.chain_of_calls_retriever_base import ChainOfCallsRetrieverBase @@ -25,7 +27,19 @@ from exploit_iq_commons.utils.source_rpm_downloader import RPMDependencyManager from vuln_analysis.utils.prompt_factory import FL_EXAMPLES -logger = LoggingFactory.get_agent_logger(f"exploit-iq.{__name__}") +logger = LoggingFactory.get_agent_logger(f"morpheus.{__name__}") + +# Function Locator status hooks (Go path). Observation/memory parse these prefixes. +FL_FUNCTION_EXACT_PREFIX = "FL_FUNCTION_EXACT:" +FL_FUNCTION_STRONG_MATCH_PREFIX = "FL_FUNCTION_STRONG_MATCH:" +FL_FUNCTION_WEAK_MATCH_PREFIX = "FL_FUNCTION_WEAK_MATCH:" +FL_FUNCTION_ABSENT_PREFIX = "FL_FUNCTION_ABSENT:" +FL_SUBPACKAGE_AMBIGUOUS_PREFIX = "FL_SUBPACKAGE_AMBIGUOUS:" + +# Fuzzy match thresholds for Go function classification (SequenceMatcher.ratio). +FL_STRONG_MATCH_CUTOFF = 0.8 +FL_WEAK_MATCH_CUTOFF = 0.3 +FL_FUZZY_MATCH_LIMIT = 10 class FunctionNameLocator: @@ -63,11 +77,16 @@ def check_fuzzy_match_in_sbom(self, package: str) -> tuple[bool, str]: else: return False, None - def build_short_go_package_name(self) -> dict: - short_go_package_name = {} + def build_short_go_package_name(self) -> dict[str, list[str]]: + short_go_package_name: dict[str, list[str]] = {} for package in self.coc_retriever.supported_packages: - short_name = package.split("/")[-1] - short_go_package_name[short_name] = package + parts = package.split("/") + short_name = parts[-1] + # Go major version suffixes (v2, v3, ...) are not useful as short names; + # use the second-to-last segment instead (the actual module name). + if re.match(r'^v\d+$', short_name) and len(parts) >= 2: + short_name = parts[-2] + short_go_package_name.setdefault(short_name, []).append(package) return short_go_package_name def _resolve_go_fqdn_to_module(self, package: str) -> str | None: @@ -172,7 +191,10 @@ def python_flow_control(self, input_function: str, package_docs) -> list[str]: list_of_matching_combinations = set() for doc in package_docs: if self.lang_parser: - function_name = self.lang_parser.get_function_name(doc) + try: + function_name = self.lang_parser.get_function_name(doc) + except ValueError: + continue module_name = doc.metadata.get('source').split('/')[-1].split('.')[0] class_name = self.lang_parser.get_class_name_from_class_function(doc) if count_of_dots == 0: @@ -244,7 +266,147 @@ def common_flow_control(self,input_function: str, package_docs) -> list[str]: close_matches = difflib.get_close_matches(input_function, package_functions, n=10, cutoff=0.3) logger.info("Close matches found: %s", close_matches) return close_matches - + + @staticmethod + def _extract_go_subpackage(source_path: str, module_path: str) -> str: + """Extract the full Go import path (module + sub-package) from a doc source path. + + Example: source='vendor/google.golang.org/protobuf/encoding/protojson/decode.go', + module='google.golang.org/protobuf' + → returns 'google.golang.org/protobuf/encoding/protojson' + """ + path_no_vendor = source_path + if path_no_vendor.startswith("vendor/"): + path_no_vendor = path_no_vendor[len("vendor/"):] + dir_path = os.path.dirname(path_no_vendor) + if not dir_path: + return module_path + norm_module = module_path.rstrip("/") + if dir_path == norm_module or dir_path.startswith(norm_module + "/"): + return dir_path + return module_path + + @staticmethod + def _rank_fuzzy_matches( + input_function: str, candidates: list[str], cutoff: float + ) -> list[tuple[str, float]]: + """Return (name, ratio) pairs for candidates at or above cutoff, best first.""" + scored: list[tuple[str, float]] = [] + for name in candidates: + ratio = difflib.SequenceMatcher(None, input_function, name).ratio() + if ratio >= cutoff: + scored.append((name, ratio)) + scored.sort(key=lambda item: item[1], reverse=True) + return scored[:FL_FUZZY_MATCH_LIMIT] + + def _subpackages_for_functions( + self, func_names: list[str], func_to_docs: dict, module_path: str + ) -> dict[str, set[str]]: + """Map sub-package import path → function names found under that path.""" + subpkg_to_funcs: dict[str, set[str]] = defaultdict(set) + for func_name in func_names: + for doc in func_to_docs[func_name]: + source = doc.metadata.get("source", "") + subpkg = self._extract_go_subpackage(source, module_path) + subpkg_to_funcs[subpkg].add(func_name) + return subpkg_to_funcs + + def _format_subpackage_ambiguous( + self, matched_name: str, subpkg_to_funcs: dict[str, set[str]], module_path: str + ) -> list[str]: + """Emit FL_SUBPACKAGE_AMBIGUOUS + CCA guidance when one name spans multiple sub-packages.""" + lines = [] + for subpkg in sorted(subpkg_to_funcs): + funcs_in_subpkg = sorted(subpkg_to_funcs[subpkg]) + lines.append(f" {subpkg}: {', '.join(funcs_in_subpkg)}") + logger.info( + "Go sub-package disambiguation for '%s': matched functions exist in multiple sub-packages:\n%s", + module_path, + "\n".join(lines), + ) + example_subpkg = next(iter(sorted(subpkg_to_funcs))) + subpkg_list = ", ".join(sorted(subpkg_to_funcs.keys())) + return [ + f"{FL_SUBPACKAGE_AMBIGUOUS_PREFIX} {matched_name} in {subpkg_list}", + ( + f"INFO: For CCA use {example_subpkg},{matched_name} " + f"(or another listed sub-package)" + ), + ] + + def _go_subpackage_flow_control(self, input_function: str, package_docs, module_path: str) -> list[str]: + """Go-specific flow: classify exact/strong/weak matches; gate sub-package INFO on exact/strong only.""" + func_to_docs = defaultdict(list) + for doc in package_docs: + if not self.lang_parser: + continue + try: + function_name = self.lang_parser.get_function_name(doc) + except ValueError: + continue + if function_name: + func_to_docs[function_name].append(doc) + + all_functions = list(func_to_docs.keys()) + + # Exact name present in indexed sources — safe for FCF/CCA. + if input_function in func_to_docs: + logger.info("Exact function match found: %s", input_function) + result = [f"{FL_FUNCTION_EXACT_PREFIX} {input_function} found"] + subpkg_to_funcs = self._subpackages_for_functions( + [input_function], func_to_docs, module_path + ) + if len(subpkg_to_funcs) > 1: + result.extend( + self._format_subpackage_ambiguous(input_function, subpkg_to_funcs, module_path) + ) + return result + + strong_matches = self._rank_fuzzy_matches( + input_function, all_functions, FL_STRONG_MATCH_CUTOFF + ) + if strong_matches: + names = [name for name, _ in strong_matches] + logger.info("Strong fuzzy matches found: %s", strong_matches) + result = [ + ( + f"{FL_FUNCTION_STRONG_MATCH_PREFIX} suggestions=[{', '.join(names)}] " + f"(high-confidence fuzzy; not the queried function)" + ) + ] + # Ambiguity only for the best strong name when it spans multiple sub-packages. + best_name = names[0] + subpkg_to_funcs = self._subpackages_for_functions( + [best_name], func_to_docs, module_path + ) + if len(subpkg_to_funcs) > 1: + result.extend( + self._format_subpackage_ambiguous(best_name, subpkg_to_funcs, module_path) + ) + return result + + weak_matches = self._rank_fuzzy_matches( + input_function, all_functions, FL_WEAK_MATCH_CUTOFF + ) + result = [ + ( + f"{FL_FUNCTION_ABSENT_PREFIX} {input_function} not found in indexed sources " + f"for {module_path}" + ) + ] + if weak_matches: + names = [name for name, _ in weak_matches] + logger.info("Weak fuzzy matches only (queried function absent): %s", weak_matches) + result.append( + ( + f"{FL_FUNCTION_WEAK_MATCH_PREFIX} suggestions=[{', '.join(names)}] " + f"(fuzzy only; not the queried function)" + ) + ) + else: + logger.info("No fuzzy matches for '%s' in module '%s'", input_function, module_path) + return result + def search_in_third_party_packages(self, package: str) -> bool: """ @@ -398,23 +560,31 @@ async def locate_functions(self, query: str) -> list[str]: # Package is valid, proceed with function search self.is_package_valid = True package_docs = [] + resolved_module_path = package if (self.coc_retriever.ecosystem and self.coc_retriever.ecosystem.value == Ecosystem.C_CPP.value): package_docs = self.coc_retriever.sort_docs[package] else: # package_docs = self.coc_retriever.documents_of_functions + packages_to_search = [] if package in self.short_go_package_name: - package = self.short_go_package_name[package] + packages_to_search = self.short_go_package_name[package] + if packages_to_search: + resolved_module_path = packages_to_search[0] elif self.coc_retriever.ecosystem.value == Ecosystem.GO.value: resolved = self._resolve_go_fqdn_to_module(package) if resolved: - package = resolved + packages_to_search = [resolved] + resolved_module_path = resolved + if not packages_to_search: + packages_to_search = [package] package_docs = [ doc for doc in (self.coc_retriever.documents_of_functions or []) if ( self.lang_parser - and self.lang_parser.get_package_name(doc, package) + and any(self.lang_parser.get_package_name(doc, p) + for p in packages_to_search) ) ] @@ -432,7 +602,7 @@ async def locate_functions(self, query: str) -> list[str]: case Ecosystem.C_CPP.value: return self.common_flow_control(function, package_docs) case Ecosystem.GO.value: - return self.common_flow_control(function, package_docs) + return self._go_subpackage_flow_control(function, package_docs, resolved_module_path) case Ecosystem.JAVA.value: return self.java_flow_control(function, package_docs) case _: @@ -467,7 +637,7 @@ async def quick_standard_lib_check(package_name: str, ecosystem: Ecosystem) -> t True if package is standard library, False otherwise """ try: - search = ExploitIqSerpAPIWrapper(max_retries=2) + search = ExploitIqSerpAPIWrapper() result = await search.arun(f"Is '{package_name}' part of the {ecosystem.value} standard library?") logger.info("quick_standard_lib_check Standard library check result: %s", result) text = str(result).lower() diff --git a/src/vuln_analysis/utils/intel_retriever.py b/src/vuln_analysis/utils/intel_retriever.py index 7c67dcb50..e1f77b8e8 100644 --- a/src/vuln_analysis/utils/intel_retriever.py +++ b/src/vuln_analysis/utils/intel_retriever.py @@ -78,8 +78,9 @@ def __init__(self, self._rhsa_client = RHSAClient(session=self._session, retry_count=max_retries, retry_on_client_errors=retry_on_client_errors) + # Ubuntu's CVE API is often flaky (503/504); fail faster than other intel sources self._ubuntu_client = UbuntuClient(session=self._session, - retry_count=max_retries, + retry_count=5, retry_on_client_errors=retry_on_client_errors) self._osidb_client = None diff --git a/src/vuln_analysis/utils/intel_utils.py b/src/vuln_analysis/utils/intel_utils.py index 5c998a615..fae9bde9f 100644 --- a/src/vuln_analysis/utils/intel_utils.py +++ b/src/vuln_analysis/utils/intel_utils.py @@ -14,6 +14,7 @@ # limitations under the License. import ipaddress +import os import re from pathlib import Path from urllib.parse import urlparse @@ -525,6 +526,31 @@ def _match_func_name(text: str, patterns: list[re.Pattern], out: set[str]) -> No break +def extract_go_subpackages_from_patch(parsed_patch, module_path: str) -> set[str]: + """Extract Go sub-package paths from patched file paths. + + For a module like 'google.golang.org/protobuf' and a patched file + 'encoding/protojson/decode.go', returns the full sub-package path + 'google.golang.org/protobuf/encoding/protojson'. + Only considers .go source files. + """ + subpackages: set[str] = set() + for pf in parsed_patch.files: + if TEST_FILE_RE.search(pf.target_path): + continue + path = pf.target_path + if path.startswith("b/"): + path = path[2:] + if path.startswith("a/"): + path = path[2:] + if not path.endswith(".go"): + continue + dir_path = os.path.dirname(path) + if dir_path and dir_path != ".": + subpackages.add(f"{module_path}/{dir_path}") + return subpackages + + def extract_functions_from_parsed_patch(parsed_patch, ecosystem: str = "") -> set[str]: """Extract function names from a ParsedPatch object. @@ -545,7 +571,8 @@ def extract_functions_from_parsed_patch(parsed_patch, ecosystem: str = "") -> se return funcs -def _append_enrichment_to_context(all_funcs: set[str], critical_context: list[str], source_label: str) -> None: +def _append_enrichment_to_context(all_funcs: set[str], critical_context: list[str], source_label: str, + go_subpackages: set[str] | None = None) -> None: """Append extracted function names to critical_context.""" critical_context.append( f"Vulnerable functions (remediation patch hint): {', '.join(sorted(all_funcs))}" @@ -554,15 +581,37 @@ def _append_enrichment_to_context(all_funcs: set[str], critical_context: list[st if short_names: unique = list(dict.fromkeys(list(all_funcs) + short_names)) critical_context.append(f"Search keywords: {', '.join(unique)}") + if go_subpackages: + critical_context.append( + f"Vulnerable Go sub-packages (from patch): {', '.join(sorted(go_subpackages))}. " + f"Use these specific sub-package paths with Function Locator and Call Chain Analyzer." + ) logger.info("Patch enrichment extracted functions from %s: %s", source_label, sorted(all_funcs)) +def _extract_go_subpackages_from_result(parsed_patch, candidate_packages: list[dict] | None) -> set[str] | None: + """Extract Go sub-packages from a parsed patch when candidate_packages has a Go module.""" + if not candidate_packages or not parsed_patch: + return None + go_module = None + for cp in candidate_packages: + eco = cp.get("ecosystem", "") + if eco.lower() == "go" and cp.get("name"): + go_module = cp["name"] + break + if not go_module: + return None + subpkgs = extract_go_subpackages_from_patch(parsed_patch, go_module) + return subpkgs if subpkgs else None + + async def enrich_vulnerable_functions_from_patch( cve_intel_list, critical_context: list[str], vulnerable_functions: set[str], ecosystem: str = "", patch_result=None, + candidate_packages: list[dict] | None = None, ) -> None: """Extract vulnerable function names from remediation patch commits. @@ -582,7 +631,10 @@ async def enrich_vulnerable_functions_from_patch( if patch_result and patch_result.parsed_patch: all_funcs = extract_functions_from_parsed_patch(patch_result.parsed_patch, ecosystem) if all_funcs: - _append_enrichment_to_context(all_funcs, critical_context, f"pre-fetched patch ({patch_result.source})") + go_subpkgs = _extract_go_subpackages_from_result(patch_result.parsed_patch, candidate_packages) + _append_enrichment_to_context(all_funcs, critical_context, + f"pre-fetched patch ({patch_result.source})", + go_subpackages=go_subpkgs) return try: @@ -613,7 +665,10 @@ async def enrich_vulnerable_functions_from_patch( if result and result.parsed_patch: all_funcs = extract_functions_from_parsed_patch(result.parsed_patch, ecosystem) if all_funcs: - _append_enrichment_to_context(all_funcs, critical_context, f"fetched patch ({result.source})") + go_subpkgs = _extract_go_subpackages_from_result(result.parsed_patch, candidate_packages) + _append_enrichment_to_context(all_funcs, critical_context, + f"fetched patch ({result.source})", + go_subpackages=go_subpkgs) return else: logger.info("Patch fetched for %s but no functions extracted", vuln_id) diff --git a/src/vuln_analysis/utils/tests/test_function_name_locator_go.py b/src/vuln_analysis/utils/tests/test_function_name_locator_go.py index 755c5a3af..98b0c5c61 100644 --- a/src/vuln_analysis/utils/tests/test_function_name_locator_go.py +++ b/src/vuln_analysis/utils/tests/test_function_name_locator_go.py @@ -155,3 +155,116 @@ def test_single_segment_package_no_match(self): ["protobuf-lite"], ) assert result is None + + +class TestFLGoSubpackageFlowHooks: + """Phase 1: Go FL emits exact/strong/weak/absent hooks; no false sub-package claims.""" + + MODULE = "github.com/docker/docker" + + def _make_locator(self): + parser = MagicMock() + retriever = MagicMock() + retriever.language_parser = parser + retriever.ecosystem = Ecosystem.GO + retriever.supported_packages = [self.MODULE] + retriever.documents_of_functions = [] + locator = FunctionNameLocator(retriever) + return locator, parser + + @staticmethod + def _doc(source: str, func_name: str): + doc = MagicMock() + doc.metadata = {"source": source} + doc._func_name = func_name + return doc + + def _wire_parser(self, parser, docs): + name_by_id = {id(d): d._func_name for d in docs} + parser.get_function_name.side_effect = lambda doc: name_by_id[id(doc)] + + def test_exact_match_emits_fl_function_exact(self): + locator, parser = self._make_locator() + docs = [ + self._doc(f"vendor/{self.MODULE}/oci/defaults.go", "DefaultLinuxSpec"), + ] + self._wire_parser(parser, docs) + + result = locator._go_subpackage_flow_control( + "DefaultLinuxSpec", docs, self.MODULE + ) + + assert result[0].startswith("FL_FUNCTION_EXACT:") + assert "DefaultLinuxSpec" in result[0] + assert not any(line.startswith("FL_SUBPACKAGE_AMBIGUOUS:") for line in result) + assert not any(line.startswith("FL_FUNCTION_ABSENT:") for line in result) + + def test_exact_match_multi_subpackage_emits_ambiguous(self): + locator, parser = self._make_locator() + docs = [ + self._doc(f"vendor/{self.MODULE}/a/foo.go", "Foo"), + self._doc(f"vendor/{self.MODULE}/b/foo.go", "Foo"), + ] + self._wire_parser(parser, docs) + + result = locator._go_subpackage_flow_control("Foo", docs, self.MODULE) + + assert result[0].startswith("FL_FUNCTION_EXACT:") + assert any(line.startswith("FL_SUBPACKAGE_AMBIGUOUS:") for line in result) + assert f"{self.MODULE}/a" in " ".join(result) + assert f"{self.MODULE}/b" in " ".join(result) + + def test_cve_absent_weak_no_false_subpackage_claim(self): + """CVE-2022-24769 shape: module present, queried name absent, weak near-names + across utility sub-packages must NOT claim DefaultLinuxSpec was found there.""" + locator, parser = self._make_locator() + docs = [ + self._doc(f"vendor/{self.MODULE}/api/types/container/config.go", "Deadline"), + self._doc(f"vendor/{self.MODULE}/errdefs/helpers.go", "IsDefault"), + self._doc(f"vendor/{self.MODULE}/pkg/archive/archive.go", "DefaultPathEnv"), + ] + self._wire_parser(parser, docs) + + result = locator._go_subpackage_flow_control( + "DefaultLinuxSpec", docs, self.MODULE + ) + + assert any(line.startswith("FL_FUNCTION_ABSENT:") for line in result) + assert "DefaultLinuxSpec" in result[0] + assert any(line.startswith("FL_FUNCTION_WEAK_MATCH:") for line in result) + assert not any(line.startswith("FL_FUNCTION_EXACT:") for line in result) + assert not any(line.startswith("FL_SUBPACKAGE_AMBIGUOUS:") for line in result) + assert not any(line.startswith("INFO:") and "DefaultLinuxSpec" in line for line in result) + joined = " ".join(result) + assert "Functions found in multiple sub-packages" not in joined + + def test_strong_match_not_exact(self): + locator, parser = self._make_locator() + # Ratio DefaultLinuxSpec ↔ DefaultLinuxSpecs is well above 0.8 + docs = [ + self._doc(f"vendor/{self.MODULE}/oci/defaults.go", "DefaultLinuxSpecs"), + ] + self._wire_parser(parser, docs) + + result = locator._go_subpackage_flow_control( + "DefaultLinuxSpec", docs, self.MODULE + ) + + assert result[0].startswith("FL_FUNCTION_STRONG_MATCH:") + assert "DefaultLinuxSpecs" in result[0] + assert not any(line.startswith("FL_FUNCTION_EXACT:") for line in result) + assert not any(line.startswith("FL_FUNCTION_ABSENT:") for line in result) + + def test_no_matches_emits_absent_only(self): + locator, parser = self._make_locator() + docs = [ + self._doc(f"vendor/{self.MODULE}/pkg/util.go", "CompletelyDifferentName"), + ] + self._wire_parser(parser, docs) + + result = locator._go_subpackage_flow_control( + "DefaultLinuxSpec", docs, self.MODULE + ) + + assert len(result) == 1 + assert result[0].startswith("FL_FUNCTION_ABSENT:") diff --git a/tests/test_fl_observation_hooks.py b/tests/test_fl_observation_hooks.py new file mode 100644 index 000000000..e9f066f7d --- /dev/null +++ b/tests/test_fl_observation_hooks.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Phase 2: Function Locator observation hooks — deterministic parse + memory categories.""" + +from vuln_analysis.functions.base_graph_agent import parse_function_locator_findings +from vuln_analysis.functions.react_internals import ( + COMPREHENSION_PROMPT, + MEMORY_UPDATE_PROMPT, + parse_to_memory_item, +) +from vuln_analysis.tools.tool_names import ToolNames + + +CVE_FL_PAYLOAD = ( + "{'ecosystem': 'GO', " + "'package_msg': 'FL_PACKAGE_PRESENT: github.com/docker/docker', " + "'status': 'present', " + "'result': [" + "'FL_FUNCTION_ABSENT: DefaultLinuxSpec not found in indexed sources " + "for github.com/docker/docker', " + "'FL_FUNCTION_WEAK_MATCH: suggestions=[Deadline, IsDefault, DefaultPathEnv] " + "(fuzzy only; not the queried function)'" + "]}" +) + +CVE_TOOL_INPUT = "github.com/docker/docker,DefaultLinuxSpec" + +EXACT_FL_PAYLOAD = ( + "{'package_msg': 'FL_PACKAGE_PRESENT: github.com/lib/foo', " + "'result': [" + "'FL_FUNCTION_EXACT: Foo found', " + "'FL_SUBPACKAGE_AMBIGUOUS: Foo in github.com/lib/foo/a, github.com/lib/foo/b'" + "]}" +) + + +class TestParseFunctionLocatorFindings: + """Deterministic FL → CodeFindings (no LLM).""" + + def test_cve_absent_payload_no_false_validated(self): + findings = parse_function_locator_findings(CVE_FL_PAYLOAD, CVE_TOOL_INPUT) + assert findings is not None + texts = findings.findings + assert any(f.startswith("FL_PACKAGE_PRESENT:") for f in texts) + assert any(f.startswith("FL_FUNCTION_ABSENT:") for f in texts) + assert any(f.startswith("FL_FUNCTION_WEAK_MATCH:") for f in texts) + assert not any(f.startswith("VALIDATED:") for f in texts) + assert "DefaultLinuxSpec exists" not in " ".join(texts) + + def test_exact_payload_emits_validated_alias(self): + findings = parse_function_locator_findings( + EXACT_FL_PAYLOAD, "github.com/lib/foo,Foo" + ) + assert findings is not None + assert any(f.startswith("FL_FUNCTION_EXACT:") for f in findings.findings) + assert "VALIDATED: github.com/lib/foo,Foo exists" in findings.findings + + def test_no_hooks_returns_none_for_llm_fallback(self): + assert parse_function_locator_findings("['Deadline', 'IsDefault']", "a,b") is None + + def test_empty_output_returns_none(self): + assert parse_function_locator_findings("", "a,b") is None + assert parse_function_locator_findings(None, "a,b") is None # type: ignore[arg-type] + + +class TestFlMemoryCategories: + """CATEGORY_PATTERNS recognize FL_* prefixes.""" + + def test_fl_absent_category(self): + item = parse_to_memory_item( + "FL_FUNCTION_ABSENT: DefaultLinuxSpec not found in indexed sources " + "for github.com/docker/docker" + ) + assert item.category == "fl_absent" + assert item.priority == 3 + + def test_fl_weak_category(self): + item = parse_to_memory_item( + "FL_FUNCTION_WEAK_MATCH: suggestions=[Deadline] (fuzzy only; not the queried function)" + ) + assert item.category == "fl_weak" + assert item.priority == 1 + + def test_fl_exact_category(self): + item = parse_to_memory_item("FL_FUNCTION_EXACT: Foo found") + assert item.category == "fl_exact" + assert item.priority == 3 + + def test_fl_package_present_category(self): + item = parse_to_memory_item("FL_PACKAGE_PRESENT: github.com/docker/docker") + assert item.category == "fl_package_present" + + def test_cve_findings_map_without_validated(self): + findings = parse_function_locator_findings(CVE_FL_PAYLOAD, CVE_TOOL_INPUT) + categories = {parse_to_memory_item(f).category for f in findings.findings} + assert "fl_absent" in categories + assert "fl_weak" in categories + assert "fl_package_present" in categories + assert "validated" not in categories + + +class TestFlPromptGuards: + """Comprehension/memory prompts must not invent VALIDATED for non-FL tools.""" + + def test_comprehension_fl_rules_are_tool_gated(self): + assert 'When TOOL USED is "Function Locator"' in COMPREHENSION_PROMPT + assert "For Function Locator results, include: \"VALIDATED:" not in COMPREHENSION_PROMPT + + def test_memory_prompt_does_not_blanket_validated(self): + assert "For Function Locator validation, record:" not in MEMORY_UPDATE_PROMPT + assert "ONLY when NEW FINDINGS include FL_FUNCTION_EXACT" in MEMORY_UPDATE_PROMPT + + def test_non_fl_tool_name_constant_unchanged(self): + # Negative: Keyword Search / CCA are distinct tools; parser is FL-gated by caller. + assert ToolNames.CODE_KEYWORD_SEARCH != ToolNames.FUNCTION_LOCATOR + assert ToolNames.CALL_CHAIN_ANALYZER != ToolNames.FUNCTION_LOCATOR + # Parser itself has no tool gate; observation_node gates on ToolNames.FUNCTION_LOCATOR. + keyword_output = "Main application: rate_limit found in app/main.go" + assert parse_function_locator_findings(keyword_output, "rate_limit") is None diff --git a/tests/test_go_subpackage_fixes.py b/tests/test_go_subpackage_fixes.py new file mode 100644 index 000000000..e3972c676 --- /dev/null +++ b/tests/test_go_subpackage_fixes.py @@ -0,0 +1,460 @@ +"""Tests for Go sub-package granularity fixes (FL feedback, CCA filtering, intel enrichment).""" + +import os +from collections import defaultdict +from unittest.mock import MagicMock + +import pytest +from langchain_core.documents import Document + +from vuln_analysis.utils.function_name_locator import FunctionNameLocator +from exploit_iq_commons.utils.functions_parsers.golang_functions_parsers import GoLanguageFunctionsParser +from exploit_iq_commons.utils.functions_parsers.lang_functions_parsers import LanguageFunctionsParser +from vuln_analysis.utils.intel_utils import ( + extract_go_subpackages_from_patch, + _extract_go_subpackages_from_result, + _append_enrichment_to_context, +) + + +# --------------------------------------------------------------------------- +# FL sub-package feedback +# --------------------------------------------------------------------------- + + +class TestExtractGoSubpackage: + + def test_vendor_prefixed_path(self): + result = FunctionNameLocator._extract_go_subpackage( + "vendor/google.golang.org/protobuf/encoding/protojson/decode.go", + "google.golang.org/protobuf", + ) + assert result == "google.golang.org/protobuf/encoding/protojson" + + def test_no_vendor_prefix(self): + result = FunctionNameLocator._extract_go_subpackage( + "google.golang.org/protobuf/encoding/protojson/decode.go", + "google.golang.org/protobuf", + ) + assert result == "google.golang.org/protobuf/encoding/protojson" + + def test_file_at_module_root(self): + result = FunctionNameLocator._extract_go_subpackage( + "vendor/google.golang.org/protobuf/proto.go", + "google.golang.org/protobuf", + ) + assert result == "google.golang.org/protobuf" + + def test_deep_subpackage(self): + result = FunctionNameLocator._extract_go_subpackage( + "vendor/github.com/quic-go/quic-go/internal/wire/frame.go", + "github.com/quic-go/quic-go", + ) + assert result == "github.com/quic-go/quic-go/internal/wire" + + def test_unrelated_path_returns_module(self): + result = FunctionNameLocator._extract_go_subpackage( + "vendor/some/other/package/file.go", + "google.golang.org/protobuf", + ) + assert result == "google.golang.org/protobuf" + + +class TestGoSubpackageFlowControl: + + def _make_doc(self, func_name: str, source: str) -> Document: + return Document( + page_content=f"func {func_name}() {{}}", + metadata={"source": source, "func_name": func_name}, + ) + + def _make_locator(self) -> FunctionNameLocator: + parser = GoLanguageFunctionsParser() + mock_retriever = MagicMock() + mock_retriever.ecosystem = MagicMock() + mock_retriever.ecosystem.value = "go" + mock_retriever.documents_of_functions = [] + mock_retriever.supported_packages = [] + locator = FunctionNameLocator.__new__(FunctionNameLocator) + locator.lang_parser = parser + locator.coc_retriever = mock_retriever + return locator + + def test_single_subpackage_no_guidance(self): + locator = self._make_locator() + docs = [ + self._make_doc("Unmarshal", "vendor/google.golang.org/protobuf/proto/decode.go"), + self._make_doc("Marshal", "vendor/google.golang.org/protobuf/proto/encode.go"), + ] + result = locator._go_subpackage_flow_control("Unmarshal", docs, "google.golang.org/protobuf") + assert result[0].startswith("FL_FUNCTION_EXACT:") + assert "Unmarshal" in result[0] + assert not any("INFO:" in r for r in result) + assert not any(r.startswith("FL_SUBPACKAGE_AMBIGUOUS:") for r in result) + + def test_multiple_subpackages_adds_guidance(self): + locator = self._make_locator() + docs = [ + self._make_doc("Unmarshal", "vendor/google.golang.org/protobuf/proto/decode.go"), + self._make_doc("Unmarshal", "vendor/google.golang.org/protobuf/encoding/protojson/decode.go"), + self._make_doc("Unmarshal", "vendor/google.golang.org/protobuf/encoding/prototext/decode.go"), + ] + result = locator._go_subpackage_flow_control("Unmarshal", docs, "google.golang.org/protobuf") + assert result[0].startswith("FL_FUNCTION_EXACT:") + assert "Unmarshal" in result[0] + ambiguous = [r for r in result if r.startswith("FL_SUBPACKAGE_AMBIGUOUS:")] + assert len(ambiguous) == 1 + assert "encoding/protojson" in ambiguous[0] + info_lines = [r for r in result if "INFO:" in r] + assert len(info_lines) == 1 + assert "For CCA use" in info_lines[0] + assert "encoding/protojson" in info_lines[0] + + def test_no_matches_returns_empty(self): + locator = self._make_locator() + docs = [ + self._make_doc("SomethingElse", "vendor/google.golang.org/protobuf/proto/other.go"), + ] + result = locator._go_subpackage_flow_control("Unmarshal", docs, "google.golang.org/protobuf") + assert len(result) == 1 + assert result[0].startswith("FL_FUNCTION_ABSENT:") + assert "Unmarshal" in result[0] + + +# --------------------------------------------------------------------------- +# CCA sub-package filtering (Go parser method) +# --------------------------------------------------------------------------- + + +class TestResolveSubpackageToModule: + + def test_subpackage_returns_suffix(self): + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "google.golang.org/protobuf/encoding/protojson", + "google.golang.org/protobuf", + ) + assert result == "/encoding/protojson" + + def test_exact_match_returns_none(self): + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "google.golang.org/protobuf", + "google.golang.org/protobuf", + ) + assert result is None + + def test_unrelated_returns_none(self): + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "github.com/some/other", + "google.golang.org/protobuf", + ) + assert result is None + + def test_case_insensitive(self): + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "Google.Golang.Org/Protobuf/Encoding/Protojson", + "google.golang.org/protobuf", + ) + assert result == "/Encoding/Protojson" + + def test_deep_subpackage(self): + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "github.com/quic-go/quic-go/internal/wire", + "github.com/quic-go/quic-go", + ) + assert result == "/internal/wire" + + def test_partial_overlap_not_matched(self): + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "google.golang.org/protobuf-extra/encoding", + "google.golang.org/protobuf", + ) + assert result is None + + +# --------------------------------------------------------------------------- +# Intel sub-package enrichment from patches +# --------------------------------------------------------------------------- + + +class _FakePatchFile: + def __init__(self, target_path: str): + self.target_path = target_path + + +class _FakeParsedPatch: + def __init__(self, files: list[_FakePatchFile]): + self.files = files + + +class TestExtractGoSubpackagesFromPatch: + + def test_extracts_subpackages(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/encoding/protojson/decode.go"), + _FakePatchFile("b/encoding/prototext/decode.go"), + _FakePatchFile("b/proto/decode.go"), + ]) + result = extract_go_subpackages_from_patch(patch, "google.golang.org/protobuf") + assert result == { + "google.golang.org/protobuf/encoding/protojson", + "google.golang.org/protobuf/encoding/prototext", + "google.golang.org/protobuf/proto", + } + + def test_skips_test_files(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/encoding/protojson/decode.go"), + _FakePatchFile("b/encoding/protojson/decode_test.go"), + ]) + result = extract_go_subpackages_from_patch(patch, "google.golang.org/protobuf") + assert result == {"google.golang.org/protobuf/encoding/protojson"} + + def test_root_file_not_included(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/proto.go"), + ]) + result = extract_go_subpackages_from_patch(patch, "google.golang.org/protobuf") + assert result == set() + + def test_no_b_prefix(self): + patch = _FakeParsedPatch([ + _FakePatchFile("encoding/protojson/decode.go"), + ]) + result = extract_go_subpackages_from_patch(patch, "google.golang.org/protobuf") + assert result == {"google.golang.org/protobuf/encoding/protojson"} + + +class TestExtractGoSubpackagesFromResult: + + def test_with_go_candidate(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/encoding/protojson/decode.go"), + ]) + candidates = [{"name": "google.golang.org/protobuf", "source": "ghsa", "ecosystem": "Go"}] + result = _extract_go_subpackages_from_result(patch, candidates) + assert result == {"google.golang.org/protobuf/encoding/protojson"} + + def test_without_go_candidate_returns_none(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/src/decode.c"), + ]) + candidates = [{"name": "libxml2", "source": "ghsa", "ecosystem": "C"}] + result = _extract_go_subpackages_from_result(patch, candidates) + assert result is None + + def test_no_candidates_returns_none(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/encoding/protojson/decode.go"), + ]) + result = _extract_go_subpackages_from_result(patch, None) + assert result is None + + def test_empty_subpackages_returns_none(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/proto.go"), + ]) + candidates = [{"name": "google.golang.org/protobuf", "source": "ghsa", "ecosystem": "Go"}] + result = _extract_go_subpackages_from_result(patch, candidates) + assert result is None + + +# --------------------------------------------------------------------------- +# Additional edge cases +# --------------------------------------------------------------------------- + + +class TestExtractGoSubpackageEdgeCases: + + def test_versioned_module_path(self): + result = FunctionNameLocator._extract_go_subpackage( + "vendor/github.com/quic-go/quic-go/v3/internal/wire/frame.go", + "github.com/quic-go/quic-go/v3", + ) + assert result == "github.com/quic-go/quic-go/v3/internal/wire" + + def test_empty_source_path(self): + result = FunctionNameLocator._extract_go_subpackage("", "google.golang.org/protobuf") + assert result == "google.golang.org/protobuf" + + def test_trailing_slash_on_module(self): + result = FunctionNameLocator._extract_go_subpackage( + "vendor/google.golang.org/protobuf/proto/decode.go", + "google.golang.org/protobuf/", + ) + assert result == "google.golang.org/protobuf/proto" + + +class TestGoSubpackageFlowControlEdgeCases: + + def _make_doc(self, func_name: str, source: str) -> Document: + return Document( + page_content=f"func {func_name}() {{}}", + metadata={"source": source, "func_name": func_name}, + ) + + def _make_locator(self) -> FunctionNameLocator: + parser = GoLanguageFunctionsParser() + mock_retriever = MagicMock() + mock_retriever.ecosystem = MagicMock() + mock_retriever.ecosystem.value = "go" + locator = FunctionNameLocator.__new__(FunctionNameLocator) + locator.lang_parser = parser + locator.coc_retriever = mock_retriever + return locator + + def test_fuzzy_match_cutoff_03_fallback(self): + locator = self._make_locator() + docs = [ + self._make_doc("UnmarshalJSON", "vendor/google.golang.org/protobuf/proto/decode.go"), + ] + result = locator._go_subpackage_flow_control("Unmarshal", docs, "google.golang.org/protobuf") + assert result[0].startswith("FL_FUNCTION_STRONG_MATCH:") + assert "UnmarshalJSON" in result[0] + assert not any(r.startswith("FL_FUNCTION_EXACT:") for r in result) + + def test_same_function_two_subpackages_shows_both(self): + locator = self._make_locator() + docs = [ + self._make_doc("Reset", "vendor/google.golang.org/protobuf/proto/reset.go"), + self._make_doc("Reset", "vendor/google.golang.org/protobuf/internal/impl/reset.go"), + ] + result = locator._go_subpackage_flow_control("Reset", docs, "google.golang.org/protobuf") + assert result[0].startswith("FL_FUNCTION_EXACT:") + assert "Reset" in result[0] + ambiguous = [r for r in result if r.startswith("FL_SUBPACKAGE_AMBIGUOUS:")] + assert len(ambiguous) == 1 + assert "proto" in ambiguous[0] + assert "internal/impl" in ambiguous[0] + info_lines = [r for r in result if "INFO:" in r] + assert len(info_lines) == 1 + assert "For CCA use" in info_lines[0] + assert "Reset" in info_lines[0] + # INFO lists one example sub-package; both appear on FL_SUBPACKAGE_AMBIGUOUS + assert ("proto" in info_lines[0]) or ("internal/impl" in info_lines[0]) + + +# --------------------------------------------------------------------------- +# Additional edge cases for resolve_subpackage_to_module +# --------------------------------------------------------------------------- + + +class TestResolveSubpackageToModuleEdgeCases: + + def test_base_class_returns_none(self): + """Non-Go ecosystems should always return None.""" + from exploit_iq_commons.utils.functions_parsers.c_lang_function_parsers import CLanguageFunctionsParser + parser = CLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module("libxml2/encoding", "libxml2") + assert result is None + + def test_versioned_module_v2(self): + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "github.com/golang-jwt/jwt/v5/parser", + "github.com/golang-jwt/jwt/v5", + ) + assert result == "/parser" + + def test_prefix_match_without_slash_boundary(self): + """'protobuf-extra' should NOT match 'protobuf'.""" + parser = GoLanguageFunctionsParser() + result = parser.resolve_subpackage_to_module( + "google.golang.org/protobufextra/encoding", + "google.golang.org/protobuf", + ) + assert result is None + + +# --------------------------------------------------------------------------- +# _append_enrichment_to_context with go_subpackages +# --------------------------------------------------------------------------- + + +class TestAppendEnrichmentGoSubpackages: + + def test_appends_subpackage_guidance(self): + ctx = [] + _append_enrichment_to_context( + {"Unmarshal", "UnmarshalOptions"}, + ctx, + "test-source", + go_subpackages={"google.golang.org/protobuf/encoding/protojson"}, + ) + subpkg_lines = [c for c in ctx if "Go sub-packages" in c] + assert len(subpkg_lines) == 1 + assert "encoding/protojson" in subpkg_lines[0] + assert "Function Locator" in subpkg_lines[0] + + def test_no_subpackages_no_guidance(self): + ctx = [] + _append_enrichment_to_context({"Unmarshal"}, ctx, "test-source") + subpkg_lines = [c for c in ctx if "Go sub-packages" in c] + assert len(subpkg_lines) == 0 + + def test_none_subpackages_no_guidance(self): + ctx = [] + _append_enrichment_to_context({"Unmarshal"}, ctx, "test-source", go_subpackages=None) + subpkg_lines = [c for c in ctx if "Go sub-packages" in c] + assert len(subpkg_lines) == 0 + + def test_empty_set_subpackages_no_guidance(self): + ctx = [] + _append_enrichment_to_context({"Unmarshal"}, ctx, "test-source", go_subpackages=set()) + subpkg_lines = [c for c in ctx if "Go sub-packages" in c] + assert len(subpkg_lines) == 0 + + def test_multiple_subpackages_sorted(self): + ctx = [] + _append_enrichment_to_context( + {"Unmarshal"}, + ctx, + "test-source", + go_subpackages={ + "google.golang.org/protobuf/encoding/prototext", + "google.golang.org/protobuf/encoding/protojson", + }, + ) + subpkg_lines = [c for c in ctx if "Go sub-packages" in c] + assert len(subpkg_lines) == 1 + protojson_pos = subpkg_lines[0].index("protojson") + prototext_pos = subpkg_lines[0].index("prototext") + assert protojson_pos < prototext_pos + + +# --------------------------------------------------------------------------- +# Patch with mixed Go and non-Go files +# --------------------------------------------------------------------------- + + +class TestExtractGoSubpackagesMixedFiles: + + def test_ignores_non_go_files(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/encoding/protojson/decode.go"), + _FakePatchFile("b/README.md"), + _FakePatchFile("b/.github/workflows/ci.yml"), + ]) + result = extract_go_subpackages_from_patch(patch, "google.golang.org/protobuf") + assert result == {"google.golang.org/protobuf/encoding/protojson"} + + def test_deduplicates_subpackages(self): + patch = _FakeParsedPatch([ + _FakePatchFile("b/encoding/protojson/decode.go"), + _FakePatchFile("b/encoding/protojson/encode.go"), + _FakePatchFile("b/encoding/protojson/well_known_types.go"), + ]) + result = extract_go_subpackages_from_patch(patch, "google.golang.org/protobuf") + assert result == {"google.golang.org/protobuf/encoding/protojson"} + + def test_a_prefix_stripped(self): + patch = _FakeParsedPatch([ + _FakePatchFile("a/encoding/protojson/decode.go"), + ]) + result = extract_go_subpackages_from_patch(patch, "google.golang.org/protobuf") + assert result == {"google.golang.org/protobuf/encoding/protojson"} \ No newline at end of file