diff --git a/dissect/cstruct/lexer.py b/dissect/cstruct/lexer.py index 30a03ec..c6dc82e 100644 --- a/dissect/cstruct/lexer.py +++ b/dissect/cstruct/lexer.py @@ -390,6 +390,41 @@ def _read_string(self) -> str: return result + def _read_define_value(self) -> str: + """Read value of a define, ignoring comments and unecessary white space.""" + text_offsets: list[tuple[int, int]] = [] + + while not self.eof: + char = self._current() + if char == "\n": + break + + if char in ("\"'"): + sof = self._pos + self._read_string() + text_offsets.append((sof, self._pos)) + continue + + if char not in " \t\r": + sof = self._pos + self._read_until(" \t\r\n") + text_offsets.append((sof, self._pos)) + continue + + # Skip whitespace + self._read_while(" \t\r", or_eof=True) + if self.eof: + break + self._skip_comment() + if self.eof: + break + + result: list[str] = [] + for start, end in text_offsets: + self._pos = start + result.append(self._take(end - start)) + return " ".join(result) + def _read_preprocessor(self) -> None: """Read a preprocessor directive starting with ``#``.""" line = self._line @@ -428,7 +463,7 @@ def _read_preprocessor(self) -> None: # No value, just a simple macro definition return - if (value := self._read_until("\n")).strip(): + if value := self._read_define_value(): self._emit(TokenType.STRING, value, line) elif token_type == TokenType.PP_INCLUDE: diff --git a/tests/test_lexer.py b/tests/test_lexer.py index f08a584..64d13c6 100644 --- a/tests/test_lexer.py +++ b/tests/test_lexer.py @@ -96,7 +96,7 @@ + 2) """, [TokenType.PP_DEFINE, TokenType.IDENTIFIER, TokenType.STRING], - ["define", "FOO", "(1 + 2)"], + ["define", "FOO", "(1 + 2)"], ), ("#undef", [TokenType.PP_UNDEF], ["undef"]), ("#ifdef", [TokenType.PP_IFDEF], ["ifdef"]), diff --git a/tests/test_parser.py b/tests/test_parser.py index 54273f0..d3c45f1 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -251,11 +251,13 @@ def test_define(cs: cstruct) -> None: #define MULTILINE (1 + \ 2 + \ 3) - #define QUOTES "\'\"a'b\"" + #define QUOTES "'\\"a'b\\"" #define ESCAPE "\\'\\"a'b\\"\\n" #define BYTES_ESCAPE b"`\\n" #define FUNC(x) ( x == 0 ) #define TERNARY(x) ( x ? 1 : 0 ) + #define MULTI_LINE_QUOTE "\ + #define DEFINITION " """ cs.load(cdef) @@ -275,6 +277,7 @@ def test_define(cs: cstruct) -> None: # We don't evaluate function-like macros yet, so they should be stored as their raw string representation assert cs.consts["FUNC"] == "(x) ( x == 0 )" assert cs.consts["TERNARY"] == "(x) ( x ? 1 : 0 )" + assert cs.consts["MULTI_LINE_QUOTE"] == " #define DEFINITION " def test_define_flag_value(cs: cstruct) -> None: @@ -515,6 +518,34 @@ def test_preprocessor_in_struct_body(cs: cstruct) -> None: assert cs.test.fields["bonus"].type == cs.uint64 +def test_preprocessor_define_with_comments(cs: cstruct) -> None: + cdef = """ + #define IGNORE_SCOPED_COMMENT data1 /* ... */ + #define IGNORE_LINE_COMMENT data2 // + #define IGNORE_MULTILINE_COMMENT data3 /* Multiline + comment + */ + #define IGNORE_DEFINE_IN_COMMENT data4 /* define + #define TEST_NO data_no inside + comment + */ + #define TEXT_CONTAINING_COMMENT "text with comments /* data5 */" + #define IGNORE_STRING_AFTER_COMMENT "data6" // \ + "string that cannot be reached" + #define NO_COMMENT_BETWEEN_CONCATINATION "data7" /* Comment */ \ + "should get concatinated" + """ + + cs.load(cdef) + assert cs.consts["IGNORE_SCOPED_COMMENT"] == "data1" + assert cs.consts["IGNORE_LINE_COMMENT"] == "data2" + assert cs.consts["IGNORE_MULTILINE_COMMENT"] == "data3" + assert cs.consts["IGNORE_DEFINE_IN_COMMENT"] == "data4" + assert cs.consts["TEXT_CONTAINING_COMMENT"] == "text with comments /* data5 */" + assert cs.consts["IGNORE_STRING_AFTER_COMMENT"] == "data6" + assert cs.consts["NO_COMMENT_BETWEEN_CONCATINATION"] == 'data7" "should get concatinated' + + def test_preprocessor_define_from_enum_in_struct(cs: cstruct) -> None: """Test #define referencing enum values used for conditional fields and array sizes.""" cdef = """