cros_generate_update_payload: Move utility funtions to paygen
There are a few utility functions in the cros_generate_update_payload that are
more suited to be moved to chromite/lib/paygen which they belong. Eventually the
cros_generate_update_payload will just call a function like
GenerateUpdatePayload() from lib/paygen instead of directly forking the
delta_generator.
The utility functions are mostly related to extracting partitions out of an
image. They are moved to a new file lib/paygen/partition_lib.py. Appropriate
unittests are also moved along. In addition these functions have been improved
and udpated. Specifically, they don't return the (potentially temporary file)
path to extracted partition anymore. The path needs to be passed to them.
BUG=chromium:862679
TEST=cros flash inside and outside chroot
TEST=unittests
Change-Id: Ic8cd2eb58f4e451b74250eea6789e5f2a415354a
Reviewed-on: https://chromium-review.googlesource.com/1278089
Commit-Ready: Amin Hassani <ahassani@chromium.org>
Tested-by: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Don Garrett <dgarrett@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/lib/paygen/filelib.py b/lib/paygen/filelib.py
index 18e5f86..5f5ffe6 100644
--- a/lib/paygen/filelib.py
+++ b/lib/paygen/filelib.py
@@ -99,3 +99,24 @@
sha256_hex = base64.b64encode(sha256.digest())
return sha1_hex, sha256_hex
+
+
+def CopyFileSegment(in_file, in_mode, in_len, out_file, out_mode, in_seek=0):
+ """Simulates a `dd` operation with seeks.
+
+ Args:
+ in_file: The input file
+ in_mode: The mode to open the input file
+ in_len: The length to copy
+ out_file: The output file
+ out_mode: The mode to open the output file
+ in_seek: How many bytes to seek from the |in_file|
+ """
+ with open(in_file, in_mode) as in_stream, \
+ open(out_file, out_mode) as out_stream:
+ in_stream.seek(in_seek)
+ remaining = in_len
+ while remaining:
+ chunk = in_stream.read(min(8192 * 1024, remaining))
+ remaining -= len(chunk)
+ out_stream.write(chunk)
diff --git a/lib/paygen/filelib_unittest.py b/lib/paygen/filelib_unittest.py
index 17cf5ba..76acd66 100644
--- a/lib/paygen/filelib_unittest.py
+++ b/lib/paygen/filelib_unittest.py
@@ -89,3 +89,12 @@
filelib.Copy(path1, relative_path)
self.assertExists(path2)
+
+ def testCopyFileSegment(self):
+ """Test copying on a simple file segment."""
+ a = os.path.join(self.tempdir, 'a.txt')
+ osutils.WriteFile(a, '789')
+ b = os.path.join(self.tempdir, 'b.txt')
+ osutils.WriteFile(b, '0123')
+ filelib.CopyFileSegment(a, 'r', 2, b, 'r+', in_seek=1)
+ self.assertEqual(osutils.ReadFile(b), '8923')
diff --git a/lib/paygen/partition_lib.py b/lib/paygen/partition_lib.py
new file mode 100644
index 0000000..dd9aa2c
--- /dev/null
+++ b/lib/paygen/partition_lib.py
@@ -0,0 +1,116 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Library for handling Chrome OS partition."""
+
+from __future__ import print_function
+
+import os
+import shutil
+import tempfile
+
+from chromite.lib import constants
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_logging as logging
+
+from chromite.lib.paygen import filelib
+
+
+def ExtractPartition(filename, partition, out_part):
+ """Extracts partition from an image file.
+
+ Args:
+ filename: The image file.
+ partition: The partition name. e.g. ROOT or KERNEL.
+ out_part: The output partition file.
+ """
+ parts = cros_build_lib.GetImageDiskPartitionInfo(filename)
+ part_info = [p for p in parts if p.name == partition][0]
+
+ offset = int(part_info.start)
+ length = int(part_info.size)
+
+ filelib.CopyFileSegment(filename, 'r', length, out_part, 'w', in_seek=offset)
+
+
+def Ext2FileSystemSize(ext2_file):
+ """Return the size of an ext2 filesystem in bytes.
+
+ Args:
+ ext2_file: The path to the ext2 file.
+ """
+ # dumpe2fs is normally installed in /sbin but doesn't require root.
+ dump = cros_build_lib.RunCommand(['/sbin/dumpe2fs', '-h', ext2_file],
+ print_cmd=False, capture_output=True).output
+ fs_blocks = 0
+ fs_blocksize = 0
+ for line in dump.split('\n'):
+ if not line:
+ continue
+ label, data = line.split(':')[:2]
+ if label == 'Block count':
+ fs_blocks = int(data)
+ elif label == 'Block size':
+ fs_blocksize = int(data)
+
+ return fs_blocks * fs_blocksize
+
+
+def PatchKernel(image, kern_file):
+ """Patches a kernel with vblock from a stateful partition.
+
+ Args:
+ image: The stateful partition image.
+ kern_file: The kernel file.
+ """
+
+ with tempfile.NamedTemporaryFile(prefix='stateful') as state_out, \
+ tempfile.NamedTemporaryFile(prefix='vmlinuz_hd.vblock') as vblock:
+ ExtractPartition(image, constants.PART_STATE, state_out)
+ cros_build_lib.RunCommand(
+ ['e2cp', '%s:/vmlinuz_hd.vblock' % state_out, vblock])
+ filelib.CopyFileSegment(
+ vblock, 'r', os.path.getsize(vblock), kern_file, 'r+')
+
+
+def ExtractKernel(image, kern_out):
+ """Extracts the kernel from the given image.
+
+ Args:
+ image: The image containing the kernel partition.
+ kern_out: The output kernel file.
+ """
+ ExtractPartition(image, constants.PART_KERN_B, kern_out)
+ with open(kern_out, 'r') as kern:
+ if not any(kern.read(65536)):
+ logging.warn('%s: Kernel B is empty, patching kernel A.', image)
+ ExtractPartition(image, constants.PART_KERN_A, kern_out)
+ PatchKernel(image, kern_out)
+
+
+def ExtractRoot(image, root_out, root_pretruncate=None):
+ """Extract the rootfs partition from a gpt image.
+
+ Args:
+ image: The input image file.
+ root_out: The output root partition file.
+ root_pretruncate: The output root partition before being truncated.
+ """
+ ExtractPartition(image, constants.PART_ROOT_A, root_out)
+
+ if root_pretruncate:
+ logging.info('Saving pre-truncated root as %s.', root_pretruncate)
+ shutil.copyfile(root_out, root_pretruncate)
+
+ # We only update the filesystem part of the partition, which is stored in the
+ # gpt script.
+ root_out_size = Ext2FileSystemSize(root_out)
+ if root_out_size:
+ with open(root_out, 'a') as root:
+ logging.info('Root size currently %d bytes.', os.path.getsize(root_out))
+ root.truncate(root_out_size)
+ logging.info('Truncated root to %d bytes.', root_out_size)
+ else:
+ raise IOError('Error truncating the rootfs to filesystem size.')
diff --git a/lib/paygen/partition_lib_unittest b/lib/paygen/partition_lib_unittest
new file mode 120000
index 0000000..ef3e37b
--- /dev/null
+++ b/lib/paygen/partition_lib_unittest
@@ -0,0 +1 @@
+../../scripts/wrapper.py
\ No newline at end of file
diff --git a/lib/paygen/partition_lib_unittest.py b/lib/paygen/partition_lib_unittest.py
new file mode 100644
index 0000000..3b49558
--- /dev/null
+++ b/lib/paygen/partition_lib_unittest.py
@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Test the partition_lib module."""
+
+from __future__ import print_function
+
+import os
+
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_test_lib
+from chromite.lib import osutils
+
+from chromite.lib.paygen import partition_lib
+
+class ExtractRootTest(cros_test_lib.MockTempDirTestCase):
+ """Test partition_lib.ExtractRoot"""
+ def testTruncate(self):
+ """Test truncating on extraction."""
+ root = os.path.join(self.tempdir, 'root.bin')
+ root_pretruncate = os.path.join(self.tempdir, 'root_pretruncate.bin')
+
+ content = '0123456789'
+ osutils.WriteFile(root, content)
+ self.PatchObject(partition_lib, 'ExtractPartition',
+ return_value=root)
+ self.PatchObject(partition_lib, 'Ext2FileSystemSize',
+ return_value=2)
+
+ partition_lib.ExtractRoot(None, root, root_pretruncate)
+ self.assertEqual(osutils.ReadFile(root), content[:2])
+ self.assertEqual(osutils.ReadFile(root_pretruncate), content)
+
+
+class Ext2FileSystemSizeTest(cros_test_lib.MockTestCase):
+ """Test partition_lib.Ext2FileSystemSize"""
+ def testExt2FileSystemSize(self):
+ """Test on simple output."""
+ block_size = 4096
+ block_count = 123
+
+ self.PatchObject(cros_build_lib, 'RunCommand',
+ return_value=cros_build_lib.CommandResult(output='''
+Block size: %d
+Other thing: 123456798
+Not an integer: cdsad132csda
+Block count: %d
+''' % (block_size, block_count)))
+
+ size = partition_lib.Ext2FileSystemSize('/dev/null')
+ self.assertEqual(size, block_size * block_count)
+
+
+class ExtractPartitionTest(cros_test_lib.MockTempDirTestCase):
+ """Test partition_lib.ExtractPartition"""
+ def testExtractPartition(self):
+ """Tests extraction on a simple image."""
+ part_a_bin = '0123'
+ part_b_bin = '4567'
+ image_bin = part_a_bin + part_b_bin
+
+ image = os.path.join(self.tempdir, 'image.bin')
+ osutils.WriteFile(image, image_bin)
+ part_a = os.path.join(self.tempdir, 'a.bin')
+
+ fake_partitions = (
+ cros_build_lib.PartitionInfo(1, 0, 4, 4, 'fs', 'PART-A', ''),
+ cros_build_lib.PartitionInfo(2, 4, 8, 4, 'fs', 'PART-B', ''),
+ )
+ self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
+ return_value=fake_partitions)
+
+ partition_lib.ExtractPartition(image, 'PART-A', part_a)
+ self.assertEqual(osutils.ReadFile(part_a), part_a_bin)
diff --git a/scripts/cros_generate_update_payload.py b/scripts/cros_generate_update_payload.py
index aeb0431..a78f479 100644
--- a/scripts/cros_generate_update_payload.py
+++ b/scripts/cros_generate_update_payload.py
@@ -10,7 +10,6 @@
from __future__ import print_function
-import os
import shutil
import tempfile
@@ -20,118 +19,12 @@
from chromite.lib import cros_logging as logging
from chromite.lib import osutils
+from chromite.lib.paygen import partition_lib
+
_DELTA_GENERATOR = 'delta_generator'
-# TODO(tbrindus): move this to paygen/filelib.py.
-def CopyFileSegment(in_file, in_mode, in_len, out_file, out_mode, in_seek=0):
- """Simulates a `dd` operation with seeks."""
- with open(in_file, in_mode) as in_stream, \
- open(out_file, out_mode) as out_stream:
- in_stream.seek(in_seek)
- remaining = in_len
- while remaining:
- chunk = in_stream.read(min(8192 * 1024, remaining))
- remaining -= len(chunk)
- out_stream.write(chunk)
-
-
-# TODO(tbrindus): move this to paygen/filelib.py.
-def ExtractPartitionToTempFile(filename, partition, temp_file=None):
- """Extracts |partition| from |filename| into |temp_file|.
-
- If |temp_file| is not specified, an arbitrary file is used.
-
- Returns the location of the extracted partition.
- """
- if temp_file is None:
- temp_file = tempfile.mktemp(prefix='cros_generate_update_payload')
-
- parts = cros_build_lib.GetImageDiskPartitionInfo(filename)
- part_info = [p for p in parts if p.name == partition][0]
-
- offset = int(part_info.start)
- length = int(part_info.size)
-
- CopyFileSegment(filename, 'r', length, temp_file, 'w', in_seek=offset)
-
- return temp_file
-
-
-def PatchKernel(image, kern_file):
- """Patches kernel |kern_file| with vblock from |image| stateful partition."""
- state_out = ExtractPartitionToTempFile(image, constants.PART_STATE)
- vblock = tempfile.mktemp(prefix='vmlinuz_hd.vblock')
- cros_build_lib.RunCommand(['e2cp', '%s:/vmlinuz_hd.vblock' % state_out,
- vblock])
-
- CopyFileSegment(vblock, 'r', os.path.getsize(vblock), kern_file, 'r+')
-
- osutils.SafeUnlink(state_out)
- osutils.SafeUnlink(vblock)
-
-
-def ExtractKernel(bin_file, kern_out):
- """Extracts the kernel from the given |bin_file|, into |kern_out|."""
- kern_out = ExtractPartitionToTempFile(bin_file, constants.PART_KERN_B,
- kern_out)
- with open(kern_out, 'r') as kern:
- if not any(kern.read(65536)):
- logging.warn('%s: Kernel B is empty, patching kernel A.', bin_file)
- ExtractPartitionToTempFile(bin_file, constants.PART_KERN_A, kern_out)
- PatchKernel(bin_file, kern_out)
-
- return kern_out
-
-
-def Ext2FileSystemSize(rootfs):
- """Return the size in bytes of the ext2 filesystem passed in |rootfs|."""
- # dumpe2fs is normally installed in /sbin but doesn't require root.
- dump = cros_build_lib.RunCommand(['/sbin/dumpe2fs', '-h', rootfs],
- print_cmd=False, capture_output=True).output
- fs_blocks = 0
- fs_blocksize = 0
- for line in dump.split('\n'):
- if not line:
- continue
- label, data = line.split(':')[:2]
- if label == 'Block count':
- fs_blocks = int(data)
- elif label == 'Block size':
- fs_blocksize = int(data)
-
- return fs_blocks * fs_blocksize
-
-
-def ExtractRoot(bin_file, root_out, root_pretruncate=None):
- """Extract the rootfs partition from the gpt image |bin_file|.
-
- Stores it in |root_out|. If |root_out| is empty, a new temp file will be used.
- If |root_pretruncate| is non-empty, saves the pretruncated rootfs partition
- there.
- """
- root_out = ExtractPartitionToTempFile(bin_file, constants.PART_ROOT_A,
- root_out)
-
- if root_pretruncate:
- logging.info('Saving pre-truncated root as %s.', root_pretruncate)
- shutil.copyfile(root_out, root_pretruncate)
-
- # We only update the filesystem part of the partition, which is stored in the
- # gpt script.
- root_out_size = Ext2FileSystemSize(root_out)
- if root_out_size:
- with open(root_out, 'a') as root:
- logging.info('Root size currently %d bytes.', os.path.getsize(root_out))
- root.truncate(root_out_size)
- logging.info('Truncated root to %d bytes.', root_out_size)
- else:
- raise IOError('Error truncating the rootfs to filesystem size.')
-
- return root_out
-
-
def GenerateUpdatePayload(opts):
"""Generates the output files for the given commandline |opts|."""
# TODO(tbrindus): we can support calling outside of chroot easily with
@@ -150,13 +43,14 @@
if opts.src_image:
src_kernel_path = opts.src_kern_path or 'old_kern.dat'
src_root_path = opts.src_root_path or 'old_root.dat'
- ExtractKernel(opts.src_image, src_kernel_path)
- ExtractRoot(opts.src_image, src_root_path)
+ partition_lib.ExtractKernel(opts.src_image, src_kernel_path)
+ partition_lib.ExtractRoot(opts.src_image, src_root_path)
if opts.image:
dst_kernel_path = opts.kern_path or 'new_kern.dat'
dst_root_path = opts.root_path or 'new_root.dat'
- ExtractKernel(opts.image, dst_kernel_path)
- ExtractRoot(opts.image, dst_root_path, opts.root_pretruncate_path)
+ partition_lib.ExtractKernel(opts.image, dst_kernel_path)
+ partition_lib.ExtractRoot(opts.image, dst_root_path,
+ opts.root_pretruncate_path)
logging.info('Done extracting kernel/root')
return
@@ -164,17 +58,24 @@
logging.info('Generating %s update', payload_type)
+ src_kernel_path = ''
+ src_root_path = ''
if delta:
if not opts.full_kernel:
- src_kernel_path = ExtractKernel(opts.src_image, opts.src_kern_path)
+ src_kernel_path = (opts.src_kern_path or
+ tempfile.NamedTemporaryFile().name)
+ partition_lib.ExtractKernel(opts.src_image, src_kernel_path)
else:
logging.info('Generating full kernel update')
- src_root_path = ExtractRoot(opts.src_image, opts.src_root_path)
+ src_root_path = opts.src_root_path or tempfile.NamedTemporaryFile().name
+ partition_lib.ExtractRoot(opts.src_image, src_root_path)
- dst_kernel_path = ExtractKernel(opts.image, opts.kern_path)
- dst_root_path = ExtractRoot(opts.image, opts.root_path,
- opts.root_pretruncate_path)
+ dst_kernel_path = opts.kern_path or tempfile.NamedTemporaryFile().name
+ dst_root_path = opts.root_path or tempfile.NamedTemporaryFile().name
+ partition_lib.ExtractKernel(opts.image, dst_kernel_path)
+ partition_lib.ExtractRoot(opts.image, dst_root_path,
+ opts.root_pretruncate_path)
# In major version 2 we need to explicity mark the postinst on the root
# partition to run.
diff --git a/scripts/cros_generate_update_payload_unittest.py b/scripts/cros_generate_update_payload_unittest.py
index a3b935a..219af01 100644
--- a/scripts/cros_generate_update_payload_unittest.py
+++ b/scripts/cros_generate_update_payload_unittest.py
@@ -13,98 +13,17 @@
from chromite.lib import cros_build_lib
from chromite.lib import cros_test_lib
from chromite.lib import osutils
+from chromite.lib.paygen import partition_lib
from chromite.scripts import cros_generate_update_payload
-
-class CopyFileSegmentTest(cros_test_lib.TempDirTestCase):
- """Test cros_generate_update_payload.CopyFileSegment"""
- def testCopyFileSegment(self):
- """Test copying on a simple file."""
- a = os.path.join(self.tempdir, 'a.txt')
- osutils.WriteFile(a, '789')
- b = os.path.join(self.tempdir, 'b.txt')
- osutils.WriteFile(b, '0123')
- cros_generate_update_payload.CopyFileSegment(a, 'r', 2, b, 'r+', in_seek=1)
- self.assertEqual(osutils.ReadFile(b), '8923')
-
-
-class ExtractRootTest(cros_test_lib.MockTempDirTestCase):
- """Test cros_generate_update_payload.ExtractRoot"""
- def testTruncate(self):
- """Test truncating on extraction."""
- root = os.path.join(self.tempdir, 'root.bin')
- root_pretruncate = os.path.join(self.tempdir, 'root_pretruncate.bin')
-
- content = '0123456789'
- osutils.WriteFile(root, content)
- self.PatchObject(cros_generate_update_payload, 'ExtractPartitionToTempFile',
- return_value=root)
- self.PatchObject(cros_generate_update_payload, 'Ext2FileSystemSize',
- return_value=2)
-
- cros_generate_update_payload.ExtractRoot(None, root, root_pretruncate)
- self.assertEqual(osutils.ReadFile(root), content[:2])
- self.assertEqual(osutils.ReadFile(root_pretruncate), content)
-
-
-class Ext2FileSystemSizeTest(cros_test_lib.MockTestCase):
- """Test cros_generate_update_payload.Ext2FileSystemSize"""
- def testExt2FileSystemSize(self):
- """Test on simple output."""
- block_size = 4096
- block_count = 123
-
- self.PatchObject(cros_build_lib, 'RunCommand',
- return_value=cros_build_lib.CommandResult(output='''
-Block size: %d
-Other thing: 123456798
-Not an integer: cdsad132csda
-Block count: %d
-''' % (block_size, block_count)))
-
- size = cros_generate_update_payload.Ext2FileSystemSize('/dev/null')
- self.assertEqual(size, block_size * block_count)
-
-
-class ExtractPartitionToTempFileTest(cros_test_lib.MockTempDirTestCase):
- """Test cros_generate_update_payload.ExtractPartitionToTempFile"""
- def testExtractPartitionToTempFile(self):
- """Tests extraction on a simple image."""
- part_a_bin = '0123'
- part_b_bin = '4567'
- image_bin = part_a_bin + part_b_bin
-
- image = os.path.join(self.tempdir, 'image.bin')
- osutils.WriteFile(image, image_bin)
- part_a = os.path.join(self.tempdir, 'a.bin')
-
- fake_partitions = (
- cros_build_lib.PartitionInfo(1, 0, 4, 4, 'fs', 'PART-A', ''),
- cros_build_lib.PartitionInfo(2, 4, 8, 4, 'fs', 'PART-B', ''),
- )
- self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
- return_value=fake_partitions)
-
- cros_generate_update_payload.ExtractPartitionToTempFile(image, 'PART-A',
- temp_file=part_a)
- self.assertEqual(osutils.ReadFile(part_a), part_a_bin)
-
- # Make sure we correctly generate new temp files.
- tmp = cros_generate_update_payload.ExtractPartitionToTempFile(image,
- 'PART-B')
- self.assertEqual(osutils.ReadFile(tmp), part_b_bin)
-
-
class DeltaGeneratorTest(cros_test_lib.RunCommandTempDirTestCase):
"""Test correct arguments passed to delta_generator."""
def testDeltaGenerator(self):
"""Test correct arguments propagated to delta_generator call."""
temp = os.path.join(self.tempdir, 'temp.bin')
osutils.WriteFile(temp, '0123456789')
- self.PatchObject(cros_generate_update_payload, 'ExtractPartitionToTempFile',
- return_value=temp)
- self.PatchObject(cros_generate_update_payload, 'Ext2FileSystemSize',
- return_value=4)
+ self.PatchObject(partition_lib, 'ExtractPartition', return_value=temp)
+ self.PatchObject(partition_lib, 'Ext2FileSystemSize', return_value=4)
fake_partitions = (
cros_build_lib.PartitionInfo(3, 0, 4, 4, 'fs', 'ROOT-A', ''),
