Allow comments in prototypes of unit test functions

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>

tests/scripts/generate_test_code.py

@@ -519,6 +519,41 @@ def generate_function_code(name, code, local_vars, args_dispatch,
         gen_dependencies(dependencies)
     return preprocessor_check_start + code + preprocessor_check_end
 
+COMMENT_START_REGEX = re.compile(r'/[*/]')
+
+def skip_comments(line, stream):
+    """Remove comments in line.
+
+    If the line contains an unfinished comment, read more lines from stream
+    until the line that contains the comment.
+
+    :return: The original line with inner comments replaced by spaces.
+             Trailing comments and whitespace may be removed completely.
+    """
+    pos = 0
+    while True:
+        opening = COMMENT_START_REGEX.search(line, pos)
+        if not opening:
+            break
+        if line[opening.start(0) + 1] == '/': # //...
+            continuation = line
+            while continuation.endswith('\\\n'):
+                # This errors out if the file ends with an unfinished line
+                # comment. That's acceptable to not complicat the code further.
+                continuation = next(stream)
+            return line[:opening.start(0)].rstrip() + '\n'
+        # Parsing /*...*/, looking for the end
+        closing = line.find('*/', opening.end(0))
+        while closing == -1:
+            # This errors out if the file ends with an unfinished block
+            # comment. That's acceptable to not complicat the code further.
+            line += next(stream)
+            closing = line.find('*/', opening.end(0))
+        pos = closing + 2
+        line = (line[:opening.start(0)] +
+                ' ' * (pos - opening.start(0)) +
+                line[pos:])
+    return re.sub(r' +(\n|\Z)', r'\1', line)
 
 def parse_function_code(funcs_f, dependencies, suite_dependencies):
     """
@@ -538,6 +573,7 @@ def parse_function_code(funcs_f, dependencies, suite_dependencies):
         # across multiple lines. Here we try to find the start of
         # arguments list, then remove '\n's and apply the regex to
         # detect function start.
+        line = skip_comments(line, funcs_f)
         up_to_arg_list_start = code + line[:line.find('(') + 1]
         match = re.match(TEST_FUNCTION_VALIDATION_REGEX,
                          up_to_arg_list_start.replace('\n', ' '), re.I)
@@ -546,7 +582,7 @@ def parse_function_code(funcs_f, dependencies, suite_dependencies):
             name = match.group('func_name')
             if not re.match(FUNCTION_ARG_LIST_END_REGEX, line):
                 for lin in funcs_f:
-                    line += lin
+                    line += skip_comments(lin, funcs_f)
                     if re.search(FUNCTION_ARG_LIST_END_REGEX, line):
                         break
             args, local_vars, args_dispatch = parse_function_arguments(

tests/scripts/test_generate_test_code.py

@@ -724,6 +724,102 @@ exit:
     yes sir yes sir
     3 bags full
 }
+'''
+        self.assertEqual(code, expected)
+
+    @patch("generate_test_code.gen_dispatch")
+    @patch("generate_test_code.gen_dependencies")
+    @patch("generate_test_code.gen_function_wrapper")
+    @patch("generate_test_code.parse_function_arguments")
+    def test_case_starting_with_comment(self, parse_function_arguments_mock,
+                                        gen_function_wrapper_mock,
+                                        gen_dependencies_mock,
+                                        gen_dispatch_mock):
+        """
+        Test with comments before the function signature
+        :return:
+        """
+        parse_function_arguments_mock.return_value = ([], '', [])
+        gen_function_wrapper_mock.return_value = ''
+        gen_dependencies_mock.side_effect = gen_dependencies
+        gen_dispatch_mock.side_effect = gen_dispatch
+        data = '''/* comment */
+/* more
+ * comment */
+// this is\
+still \
+a comment
+void func()
+{
+    ba ba black sheep
+    have you any wool
+exit:
+    yes sir yes sir
+    3 bags full
+}
+/* END_CASE */
+'''
+        stream = StringIOWrapper('test_suite_ut.function', data)
+        _, _, code, _ = parse_function_code(stream, [], [])
+
+        expected = '''#line 1 "test_suite_ut.function"
+
+
+
+void test_func()
+{
+    ba ba black sheep
+    have you any wool
+exit:
+    yes sir yes sir
+    3 bags full
+}
+'''
+        self.assertEqual(code, expected)
+
+    @patch("generate_test_code.gen_dispatch")
+    @patch("generate_test_code.gen_dependencies")
+    @patch("generate_test_code.gen_function_wrapper")
+    @patch("generate_test_code.parse_function_arguments")
+    def test_comment_in_prototype(self, parse_function_arguments_mock,
+                                  gen_function_wrapper_mock,
+                                  gen_dependencies_mock,
+                                  gen_dispatch_mock):
+        """
+        Test with comments in the function prototype
+        :return:
+        """
+        parse_function_arguments_mock.return_value = ([], '', [])
+        gen_function_wrapper_mock.return_value = ''
+        gen_dependencies_mock.side_effect = gen_dependencies
+        gen_dispatch_mock.side_effect = gen_dispatch
+        data = '''
+void func( int x, // (line \\
+                     comment)
+           int y /* lone closing parenthesis) */ )
+{
+    ba ba black sheep
+    have you any wool
+exit:
+    yes sir yes sir
+    3 bags full
+}
+/* END_CASE */
+'''
+        stream = StringIOWrapper('test_suite_ut.function', data)
+        _, _, code, _ = parse_function_code(stream, [], [])
+
+        expected = '''#line 1 "test_suite_ut.function"
+
+void test_func( int x,
+           int y                                 )
+{
+    ba ba black sheep
+    have you any wool
+exit:
+    yes sir yes sir
+    3 bags full
+}
 '''
         self.assertEqual(code, expected)
 

tests/suites/test_suite_psa_crypto.function

@@ -1412,6 +1412,7 @@ exit:
 /* END_CASE */
 
 /* BEGIN_CASE */
+/* Construct and attempt to import a large unstructured key. */
 void import_large_key( int type_arg, int byte_size_arg,
                        int expected_status_arg )
 {
@@ -1468,6 +1469,9 @@ exit:
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_ASN1_WRITE_C */
+/* Import an RSA key with a valid structure (but not valid numbers
+ * inside, beyond having sensible size and parity). This is expected to
+ * fail for large keys. */
 void import_rsa_made_up( int bits_arg, int keypair, int expected_status_arg )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
@@ -1513,6 +1517,7 @@ void import_export( data_t *data,
                     int expected_bits,
                     int export_size_delta,
                     int expected_export_status_arg,
+                    /*whether reexport must give the original input exactly*/
                     int canonical_input )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
@@ -1617,7 +1622,7 @@ exit:
 
 /* BEGIN_CASE */
 void import_export_public_key( data_t *data,
-                               int type_arg,
+                               int type_arg, // key pair or public key
                                int alg_arg,
                                int lifetime_arg,
                                int export_size_delta,