1# Copyright 2020 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import logging
6import os
7
8from autotest_lib.client.common_lib import error
9from autotest_lib.server.cros.faft.firmware_test import FirmwareTest
10from autotest_lib.server.cros.faft.firmware_test import ConnectionError
11
12
13BIOS = 'bios'
14EC = 'ec'
15
16
17class firmware_WriteProtectFunc(FirmwareTest):
18    """
19    This test checks whether the SPI flash write-protection functionally works
20    """
21    version = 1
22
23    def initialize(self, host, cmdline_args, dev_mode=False):
24        """Initialize the test"""
25        super(firmware_WriteProtectFunc, self).initialize(host, cmdline_args)
26        self.switcher.setup_mode('dev' if dev_mode else 'normal',
27                                 allow_gbb_force=True)
28        if self.faft_config.chrome_ec:
29            self._targets = (BIOS, EC)
30        else:
31            self._targets = (BIOS, )
32        self._rpcs = {BIOS: self.faft_client.bios,
33                EC: self.faft_client.ec}
34        self._flashrom_targets = {BIOS: 'host', EC: 'ec'}
35        self._original_sw_wps = {}
36        for target in self._targets:
37            sw_wp_dict = self._rpcs[target].get_write_protect_status()
38            logging.debug("self._rpcs[%s].get_write_protect_status() = %s",
39                          target, sw_wp_dict)
40            self._original_sw_wps[target] = sw_wp_dict['enabled']
41        self._original_hw_wp = 'on' in self.servo.get('fw_wp_state')
42        self.backup_firmware()
43        self.work_path = self.faft_client.system.create_temp_dir(
44                'flashrom_', '/mnt/stateful_partition/')
45
46    def cleanup(self):
47        """Cleanup the test"""
48        try:
49            if self.is_firmware_saved():
50                self.restore_firmware()
51        except ConnectionError:
52            logging.error("ERROR: DUT did not come up after firmware restore!")
53
54        try:
55            # Recover SW WP status.
56            if hasattr(self, '_original_sw_wps'):
57                # If HW WP is enabled, we have to disable it first so that
58                # SW WP can be changed.
59                current_hw_wp = 'on' in self.servo.get('fw_wp_state')
60                if current_hw_wp:
61                    self.set_ap_write_protect_and_reboot(False)
62                for target, original_sw_wp in self._original_sw_wps.items():
63                    self._set_write_protect(target, original_sw_wp)
64                self.set_ap_write_protect_and_reboot(current_hw_wp)
65            # Recover HW WP status.
66            if hasattr(self, '_original_hw_wp'):
67                self.set_ap_write_protect_and_reboot(self._original_hw_wp)
68        except Exception as e:
69            logging.error('Caught exception: %s', str(e))
70
71        self.faft_client.system.remove_dir(self.work_path)
72        super(firmware_WriteProtectFunc, self).cleanup()
73
74    def _set_write_protect(self, target, enable):
75        """
76        Set write_protect to `enable` for the specified target.
77
78        @param target: Which firmware to toggle the write-protect for,
79                       either 'bios' or 'ec'
80        @type target: string
81        @param enable: Whether to enable or disable write-protect
82        @type enable: bool
83        """
84        assert target in (BIOS, EC)
85        if target == BIOS:
86            # Unlock registers to alter the region/range
87            self.set_ap_write_protect_and_reboot(False)
88            self.faft_client.bios.set_write_protect_region('WP_RO', enable)
89            if enable:
90                self.set_ap_write_protect_and_reboot(True)
91        elif target == EC:
92            self.switcher.mode_aware_reboot('custom',
93                    lambda:self.set_ec_write_protect_and_reboot(enable))
94
95    def _get_relative_path(self, target):
96        """
97        Send an RPC.updater call to get the relative path for the target.
98
99        @param target: Which firmware to get the relative path to,
100                       either 'bios' or 'ec'.
101        @type target: string
102        @return: The relative path of the bios/ec image in the shellball.
103        """
104        assert target in (BIOS, EC)
105        if target == BIOS:
106            return self.faft_client.updater.get_bios_relative_path()
107        elif target == EC:
108            return self.faft_client.updater.get_ec_relative_path()
109
110    def run_cmd(self, command, checkfor=''):
111        """
112        Log and execute command and return the output.
113
114        @param command: Command to execute on device.
115        @param checkfor: If not empty, make the test fail when this param
116            is not found in the command output.
117        @returns the output of command.
118        """
119        command = command + ' 2>&1'
120        logging.info('Execute %s', command)
121        output = self.faft_client.system.run_shell_command_get_output(command)
122        logging.info('Output >>> %s <<<', output)
123        if checkfor and checkfor not in '\n'.join(output):
124            raise error.TestFail('Expect %s in output of cmd <%s>:\n\t%s' %
125                                 (checkfor, command, '\n\t'.join(output)))
126        return output
127
128    def get_wp_ro_firmware_section(self, firmware_file, wp_ro_firmware_file):
129        """
130        Read out WP_RO section from the firmware file.
131
132        @param firmware_file: The AP or EC firmware binary to be parsed.
133        @param wp_ro_firmware_file: The file path for the WP_RO section
134            dumped from the firmware_file.
135        @returns the output of the dd command.
136        """
137        cmd_output = self.run_cmd(
138                'futility dump_fmap -p %s WP_RO'% firmware_file)
139        if cmd_output:
140            unused_name, offset, size = cmd_output[0].split()
141
142        return self.run_cmd('dd bs=1 skip=%s count=%s if=%s of=%s' %
143                            (offset, size, firmware_file, wp_ro_firmware_file))
144
145    def run_once(self):
146        """Runs a single iteration of the test."""
147        # Enable WP
148        for target in self._targets:
149            self._set_write_protect(target, True)
150
151        # Check WP is properly enabled at the start
152        for target in self._targets:
153            sw_wp_dict = self._rpcs[target].get_write_protect_status()
154            logging.debug("self._rpcs[%s].get_write_protect_status() = %s",
155                          target, sw_wp_dict)
156            if not sw_wp_dict['enabled']:
157                raise error.TestFail('Failed to enable %s SW WP at '
158                                     'test start' % target.upper())
159
160        reboots = (('shutdown cmd', lambda:self.run_shutdown_process(
161                                        lambda:self.run_shutdown_cmd())),
162                   ('reboot cmd', lambda:self.run_cmd('reboot')),
163                   ('power button', lambda:self.full_power_off_and_on()))
164
165        if self.faft_config.chrome_ec:
166            reboots += (('ec reboot', lambda:self.sync_and_ec_reboot('hard')), )
167
168        # Check if enabled SW WP can stay preserved across reboots.
169        for (reboot_name, reboot_method) in reboots:
170            self.switcher.mode_aware_reboot('custom', reboot_method)
171            for target in self._targets:
172                sw_wp_dict = self._rpcs[target].get_write_protect_status()
173                if not sw_wp_dict['enabled']:
174                    raise error.TestFail('%s SW WP can not stay preserved '
175                                         'accross %s' %
176                                         (target.upper(), reboot_name))
177
178        work_path = self.work_path
179        # Check if RO FW really can't be overwritten when WP is enabled.
180        for target in self._targets:
181            # Current firmware image as read from flash
182            ro_before = os.path.join(work_path, '%s_ro_before.bin' % target)
183            # Current firmware image with modification to test writing
184            ro_test = os.path.join(work_path, '%s_ro_test.bin' % target)
185            # Firmware as read after writing flash
186            ro_after = os.path.join(work_path, '%s_ro_after.bin' % target)
187
188            # Fetch firmware from flash. This serves as the base of ro_test
189            self.run_cmd(
190                    'flashrom -p %s -r -i WP_RO:%s ' %
191                    (self._flashrom_targets[target], ro_before), 'SUCCESS')
192
193            lines = self.run_cmd('dump_fmap -p %s' % ro_before)
194            FMAP_AREA_NAMES = ['name', 'offset', 'size']
195
196            modified = False
197            wpro_offset = -1
198            for line in lines:
199                region = dict(zip(FMAP_AREA_NAMES, line.split()))
200                if region['name'] == 'WP_RO':
201                    wpro_offset = int(region['offset'])
202            if wpro_offset == -1:
203                raise error.TestFail('WP_RO not found in fmap')
204            for line in lines:
205                region = dict(zip(FMAP_AREA_NAMES, line.split()))
206                if region['name'] == 'RO_FRID':
207                    modified = True
208                    self.run_cmd('cp %s %s' % (ro_before, ro_test))
209                    self.run_cmd(
210                            'dd if=%s bs=1 count=%d skip=%d '
211                            '| tr "[a-zA-Z]" "[A-Za-z]" '
212                            '| dd of=%s bs=1 count=%d seek=%d conv=notrunc' %
213                            (ro_test, int(region['size']),
214                             int(region['offset']) - wpro_offset, ro_test,
215                             int(region['size']),
216                             int(region['offset']) - wpro_offset))
217
218            if not modified:
219                raise error.TestFail('Could not find RO_FRID in %s' %
220                                     target.upper())
221
222            # Writing WP_RO section is expected to fail.
223            self.run_cmd('flashrom -p %s -w -i WP_RO:%s' %
224                    (self._flashrom_targets[target], ro_test),
225                    'FAIL')
226            self.run_cmd('flashrom -p %s -r -i WP_RO:%s' %
227                    (self._flashrom_targets[target], ro_after),
228                    'SUCCESS')
229
230            self.switcher.mode_aware_reboot(reboot_type='cold')
231
232            # The WP_RO section on the DUT should not change.
233            cmp_output = self.run_cmd('cmp %s %s' % (ro_before, ro_after))
234            if ''.join(cmp_output) != '':
235                raise error.TestFail('%s RO changes when WP is on!' %
236                        target.upper())
237
238        # Disable WP
239        for target in self._targets:
240            self._set_write_protect(target, False)
241
242        # Check if RO FW can be overwritten when WP is disabled.
243        for target in self._targets:
244            ro_after = os.path.join(work_path, '%s_ro_after.bin' % target)
245            ro_test = os.path.join(work_path, '%s_ro_test.bin' % target)
246
247            # Writing WP_RO section is expected to succeed.
248            self.run_cmd('flashrom -p %s -w -i WP_RO:%s' %
249                    (self._flashrom_targets[target], ro_test),
250                    'SUCCESS')
251            self.run_cmd('flashrom -p %s -r -i WP_RO:%s' %
252                    (self._flashrom_targets[target], ro_after),
253                    'SUCCESS')
254
255            # The DUT's WP_RO section should be the same as the test firmware.
256            cmp_output = self.run_cmd('cmp %s %s' % (ro_test, ro_after))
257            if ''.join(cmp_output) != '':
258                raise error.TestFail('%s RO is not flashed correctly'
259                                     'when WP is off!' % target.upper())
260