Call delta_generator to get the signature size

The signature size is needed during payload hashing and signing.
We used to sign the update_engine's payload with RSA keys only. In
this case, the signature size always equals the key size. But the
assumption is no longer true for EC keys, whose DER-encoded signature
size is a variant with a maximum size.

Therefore, we always give the maximum signature size to the delta
generator, who then add paddings to the real signature if necessary.
The maximum signature size is calculated by calling the delta_generator
with the new option '--calculate_signature_size'. For custom payload
signers, we also deprecate the '--payload_signer_key_size' and replace
it with '--payload_signer_maximum_signature_size'.

The EC key in the test is generated with:
The EC key in the unittest is generated with the command:
openssl ecparam -name prime256v1 -genkey -noout -out prime256v1-key.pem
openssl pkey -in prime256v1-key.pem -out testkey_EC.key

Bug: 141244025
Test: sign and verify a payload
Change-Id: Ife6e269d8aa3d870405aca20086330f1795e167f
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 1e7bb3a..dfcfb49 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -171,8 +171,16 @@
   --payload_signer_args <args>
       Specify the arguments needed for payload signer.
 
+  --payload_signer_maximum_signature_size <signature_size>
+      The maximum signature size (in bytes) that would be generated by the given
+      payload signer. Only meaningful when custom payload signer is specified
+      via '--payload_signer'.
+      If the signer uses a RSA key, this should be the number of bytes to
+      represent the modulus. If it uses an EC key, this is the size of a
+      DER-encoded ECDSA signature.
+
   --payload_signer_key_size <key_size>
-      Specify the key size in bytes of the payload signer.
+      Deprecated. Use the '--payload_signer_maximum_signature_size' instead.
 
   --skip_postinstall
       Skip the postinstall hooks when generating an A/B OTA package (default:
@@ -231,7 +239,7 @@
 OPTIONS.log_diff = None
 OPTIONS.payload_signer = None
 OPTIONS.payload_signer_args = []
-OPTIONS.payload_signer_key_size = None
+OPTIONS.payload_signer_maximum_signature_size = None
 OPTIONS.extracted_input = None
 OPTIONS.key_passwords = []
 OPTIONS.skip_postinstall = False
@@ -288,35 +296,31 @@
       self.signer = "openssl"
       self.signer_args = ["pkeyutl", "-sign", "-inkey", signing_key,
                           "-pkeyopt", "digest:sha256"]
-      self.key_size = self._GetKeySizeInBytes(signing_key)
+      self.maximum_signature_size = self._GetMaximumSignatureSizeInBytes(
+          signing_key)
     else:
       self.signer = OPTIONS.payload_signer
       self.signer_args = OPTIONS.payload_signer_args
-      if OPTIONS.payload_signer_key_size:
-        self.key_size = int(OPTIONS.payload_signer_key_size)
-        assert self.key_size == 256 or self.key_size == 512, \
-            "Unsupported key size {}".format(OPTIONS.payload_signer_key_size)
+      if OPTIONS.payload_signer_maximum_signature_size:
+        self.maximum_signature_size = int(
+            OPTIONS.payload_signer_maximum_signature_size)
       else:
-        self.key_size = 256
+        # The legacy config uses RSA2048 keys.
+        logger.warning("The maximum signature size for payload signer is not"
+                       " set, default to 256 bytes.")
+        self.maximum_signature_size = 256
 
   @staticmethod
-  def _GetKeySizeInBytes(signing_key):
-    modulus_file = common.MakeTempFile(prefix="modulus-")
-    cmd = ["openssl", "rsa", "-inform", "PEM", "-in", signing_key, "-modulus",
-           "-noout", "-out", modulus_file]
-    common.RunAndCheckOutput(cmd, verbose=False)
-
-    with open(modulus_file) as f:
-      modulus_string = f.read()
-    # The modulus string has the format "Modulus=$data", where $data is the
-    # concatenation of hex dump of the modulus.
-    MODULUS_PREFIX = "Modulus="
-    assert modulus_string.startswith(MODULUS_PREFIX)
-    modulus_string = modulus_string[len(MODULUS_PREFIX):]
-    key_size = len(modulus_string) // 2
-    assert key_size == 256 or key_size == 512, \
-        "Unsupported key size {}".format(key_size)
-    return key_size
+  def _GetMaximumSignatureSizeInBytes(signing_key):
+    out_signature_size_file = common.MakeTempFile("signature_size")
+    cmd = ["delta_generator", "--out_maximum_signature_size_file={}".format(
+        out_signature_size_file), "--private_key={}".format(signing_key)]
+    common.RunAndCheckOutput(cmd)
+    with open(out_signature_size_file) as f:
+      signature_size = f.read().rstrip()
+    logger.info("% outputs the maximum signature size: %", cmd[0],
+                signature_size)
+    return int(signature_size)
 
   def Sign(self, in_file):
     """Signs the given input file. Returns the output filename."""
@@ -396,7 +400,7 @@
     metadata_sig_file = common.MakeTempFile(prefix="sig-", suffix=".bin")
     cmd = ["brillo_update_payload", "hash",
            "--unsigned_payload", self.payload_file,
-           "--signature_size", str(payload_signer.key_size),
+           "--signature_size", str(payload_signer.maximum_signature_size),
            "--metadata_hash_file", metadata_sig_file,
            "--payload_hash_file", payload_sig_file]
     self._Run(cmd)
@@ -411,7 +415,7 @@
     cmd = ["brillo_update_payload", "sign",
            "--unsigned_payload", self.payload_file,
            "--payload", signed_payload_file,
-           "--signature_size", str(payload_signer.key_size),
+           "--signature_size", str(payload_signer.maximum_signature_size),
            "--metadata_signature_file", signed_metadata_sig_file,
            "--payload_signature_file", signed_payload_sig_file]
     self._Run(cmd)
@@ -2005,8 +2009,13 @@
       OPTIONS.payload_signer = a
     elif o == "--payload_signer_args":
       OPTIONS.payload_signer_args = shlex.split(a)
+    elif o == "--payload_signer_maximum_signature_size":
+      OPTIONS.payload_signer_maximum_signature_size = a
     elif o == "--payload_signer_key_size":
-      OPTIONS.payload_signer_key_size = a
+      # TODO(Xunchang) remove this option after cleaning up the callers.
+      logger.warning("The option '--payload_signer_key_size' is deprecated."
+                     " Use '--payload_signer_maximum_signature_size' instead.")
+      OPTIONS.payload_signer_maximum_signature_size = a
     elif o == "--extracted_input_target_files":
       OPTIONS.extracted_input = a
     elif o == "--skip_postinstall":
@@ -2047,6 +2056,7 @@
                                  "log_diff=",
                                  "payload_signer=",
                                  "payload_signer_args=",
+                                 "payload_signer_maximum_signature_size=",
                                  "payload_signer_key_size=",
                                  "extracted_input_target_files=",
                                  "skip_postinstall",
diff --git a/tools/releasetools/test_ota_from_target_files.py b/tools/releasetools/test_ota_from_target_files.py
index c3021a1..38faf64 100644
--- a/tools/releasetools/test_ota_from_target_files.py
+++ b/tools/releasetools/test_ota_from_target_files.py
@@ -885,10 +885,28 @@
       payload_offset, metadata_total = (
           property_files._GetPayloadMetadataOffsetAndSize(input_zip))
 
-    # Read in the metadata signature directly.
+    # The signature proto has the following format (details in
+    #  /platform/system/update_engine/update_metadata.proto):
+    #  message Signature {
+    #    optional uint32 version = 1;
+    #    optional bytes data = 2;
+    #    optional fixed32 unpadded_signature_size = 3;
+    #  }
+    #
+    # According to the protobuf encoding, the tail of the signature message will
+    # be [signature string(256 bytes) + encoding of the fixed32 number 256]. And
+    # 256 is encoded as 'x1d\x00\x01\x00\x00':
+    # [3 (field number) << 3 | 5 (type) + byte reverse of 0x100 (256)].
+    # Details in (https://developers.google.com/protocol-buffers/docs/encoding)
+    signature_tail_length = self.SIGNATURE_SIZE + 5
+    self.assertGreater(metadata_total, signature_tail_length)
     with open(output_file, 'rb') as verify_fp:
-      verify_fp.seek(payload_offset + metadata_total - self.SIGNATURE_SIZE)
-      metadata_signature = verify_fp.read(self.SIGNATURE_SIZE)
+      verify_fp.seek(payload_offset + metadata_total - signature_tail_length)
+      metadata_signature_proto_tail = verify_fp.read(signature_tail_length)
+
+    self.assertEqual(b'\x1d\x00\x01\x00\x00',
+                     metadata_signature_proto_tail[-5:])
+    metadata_signature = metadata_signature_proto_tail[:-5]
 
     # Now we extract the metadata hash via brillo_update_payload script, which
     # will serve as the oracle result.
@@ -1050,11 +1068,13 @@
     with open(file1, 'rb') as fp1, open(file2, 'rb') as fp2:
       self.assertEqual(fp1.read(), fp2.read())
 
+  @test_utils.SkipIfExternalToolsUnavailable()
   def test_init(self):
     payload_signer = PayloadSigner()
     self.assertEqual('openssl', payload_signer.signer)
-    self.assertEqual(256, payload_signer.key_size)
+    self.assertEqual(256, payload_signer.maximum_signature_size)
 
+  @test_utils.SkipIfExternalToolsUnavailable()
   def test_init_withPassword(self):
     common.OPTIONS.package_key = os.path.join(
         self.testdata_dir, 'testkey_with_passwd')
@@ -1067,18 +1087,27 @@
   def test_init_withExternalSigner(self):
     common.OPTIONS.payload_signer = 'abc'
     common.OPTIONS.payload_signer_args = ['arg1', 'arg2']
-    common.OPTIONS.payload_signer_key_size = '512'
+    common.OPTIONS.payload_signer_maximum_signature_size = '512'
     payload_signer = PayloadSigner()
     self.assertEqual('abc', payload_signer.signer)
     self.assertEqual(['arg1', 'arg2'], payload_signer.signer_args)
-    self.assertEqual(512, payload_signer.key_size)
+    self.assertEqual(512, payload_signer.maximum_signature_size)
 
-  def test_GetKeySizeInBytes_512Bytes(self):
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_GetMaximumSignatureSizeInBytes_512Bytes(self):
     signing_key = os.path.join(self.testdata_dir, 'testkey_RSA4096.key')
     # pylint: disable=protected-access
-    key_size = PayloadSigner._GetKeySizeInBytes(signing_key)
-    self.assertEqual(512, key_size)
+    signature_size = PayloadSigner._GetMaximumSignatureSizeInBytes(signing_key)
+    self.assertEqual(512, signature_size)
 
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_GetMaximumSignatureSizeInBytes_ECKey(self):
+    signing_key = os.path.join(self.testdata_dir, 'testkey_EC.key')
+    # pylint: disable=protected-access
+    signature_size = PayloadSigner._GetMaximumSignatureSizeInBytes(signing_key)
+    self.assertEqual(72, signature_size)
+
+  @test_utils.SkipIfExternalToolsUnavailable()
   def test_Sign(self):
     payload_signer = PayloadSigner()
     input_file = os.path.join(self.testdata_dir, self.SIGFILE)
diff --git a/tools/releasetools/testdata/testkey_EC.key b/tools/releasetools/testdata/testkey_EC.key
new file mode 100644
index 0000000..9e65a68
--- /dev/null
+++ b/tools/releasetools/testdata/testkey_EC.key
@@ -0,0 +1,5 @@
+-----BEGIN PRIVATE KEY-----
+MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgGaguGj8Yb1KkqKHd
+ISblUsjtOCbzAuVpX81i02sm8FWhRANCAARBnuotwKOsuvjH6iwTDhOAi7Q5pLWz
+xDkZjg2pcfbfi9FFTvLYETas7B2W6fx9PUezUmHTFTDV2JZuMYYFdZOw
+-----END PRIVATE KEY-----