1#!/usr/bin/python
2
3#
4# Copyright (C) 2014 The Android Open Source Project
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10#      http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17#
18
19"""Unit tests for the TPM 2.0 code generator."""
20
21from __future__ import print_function
22
23import StringIO
24import unittest
25
26import generator
27
28
29class TestGenerators(unittest.TestCase):
30  """Test generator classes."""
31
32  def testTypedef(self):
33    """Test generation of typedefs and dependencies."""
34    typedef = generator.Typedef('int', 'INT')
35    defined_types = set(['int'])
36    typemap = {}
37    out_file = StringIO.StringIO()
38    # Expect this to just write the typedef.
39    typedef.OutputForward(out_file, defined_types, typemap)
40    # Expect this to know it has already been written.
41    typedef.Output(out_file, defined_types, typemap)
42    self.assertEqual(out_file.getvalue(), 'typedef int INT;\n')
43    self.assertIn('INT', defined_types)
44    typedef2 = generator.Typedef('TYPE1', 'TYPE2')
45    typemap = {'TYPE1': generator.Structure('TYPE1', False)}
46    defined_types = set([])
47    out_file2 = StringIO.StringIO()
48    # Expect this to write first TYPE1 forward then TYPE2 typedef.
49    typedef2.Output(out_file2, defined_types, typemap)
50    output_re = r'struct TYPE1;\s+typedef TYPE1 TYPE2;\s+'
51    self.assertRegexpMatches(out_file2.getvalue(), output_re)
52    self.assertIn('TYPE2', defined_types)
53    out_file.close()
54    out_file2.close()
55
56  def testTypedefSerialize(self):
57    """Test generation of serialization code for typedefs."""
58    serialized_types = set(['int'])
59    typedef = generator.Typedef('int', 'INT')
60    typedef2 = generator.Typedef('INT', 'INT2')
61    typemap = {'INT': typedef}
62    out_file = StringIO.StringIO()
63    typedef2.OutputSerialize(out_file, serialized_types, typemap)
64    self.assertIn('INT', serialized_types)
65    self.assertIn('INT2', serialized_types)
66    out_file.close()
67
68  def testConstant(self):
69    """Test generation of constant definitions and type dependencies."""
70    constant = generator.Constant('INT', 'test', '1')
71    typemap = {'INT': generator.Structure('INT', False)}
72    defined_types = set([])
73    out_file = StringIO.StringIO()
74    constant.Output(out_file, defined_types, typemap)
75    output_re = r'struct INT;\s+const INT test = 1;\s+'
76    self.assertRegexpMatches(out_file.getvalue(), output_re)
77    out_file.close()
78
79  def testStructure(self):
80    """Test generation of structure declarations and field dependencies."""
81    struct = generator.Structure('STRUCT', False)
82    struct.AddField('int', 'i')
83    struct.AddDependency('DEPEND')
84    union = generator.Structure('UNION', True)
85    union.AddField('STRUCT', 'inner')
86    depend = generator.Structure('DEPEND', False)
87    defined_types = set(['int'])
88    out_file = StringIO.StringIO()
89    typemap = {'STRUCT': struct, 'DEPEND': depend}
90    # Only output |union|, this will test the dependency logic.
91    union.OutputForward(out_file, defined_types, typemap)
92    union.OutputForward(out_file, defined_types, typemap)
93    union.Output(out_file, defined_types, typemap)
94    output_re = r'union UNION;\s+struct DEPEND {\s+};\s+'
95    output_re += r'struct STRUCT {\s+int i;\s+};\s+'
96    output_re += r'union UNION {\s+STRUCT inner;\s+};\s+'
97    self.assertRegexpMatches(out_file.getvalue(), output_re)
98    for t in ('STRUCT', 'DEPEND', 'UNION'):
99      self.assertIn(t, defined_types)
100    # Test serialize / parse code generation.
101    out_file.close()
102
103  def testStructSerialize(self):
104    """Test generation of serialization code for typedefs."""
105    serialized_types = set(['int', 'FOO', 'BAR', 'TPMI_ALG_SYM_OBJECT'])
106    struct = generator.Structure('TEST_STRUCT', False)
107    struct.fields = [('TPMI_ALG_SYM_OBJECT', 'selector'),
108                     ('TPMU_SYM_MODE', 'mode'),
109                     ('int', 'sizeOfFoo'),
110                     ('int', 'foo[FOO_MAX]')]
111    # Choose TPMU_SYM_MODE because it exists in the selectors definition and it
112    # has few fields.
113    union = generator.Structure('TPMU_SYM_MODE', True)
114    union.fields = [('FOO', 'aes'), ('BAR', 'sm4')]
115    typemap = {'TPMU_SYM_MODE': union}
116    out_file = StringIO.StringIO()
117    struct.OutputSerialize(out_file, serialized_types, typemap)
118    self.assertIn('TPMU_SYM_MODE', serialized_types)
119    self.assertIn('TEST_STRUCT', serialized_types)
120    out_file.close()
121
122  def testDefine(self):
123    """Test generation of preprocessor defines."""
124    define = generator.Define('name', 'value')
125    out_file = StringIO.StringIO()
126    define.Output(out_file)
127    output_re = r'#if !defined\(name\)\s+#define name value\s+#endif\s+'
128    self.assertRegexpMatches(out_file.getvalue(), output_re)
129    out_file.close()
130
131  def _MakeArg(self, arg_type, arg_name):
132    return {'type': arg_type,
133            'name': arg_name,
134            'command_code': None,
135            'description': None}
136
137  def testCommand(self):
138    """Test generation of command methods and callbacks."""
139    command = generator.Command('TPM2_Test')
140    command.request_args = [self._MakeArg('int', 'input')]
141    command.response_args = [self._MakeArg('char', 'output')]
142    out_file = StringIO.StringIO()
143    command.OutputDeclarations(out_file)
144    expected_callback = """typedef base::Callback<void(
145      TPM_RC response_code,
146      const char& output)> TestResponse;"""
147    self.assertIn(expected_callback, out_file.getvalue())
148    expected_serialize = """static TPM_RC SerializeCommand_Test(
149      const int& input,
150      std::string* serialized_command,
151      AuthorizationDelegate* authorization_delegate);"""
152    self.assertIn(expected_serialize, out_file.getvalue())
153    expected_parse = """static TPM_RC ParseResponse_Test(
154      const std::string& response,
155      char* output,
156      AuthorizationDelegate* authorization_delegate);"""
157    self.assertIn(expected_parse, out_file.getvalue())
158    expected_async = """virtual void Test(
159      const int& input,
160      AuthorizationDelegate* authorization_delegate,
161      const TestResponse& callback);"""
162    self.assertIn(expected_async, out_file.getvalue())
163    expected_sync = """virtual TPM_RC TestSync(
164      const int& input,
165      char* output,
166      AuthorizationDelegate* authorization_delegate);"""
167    self.assertIn(expected_sync, out_file.getvalue())
168    out_file.close()
169
170
171class TestParsers(unittest.TestCase):
172  """Test parser classes."""
173
174  FAKE_TYPEDEF = '_BEGIN_TYPES\n_OLD_TYPE type1\n_NEW_TYPE type2\n_END\n'
175  FAKE_CONSTANT = ('_BEGIN_CONSTANTS\n_CONSTANTS (base_type) const_type\n'
176                   '_TYPE base_type\n_NAME const_name\n_VALUE const_value\n'
177                   '_END\n')
178  FAKE_STRUCTURE = ('_BEGIN_STRUCTURES\n_STRUCTURE struct_type\n'
179                    '_TYPE field_type\n'
180                    '_NAME field_name[sizeof(depend_type)]\n_END\n')
181  FAKE_DEFINE = '_BEGIN_DEFINES\n_NAME define_name\n_VALUE define_value\n_END'
182  FAKE_COMMAND = ('_BEGIN\n_INPUT_START TPM2_Test\n'
183                  '_TYPE UINT32\n_NAME commandSize\n'
184                  '_TYPE TPM_CC\n_NAME commandCode\n_COMMENT TPM_CC_Test\n'
185                  '_TYPE UINT16\n_NAME input\n'
186                  '_OUTPUT_START TPM2_Test\n_END\n')
187
188  def testStructureParserWithBadData(self):
189    """Test the structure parser with invalid data."""
190    input_data = 'bad_data'
191    in_file = StringIO.StringIO(input_data)
192    parser = generator.StructureParser(in_file)
193    types, constants, structs, defines, typemap = parser.Parse()
194    self.assertIsNotNone(types)
195    self.assertIsNotNone(constants)
196    self.assertIsNotNone(structs)
197    self.assertIsNotNone(defines)
198    self.assertIsNotNone(typemap)
199
200  def testStructureParser(self):
201    """Test the structure parser with valid data."""
202    input_data = (self.FAKE_TYPEDEF + self.FAKE_CONSTANT + self.FAKE_STRUCTURE +
203                  self.FAKE_DEFINE)
204    in_file = StringIO.StringIO(input_data)
205    parser = generator.StructureParser(in_file)
206    types, constants, structs, defines, typemap = parser.Parse()
207    # Be flexible on these counts because the parser may add special cases.
208    self.assertGreaterEqual(len(types), 2)
209    self.assertGreaterEqual(len(constants), 1)
210    self.assertGreaterEqual(len(structs), 1)
211    self.assertGreaterEqual(len(defines), 1)
212    self.assertGreaterEqual(len(typemap), 3)
213    self.assertEqual(types[0].old_type, 'type1')
214    self.assertEqual(types[0].new_type, 'type2')
215    self.assertEqual(types[1].old_type, 'base_type')
216    self.assertEqual(types[1].new_type, 'const_type')
217    self.assertEqual(constants[0].const_type, 'const_type')
218    self.assertEqual(constants[0].name, 'const_name')
219    self.assertEqual(constants[0].value, 'const_value')
220    self.assertEqual(structs[0].name, 'struct_type')
221    self.assertEqual(structs[0].is_union, False)
222    self.assertEqual(len(structs[0].fields), 1)
223    self.assertEqual(structs[0].fields[0][0], 'field_type')
224    self.assertEqual(structs[0].fields[0][1], 'field_name[sizeof(depend_type)]')
225    self.assertEqual(len(structs[0].depends_on), 1)
226    self.assertEqual(structs[0].depends_on[0], 'depend_type')
227    self.assertEqual(defines[0].name, 'define_name')
228    self.assertEqual(defines[0].value, 'define_value')
229
230  def testCommandParserWithBadData(self):
231    """Test the command parser with invalid data."""
232    input_data = 'bad_data'
233    in_file = StringIO.StringIO(input_data)
234    parser = generator.CommandParser(in_file)
235    commands = parser.Parse()
236    self.assertIsNotNone(commands)
237
238  def testCommandParser(self):
239    """Test the command parser with valid data."""
240    input_data = self.FAKE_COMMAND
241    in_file = StringIO.StringIO(input_data)
242    parser = generator.CommandParser(in_file)
243    commands = parser.Parse()
244    self.assertEqual(len(commands), 1)
245    self.assertEqual(commands[0].name, 'TPM2_Test')
246    self.assertEqual(commands[0].command_code, 'TPM_CC_Test')
247    # We expect the 'commandSize' and 'commandCode' args to be filtered out.
248    self.assertEqual(len(commands[0].request_args), 1)
249    self.assertEqual(commands[0].request_args[0]['type'], 'UINT16')
250    self.assertEqual(commands[0].request_args[0]['name'], 'input')
251    self.assertIsNotNone(commands[0].response_args)
252    self.assertFalse(commands[0].response_args)
253
254
255if __name__ == '__main__':
256  unittest.main()
257