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