mirror of
https://github.com/chenxiaolong/avbroot.git
synced 2026-07-03 14:05:11 +02:00
ota: Reimplement zip signing ourselves more efficiently
The original AOSP ota_utils implementation is quite disk and memory heavy: * ota_utils.FinalizeMetadata() performs metadata file replacement as separate delete + add operations, which require an extra intermediate zip to be written. * ota_utils.FinalizeMetadata() keeps unneeded intermediate temporary files around until common.Cleanup() is called. * common.ZipDelete() deletes files by excluding entries when copying from one zip to another. The data copy reads entire entries (including the huge payload) into memory instead of copying chunks. Additionally, ota_utils is a bit exception/crash-unsafe in general. There are several places where temp files won't be cleaned up unless everything is following the happy path. This commit's reimplementation tries to use streaming and chunked operations where possible and only keeps temporary files around as long as needed. This is about as efficient as we can get while still relying on AOSP's signapk. (And we can finally get rid of the common.OPTIONS hackery!) Signed-off-by: Andrew Gunnerson <accounts+github@chiller3.com>
This commit is contained in:
+15
-31
@@ -18,14 +18,8 @@ from avbroot import util
|
||||
from avbroot import vbmeta
|
||||
|
||||
|
||||
PATH_METADATA_PB = 'META-INF/com/android/metadata.pb'
|
||||
PATH_PAYLOAD = 'payload.bin'
|
||||
PATH_PROPERTIES = 'payload_properties.txt'
|
||||
SKIP_PATHS = (
|
||||
PATH_METADATA_PB,
|
||||
'META-INF/com/android/metadata',
|
||||
'META-INF/com/android/otacert',
|
||||
)
|
||||
|
||||
|
||||
def print_status(*args, **kwargs):
|
||||
@@ -130,7 +124,7 @@ def patch_ota_zip(f_zip_in, f_zip_out, magisk, privkey_avb, privkey_ota,
|
||||
zipfile.ZipFile(f_zip_out, 'w') as z_out,
|
||||
):
|
||||
infolist = z_in.infolist()
|
||||
missing = {PATH_METADATA_PB, PATH_PAYLOAD, PATH_PROPERTIES}
|
||||
missing = {ota.PATH_METADATA_PB, PATH_PAYLOAD, PATH_PROPERTIES}
|
||||
i_payload = -1
|
||||
i_properties = -1
|
||||
|
||||
@@ -159,15 +153,10 @@ def patch_ota_zip(f_zip_in, f_zip_out, magisk, privkey_avb, privkey_ota,
|
||||
|
||||
for info in z_in.infolist():
|
||||
# The existing metadata is needed to generate a new signed zip
|
||||
if info.filename == PATH_METADATA_PB:
|
||||
if info.filename == ota.PATH_METADATA_PB:
|
||||
with z_in.open(info, 'r') as f_in:
|
||||
metadata = f_in.read()
|
||||
|
||||
# Skip files that are created during zip signing
|
||||
if info.filename in SKIP_PATHS:
|
||||
print_status('Skipping', info.filename)
|
||||
continue
|
||||
|
||||
# Copy other files, patching if needed
|
||||
with (
|
||||
z_in.open(info, 'r') as f_in,
|
||||
@@ -222,39 +211,34 @@ def patch_subcommand(args):
|
||||
dec_privkey_avb = os.path.join(key_dir, 'avb.key')
|
||||
openssl.decrypt_key(args.privkey_avb, dec_privkey_avb)
|
||||
|
||||
# AOSP's OTA utils require a DER-encoded private key with the `.pk8`
|
||||
# extension and a PEM-encoded certificate with the `.x509.pem` extension
|
||||
dec_privkey_ota = os.path.join(key_dir, 'ota.pk8')
|
||||
key_prefix_ota = dec_privkey_ota[:-4]
|
||||
# signapk requires a DER-encoded private key
|
||||
dec_privkey_ota = os.path.join(key_dir, 'ota.key')
|
||||
openssl.decrypt_key(args.privkey_ota, dec_privkey_ota, out_form='DER')
|
||||
|
||||
cert_ota = os.path.join(key_dir, 'ota.x509.pem')
|
||||
shutil.copyfile(args.cert_ota, cert_ota)
|
||||
|
||||
# Ensure that the certificate matches the private key
|
||||
if not openssl.cert_matches_key(cert_ota, dec_privkey_ota):
|
||||
if not openssl.cert_matches_key(args.cert_ota, dec_privkey_ota):
|
||||
raise Exception('OTA certificate does not match private key')
|
||||
|
||||
start = time.perf_counter_ns()
|
||||
|
||||
with tempfile.NamedTemporaryFile() as temp_unsigned:
|
||||
with util.open_output_file(output) as temp:
|
||||
metadata = patch_ota_zip(
|
||||
args.input,
|
||||
temp_unsigned,
|
||||
temp,
|
||||
args.magisk,
|
||||
dec_privkey_avb,
|
||||
dec_privkey_ota,
|
||||
cert_ota,
|
||||
args.cert_ota,
|
||||
)
|
||||
|
||||
print_status('Signing OTA zip')
|
||||
with util.open_output_file(output) as temp_signed:
|
||||
ota.sign_zip(
|
||||
temp_unsigned.name,
|
||||
temp_signed.name,
|
||||
key_prefix_ota,
|
||||
metadata,
|
||||
)
|
||||
ota.sign_zip(
|
||||
temp.name,
|
||||
temp.name,
|
||||
dec_privkey_ota,
|
||||
args.cert_ota,
|
||||
metadata,
|
||||
)
|
||||
|
||||
# Excluding the time it takes for the user to type in the passwords
|
||||
elapsed = time.perf_counter_ns() - start
|
||||
|
||||
+128
-37
@@ -2,19 +2,22 @@ import base64
|
||||
import binascii
|
||||
import bz2
|
||||
import hashlib
|
||||
import itertools
|
||||
import lzma
|
||||
import os
|
||||
import shutil
|
||||
import struct
|
||||
import sys
|
||||
import subprocess
|
||||
import zipfile
|
||||
|
||||
# Silence undesired warning
|
||||
orig_argv0 = sys.argv[0]
|
||||
sys.argv[0] = sys.argv[0].removesuffix('.py')
|
||||
import common
|
||||
import ota_utils
|
||||
sys.argv[0] = orig_argv0
|
||||
|
||||
import ota_metadata_pb2
|
||||
import ota_utils
|
||||
import update_metadata_pb2
|
||||
|
||||
from . import openssl
|
||||
@@ -23,6 +26,9 @@ from . import util
|
||||
|
||||
OTA_MAGIC = b'CrAU'
|
||||
|
||||
PATH_METADATA = 'META-INF/com/android/metadata'
|
||||
PATH_METADATA_PB = f'{PATH_METADATA}.pb'
|
||||
|
||||
|
||||
def parse_payload(f):
|
||||
'''
|
||||
@@ -77,12 +83,12 @@ def _extract_image(f_payload, f_out, block_size, blob_offset, partition):
|
||||
hasher=h_data)
|
||||
elif op.type == Type.REPLACE_BZ:
|
||||
decompressor = bz2.BZ2Decompressor()
|
||||
util.decompress_n(decompressor, f_payload, f_out, op.data_length,
|
||||
hasher=h_data)
|
||||
util.decompress_n(decompressor, f_payload, f_out,
|
||||
op.data_length, hasher=h_data)
|
||||
elif op.type == Type.REPLACE_XZ:
|
||||
decompressor = lzma.LZMADecompressor()
|
||||
util.decompress_n(decompressor, f_payload, f_out, op.data_length,
|
||||
hasher=h_data)
|
||||
util.decompress_n(decompressor, f_payload, f_out,
|
||||
op.data_length, hasher=h_data)
|
||||
elif op.type == Type.ZERO or op.type == Type.DISCARD:
|
||||
util.zero_n(f_out, extent.num_blocks * block_size)
|
||||
else:
|
||||
@@ -258,7 +264,7 @@ def patch_payload(f_in, f_out, version, manifest, blob_offset, temp_dir,
|
||||
for name, path in patched.items():
|
||||
# Find the partition in the manifest
|
||||
partition = next((p for p in manifest.partitions
|
||||
if p.partition_name == name), None)
|
||||
if p.partition_name == name), None)
|
||||
if partition is None:
|
||||
raise Exception(f'Partition {name} not found in manifest')
|
||||
|
||||
@@ -341,7 +347,7 @@ def patch_payload(f_in, f_out, version, manifest, blob_offset, temp_dir,
|
||||
blob_size = manifest.signatures_offset + manifest.signatures_size
|
||||
new_file_size = metadata_size + dummy_sig_size + blob_size
|
||||
|
||||
b64 = lambda d: base64.b64encode(d)
|
||||
def b64(d): return base64.b64encode(d)
|
||||
props = [
|
||||
b'FILE_HASH=%s\n' % b64(h_full.digest()),
|
||||
b'FILE_SIZE=%d\n' % new_file_size,
|
||||
@@ -352,12 +358,62 @@ def patch_payload(f_in, f_out, version, manifest, blob_offset, temp_dir,
|
||||
return b''.join(props)
|
||||
|
||||
|
||||
def sign_zip(input_path, output_path, key_prefix, metadata_raw):
|
||||
def _replace_metadata(z_in, z_out, metadata):
|
||||
'''
|
||||
Copy all entries from the input to the output, replacing the existing
|
||||
metadata files (legacy and protobuf) with the data serialized from the
|
||||
given metadata instance.
|
||||
'''
|
||||
|
||||
legacy_metadata = ota_utils.BuildLegacyOtaMetadata(metadata)
|
||||
legacy_metadata_str = "".join([f'{k}={v}\n' for k, v in
|
||||
sorted(legacy_metadata.items())])
|
||||
metadata_bytes = metadata.SerializeToString()
|
||||
|
||||
for info in z_in.infolist():
|
||||
with (
|
||||
z_in.open(info, 'r') as f_in,
|
||||
z_out.open(info, 'w') as f_out,
|
||||
):
|
||||
if info.filename == PATH_METADATA:
|
||||
f_out.write(legacy_metadata_str.encode('UTF-8'))
|
||||
elif info.filename == PATH_METADATA_PB:
|
||||
f_out.write(metadata_bytes)
|
||||
else:
|
||||
shutil.copyfileobj(f_in, f_out)
|
||||
|
||||
|
||||
def _sign_file(input_path, output_path, privkey, cert):
|
||||
'''
|
||||
Sign an OTA zip with AOSP's signapk tool. This may result in the zip
|
||||
entries being reordered.
|
||||
'''
|
||||
|
||||
subprocess.check_call([
|
||||
'java',
|
||||
'-Xmx4096m',
|
||||
'-jar',
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
'..', 'signapk', 'build', 'libs', 'signapk-all.jar',
|
||||
),
|
||||
'-w',
|
||||
cert,
|
||||
privkey,
|
||||
input_path,
|
||||
output_path,
|
||||
])
|
||||
|
||||
|
||||
def sign_zip(input_path, output_path, privkey_ota, cert_ota, metadata_raw):
|
||||
'''
|
||||
Sign an unsigned OTA zip. <metadata_raw> should be the serialized OTA
|
||||
metadata protobuf struct from the original OTA. The property files contained
|
||||
within the metadata that reference stored zip entries will be deleted and
|
||||
recreated during signing.
|
||||
metadata protobuf struct from the original OTA. The property files
|
||||
contained within the metadata that reference stored zip entries will be
|
||||
deleted and recreated during signing.
|
||||
|
||||
This is a reimplementation of ota_utils.FinalizeMetadata() because that
|
||||
function is quite inefficient, disk space and memory wise.
|
||||
'''
|
||||
|
||||
metadata = ota_metadata_pb2.OtaMetadata()
|
||||
@@ -365,32 +421,67 @@ def sign_zip(input_path, output_path, key_prefix, metadata_raw):
|
||||
|
||||
metadata.property_files.clear()
|
||||
|
||||
# We can't replace common.OPTIONS itself because ota_utils holds a reference
|
||||
# to the original instance
|
||||
attrs = (
|
||||
'search_path',
|
||||
'signapk_shared_library_path',
|
||||
'signapk_path',
|
||||
)
|
||||
orig_attrs = {a: getattr(common.OPTIONS, a) for a in attrs}
|
||||
props = [
|
||||
ota_utils.AbOtaPropertyFiles(),
|
||||
ota_utils.StreamingPropertyFiles(),
|
||||
]
|
||||
|
||||
try:
|
||||
common.OPTIONS.search_path = '/var/empty'
|
||||
common.OPTIONS.signapk_shared_library_path = '/var/empty'
|
||||
common.OPTIONS.signapk_path = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
'..', 'signapk', 'build', 'libs', 'signapk-all.jar',
|
||||
)
|
||||
# The first signing attempt uses the specified input file initially and
|
||||
# writes to the specified output file. If more signing attempts are needed,
|
||||
# the now-intermediate output file will be used as the input for the next
|
||||
# attempt.
|
||||
attempt_input_path = input_path
|
||||
|
||||
ota_utils.FinalizeMetadata(
|
||||
metadata,
|
||||
input_path,
|
||||
output_path,
|
||||
package_key=key_prefix,
|
||||
)
|
||||
# This is set to true once we know that the metadata entries can fit in the
|
||||
# reserved space within the property files.
|
||||
reserved_ok = False
|
||||
|
||||
finally:
|
||||
common.Cleanup()
|
||||
for attempt in itertools.count():
|
||||
# Compute initial property files with reserved space as placeholders
|
||||
# to store the metadata entries later
|
||||
if not reserved_ok:
|
||||
with zipfile.ZipFile(attempt_input_path, 'r') as z:
|
||||
for p in props:
|
||||
metadata.property_files[p.name] = p.Compute(z)
|
||||
|
||||
for k, v in orig_attrs.items():
|
||||
setattr(common.OPTIONS, k, v)
|
||||
# Replace the metadata files
|
||||
with (
|
||||
util.open_output_file(output_path) as f_replace,
|
||||
zipfile.ZipFile(attempt_input_path, 'r') as z_in,
|
||||
zipfile.ZipFile(f_replace, 'w') as z_out,
|
||||
):
|
||||
_replace_metadata(z_in, z_out, metadata)
|
||||
|
||||
# The newly generated zip is now the input for further attempts
|
||||
attempt_input_path = output_path
|
||||
|
||||
# Sign the zip "in place"
|
||||
with util.open_output_file(output_path) as f_replace:
|
||||
_sign_file(output_path, f_replace.name, privkey_ota, cert_ota)
|
||||
|
||||
with zipfile.ZipFile(output_path, 'r') as z:
|
||||
# Compute the final property files, including metadata entries
|
||||
try:
|
||||
for p in props:
|
||||
metadata.property_files[p.name] = \
|
||||
p.Finalize(z, len(metadata.property_files[p.name]))
|
||||
|
||||
reserved_ok = True
|
||||
except ota_utils.PropertyFiles.InsufficientSpaceException:
|
||||
# There is not enough space in the placeholders to store the
|
||||
# additional metadata entries and the new file offsets after
|
||||
# potential reordering by signapk. Try again since reordering
|
||||
# should no longer happen.
|
||||
if attempt == 0:
|
||||
continue
|
||||
else:
|
||||
raise
|
||||
|
||||
# Attempt number 0 only has placeholders, so don't bother verifying
|
||||
if attempt > 0:
|
||||
# Verify that the zip entry offsets in the property files are
|
||||
# still valid after signing
|
||||
for p in props:
|
||||
p.Verify(z, metadata.property_files[p.name].strip())
|
||||
|
||||
break
|
||||
|
||||
Reference in New Issue
Block a user