blob: a9867f5deb5bb4660d06980f2a736780d32ccab6 [file] [log] [blame]
Amin Hassani92f6c4a2021-02-20 17:36:09 -08001# Copyright 2021 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
5"""Library containing functions to install an image on a Chromium OS device."""
6
Amin Hassanid4b3ff82021-02-20 23:05:14 -08007import abc
Amin Hassani92f6c4a2021-02-20 17:36:09 -08008import enum
Chris McDonald14ac61d2021-07-21 11:49:56 -06009import logging
Amin Hassani92f6c4a2021-02-20 17:36:09 -080010import os
11import re
Amin Hassanid4b3ff82021-02-20 23:05:14 -080012import tempfile
13import threading
Amin Hassani55970562021-02-22 20:49:13 -080014import time
Chris McDonald14ac61d2021-07-21 11:49:56 -060015from typing import Dict, Tuple
Amin Hassani92f6c4a2021-02-20 17:36:09 -080016
Amin Hassani55970562021-02-22 20:49:13 -080017from chromite.cli import command
Amin Hassanicf8f0042021-03-12 10:42:13 -080018from chromite.cli import flash
Jae Hoon Kimcc723e02021-08-16 21:03:21 +000019from chromite.lib import cgpt
Amin Hassanid4b3ff82021-02-20 23:05:14 -080020from chromite.lib import constants
Amin Hassani92f6c4a2021-02-20 17:36:09 -080021from chromite.lib import cros_build_lib
Amin Hassani92f6c4a2021-02-20 17:36:09 -080022from chromite.lib import gs
Amin Hassanid4b3ff82021-02-20 23:05:14 -080023from chromite.lib import image_lib
Amin Hassani55970562021-02-22 20:49:13 -080024from chromite.lib import operation
Amin Hassanid4b3ff82021-02-20 23:05:14 -080025from chromite.lib import osutils
Amin Hassani92f6c4a2021-02-20 17:36:09 -080026from chromite.lib import parallel
27from chromite.lib import remote_access
28from chromite.lib import retry_util
Amin Hassani74403082021-02-22 11:40:09 -080029from chromite.lib import stateful_updater
Amin Hassani75c5f942021-02-20 23:56:53 -080030from chromite.lib.paygen import partition_lib
Amin Hassani74403082021-02-22 11:40:09 -080031from chromite.lib.paygen import paygen_stateful_payload_lib
Amin Hassani92f6c4a2021-02-20 17:36:09 -080032from chromite.lib.xbuddy import devserver_constants
33from chromite.lib.xbuddy import xbuddy
34
35
Amin Hassani92f6c4a2021-02-20 17:36:09 -080036class Error(Exception):
37 """Thrown when there is a general Chromium OS-specific flash error."""
38
39
40class ImageType(enum.Enum):
41 """Type of the image that is used for flashing the device."""
42
43 # The full image on disk (e.g. chromiumos_test_image.bin).
44 FULL = 0
45 # The remote directory path
46 # (e.g gs://chromeos-image-archive/eve-release/R90-x.x.x)
47 REMOTE_DIRECTORY = 1
48
49
50class Partition(enum.Enum):
51 """An enum for partition types like kernel and rootfs."""
52 KERNEL = 0
53 ROOTFS = 1
Jae Hoon Kimcc723e02021-08-16 21:03:21 +000054 MINIOS = 2
Amin Hassani92f6c4a2021-02-20 17:36:09 -080055
56
57class DeviceImager(object):
58 """A class to flash a Chromium OS device.
59
60 This utility uses parallelism as much as possible to achieve its goal as fast
61 as possible. For example, it uses parallel compressors, parallel transfers,
62 and simultaneous pipes.
63 """
64
65 # The parameters of the kernel and rootfs's two main partitions.
66 A = {Partition.KERNEL: 2, Partition.ROOTFS: 3}
67 B = {Partition.KERNEL: 4, Partition.ROOTFS: 5}
68
Jae Hoon Kimcc723e02021-08-16 21:03:21 +000069 MINIOS_A = {Partition.MINIOS: 9}
70 MINIOS_B = {Partition.MINIOS: 10}
71
Amin Hassani92f6c4a2021-02-20 17:36:09 -080072 def __init__(self, device, image: str,
Amin Hassanicf8f0042021-03-12 10:42:13 -080073 board: str = None,
74 version: str = None,
Amin Hassani92f6c4a2021-02-20 17:36:09 -080075 no_rootfs_update: bool = False,
76 no_stateful_update: bool = False,
Jae Hoon Kimcc723e02021-08-16 21:03:21 +000077 no_minios_update: bool = False,
Amin Hassani92f6c4a2021-02-20 17:36:09 -080078 no_reboot: bool = False,
79 disable_verification: bool = False,
80 clobber_stateful: bool = False,
81 clear_tpm_owner: bool = False):
82 """Initialize DeviceImager for flashing a Chromium OS device.
83
84 Args:
85 device: The ChromiumOSDevice to be updated.
86 image: The target image path (can be xBuddy path).
Amin Hassanicf8f0042021-03-12 10:42:13 -080087 board: Board to use.
88 version: Image version to use.
Amin Hassani92f6c4a2021-02-20 17:36:09 -080089 no_rootfs_update: Whether to do rootfs partition update.
90 no_stateful_update: Whether to do stateful partition update.
Jae Hoon Kimcc723e02021-08-16 21:03:21 +000091 no_minios_update: Whether to do minios partition update.
Amin Hassani92f6c4a2021-02-20 17:36:09 -080092 no_reboot: Whether to reboot device after update. The default is True.
93 disable_verification: Whether to disabling rootfs verification on the
94 device.
95 clobber_stateful: Whether to do a clean stateful partition.
96 clear_tpm_owner: If true, it will clear the TPM owner on reboot.
97 """
98
99 self._device = device
100 self._image = image
Amin Hassanicf8f0042021-03-12 10:42:13 -0800101 self._board = board
102 self._version = version
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800103 self._no_rootfs_update = no_rootfs_update
104 self._no_stateful_update = no_stateful_update
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000105 self._no_minios_update = no_minios_update
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800106 self._no_reboot = no_reboot
107 self._disable_verification = disable_verification
108 self._clobber_stateful = clobber_stateful
109 self._clear_tpm_owner = clear_tpm_owner
110
Amin Hassanib1993eb2021-04-28 12:00:11 -0700111 self._image_type = None
Amin Hassani0972d392021-03-31 19:04:19 -0700112 self._compression = cros_build_lib.COMP_GZIP
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800113 self._inactive_state = None
114
115 def Run(self):
116 """Update the device with image of specific version."""
Amin Hassanib1993eb2021-04-28 12:00:11 -0700117 self._LocateImage()
118 logging.notice('Preparing to update the remote device %s with image %s',
119 self._device.hostname, self._image)
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800120
121 try:
Amin Hassani55970562021-02-22 20:49:13 -0800122 if command.UseProgressBar():
123 op = DeviceImagerOperation()
124 op.Run(self._Run)
125 else:
126 self._Run()
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800127 except Exception as e:
128 raise Error(f'DeviceImager Failed with error: {e}')
129
Amin Hassani55970562021-02-22 20:49:13 -0800130 # DeviceImagerOperation will look for this log.
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800131 logging.info('DeviceImager completed.')
132
133 def _Run(self):
134 """Runs the various operations to install the image on device."""
Amin Hassanib1993eb2021-04-28 12:00:11 -0700135 # Override the compression as remote quick provision images are gzip
136 # compressed only.
137 if self._image_type == ImageType.REMOTE_DIRECTORY:
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800138 self._compression = cros_build_lib.COMP_GZIP
139
Amin Hassanib1993eb2021-04-28 12:00:11 -0700140 self._InstallPartitions()
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800141
142 if self._clear_tpm_owner:
143 self._device.ClearTpmOwner()
144
145 if not self._no_reboot:
146 self._Reboot()
147 self._VerifyBootExpectations()
148
149 if self._disable_verification:
150 self._device.DisableRootfsVerification()
151
Amin Hassanib1993eb2021-04-28 12:00:11 -0700152 def _LocateImage(self):
153 """Locates the path to the final image(s) that need to be installed.
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800154
155 If the paths is local, the image should be the Chromium OS GPT image
156 (e.g. chromiumos_test_image.bin). If the path is remote, it should be the
157 remote directory where we can find the quick-provision and stateful update
158 files (e.g. gs://chromeos-image-archive/eve-release/R90-x.x.x).
159
160 NOTE: At this point there is no caching involved. Hence we always download
161 the partition payloads or extract them from the Chromium OS image.
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800162 """
163 if os.path.isfile(self._image):
Amin Hassanib1993eb2021-04-28 12:00:11 -0700164 self._image_type = ImageType.FULL
165 return
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800166
167 # TODO(b/172212406): We could potentially also allow this by searching
168 # through the directory to see whether we have quick-provision and stateful
169 # payloads. This only makes sense when a user has their workstation at home
170 # and doesn't want to incur the bandwidth cost of downloading the same
171 # image multiple times. For that, they can simply download the GPT image
172 # image first and flash that instead.
173 if os.path.isdir(self._image):
174 raise ValueError(
175 f'{self._image}: input must be a disk image, not a directory.')
176
177 if gs.PathIsGs(self._image):
178 # TODO(b/172212406): Check whether it is a directory. If it wasn't a
179 # directory download the image into some temp location and use it instead.
Amin Hassanib1993eb2021-04-28 12:00:11 -0700180 self._image_type = ImageType.REMOTE_DIRECTORY
181 return
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800182
183 # Assuming it is an xBuddy path.
Amin Hassanicf8f0042021-03-12 10:42:13 -0800184 board = cros_build_lib.GetBoard(
185 device_board=self._device.board or flash.GetDefaultBoard(),
186 override_board=self._board, force=True)
187
Amin Hassania20d20d2021-04-28 10:18:18 -0700188 xb = xbuddy.XBuddy(board=board, version=self._version)
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800189 build_id, local_file = xb.Translate([self._image])
190 if build_id is None:
191 raise Error(f'{self._image}: unable to find matching xBuddy path.')
192 logging.info('XBuddy path translated to build ID %s', build_id)
193
194 if local_file:
Amin Hassanib1993eb2021-04-28 12:00:11 -0700195 self._image = local_file
196 self._image_type = ImageType.FULL
197 return
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800198
Amin Hassanib1993eb2021-04-28 12:00:11 -0700199 self._image = f'{devserver_constants.GS_IMAGE_DIR}/{build_id}'
200 self._image_type = ImageType.REMOTE_DIRECTORY
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800201
202 def _SplitDevPath(self, path: str) -> Tuple[str, int]:
203 """Splits the given /dev/x path into prefix and the dev number.
204
205 Args:
206 path: The path to a block dev device.
207
208 Returns:
209 A tuple of representing the prefix and the index of the dev path.
210 e.g.: '/dev/mmcblk0p1' -> ['/dev/mmcblk0p', 1]
211 """
212 match = re.search(r'(.*)([0-9]+)$', path)
213 if match is None:
214 raise Error(f'{path}: Could not parse root dev path.')
215
216 return match.group(1), int(match.group(2))
217
218 def _GetKernelState(self, root_num: int) -> Tuple[Dict, Dict]:
219 """Returns the kernel state.
220
221 Returns:
222 A tuple of two dictionaries: The current active kernel state and the
223 inactive kernel state. (Look at A and B constants in this class.)
224 """
225 if root_num == self.A[Partition.ROOTFS]:
226 return self.A, self.B
227 elif root_num == self.B[Partition.ROOTFS]:
228 return self.B, self.A
229 else:
230 raise Error(f'Invalid root partition number {root_num}')
231
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000232 def _GetMiniOSState(self, minios_num: int) -> Tuple[Dict, Dict]:
233 """Returns the miniOS state.
234
235 Returns:
236 A tuple of dictionaries: The current active miniOS state and the inactive
237 miniOS state.
238 """
239 if minios_num == self.MINIOS_A[Partition.MINIOS]:
240 return self.MINIOS_A, self.MINIOS_B
241 elif minios_num == self.MINIOS_B[Partition.MINIOS]:
242 return self.MINIOS_B, self.MINIOS_A
243 else:
244 raise Error(f'Invalid minios partition number {minios_num}')
245
Amin Hassanib1993eb2021-04-28 12:00:11 -0700246 def _InstallPartitions(self):
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800247 """The main method that installs the partitions of a Chrome OS device.
248
249 It uses parallelism to install the partitions as fast as possible.
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800250 """
Amin Hassanid684e982021-02-26 11:10:58 -0800251 prefix, root_num = self._SplitDevPath(self._device.root_dev)
252 active_state, self._inactive_state = self._GetKernelState(root_num)
253
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800254 updaters = []
Amin Hassanid684e982021-02-26 11:10:58 -0800255 if not self._no_rootfs_update:
Amin Hassani75c5f942021-02-20 23:56:53 -0800256 current_root = prefix + str(active_state[Partition.ROOTFS])
257 target_root = prefix + str(self._inactive_state[Partition.ROOTFS])
Amin Hassanib1993eb2021-04-28 12:00:11 -0700258 updaters.append(RootfsUpdater(current_root, self._device, self._image,
259 self._image_type, target_root,
260 self._compression))
Amin Hassani75c5f942021-02-20 23:56:53 -0800261
Amin Hassanid684e982021-02-26 11:10:58 -0800262 target_kernel = prefix + str(self._inactive_state[Partition.KERNEL])
Amin Hassanib1993eb2021-04-28 12:00:11 -0700263 updaters.append(KernelUpdater(self._device, self._image, self._image_type,
Amin Hassanid684e982021-02-26 11:10:58 -0800264 target_kernel, self._compression))
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800265
Amin Hassani74403082021-02-22 11:40:09 -0800266 if not self._no_stateful_update:
267 updaters.append(StatefulUpdater(self._clobber_stateful, self._device,
Amin Hassanib1993eb2021-04-28 12:00:11 -0700268 self._image, self._image_type, None,
269 None))
Amin Hassani74403082021-02-22 11:40:09 -0800270
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000271 if not self._no_minios_update:
Jae Hoon Kimb88b7962021-10-18 11:08:38 -0700272 minios_priority = self._device.run(
273 ['crossystem', constants.MINIOS_PRIORITY]).stdout
274 if minios_priority not in ['A', 'B']:
275 logging.warning('Skipping miniOS flash due to missing priority.')
276 else:
277 # Reference disk_layout_v3 for partition numbering.
278 _, inactive_minios_state = self._GetMiniOSState(
279 9 if minios_priority == 'A' else 10)
280 target_minios = prefix + str(inactive_minios_state[Partition.MINIOS])
281 minios_updater = MiniOSUpdater(self._device, self._image,
282 self._image_type, target_minios,
283 self._compression)
284 updaters.append(minios_updater)
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000285
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800286 # Retry the partitions updates that failed, in case a transient error (like
287 # SSH drop, etc) caused the error.
288 num_retries = 1
289 try:
290 retry_util.RetryException(Error, num_retries,
291 parallel.RunParallelSteps,
292 (x.Run for x in updaters if not x.IsFinished()),
293 halt_on_error=True)
294 except Exception:
295 # If one of the partitions failed to be installed, revert all partitions.
296 parallel.RunParallelSteps(x.Revert for x in updaters)
297 raise
298
299 def _Reboot(self):
300 """Reboots the device."""
301 try:
Amin Hassani9281b682021-03-08 16:38:25 -0800302 self._device.Reboot(timeout_sec=300)
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800303 except remote_access.RebootError:
304 raise Error('Could not recover from reboot. Once example reason'
305 ' could be the image provided was a non-test image'
306 ' or the system failed to boot after the update.')
307 except Exception as e:
308 raise Error(f'Failed to reboot to the device with error: {e}')
309
310 def _VerifyBootExpectations(self):
311 """Verify that we fully booted into the expected kernel state."""
312 # Discover the newly active kernel.
313 _, root_num = self._SplitDevPath(self._device.root_dev)
314 active_state, _ = self._GetKernelState(root_num)
315
316 # If this happens, we should rollback.
317 if active_state != self._inactive_state:
318 raise Error('The expected kernel state after update is invalid.')
319
320 logging.info('Verified boot expectations.')
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800321
322
323class ReaderBase(threading.Thread):
324 """The base class for reading different inputs and writing into output.
325
326 This class extends threading.Thread, so it will be run on its own thread. Also
327 it can be used as a context manager. Internally, it opens necessary files for
328 writing to and reading from. This class cannot be instantiated, it needs to be
329 sub-classed first to provide necessary function implementations.
330 """
331
332 def __init__(self, use_named_pipes: bool = False):
333 """Initializes the class.
334
335 Args:
336 use_named_pipes: Whether to use a named pipe or anonymous file
337 descriptors.
338 """
339 super().__init__()
340 self._use_named_pipes = use_named_pipes
341 self._pipe_target = None
342 self._pipe_source = None
343
344 def __del__(self):
345 """Destructor.
346
347 Make sure to clean up any named pipes we might have created.
348 """
349 if self._use_named_pipes:
350 osutils.SafeUnlink(self._pipe_target)
351
352 def __enter__(self):
353 """Enters the context manager"""
354 if self._use_named_pipes:
355 # There is no need for the temp file, we only need its path. So the named
356 # pipe is created after this temp file is deleted.
357 with tempfile.NamedTemporaryFile(prefix='chromite-device-imager') as fp:
358 self._pipe_target = self._pipe_source = fp.name
359 os.mkfifo(self._pipe_target)
360 else:
361 self._pipe_target, self._pipe_source = os.pipe()
362
363 self.start()
364 return self
365
366 def __exit__(self, *args, **kwargs):
367 """Exits the context manager."""
368 self.join()
369
370 def _Source(self):
371 """Returns the source pipe to write data into.
372
373 Sub-classes can use this function to determine where to write their data
374 into.
375 """
376 return self._pipe_source
377
378 def _CloseSource(self):
379 """Closes the source pipe.
380
381 Sub-classes should use this function to close the pipe after they are done
382 writing into it. Failure to do so may result reader of the data to hang
383 indefinitely.
384 """
385 if not self._use_named_pipes:
386 os.close(self._pipe_source)
387
388 def Target(self):
389 """Returns the target pipe to read data from.
390
391 Users of this class can use this path to read data from.
392 """
393 return self._pipe_target
394
395 def CloseTarget(self):
396 """Closes the target pipe.
397
398 Users of this class should use this function to close the pipe after they
399 are done reading from it.
400 """
401 if self._use_named_pipes:
402 os.remove(self._pipe_target)
403 else:
404 os.close(self._pipe_target)
405
406
407class PartialFileReader(ReaderBase):
408 """A class to read specific offset and length from a file and compress it.
409
410 This class can be used to read from specific location and length in a file
411 (e.g. A partition in a GPT image). Then it compresses the input and writes it
412 out (to a pipe). Look at the base class for more information.
413 """
414
415 # The offset of different partitions in a Chromium OS image does not always
416 # align to larger values like 4096. It seems that 512 is the maximum value to
417 # be divisible by partition offsets. This size should not be increased just
418 # for 'performance reasons'. Since we are doing everything in parallel, in
419 # practice there is not much difference between this and larger block sizes as
420 # parallelism hides the possible extra latency provided by smaller block
421 # sizes.
422 _BLOCK_SIZE = 512
423
424 def __init__(self, image: str, offset: int, length: int, compression):
425 """Initializes the class.
426
427 Args:
428 image: The path to an image (local or remote directory).
429 offset: The offset (in bytes) to read from the image.
430 length: The length (in bytes) to read from the image.
431 compression: The compression type (see cros_build_lib.COMP_XXX).
432 """
433 super().__init__()
434
435 self._image = image
436 self._offset = offset
437 self._length = length
438 self._compression = compression
439
440 def run(self):
441 """Runs the reading and compression."""
442 cmd = [
443 'dd',
444 'status=none',
445 f'if={self._image}',
446 f'ibs={self._BLOCK_SIZE}',
447 f'skip={int(self._offset/self._BLOCK_SIZE)}',
448 f'count={int(self._length/self._BLOCK_SIZE)}',
449 '|',
450 cros_build_lib.FindCompressor(self._compression),
451 ]
452
453 try:
454 cros_build_lib.run(' '.join(cmd), stdout=self._Source(), shell=True)
455 finally:
456 self._CloseSource()
457
458
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800459class GsFileCopier(ReaderBase):
460 """A class for downloading gzip compressed file from GS bucket into a pipe."""
461
462 def __init__(self, image: str):
463 """Initializes the class.
464
465 Args:
466 image: The path to an image (local or remote directory).
467 """
468 super().__init__(use_named_pipes=True)
469 self._image = image
470
471 def run(self):
472 """Runs the download and write into the output pipe."""
473 try:
474 gs.GSContext().Copy(self._image, self._Source())
475 finally:
476 self._CloseSource()
477
478
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800479class PartitionUpdaterBase(object):
480 """A base abstract class to use for installing an image into a partition.
481
482 Sub-classes should implement the abstract methods to provide the core
483 functionality.
484 """
485 def __init__(self, device, image: str, image_type, target: str, compression):
486 """Initializes this base class with values that most sub-classes will need.
487
488 Args:
489 device: The ChromiumOSDevice to be updated.
490 image: The target image path for the partition update.
491 image_type: The type of the image (ImageType).
492 target: The target path (e.g. block dev) to install the update.
493 compression: The compression used for compressing the update payload.
494 """
495 self._device = device
496 self._image = image
497 self._image_type = image_type
498 self._target = target
499 self._compression = compression
500 self._finished = False
501
502 def Run(self):
503 """The main function that does the partition update job."""
504 with cros_build_lib.TimedSection() as timer:
505 try:
506 self._Run()
507 finally:
508 self._finished = True
509
510 logging.debug('Completed %s in %s', self.__class__.__name__, timer.delta)
511
512 @abc.abstractmethod
513 def _Run(self):
514 """The method that need to be implemented by sub-classes."""
515 raise NotImplementedError('Sub-classes need to implement this.')
516
517 def IsFinished(self):
518 """Returns whether the partition update has been successful."""
519 return self._finished
520
521 @abc.abstractmethod
522 def Revert(self):
523 """Reverts the partition update.
524
525 Sub-classes need to implement this function to provide revert capability.
526 """
527 raise NotImplementedError('Sub-classes need to implement this.')
528
529
530class RawPartitionUpdater(PartitionUpdaterBase):
531 """A class to update a raw partition on a Chromium OS device."""
532
533 def _Run(self):
534 """The function that does the job of kernel partition update."""
535 if self._image_type == ImageType.FULL:
536 self._CopyPartitionFromImage(self._GetPartitionName())
537 elif self._image_type == ImageType.REMOTE_DIRECTORY:
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800538 self._RedirectPartition(self._GetRemotePartitionName())
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800539 else:
540 raise ValueError(f'Invalid image type {self._image_type}')
541
542 def _GetPartitionName(self):
543 """Returns the name of the partition in a Chromium OS GPT layout.
544
545 Subclasses should override this function to return correct name.
546 """
547 raise NotImplementedError('Subclasses need to implement this.')
548
549 def _CopyPartitionFromImage(self, part_name: str):
550 """Updates the device's partition from a local Chromium OS image.
551
552 Args:
553 part_name: The name of the partition in the source image that needs to be
554 extracted.
555 """
556 cmd = self._GetWriteToTargetCommand()
557
558 offset, length = self._GetPartLocation(part_name)
559 offset, length = self._OptimizePartLocation(offset, length)
560 with PartialFileReader(self._image, offset, length,
561 self._compression) as generator:
562 try:
563 self._device.run(cmd, input=generator.Target(), shell=True)
564 finally:
565 generator.CloseTarget()
566
567 def _GetWriteToTargetCommand(self):
568 """Returns a write to target command to run on a Chromium OS device.
569
570 Returns:
571 A string command to run on a device to read data from stdin, uncompress it
572 and write it to the target partition.
573 """
574 cmd = self._device.GetDecompressor(self._compression)
575 # Using oflag=direct to tell the OS not to cache the writes (faster).
576 cmd += ['|', 'dd', 'bs=1M', 'oflag=direct', f'of={self._target}']
577 return ' '.join(cmd)
578
579 def _GetPartLocation(self, part_name: str):
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000580 """Extracts the location and size of the raw partition from the image.
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800581
582 Args:
583 part_name: The name of the partition in the source image that needs to be
584 extracted.
585
586 Returns:
587 A tuple of offset and length (in bytes) from the image.
588 """
589 try:
590 parts = image_lib.GetImageDiskPartitionInfo(self._image)
591 part_info = [p for p in parts if p.name == part_name][0]
592 except IndexError:
593 raise Error(f'No partition named {part_name} found.')
594
595 return int(part_info.start), int(part_info.size)
596
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800597 def _GetRemotePartitionName(self):
598 """Returns the name of the quick-provision partition file.
599
600 Subclasses should override this function to return correct name.
601 """
602 raise NotImplementedError('Subclasses need to implement this.')
603
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800604 def _OptimizePartLocation(self, offset: int, length: int):
605 """Optimizes the offset and length of the partition.
606
607 Subclasses can override this to provide better offset/length than what is
608 defined in the PGT partition layout.
609
610 Args:
611 offset: The offset (in bytes) of the partition in the image.
612 length: The length (in bytes) of the partition.
613
614 Returns:
615 A tuple of offset and length (in bytes) from the image.
616 """
617 return offset, length
Amin Hassanid684e982021-02-26 11:10:58 -0800618
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800619 def _RedirectPartition(self, file_name: str):
620 """Downloads the partition from a remote path and writes it into target.
621
622 Args:
623 file_name: The file name in the remote directory self._image.
624 """
625 cmd = self._GetWriteToTargetCommand()
626
627 image_path = os.path.join(self._image, file_name)
628 with GsFileCopier(image_path) as generator:
629 try:
630 with open(generator.Target(), 'rb') as fp:
631 self._device.run(cmd, input=fp, shell=True)
632 finally:
633 generator.CloseTarget()
634
Amin Hassanid684e982021-02-26 11:10:58 -0800635
636class KernelUpdater(RawPartitionUpdater):
637 """A class to update the kernel partition on a Chromium OS device."""
638
639 def _GetPartitionName(self):
640 """See RawPartitionUpdater._GetPartitionName()."""
641 return constants.PART_KERN_B
642
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800643 def _GetRemotePartitionName(self):
644 """See RawPartitionUpdater._GetRemotePartitionName()."""
645 return constants.QUICK_PROVISION_PAYLOAD_KERNEL
646
Amin Hassanid684e982021-02-26 11:10:58 -0800647 def Revert(self):
648 """Reverts the kernel partition update."""
649 # There is nothing to do for reverting kernel partition.
Amin Hassani75c5f942021-02-20 23:56:53 -0800650
651
652class RootfsUpdater(RawPartitionUpdater):
653 """A class to update the root partition on a Chromium OS device."""
654
655 def __init__(self, current_root: str, *args):
656 """Initializes the class.
657
658 Args:
659 current_root: The current root device path.
660 *args: See PartitionUpdaterBase
661 """
662 super().__init__(*args)
663
664 self._current_root = current_root
665 self._ran_postinst = False
666
667 def _GetPartitionName(self):
668 """See RawPartitionUpdater._GetPartitionName()."""
669 return constants.PART_ROOT_A
670
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800671 def _GetRemotePartitionName(self):
672 """See RawPartitionUpdater._GetRemotePartitionName()."""
673 return constants.QUICK_PROVISION_PAYLOAD_ROOTFS
674
Amin Hassani75c5f942021-02-20 23:56:53 -0800675 def _Run(self):
676 """The function that does the job of rootfs partition update."""
Amin Hassani55970562021-02-22 20:49:13 -0800677 with ProgressWatcher(self._device, self._target):
678 super()._Run()
Amin Hassani75c5f942021-02-20 23:56:53 -0800679
680 self._RunPostInst()
681
682 def _OptimizePartLocation(self, offset: int, length: int):
683 """Optimizes the size of the root partition of the image.
684
685 Normally the file system does not occupy the entire partition. Furthermore
686 we don't need the verity hash tree at the end of the root file system
687 because postinst will recreate it. This function reads the (approximate)
688 superblock of the ext4 partition and extracts the actual file system size in
689 the root partition.
690 """
691 superblock_size = 4096 * 2
692 with open(self._image, 'rb') as r:
693 r.seek(offset)
694 with tempfile.NamedTemporaryFile(delete=False) as fp:
695 fp.write(r.read(superblock_size))
696 fp.close()
697 return offset, partition_lib.Ext2FileSystemSize(fp.name)
698
699 def _RunPostInst(self, on_target: bool = True):
700 """Runs the postinst process in the root partition.
701
702 Args:
703 on_target: If true the postinst is run on the target (inactive)
704 partition. This is used when doing normal updates. If false, the
705 postinst is run on the current (active) partition. This is used when
706 reverting an update.
707 """
708 try:
709 postinst_dir = '/'
710 partition = self._current_root
711 if on_target:
712 postinst_dir = self._device.run(
713 ['mktemp', '-d', '-p', self._device.work_dir],
714 capture_output=True).stdout.strip()
715 self._device.run(['mount', '-o', 'ro', self._target, postinst_dir])
716 partition = self._target
717
718 self._ran_postinst = True
719 postinst = os.path.join(postinst_dir, 'postinst')
720 result = self._device.run([postinst, partition], capture_output=True)
721
722 logging.debug('Postinst result on %s: \n%s', postinst, result.stdout)
Amin Hassani55970562021-02-22 20:49:13 -0800723 # DeviceImagerOperation will look for this log.
Amin Hassani75c5f942021-02-20 23:56:53 -0800724 logging.info('Postinstall completed.')
725 finally:
726 if on_target:
727 self._device.run(['umount', postinst_dir])
728
729 def Revert(self):
730 """Reverts the root update install."""
731 logging.info('Reverting the rootfs partition update.')
732 if self._ran_postinst:
733 # We don't have to do anything for revert if we haven't changed the kernel
734 # priorities yet.
735 self._RunPostInst(on_target=False)
Amin Hassani74403082021-02-22 11:40:09 -0800736
737
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000738class MiniOSUpdater(RawPartitionUpdater):
739 """A class to update the miniOS partition on a Chromium OS device."""
740
741 def __init__(self, *args):
742 """Initializes the class.
743
744 Args:
745 *args: See PartitionUpdaterBase
746 """
747 super().__init__(*args)
748
749 self._ran_postinst = False
750
751 def _GetPartitionName(self):
752 """See RawPartitionUpdater._GetPartitionName()."""
753 return constants.PART_MINIOS_A
754
755 def _GetRemotePartitionName(self):
756 """See RawPartitionUpdater._GetRemotePartitionName()."""
Jae Hoon Kimb88b7962021-10-18 11:08:38 -0700757 return constants.QUICK_PROVISION_PAYLOAD_MINIOS
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000758
759 def _Run(self):
760 """The function that does the job of rootfs partition update."""
Jae Hoon Kimb88b7962021-10-18 11:08:38 -0700761 if self._image_type == ImageType.FULL:
762 if self._MiniOSPartitionsExistInImage():
763 logging.info('Updating miniOS partition from local.')
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000764 super()._Run()
765 else:
Jae Hoon Kimb88b7962021-10-18 11:08:38 -0700766 logging.warning('Not updating miniOS partition as it does not exist.')
767 return
768 elif self._image_type == ImageType.REMOTE_DIRECTORY:
769 if not gs.GSContext().Exists(
770 os.path.join(self._image,
771 constants.QUICK_PROVISION_PAYLOAD_MINIOS)):
772 logging.warning('Not updating miniOS, missing remote files.')
773 return
774 elif not self._MiniOSPartitionsExist():
775 logging.warning('Not updating miniOS, missing partitions.')
776 return
777 else:
778 logging.info('Updating miniOS partition from remote.')
779 super()._Run()
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000780 else:
781 # Let super() handle this error.
782 super()._Run()
783
784 self._RunPostInstall()
785
786 def _RunPostInstall(self):
787 """The function will change the priority of the miniOS partitions."""
788 self._FlipMiniOSPriority()
789 self._ran_postinst = True
790
791 def Revert(self):
792 """Reverts the miniOS partition update."""
793 if self._ran_postinst:
794 self._FlipMiniOSPriority()
795
796 def _GetMiniOSPriority(self):
797 return self._device.run(['crossystem', constants.MINIOS_PRIORITY]).output
798
799 def _SetMiniOSPriority(self, priority: str):
800 self._device.run(
801 ['crossystem', f'{constants.MINIOS_PRIORITY}={priority}'])
802
803 def _FlipMiniOSPriority(self):
804 inactive_minios_priority = 'B' if self._GetMiniOSPriority() == 'A' else 'A'
805 logging.info('Setting miniOS priority to %s', inactive_minios_priority)
806 self._SetMiniOSPriority(inactive_minios_priority)
807
Jae Hoon Kimb88b7962021-10-18 11:08:38 -0700808 def _MiniOSPartitionsExistInImage(self):
809 """Checks if miniOS partition exists in the image."""
810 d = cgpt.Disk.FromImage(self._image)
811 try:
812 d.GetPartitionByTypeGuid(cgpt.MINIOS_TYPE_GUID)
813 return True
814 except KeyError:
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000815 return False
816
Jae Hoon Kimb88b7962021-10-18 11:08:38 -0700817 def _MiniOSPartitionsExist(self):
818 """Checks if the device has miniOS partitions."""
819 run = lambda x: self._device.run(x).stdout.strip()
820 device_drive = run(['rootdev', '-s', '-d'])
821 cmd = ['cgpt', 'show', '-t', device_drive, '-i']
822 return all((run(cmd + [p]) == cgpt.MINIOS_TYPE_GUID) for p in ('9', '10'))
823
Jae Hoon Kimcc723e02021-08-16 21:03:21 +0000824
Amin Hassani74403082021-02-22 11:40:09 -0800825class StatefulPayloadGenerator(ReaderBase):
826 """A class for generating a stateful update payload in a separate thread."""
827 def __init__(self, image: str):
828 """Initializes that class.
829
830 Args:
831 image: The path to a local Chromium OS image.
832 """
833 super().__init__()
834 self._image = image
835
836 def run(self):
837 """Generates the stateful update and writes it into the output pipe."""
838 try:
839 paygen_stateful_payload_lib.GenerateStatefulPayload(
840 self._image, self._Source())
841 finally:
842 self._CloseSource()
843
844
845class StatefulUpdater(PartitionUpdaterBase):
846 """A class to update the stateful partition on a device."""
847 def __init__(self, clobber_stateful: bool, *args):
848 """Initializes the class
849
850 Args:
851 clobber_stateful: Whether to clobber the stateful or not.
852 *args: Look at PartitionUpdaterBase.
853 """
854 super().__init__(*args)
855 self._clobber_stateful = clobber_stateful
856
857 def _Run(self):
858 """Reads/Downloads the stateful updates and writes it into the device."""
859 if self._image_type == ImageType.FULL:
860 generator_cls = StatefulPayloadGenerator
861 elif self._image_type == ImageType.REMOTE_DIRECTORY:
862 generator_cls = GsFileCopier
863 self._image = os.path.join(self._image,
864 paygen_stateful_payload_lib.STATEFUL_FILE)
865 else:
866 raise ValueError(f'Invalid image type {self._image_type}')
867
868 with generator_cls(self._image) as generator:
869 try:
870 updater = stateful_updater.StatefulUpdater(self._device)
871 updater.Update(
872 generator.Target(),
873 is_payload_on_device=False,
874 update_type=(stateful_updater.StatefulUpdater.UPDATE_TYPE_CLOBBER if
875 self._clobber_stateful else None))
876 finally:
877 generator.CloseTarget()
878
879 def Revert(self):
880 """Reverts the stateful partition update."""
881 logging.info('Reverting the stateful update.')
882 stateful_updater.StatefulUpdater(self._device).Reset()
Amin Hassani55970562021-02-22 20:49:13 -0800883
884
885class ProgressWatcher(threading.Thread):
886 """A class used for watching the progress of rootfs update."""
887
888 def __init__(self, device, target_root: str):
889 """Initializes the class.
890
891 Args:
892 device: The ChromiumOSDevice to be updated.
893 target_root: The target root partition to monitor the progress of.
894 """
895 super().__init__()
896
897 self._device = device
898 self._target_root = target_root
899 self._exit = False
900
901 def __enter__(self):
902 """Starts the thread."""
903 self.start()
904 return self
905
906 def __exit__(self, *args, **kwargs):
907 """Exists the thread."""
908 self._exit = True
909 self.join()
910
911 def _ShouldExit(self):
912 return self._exit
913
914 def run(self):
915 """Monitors the progress of the target root partitions' update.
916
917 This is done by periodically, reading the fd position of the process that is
918 writing into the target partition and reporting it back. Then the position
919 is divided by the size of the block device to report an approximate
920 progress.
921 """
922 cmd = ['blockdev', '--getsize64', self._target_root]
923 output = self._device.run(cmd, capture_output=True).stdout.strip()
924 if output is None:
925 raise Error(f'Cannot get the block device size from {output}.')
926 dev_size = int(output)
927
928 # Using lsof to find out which process is writing to the target rootfs.
929 cmd = f'lsof 2>/dev/null | grep {self._target_root}'
930 while not self._ShouldExit():
931 try:
932 output = self._device.run(cmd, capture_output=True,
933 shell=True).stdout.strip()
934 if output:
935 break
936 except cros_build_lib.RunCommandError:
937 continue
938 finally:
939 time.sleep(1)
940
941 # Now that we know which process is writing to it, we can look the fdinfo of
942 # stdout of that process to get its offset. We're assuming there will be no
943 # seek, which is correct.
944 pid = output.split()[1]
945 cmd = ['cat', f'/proc/{pid}/fdinfo/1']
946 while not self._ShouldExit():
947 try:
948 output = self._device.run(cmd, capture_output=True).stdout.strip()
949 m = re.search(r'^pos:\s*(\d+)$', output, flags=re.M)
950 if m:
951 offset = int(m.group(1))
952 # DeviceImagerOperation will look for this log.
953 logging.info('RootFS progress: %f', offset/dev_size)
954 except cros_build_lib.RunCommandError:
955 continue
956 finally:
957 time.sleep(1)
958
959
960class DeviceImagerOperation(operation.ProgressBarOperation):
961 """A class to provide a progress bar for DeviceImager operation."""
962
963 def __init__(self):
964 """Initializes the class."""
965 super().__init__()
966
967 self._progress = 0.0
968
969 def ParseOutput(self, output=None):
970 """Override function to parse the output and provide progress.
971
972 Args:
973 output: The stderr or stdout.
974 """
975 output = self._stdout.read()
976 match = re.findall(r'RootFS progress: (\d+(?:\.\d+)?)', output)
977 if match:
978 progress = float(match[0])
979 self._progress = max(self._progress, progress)
980
981 # If postinstall completes, move half of the remaining progress.
982 if re.findall(r'Postinstall completed', output):
983 self._progress += (1.0 - self._progress) / 2
984
985 # While waiting for reboot, each time, move half of the remaining progress.
986 if re.findall(r'Unable to get new boot_id', output):
987 self._progress += (1.0 - self._progress) / 2
988
989 if re.findall(r'DeviceImager completed.', output):
990 self._progress = 1.0
991
992 self.ProgressBar(self._progress)