blob: 2377091ec71a36a82b6ac7d6a804fb8f3a3a39a5 [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
9import os
10import re
11import sys
Amin Hassanid4b3ff82021-02-20 23:05:14 -080012import tempfile
13import threading
Amin Hassani55970562021-02-22 20:49:13 -080014import time
Amin Hassani92f6c4a2021-02-20 17:36:09 -080015from typing import Tuple, Dict
16
Amin Hassani55970562021-02-22 20:49:13 -080017from chromite.cli import command
Amin Hassanid4b3ff82021-02-20 23:05:14 -080018from chromite.lib import constants
Amin Hassani92f6c4a2021-02-20 17:36:09 -080019from chromite.lib import cros_build_lib
20from chromite.lib import cros_logging as logging
21from chromite.lib import gs
Amin Hassanid4b3ff82021-02-20 23:05:14 -080022from chromite.lib import image_lib
Amin Hassani55970562021-02-22 20:49:13 -080023from chromite.lib import operation
Amin Hassanid4b3ff82021-02-20 23:05:14 -080024from chromite.lib import osutils
Amin Hassani92f6c4a2021-02-20 17:36:09 -080025from chromite.lib import parallel
26from chromite.lib import remote_access
27from chromite.lib import retry_util
Amin Hassani74403082021-02-22 11:40:09 -080028from chromite.lib import stateful_updater
Amin Hassani75c5f942021-02-20 23:56:53 -080029from chromite.lib.paygen import partition_lib
Amin Hassani74403082021-02-22 11:40:09 -080030from chromite.lib.paygen import paygen_stateful_payload_lib
Amin Hassani92f6c4a2021-02-20 17:36:09 -080031from chromite.lib.xbuddy import devserver_constants
32from chromite.lib.xbuddy import xbuddy
33
34
35assert sys.version_info >= (3, 6), 'This module requires Python 3.6+'
36
37
38class Error(Exception):
39 """Thrown when there is a general Chromium OS-specific flash error."""
40
41
42class ImageType(enum.Enum):
43 """Type of the image that is used for flashing the device."""
44
45 # The full image on disk (e.g. chromiumos_test_image.bin).
46 FULL = 0
47 # The remote directory path
48 # (e.g gs://chromeos-image-archive/eve-release/R90-x.x.x)
49 REMOTE_DIRECTORY = 1
50
51
52class Partition(enum.Enum):
53 """An enum for partition types like kernel and rootfs."""
54 KERNEL = 0
55 ROOTFS = 1
56
57
58class DeviceImager(object):
59 """A class to flash a Chromium OS device.
60
61 This utility uses parallelism as much as possible to achieve its goal as fast
62 as possible. For example, it uses parallel compressors, parallel transfers,
63 and simultaneous pipes.
64 """
65
66 # The parameters of the kernel and rootfs's two main partitions.
67 A = {Partition.KERNEL: 2, Partition.ROOTFS: 3}
68 B = {Partition.KERNEL: 4, Partition.ROOTFS: 5}
69
70 def __init__(self, device, image: str,
71 no_rootfs_update: bool = False,
72 no_stateful_update: bool = False,
73 no_reboot: bool = False,
74 disable_verification: bool = False,
75 clobber_stateful: bool = False,
76 clear_tpm_owner: bool = False):
77 """Initialize DeviceImager for flashing a Chromium OS device.
78
79 Args:
80 device: The ChromiumOSDevice to be updated.
81 image: The target image path (can be xBuddy path).
82 no_rootfs_update: Whether to do rootfs partition update.
83 no_stateful_update: Whether to do stateful partition update.
84 no_reboot: Whether to reboot device after update. The default is True.
85 disable_verification: Whether to disabling rootfs verification on the
86 device.
87 clobber_stateful: Whether to do a clean stateful partition.
88 clear_tpm_owner: If true, it will clear the TPM owner on reboot.
89 """
90
91 self._device = device
92 self._image = image
93 self._no_rootfs_update = no_rootfs_update
94 self._no_stateful_update = no_stateful_update
95 self._no_reboot = no_reboot
96 self._disable_verification = disable_verification
97 self._clobber_stateful = clobber_stateful
98 self._clear_tpm_owner = clear_tpm_owner
99
100 self._compression = cros_build_lib.COMP_XZ
101 self._inactive_state = None
102
103 def Run(self):
104 """Update the device with image of specific version."""
105
106 try:
Amin Hassani55970562021-02-22 20:49:13 -0800107 if command.UseProgressBar():
108 op = DeviceImagerOperation()
109 op.Run(self._Run)
110 else:
111 self._Run()
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800112 except Exception as e:
113 raise Error(f'DeviceImager Failed with error: {e}')
114
Amin Hassani55970562021-02-22 20:49:13 -0800115 # DeviceImagerOperation will look for this log.
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800116 logging.info('DeviceImager completed.')
117
118 def _Run(self):
119 """Runs the various operations to install the image on device."""
120 image, image_type = self._GetImage()
121 logging.info('Using image %s of type %s', image, image_type )
122
123 if image_type == ImageType.REMOTE_DIRECTORY:
124 self._compression = cros_build_lib.COMP_GZIP
125
126 self._InstallPartitions(image, image_type)
127
128 if self._clear_tpm_owner:
129 self._device.ClearTpmOwner()
130
131 if not self._no_reboot:
132 self._Reboot()
133 self._VerifyBootExpectations()
134
135 if self._disable_verification:
136 self._device.DisableRootfsVerification()
137
138 def _GetImage(self) -> Tuple[str, ImageType]:
139 """Returns the path to the final image(s) that need to be installed.
140
141 If the paths is local, the image should be the Chromium OS GPT image
142 (e.g. chromiumos_test_image.bin). If the path is remote, it should be the
143 remote directory where we can find the quick-provision and stateful update
144 files (e.g. gs://chromeos-image-archive/eve-release/R90-x.x.x).
145
146 NOTE: At this point there is no caching involved. Hence we always download
147 the partition payloads or extract them from the Chromium OS image.
148
149 Returns:
150 A tuple of image path and image type.
151 """
152 if os.path.isfile(self._image):
153 return self._image, ImageType.FULL
154
155 # TODO(b/172212406): We could potentially also allow this by searching
156 # through the directory to see whether we have quick-provision and stateful
157 # payloads. This only makes sense when a user has their workstation at home
158 # and doesn't want to incur the bandwidth cost of downloading the same
159 # image multiple times. For that, they can simply download the GPT image
160 # image first and flash that instead.
161 if os.path.isdir(self._image):
162 raise ValueError(
163 f'{self._image}: input must be a disk image, not a directory.')
164
165 if gs.PathIsGs(self._image):
166 # TODO(b/172212406): Check whether it is a directory. If it wasn't a
167 # directory download the image into some temp location and use it instead.
168 return self._image, ImageType.REMOTE_DIRECTORY
169
170 # Assuming it is an xBuddy path.
171 xb = xbuddy.XBuddy(log_screen=False)
172 build_id, local_file = xb.Translate([self._image])
173 if build_id is None:
174 raise Error(f'{self._image}: unable to find matching xBuddy path.')
175 logging.info('XBuddy path translated to build ID %s', build_id)
176
177 if local_file:
178 return local_file, ImageType.FULL
179
180 return (f'{devserver_constants.GS_IMAGE_DIR}/{build_id}',
181 ImageType.REMOTE_DIRECTORY)
182
183 def _SplitDevPath(self, path: str) -> Tuple[str, int]:
184 """Splits the given /dev/x path into prefix and the dev number.
185
186 Args:
187 path: The path to a block dev device.
188
189 Returns:
190 A tuple of representing the prefix and the index of the dev path.
191 e.g.: '/dev/mmcblk0p1' -> ['/dev/mmcblk0p', 1]
192 """
193 match = re.search(r'(.*)([0-9]+)$', path)
194 if match is None:
195 raise Error(f'{path}: Could not parse root dev path.')
196
197 return match.group(1), int(match.group(2))
198
199 def _GetKernelState(self, root_num: int) -> Tuple[Dict, Dict]:
200 """Returns the kernel state.
201
202 Returns:
203 A tuple of two dictionaries: The current active kernel state and the
204 inactive kernel state. (Look at A and B constants in this class.)
205 """
206 if root_num == self.A[Partition.ROOTFS]:
207 return self.A, self.B
208 elif root_num == self.B[Partition.ROOTFS]:
209 return self.B, self.A
210 else:
211 raise Error(f'Invalid root partition number {root_num}')
212
213 def _InstallPartitions(self, image: str, image_type):
214 """The main method that installs the partitions of a Chrome OS device.
215
216 It uses parallelism to install the partitions as fast as possible.
217
218 Args:
219 image: The image path (local file or remote directory).
220 image_type: The type of the image (ImageType).
221 """
Amin Hassanid684e982021-02-26 11:10:58 -0800222 prefix, root_num = self._SplitDevPath(self._device.root_dev)
223 active_state, self._inactive_state = self._GetKernelState(root_num)
224
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800225 updaters = []
Amin Hassanid684e982021-02-26 11:10:58 -0800226 if not self._no_rootfs_update:
Amin Hassani75c5f942021-02-20 23:56:53 -0800227 current_root = prefix + str(active_state[Partition.ROOTFS])
228 target_root = prefix + str(self._inactive_state[Partition.ROOTFS])
229 updaters.append(RootfsUpdater(current_root, self._device, image,
230 image_type, target_root, self._compression))
231
Amin Hassanid684e982021-02-26 11:10:58 -0800232 target_kernel = prefix + str(self._inactive_state[Partition.KERNEL])
233 updaters.append(KernelUpdater(self._device, image, image_type,
234 target_kernel, self._compression))
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800235
Amin Hassani74403082021-02-22 11:40:09 -0800236 if not self._no_stateful_update:
237 updaters.append(StatefulUpdater(self._clobber_stateful, self._device,
238 image, image_type, None, None))
239
Amin Hassani92f6c4a2021-02-20 17:36:09 -0800240 # Retry the partitions updates that failed, in case a transient error (like
241 # SSH drop, etc) caused the error.
242 num_retries = 1
243 try:
244 retry_util.RetryException(Error, num_retries,
245 parallel.RunParallelSteps,
246 (x.Run for x in updaters if not x.IsFinished()),
247 halt_on_error=True)
248 except Exception:
249 # If one of the partitions failed to be installed, revert all partitions.
250 parallel.RunParallelSteps(x.Revert for x in updaters)
251 raise
252
253 def _Reboot(self):
254 """Reboots the device."""
255 try:
256 self._device.Reboot(timeout_sec=60)
257 except remote_access.RebootError:
258 raise Error('Could not recover from reboot. Once example reason'
259 ' could be the image provided was a non-test image'
260 ' or the system failed to boot after the update.')
261 except Exception as e:
262 raise Error(f'Failed to reboot to the device with error: {e}')
263
264 def _VerifyBootExpectations(self):
265 """Verify that we fully booted into the expected kernel state."""
266 # Discover the newly active kernel.
267 _, root_num = self._SplitDevPath(self._device.root_dev)
268 active_state, _ = self._GetKernelState(root_num)
269
270 # If this happens, we should rollback.
271 if active_state != self._inactive_state:
272 raise Error('The expected kernel state after update is invalid.')
273
274 logging.info('Verified boot expectations.')
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800275
276
277class ReaderBase(threading.Thread):
278 """The base class for reading different inputs and writing into output.
279
280 This class extends threading.Thread, so it will be run on its own thread. Also
281 it can be used as a context manager. Internally, it opens necessary files for
282 writing to and reading from. This class cannot be instantiated, it needs to be
283 sub-classed first to provide necessary function implementations.
284 """
285
286 def __init__(self, use_named_pipes: bool = False):
287 """Initializes the class.
288
289 Args:
290 use_named_pipes: Whether to use a named pipe or anonymous file
291 descriptors.
292 """
293 super().__init__()
294 self._use_named_pipes = use_named_pipes
295 self._pipe_target = None
296 self._pipe_source = None
297
298 def __del__(self):
299 """Destructor.
300
301 Make sure to clean up any named pipes we might have created.
302 """
303 if self._use_named_pipes:
304 osutils.SafeUnlink(self._pipe_target)
305
306 def __enter__(self):
307 """Enters the context manager"""
308 if self._use_named_pipes:
309 # There is no need for the temp file, we only need its path. So the named
310 # pipe is created after this temp file is deleted.
311 with tempfile.NamedTemporaryFile(prefix='chromite-device-imager') as fp:
312 self._pipe_target = self._pipe_source = fp.name
313 os.mkfifo(self._pipe_target)
314 else:
315 self._pipe_target, self._pipe_source = os.pipe()
316
317 self.start()
318 return self
319
320 def __exit__(self, *args, **kwargs):
321 """Exits the context manager."""
322 self.join()
323
324 def _Source(self):
325 """Returns the source pipe to write data into.
326
327 Sub-classes can use this function to determine where to write their data
328 into.
329 """
330 return self._pipe_source
331
332 def _CloseSource(self):
333 """Closes the source pipe.
334
335 Sub-classes should use this function to close the pipe after they are done
336 writing into it. Failure to do so may result reader of the data to hang
337 indefinitely.
338 """
339 if not self._use_named_pipes:
340 os.close(self._pipe_source)
341
342 def Target(self):
343 """Returns the target pipe to read data from.
344
345 Users of this class can use this path to read data from.
346 """
347 return self._pipe_target
348
349 def CloseTarget(self):
350 """Closes the target pipe.
351
352 Users of this class should use this function to close the pipe after they
353 are done reading from it.
354 """
355 if self._use_named_pipes:
356 os.remove(self._pipe_target)
357 else:
358 os.close(self._pipe_target)
359
360
361class PartialFileReader(ReaderBase):
362 """A class to read specific offset and length from a file and compress it.
363
364 This class can be used to read from specific location and length in a file
365 (e.g. A partition in a GPT image). Then it compresses the input and writes it
366 out (to a pipe). Look at the base class for more information.
367 """
368
369 # The offset of different partitions in a Chromium OS image does not always
370 # align to larger values like 4096. It seems that 512 is the maximum value to
371 # be divisible by partition offsets. This size should not be increased just
372 # for 'performance reasons'. Since we are doing everything in parallel, in
373 # practice there is not much difference between this and larger block sizes as
374 # parallelism hides the possible extra latency provided by smaller block
375 # sizes.
376 _BLOCK_SIZE = 512
377
378 def __init__(self, image: str, offset: int, length: int, compression):
379 """Initializes the class.
380
381 Args:
382 image: The path to an image (local or remote directory).
383 offset: The offset (in bytes) to read from the image.
384 length: The length (in bytes) to read from the image.
385 compression: The compression type (see cros_build_lib.COMP_XXX).
386 """
387 super().__init__()
388
389 self._image = image
390 self._offset = offset
391 self._length = length
392 self._compression = compression
393
394 def run(self):
395 """Runs the reading and compression."""
396 cmd = [
397 'dd',
398 'status=none',
399 f'if={self._image}',
400 f'ibs={self._BLOCK_SIZE}',
401 f'skip={int(self._offset/self._BLOCK_SIZE)}',
402 f'count={int(self._length/self._BLOCK_SIZE)}',
403 '|',
404 cros_build_lib.FindCompressor(self._compression),
405 ]
406
407 try:
408 cros_build_lib.run(' '.join(cmd), stdout=self._Source(), shell=True)
409 finally:
410 self._CloseSource()
411
412
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800413class GsFileCopier(ReaderBase):
414 """A class for downloading gzip compressed file from GS bucket into a pipe."""
415
416 def __init__(self, image: str):
417 """Initializes the class.
418
419 Args:
420 image: The path to an image (local or remote directory).
421 """
422 super().__init__(use_named_pipes=True)
423 self._image = image
424
425 def run(self):
426 """Runs the download and write into the output pipe."""
427 try:
428 gs.GSContext().Copy(self._image, self._Source())
429 finally:
430 self._CloseSource()
431
432
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800433class PartitionUpdaterBase(object):
434 """A base abstract class to use for installing an image into a partition.
435
436 Sub-classes should implement the abstract methods to provide the core
437 functionality.
438 """
439 def __init__(self, device, image: str, image_type, target: str, compression):
440 """Initializes this base class with values that most sub-classes will need.
441
442 Args:
443 device: The ChromiumOSDevice to be updated.
444 image: The target image path for the partition update.
445 image_type: The type of the image (ImageType).
446 target: The target path (e.g. block dev) to install the update.
447 compression: The compression used for compressing the update payload.
448 """
449 self._device = device
450 self._image = image
451 self._image_type = image_type
452 self._target = target
453 self._compression = compression
454 self._finished = False
455
456 def Run(self):
457 """The main function that does the partition update job."""
458 with cros_build_lib.TimedSection() as timer:
459 try:
460 self._Run()
461 finally:
462 self._finished = True
463
464 logging.debug('Completed %s in %s', self.__class__.__name__, timer.delta)
465
466 @abc.abstractmethod
467 def _Run(self):
468 """The method that need to be implemented by sub-classes."""
469 raise NotImplementedError('Sub-classes need to implement this.')
470
471 def IsFinished(self):
472 """Returns whether the partition update has been successful."""
473 return self._finished
474
475 @abc.abstractmethod
476 def Revert(self):
477 """Reverts the partition update.
478
479 Sub-classes need to implement this function to provide revert capability.
480 """
481 raise NotImplementedError('Sub-classes need to implement this.')
482
483
484class RawPartitionUpdater(PartitionUpdaterBase):
485 """A class to update a raw partition on a Chromium OS device."""
486
487 def _Run(self):
488 """The function that does the job of kernel partition update."""
489 if self._image_type == ImageType.FULL:
490 self._CopyPartitionFromImage(self._GetPartitionName())
491 elif self._image_type == ImageType.REMOTE_DIRECTORY:
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800492 self._RedirectPartition(self._GetRemotePartitionName())
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800493 else:
494 raise ValueError(f'Invalid image type {self._image_type}')
495
496 def _GetPartitionName(self):
497 """Returns the name of the partition in a Chromium OS GPT layout.
498
499 Subclasses should override this function to return correct name.
500 """
501 raise NotImplementedError('Subclasses need to implement this.')
502
503 def _CopyPartitionFromImage(self, part_name: str):
504 """Updates the device's partition from a local Chromium OS image.
505
506 Args:
507 part_name: The name of the partition in the source image that needs to be
508 extracted.
509 """
510 cmd = self._GetWriteToTargetCommand()
511
512 offset, length = self._GetPartLocation(part_name)
513 offset, length = self._OptimizePartLocation(offset, length)
514 with PartialFileReader(self._image, offset, length,
515 self._compression) as generator:
516 try:
517 self._device.run(cmd, input=generator.Target(), shell=True)
518 finally:
519 generator.CloseTarget()
520
521 def _GetWriteToTargetCommand(self):
522 """Returns a write to target command to run on a Chromium OS device.
523
524 Returns:
525 A string command to run on a device to read data from stdin, uncompress it
526 and write it to the target partition.
527 """
528 cmd = self._device.GetDecompressor(self._compression)
529 # Using oflag=direct to tell the OS not to cache the writes (faster).
530 cmd += ['|', 'dd', 'bs=1M', 'oflag=direct', f'of={self._target}']
531 return ' '.join(cmd)
532
533 def _GetPartLocation(self, part_name: str):
534 """Extracts the location and size of the kernel partition from the image.
535
536 Args:
537 part_name: The name of the partition in the source image that needs to be
538 extracted.
539
540 Returns:
541 A tuple of offset and length (in bytes) from the image.
542 """
543 try:
544 parts = image_lib.GetImageDiskPartitionInfo(self._image)
545 part_info = [p for p in parts if p.name == part_name][0]
546 except IndexError:
547 raise Error(f'No partition named {part_name} found.')
548
549 return int(part_info.start), int(part_info.size)
550
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800551 def _GetRemotePartitionName(self):
552 """Returns the name of the quick-provision partition file.
553
554 Subclasses should override this function to return correct name.
555 """
556 raise NotImplementedError('Subclasses need to implement this.')
557
Amin Hassanid4b3ff82021-02-20 23:05:14 -0800558 def _OptimizePartLocation(self, offset: int, length: int):
559 """Optimizes the offset and length of the partition.
560
561 Subclasses can override this to provide better offset/length than what is
562 defined in the PGT partition layout.
563
564 Args:
565 offset: The offset (in bytes) of the partition in the image.
566 length: The length (in bytes) of the partition.
567
568 Returns:
569 A tuple of offset and length (in bytes) from the image.
570 """
571 return offset, length
Amin Hassanid684e982021-02-26 11:10:58 -0800572
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800573 def _RedirectPartition(self, file_name: str):
574 """Downloads the partition from a remote path and writes it into target.
575
576 Args:
577 file_name: The file name in the remote directory self._image.
578 """
579 cmd = self._GetWriteToTargetCommand()
580
581 image_path = os.path.join(self._image, file_name)
582 with GsFileCopier(image_path) as generator:
583 try:
584 with open(generator.Target(), 'rb') as fp:
585 self._device.run(cmd, input=fp, shell=True)
586 finally:
587 generator.CloseTarget()
588
Amin Hassanid684e982021-02-26 11:10:58 -0800589
590class KernelUpdater(RawPartitionUpdater):
591 """A class to update the kernel partition on a Chromium OS device."""
592
593 def _GetPartitionName(self):
594 """See RawPartitionUpdater._GetPartitionName()."""
595 return constants.PART_KERN_B
596
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800597 def _GetRemotePartitionName(self):
598 """See RawPartitionUpdater._GetRemotePartitionName()."""
599 return constants.QUICK_PROVISION_PAYLOAD_KERNEL
600
Amin Hassanid684e982021-02-26 11:10:58 -0800601 def Revert(self):
602 """Reverts the kernel partition update."""
603 # There is nothing to do for reverting kernel partition.
Amin Hassani75c5f942021-02-20 23:56:53 -0800604
605
606class RootfsUpdater(RawPartitionUpdater):
607 """A class to update the root partition on a Chromium OS device."""
608
609 def __init__(self, current_root: str, *args):
610 """Initializes the class.
611
612 Args:
613 current_root: The current root device path.
614 *args: See PartitionUpdaterBase
615 """
616 super().__init__(*args)
617
618 self._current_root = current_root
619 self._ran_postinst = False
620
621 def _GetPartitionName(self):
622 """See RawPartitionUpdater._GetPartitionName()."""
623 return constants.PART_ROOT_A
624
Amin Hassani0fe49ae2021-02-21 23:41:58 -0800625 def _GetRemotePartitionName(self):
626 """See RawPartitionUpdater._GetRemotePartitionName()."""
627 return constants.QUICK_PROVISION_PAYLOAD_ROOTFS
628
Amin Hassani75c5f942021-02-20 23:56:53 -0800629 def _Run(self):
630 """The function that does the job of rootfs partition update."""
Amin Hassani55970562021-02-22 20:49:13 -0800631 with ProgressWatcher(self._device, self._target):
632 super()._Run()
Amin Hassani75c5f942021-02-20 23:56:53 -0800633
634 self._RunPostInst()
635
636 def _OptimizePartLocation(self, offset: int, length: int):
637 """Optimizes the size of the root partition of the image.
638
639 Normally the file system does not occupy the entire partition. Furthermore
640 we don't need the verity hash tree at the end of the root file system
641 because postinst will recreate it. This function reads the (approximate)
642 superblock of the ext4 partition and extracts the actual file system size in
643 the root partition.
644 """
645 superblock_size = 4096 * 2
646 with open(self._image, 'rb') as r:
647 r.seek(offset)
648 with tempfile.NamedTemporaryFile(delete=False) as fp:
649 fp.write(r.read(superblock_size))
650 fp.close()
651 return offset, partition_lib.Ext2FileSystemSize(fp.name)
652
653 def _RunPostInst(self, on_target: bool = True):
654 """Runs the postinst process in the root partition.
655
656 Args:
657 on_target: If true the postinst is run on the target (inactive)
658 partition. This is used when doing normal updates. If false, the
659 postinst is run on the current (active) partition. This is used when
660 reverting an update.
661 """
662 try:
663 postinst_dir = '/'
664 partition = self._current_root
665 if on_target:
666 postinst_dir = self._device.run(
667 ['mktemp', '-d', '-p', self._device.work_dir],
668 capture_output=True).stdout.strip()
669 self._device.run(['mount', '-o', 'ro', self._target, postinst_dir])
670 partition = self._target
671
672 self._ran_postinst = True
673 postinst = os.path.join(postinst_dir, 'postinst')
674 result = self._device.run([postinst, partition], capture_output=True)
675
676 logging.debug('Postinst result on %s: \n%s', postinst, result.stdout)
Amin Hassani55970562021-02-22 20:49:13 -0800677 # DeviceImagerOperation will look for this log.
Amin Hassani75c5f942021-02-20 23:56:53 -0800678 logging.info('Postinstall completed.')
679 finally:
680 if on_target:
681 self._device.run(['umount', postinst_dir])
682
683 def Revert(self):
684 """Reverts the root update install."""
685 logging.info('Reverting the rootfs partition update.')
686 if self._ran_postinst:
687 # We don't have to do anything for revert if we haven't changed the kernel
688 # priorities yet.
689 self._RunPostInst(on_target=False)
Amin Hassani74403082021-02-22 11:40:09 -0800690
691
692class StatefulPayloadGenerator(ReaderBase):
693 """A class for generating a stateful update payload in a separate thread."""
694 def __init__(self, image: str):
695 """Initializes that class.
696
697 Args:
698 image: The path to a local Chromium OS image.
699 """
700 super().__init__()
701 self._image = image
702
703 def run(self):
704 """Generates the stateful update and writes it into the output pipe."""
705 try:
706 paygen_stateful_payload_lib.GenerateStatefulPayload(
707 self._image, self._Source())
708 finally:
709 self._CloseSource()
710
711
712class StatefulUpdater(PartitionUpdaterBase):
713 """A class to update the stateful partition on a device."""
714 def __init__(self, clobber_stateful: bool, *args):
715 """Initializes the class
716
717 Args:
718 clobber_stateful: Whether to clobber the stateful or not.
719 *args: Look at PartitionUpdaterBase.
720 """
721 super().__init__(*args)
722 self._clobber_stateful = clobber_stateful
723
724 def _Run(self):
725 """Reads/Downloads the stateful updates and writes it into the device."""
726 if self._image_type == ImageType.FULL:
727 generator_cls = StatefulPayloadGenerator
728 elif self._image_type == ImageType.REMOTE_DIRECTORY:
729 generator_cls = GsFileCopier
730 self._image = os.path.join(self._image,
731 paygen_stateful_payload_lib.STATEFUL_FILE)
732 else:
733 raise ValueError(f'Invalid image type {self._image_type}')
734
735 with generator_cls(self._image) as generator:
736 try:
737 updater = stateful_updater.StatefulUpdater(self._device)
738 updater.Update(
739 generator.Target(),
740 is_payload_on_device=False,
741 update_type=(stateful_updater.StatefulUpdater.UPDATE_TYPE_CLOBBER if
742 self._clobber_stateful else None))
743 finally:
744 generator.CloseTarget()
745
746 def Revert(self):
747 """Reverts the stateful partition update."""
748 logging.info('Reverting the stateful update.')
749 stateful_updater.StatefulUpdater(self._device).Reset()
Amin Hassani55970562021-02-22 20:49:13 -0800750
751
752class ProgressWatcher(threading.Thread):
753 """A class used for watching the progress of rootfs update."""
754
755 def __init__(self, device, target_root: str):
756 """Initializes the class.
757
758 Args:
759 device: The ChromiumOSDevice to be updated.
760 target_root: The target root partition to monitor the progress of.
761 """
762 super().__init__()
763
764 self._device = device
765 self._target_root = target_root
766 self._exit = False
767
768 def __enter__(self):
769 """Starts the thread."""
770 self.start()
771 return self
772
773 def __exit__(self, *args, **kwargs):
774 """Exists the thread."""
775 self._exit = True
776 self.join()
777
778 def _ShouldExit(self):
779 return self._exit
780
781 def run(self):
782 """Monitors the progress of the target root partitions' update.
783
784 This is done by periodically, reading the fd position of the process that is
785 writing into the target partition and reporting it back. Then the position
786 is divided by the size of the block device to report an approximate
787 progress.
788 """
789 cmd = ['blockdev', '--getsize64', self._target_root]
790 output = self._device.run(cmd, capture_output=True).stdout.strip()
791 if output is None:
792 raise Error(f'Cannot get the block device size from {output}.')
793 dev_size = int(output)
794
795 # Using lsof to find out which process is writing to the target rootfs.
796 cmd = f'lsof 2>/dev/null | grep {self._target_root}'
797 while not self._ShouldExit():
798 try:
799 output = self._device.run(cmd, capture_output=True,
800 shell=True).stdout.strip()
801 if output:
802 break
803 except cros_build_lib.RunCommandError:
804 continue
805 finally:
806 time.sleep(1)
807
808 # Now that we know which process is writing to it, we can look the fdinfo of
809 # stdout of that process to get its offset. We're assuming there will be no
810 # seek, which is correct.
811 pid = output.split()[1]
812 cmd = ['cat', f'/proc/{pid}/fdinfo/1']
813 while not self._ShouldExit():
814 try:
815 output = self._device.run(cmd, capture_output=True).stdout.strip()
816 m = re.search(r'^pos:\s*(\d+)$', output, flags=re.M)
817 if m:
818 offset = int(m.group(1))
819 # DeviceImagerOperation will look for this log.
820 logging.info('RootFS progress: %f', offset/dev_size)
821 except cros_build_lib.RunCommandError:
822 continue
823 finally:
824 time.sleep(1)
825
826
827class DeviceImagerOperation(operation.ProgressBarOperation):
828 """A class to provide a progress bar for DeviceImager operation."""
829
830 def __init__(self):
831 """Initializes the class."""
832 super().__init__()
833
834 self._progress = 0.0
835
836 def ParseOutput(self, output=None):
837 """Override function to parse the output and provide progress.
838
839 Args:
840 output: The stderr or stdout.
841 """
842 output = self._stdout.read()
843 match = re.findall(r'RootFS progress: (\d+(?:\.\d+)?)', output)
844 if match:
845 progress = float(match[0])
846 self._progress = max(self._progress, progress)
847
848 # If postinstall completes, move half of the remaining progress.
849 if re.findall(r'Postinstall completed', output):
850 self._progress += (1.0 - self._progress) / 2
851
852 # While waiting for reboot, each time, move half of the remaining progress.
853 if re.findall(r'Unable to get new boot_id', output):
854 self._progress += (1.0 - self._progress) / 2
855
856 if re.findall(r'DeviceImager completed.', output):
857 self._progress = 1.0
858
859 self.ProgressBar(self._progress)