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:
Andrew Gunnerson
2023-01-30 19:14:05 -05:00
parent 437cfc23d0
commit df887f90e7
2 changed files with 143 additions and 68 deletions
+15 -31
View File
@@ -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
View File
@@ -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