1# Copyright 2015 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
7import re
8
9import common
10from autotest_lib.client.common_lib import error
11from autotest_lib.server import test
12
13
14# The /dev directory mapping partition names to block devices.
15_BLK_DEV_BY_NAME_DIR = '/dev/block/by-name'
16# By default, we kill and recover the active system partition.
17_DEFAULT_PART_NAME = 'system_X'
18
19
20class brillo_RecoverFromBadImage(test.test):
21    """Ensures that a Brillo device can recover from a bad image."""
22    version = 1
23
24
25    def resolve_slot(self, host, partition):
26        """Resolves a partition slot (if any).
27
28        @param host: A host object representing the DUT.
29        @param partition: The name of the partition we are using. If it ends
30                          with '_X' then we attempt to substitute it with some
31                          non-active slot.
32
33        @return A pair consisting of a fully resolved partition name and slot
34                index; the latter is None if the partition is not slotted.
35
36        @raise TestError: If a target slot could not be resolved.
37        """
38        # Check if the partition is slotted.
39        if not re.match('.+_[a-zX]$', partition):
40            return partition, None
41
42        try:
43            current_slot = int(
44                    host.run_output('bootctl get-current-slot').strip())
45            if partition[-1] == 'X':
46                # Find a non-active target slot we could use.
47                num_slots = int(
48                        host.run_output('bootctl get-number-slots').strip())
49                if num_slots < 2:
50                    raise error.TestError(
51                            'Device has no non-active slot that we can use')
52                target_slot = 0 if current_slot else 1
53                partition = partition[:-1] + chr(ord('a') + target_slot)
54                logging.info(
55                        'Current slot is %d, partition resolved to %s '
56                        '(slot %d)', current_slot, partition, target_slot)
57            else:
58                # Make sure the partition slot is different from the active one.
59                target_slot = ord(partition[-1]) - ord('a')
60                if target_slot == current_slot:
61                    target_slot = None
62                    logging.warning(
63                            'Partition %s is associated with the current boot '
64                            'slot (%d), wiping it might fail if it is mounted',
65                            partition, current_slot)
66        except error.AutoservError:
67            raise error.TestError('Error resolving device slots')
68
69        return partition, target_slot
70
71
72    def find_partition_device(self, host, partition):
73        """Returns the block device of the partition.
74
75        @param host: A host object representing the DUT.
76        @param partition: The name of the partition we are using.
77
78        @return Path to the device containing the partition.
79
80        @raise TestError: If the partition name could not be mapped to a device.
81        """
82        try:
83            cmd = 'find %s -type l' % os.path.join(_BLK_DEV_BY_NAME_DIR, '')
84            for device in host.run_output(cmd).splitlines():
85                if os.path.basename(device) == partition:
86                    logging.info('Mapped partition %s to device %s',
87                                 partition, device)
88                    return device
89        except error.AutoservError:
90            raise error.TestError(
91                    'Error finding device for partition %s' % partition)
92        raise error.TestError(
93                'No device found for partition %s' % partition)
94
95
96    def get_device_block_info(self, host, device):
97        """Returns the block size and count for a device.
98
99        @param host: A host object representing the DUT.
100        @param device: Path to a block device.
101
102        @return A pair consisting of the block size (in bytes) and the total
103                number of blocks on the device.
104
105        @raise TestError: If we failed to get the block info for the device.
106        """
107        try:
108            block_size = int(
109                    host.run_output('blockdev --getbsz %s' % device).strip())
110            device_size = int(
111                    host.run_output('blockdev --getsize64 %s' % device).strip())
112        except error.AutoservError:
113            raise error.TestError(
114                    'Failed to get block info for device %s' % device)
115        return block_size, device_size / block_size
116
117
118    def run_once(self, host=None, image_file=None, partition=_DEFAULT_PART_NAME,
119                 device=None):
120        """Runs the test.
121
122        @param host: A host object representing the DUT.
123        @param image_file: Image file to flash to the partition.
124        @param partition: Name of the partition to wipe/recover.
125        @param device: Path to the partition block device.
126
127        @raise TestError: Something went wrong while trying to execute the test.
128        @raise TestFail: The test failed.
129        """
130        # Check that the image file exists.
131        if image_file is None:
132            raise error.TestError('No image file provided')
133        if not os.path.isfile(image_file):
134            raise error.TestError('Image file %s not found' % image_file)
135
136        try:
137            # Resolve partition name and slot.
138            partition, target_slot = self.resolve_slot(host, partition)
139
140            # Figure out the partition device.
141            if device is None:
142                device = self.find_partition_device(host, partition)
143
144            # Find the block size and count for the device.
145            block_size, num_blocks = self.get_device_block_info(host, device)
146
147            # Wipe the partition.
148            logging.info('Wiping partition %s (%s)', partition, device)
149            cmd = ('dd if=/dev/zero of=%s bs=%d count=%d' %
150                   (device, block_size, num_blocks))
151            run_err = 'Failed to wipe partition using %s' % cmd
152            host.run(cmd)
153
154            # Switch to the target slot, if required.
155            if target_slot is not None:
156                run_err = 'Error setting the active boot slot'
157                host.run('bootctl set-active-boot-slot %d' % target_slot)
158
159            # Re-flash the partition with fastboot.
160            run_err = 'Failed to reboot the device into fastboot'
161            host.ensure_bootloader_mode()
162            run_err = 'Failed to flash image to partition %s' % partition
163            host.fastboot_run('flash', args=(partition, image_file))
164
165            # Reboot the device.
166            run_err = 'Failed to reboot the device after flashing image'
167            host.ensure_adb_mode()
168
169            # Make sure we've booted from the alternate slot, if required.
170            if target_slot is not None:
171                run_err = 'Error checking the current boot slot'
172                current_slot = int(
173                        host.run_output('bootctl get-current-slot').strip())
174                if current_slot != target_slot:
175                    logging.error('Rebooted from slot %d instead of %d',
176                                  current_slot, target_slot)
177                    raise error.TestError(
178                            'Device did not reboot from the expected slot')
179        except error.AutoservError:
180            raise error.TestFail(run_err)
181