Create source stamp block am: 904108a1df

Change-Id: I340e96716f2646e5223080fad678dc8964407c4d
diff --git a/src/main/java/com/android/apksig/ApkSigner.java b/src/main/java/com/android/apksig/ApkSigner.java
index ead7f6f..5057160 100644
--- a/src/main/java/com/android/apksig/ApkSigner.java
+++ b/src/main/java/com/android/apksig/ApkSigner.java
@@ -290,6 +290,14 @@
             if (mCreatedBy != null) {
                 signerEngineBuilder.setCreatedBy(mCreatedBy);
             }
+            if (mSourceStampSignerConfig != null) {
+                signerEngineBuilder.setStampSignerConfig(
+                        new DefaultApkSignerEngine.SignerConfig.Builder(
+                                        mSourceStampSignerConfig.getName(),
+                                        mSourceStampSignerConfig.getPrivateKey(),
+                                        mSourceStampSignerConfig.getCertificates())
+                                .build());
+            }
             signerEngine = signerEngineBuilder.build();
         }
 
@@ -546,7 +554,8 @@
                         outputCentralDirDataSource.size(),
                         outputCentralDirStartOffset);
 
-        // Step 11. Generate and output APK Signature Scheme v2 and/or v3 signatures, if necessary.
+        // Step 11. Generate and output APK Signature Scheme v2 and/or v3 signatures and/or
+        // SourceStamp signatures, if necessary.
         // This may insert an APK Signing Block just before the output's ZIP Central Directory
         ApkSignerEngine.OutputApkSigningBlockRequest2 outputApkSigningBlockRequest =
                 signerEngine.outputZipSections2(
@@ -570,6 +579,7 @@
         outputApkOut.consume(outputEocd);
         signerEngine.outputDone();
 
+        // Step 13. Generate and output APK Signature Scheme v4 signatures, if necessary.
         if (mV4SigningEnabled) {
             signerEngine.signV4(outputApkIn, mOutputV4File);
         }
diff --git a/src/main/java/com/android/apksig/ApkSignerEngine.java b/src/main/java/com/android/apksig/ApkSignerEngine.java
index 91069e1..59a2bd9 100644
--- a/src/main/java/com/android/apksig/ApkSignerEngine.java
+++ b/src/main/java/com/android/apksig/ApkSignerEngine.java
@@ -259,8 +259,8 @@
             DataSource zipEntries,
             DataSource zipCentralDirectory,
             DataSource zipEocd)
-                    throws IOException, ApkFormatException, NoSuchAlgorithmException,
-                            InvalidKeyException, SignatureException, IllegalStateException;
+            throws IOException, ApkFormatException, NoSuchAlgorithmException,
+            InvalidKeyException, SignatureException, IllegalStateException;
 
     /**
      * Indicates to this engine that the ZIP sections comprising the output APK have been output.
@@ -294,8 +294,8 @@
             DataSource zipEntries,
             DataSource zipCentralDirectory,
             DataSource zipEocd)
-                    throws IOException, ApkFormatException, NoSuchAlgorithmException,
-                            InvalidKeyException, SignatureException, IllegalStateException;
+            throws IOException, ApkFormatException, NoSuchAlgorithmException,
+            InvalidKeyException, SignatureException, IllegalStateException;
 
     /**
      * Indicates to this engine that the signed APK was output.
diff --git a/src/main/java/com/android/apksig/DefaultApkSignerEngine.java b/src/main/java/com/android/apksig/DefaultApkSignerEngine.java
index 989e9d7..063e3d2 100644
--- a/src/main/java/com/android/apksig/DefaultApkSignerEngine.java
+++ b/src/main/java/com/android/apksig/DefaultApkSignerEngine.java
@@ -19,7 +19,9 @@
 import com.android.apksig.apk.ApkFormatException;
 import com.android.apksig.apk.ApkUtils;
 import com.android.apksig.internal.apk.ApkSigningBlockUtils;
+import com.android.apksig.internal.apk.ContentDigestAlgorithm;
 import com.android.apksig.internal.apk.SignatureAlgorithm;
+import com.android.apksig.internal.apk.stamp.SourceStampSigner;
 import com.android.apksig.internal.apk.v1.DigestAlgorithm;
 import com.android.apksig.internal.apk.v1.V1SchemeSigner;
 import com.android.apksig.internal.apk.v1.V1SchemeVerifier;
@@ -85,6 +87,7 @@
     private final boolean mOtherSignersSignaturesPreserved;
     private final String mCreatedBy;
     private final List<SignerConfig> mSignerConfigs;
+    private final SignerConfig mSourceStampSignerConfig;
     private final int mMinSdkVersion;
     private final SigningCertificateLineage mSigningCertificateLineage;
 
@@ -146,6 +149,7 @@
 
     private DefaultApkSignerEngine(
             List<SignerConfig> signerConfigs,
+            SignerConfig sourceStampSignerConfig,
             int minSdkVersion,
             boolean v1SigningEnabled,
             boolean v2SigningEnabled,
@@ -173,6 +177,7 @@
         mOtherSignersSignaturesPreserved = otherSignersSignaturesPreserved;
         mCreatedBy = createdBy;
         mSignerConfigs = signerConfigs;
+        mSourceStampSignerConfig = sourceStampSignerConfig;
         mMinSdkVersion = minSdkVersion;
         mSigningCertificateLineage = signingCertificateLineage;
 
@@ -359,6 +364,14 @@
         return configs.get(0);
     }
 
+    private ApkSigningBlockUtils.SignerConfig createSourceStampSignerConfig()
+            throws InvalidKeyException {
+        return createSigningBlockSignerConfig(
+                mSourceStampSignerConfig,
+                /* apkSigningBlockPaddingSupported= */ true,
+                ApkSigningBlockUtils.VERSION_SOURCE_STAMP);
+    }
+
     private int getMinSdkFromV3SignatureAlgorithms(List<SignatureAlgorithm> algorithms) {
         int min = Integer.MAX_VALUE;
         for (SignatureAlgorithm algorithm : algorithms) {
@@ -427,6 +440,11 @@
                     newSignerConfig.signatureAlgorithms = new ArrayList<>();
                 }
                 break;
+            case ApkSigningBlockUtils.VERSION_SOURCE_STAMP:
+                newSignerConfig.signatureAlgorithms =
+                        Collections.singletonList(
+                                SignatureAlgorithm.VERITY_RSA_PKCS1_V1_5_WITH_SHA256);
+                break;
             default:
                 throw new IllegalArgumentException("Unknown APK Signature Scheme ID requested");
         }
@@ -798,37 +816,50 @@
         DataSource eocd = ApkSigningBlockUtils.copyWithModifiedCDOffset(beforeCentralDir, zipEocd);
 
         List<Pair<byte[], Integer>> signingSchemeBlocks = new ArrayList<>();
+        ApkSigningBlockUtils.SigningSchemeBlockAndDigests v2SigningSchemeBlockAndDigests = null;
+        ApkSigningBlockUtils.SigningSchemeBlockAndDigests v3SigningSchemeBlockAndDigests = null;
 
         // create APK Signature Scheme V2 Signature if requested
         if (mV2SigningEnabled) {
             invalidateV2Signature();
             List<ApkSigningBlockUtils.SignerConfig> v2SignerConfigs =
                     createV2SignerConfigs(apkSigningBlockPaddingSupported);
-            signingSchemeBlocks.add(
+            v2SigningSchemeBlockAndDigests =
                     V2SchemeSigner.generateApkSignatureSchemeV2Block(
-                                    mExecutor,
-                                    beforeCentralDir,
-                                    zipCentralDirectory,
-                                    eocd,
-                                    v2SignerConfigs,
-                                    mV3SigningEnabled)
-                            .signingSchemeBlock);
+                            mExecutor,
+                            beforeCentralDir,
+                            zipCentralDirectory,
+                            eocd,
+                            v2SignerConfigs,
+                            mV3SigningEnabled);
+            signingSchemeBlocks.add(v2SigningSchemeBlockAndDigests.signingSchemeBlock);
         }
         if (mV3SigningEnabled) {
             invalidateV3Signature();
             List<ApkSigningBlockUtils.SignerConfig> v3SignerConfigs =
                     createV3SignerConfigs(apkSigningBlockPaddingSupported);
-            signingSchemeBlocks.add(
+            v3SigningSchemeBlockAndDigests =
                     V3SchemeSigner.generateApkSignatureSchemeV3Block(
-                                    mExecutor,
-                                    beforeCentralDir,
-                                    zipCentralDirectory,
-                                    eocd,
-                                    v3SignerConfigs)
-                            .signingSchemeBlock);
+                            mExecutor,
+                            beforeCentralDir,
+                            zipCentralDirectory,
+                            eocd,
+                            v3SignerConfigs);
+            signingSchemeBlocks.add(v3SigningSchemeBlockAndDigests.signingSchemeBlock);
+        }
+        if (mSourceStampSignerConfig != null && (mV2SigningEnabled || mV3SigningEnabled)) {
+            ApkSigningBlockUtils.SignerConfig sourceStampSignerConfig =
+                    createSourceStampSignerConfig();
+            Map<ContentDigestAlgorithm, byte[]> digestInfo =
+                    mV3SigningEnabled
+                            ? v3SigningSchemeBlockAndDigests.digestInfo
+                            : v2SigningSchemeBlockAndDigests.digestInfo;
+            signingSchemeBlocks.add(
+                    SourceStampSigner.generateSourceStampBlock(
+                            sourceStampSignerConfig, digestInfo));
         }
 
-        // create APK Signing Block with v2 and/or v3 blocks
+        // create APK Signing Block with v2 and/or v3 and/or SourceStamp blocks
         byte[] apkSigningBlock = ApkSigningBlockUtils.generateApkSigningBlock(signingSchemeBlocks);
 
         mAddSigningBlockRequest =
@@ -1333,6 +1364,7 @@
     /** Builder of {@link DefaultApkSignerEngine} instances. */
     public static class Builder {
         private List<SignerConfig> mSignerConfigs;
+        private SignerConfig mStampSignerConfig;
         private final int mMinSdkVersion;
 
         private boolean mV1SigningEnabled = true;
@@ -1422,6 +1454,7 @@
 
             return new DefaultApkSignerEngine(
                     mSignerConfigs,
+                    mStampSignerConfig,
                     mMinSdkVersion,
                     mV1SigningEnabled,
                     mV2SigningEnabled,
@@ -1432,6 +1465,12 @@
                     mSigningCertificateLineage);
         }
 
+        /** Sets the signer configuration for the SourceStamp to be embedded in the APK. */
+        public Builder setStampSignerConfig(SignerConfig stampSignerConfig) {
+            mStampSignerConfig = stampSignerConfig;
+            return this;
+        }
+
         /**
          * Sets whether the APK should be signed using JAR signing (aka v1 signature scheme).
          *
diff --git a/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java b/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java
index 5a9175d..0662ada 100644
--- a/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java
+++ b/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java
@@ -45,8 +45,8 @@
 import com.android.apksig.util.DataSinks;
 import com.android.apksig.util.DataSource;
 import com.android.apksig.util.DataSources;
-
 import com.android.apksig.util.RunnablesExecutor;
+
 import java.io.IOException;
 import java.math.BigInteger;
 import java.nio.BufferUnderflowException;
@@ -93,6 +93,7 @@
           };
     private static final int VERITY_PADDING_BLOCK_ID = 0x42726577;
 
+    public static final int VERSION_SOURCE_STAMP = 0;
     public static final int VERSION_JAR_SIGNATURE_SCHEME = 1;
     public static final int VERSION_APK_SIGNATURE_SCHEME_V2 = 2;
     public static final int VERSION_APK_SIGNATURE_SCHEME_V3 = 3;
diff --git a/src/main/java/com/android/apksig/internal/apk/ContentDigestAlgorithm.java b/src/main/java/com/android/apksig/internal/apk/ContentDigestAlgorithm.java
index b222474..e463743 100644
--- a/src/main/java/com/android/apksig/internal/apk/ContentDigestAlgorithm.java
+++ b/src/main/java/com/android/apksig/internal/apk/ContentDigestAlgorithm.java
@@ -16,28 +16,33 @@
 
 package com.android.apksig.internal.apk;
 
-/**
- * APK Signature Scheme v2 content digest algorithm.
- */
+/** APK Signature Scheme v2 content digest algorithm. */
 public enum ContentDigestAlgorithm {
     /** SHA2-256 over 1 MB chunks. */
-    CHUNKED_SHA256("SHA-256", 256 / 8),
+    CHUNKED_SHA256(1, "SHA-256", 256 / 8),
 
     /** SHA2-512 over 1 MB chunks. */
-    CHUNKED_SHA512("SHA-512", 512 / 8),
+    CHUNKED_SHA512(2, "SHA-512", 512 / 8),
 
     /** SHA2-256 over 4 KB chunks for APK verity. */
-    VERITY_CHUNKED_SHA256("SHA-256", 256 / 8);
+    VERITY_CHUNKED_SHA256(3, "SHA-256", 256 / 8);
 
+    private final int mId;
     private final String mJcaMessageDigestAlgorithm;
     private final int mChunkDigestOutputSizeBytes;
 
     private ContentDigestAlgorithm(
-            String jcaMessageDigestAlgorithm, int chunkDigestOutputSizeBytes) {
+            int id, String jcaMessageDigestAlgorithm, int chunkDigestOutputSizeBytes) {
+        mId = id;
         mJcaMessageDigestAlgorithm = jcaMessageDigestAlgorithm;
         mChunkDigestOutputSizeBytes = chunkDigestOutputSizeBytes;
     }
 
+    /** Returns the ID of the digest algorithm used on the APK. */
+    public int getId() {
+        return mId;
+    }
+
     /**
      * Returns the {@link java.security.MessageDigest} algorithm used for computing digests of
      * chunks by this content digest algorithm.
@@ -46,10 +51,8 @@
         return mJcaMessageDigestAlgorithm;
     }
 
-    /**
-     * Returns the size (in bytes) of the digest of a chunk of content.
-     */
+    /** Returns the size (in bytes) of the digest of a chunk of content. */
     int getChunkDigestOutputSizeBytes() {
         return mChunkDigestOutputSizeBytes;
     }
-}
\ No newline at end of file
+}
diff --git a/src/main/java/com/android/apksig/internal/apk/stamp/SourceStampSigner.java b/src/main/java/com/android/apksig/internal/apk/stamp/SourceStampSigner.java
new file mode 100644
index 0000000..894d6f1
--- /dev/null
+++ b/src/main/java/com/android/apksig/internal/apk/stamp/SourceStampSigner.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.apksig.internal.apk.stamp;
+
+import static com.android.apksig.internal.apk.ApkSigningBlockUtils.encodeAsLengthPrefixedElement;
+import static com.android.apksig.internal.apk.ApkSigningBlockUtils.encodeAsSequenceOfLengthPrefixedElements;
+import static com.android.apksig.internal.apk.ApkSigningBlockUtils.encodeAsSequenceOfLengthPrefixedPairsOfIntAndLengthPrefixedBytes;
+
+import com.android.apksig.internal.apk.ApkSigningBlockUtils;
+import com.android.apksig.internal.apk.ApkSigningBlockUtils.SignerConfig;
+import com.android.apksig.internal.apk.ContentDigestAlgorithm;
+import com.android.apksig.internal.util.Pair;
+
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.security.SignatureException;
+import java.security.cert.CertificateEncodingException;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * SourceStamp signer.
+ *
+ * <p>SourceStamp improves traceability of apps with respect to unauthorized distribution.
+ *
+ * <p>The stamp is part of the APK that is protected by the signing block.
+ *
+ * <p>The APK contents hash is signed using the stamp key, and is saved as part of the signing
+ * block.
+ */
+public abstract class SourceStampSigner {
+
+    public static final int SOURCE_STAMP_BLOCK_ID = 0x2b09189e;
+
+    /** Hidden constructor to prevent instantiation. */
+    private SourceStampSigner() {}
+
+    public static Pair<byte[], Integer> generateSourceStampBlock(
+            SignerConfig sourceStampSignerConfig, Map<ContentDigestAlgorithm, byte[]> digestInfo)
+            throws SignatureException, NoSuchAlgorithmException, InvalidKeyException {
+        if (sourceStampSignerConfig.certificates.isEmpty()) {
+            throw new SignatureException("No certificates configured for signer");
+        }
+
+        List<Pair<Integer, byte[]>> digests =
+                digestInfo.entrySet().stream()
+                        .map(e -> Pair.of(e.getKey().getId(), e.getValue()))
+                        .collect(Collectors.toList());
+
+        SourceStampBlock sourceStampBlock = new SourceStampBlock();
+
+        try {
+            sourceStampBlock.stampCertificate =
+                    sourceStampSignerConfig.certificates.get(0).getEncoded();
+        } catch (CertificateEncodingException e) {
+            throw new SignatureException(
+                    "Retrieving the encoded form of the stamp certificate failed", e);
+        }
+        // TODO: Sort digests
+        sourceStampBlock.digests =
+                encodeAsSequenceOfLengthPrefixedPairsOfIntAndLengthPrefixedBytes(digests);
+        sourceStampBlock.signedDigests =
+                ApkSigningBlockUtils.generateSignaturesOverData(
+                        sourceStampSignerConfig, sourceStampBlock.digests);
+
+        // FORMAT:
+        // * length-prefixed bytes: X.509 certificate (ASN.1 DER encoded)
+        // * length-prefixed sequence of length-prefixed signatures:
+        //   * uint32: signature algorithm ID
+        //   * length-prefixed bytes: signature of signed data
+        byte[] sourceStampSignerBlock =
+                encodeAsSequenceOfLengthPrefixedElements(
+                        new byte[][] {
+                            sourceStampBlock.stampCertificate,
+                            encodeAsSequenceOfLengthPrefixedPairsOfIntAndLengthPrefixedBytes(
+                                    sourceStampBlock.signedDigests),
+                        });
+
+        // FORMAT:
+        // * length-prefixed stamp block.
+        return Pair.of(
+                encodeAsLengthPrefixedElement(sourceStampSignerBlock), SOURCE_STAMP_BLOCK_ID);
+    }
+
+    private static final class SourceStampBlock {
+        public byte[] stampCertificate;
+        public byte[] digests;
+        public List<Pair<Integer, byte[]>> signedDigests;
+    }
+}
diff --git a/src/test/java/com/android/apksig/ApkSignerTest.java b/src/test/java/com/android/apksig/ApkSignerTest.java
index 629882d..fba292c 100644
--- a/src/test/java/com/android/apksig/ApkSignerTest.java
+++ b/src/test/java/com/android/apksig/ApkSignerTest.java
@@ -30,6 +30,7 @@
 import com.android.apksig.apk.ApkUtils;
 import com.android.apksig.internal.apk.ApkSigningBlockUtils;
 import com.android.apksig.internal.apk.SignatureInfo;
+import com.android.apksig.internal.apk.stamp.SourceStampSigner;
 import com.android.apksig.internal.apk.v1.V1SchemeVerifier;
 import com.android.apksig.internal.apk.v2.V2SchemeSigner;
 import com.android.apksig.internal.apk.v3.V3SchemeSigner;
@@ -43,6 +44,7 @@
 import com.android.apksig.util.DataSource;
 import com.android.apksig.util.DataSources;
 import com.android.apksig.util.ReadableDataSink;
+import com.android.apksig.zip.ZipFormatException;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -978,6 +980,109 @@
         assertArrayEquals(expectedStampCertificateDigest, actualStampCertificateDigest);
     }
 
+    @Test
+    public void testSignApk_stampBlock_noStampGenerated() throws Exception {
+        List<ApkSigner.SignerConfig> signersList =
+                Collections.singletonList(
+                        getDefaultSignerConfigFromResources(FIRST_RSA_2048_SIGNER_RESOURCE_NAME));
+
+        DataSource signedApk =
+                sign(
+                        "original.apk",
+                        new ApkSigner.Builder(signersList)
+                                .setV1SigningEnabled(true)
+                                .setV2SigningEnabled(true)
+                                .setV3SigningEnabled(true));
+
+        ApkUtils.ZipSections zipSections = ApkUtils.findZipSections(signedApk);
+        ApkSigningBlockUtils.Result result =
+                new ApkSigningBlockUtils.Result(ApkSigningBlockUtils.VERSION_SOURCE_STAMP);
+        assertThrows(
+                ApkSigningBlockUtils.SignatureNotFoundException.class,
+                () ->
+                        ApkSigningBlockUtils.findSignature(
+                                signedApk,
+                                zipSections,
+                                ApkSigningBlockUtils.VERSION_SOURCE_STAMP,
+                                result));
+    }
+
+    @Test
+    public void testSignApk_stampBlock_whenNoV2V3SignaturePresent() throws Exception {
+        List<ApkSigner.SignerConfig> signersList =
+                Collections.singletonList(
+                        getDefaultSignerConfigFromResources(FIRST_RSA_2048_SIGNER_RESOURCE_NAME));
+        ApkSigner.SignerConfig sourceStampSigner =
+                getDefaultSignerConfigFromResources(SECOND_RSA_2048_SIGNER_RESOURCE_NAME);
+
+        DataSource signedApk =
+                sign(
+                        "original.apk",
+                        new ApkSigner.Builder(signersList)
+                                .setV1SigningEnabled(true)
+                                .setSourceStampSignerConfig(sourceStampSigner));
+
+        ApkUtils.ZipSections zipSections = ApkUtils.findZipSections(signedApk);
+        ApkSigningBlockUtils.Result result =
+                new ApkSigningBlockUtils.Result(ApkSigningBlockUtils.VERSION_SOURCE_STAMP);
+        assertThrows(
+                ApkSigningBlockUtils.SignatureNotFoundException.class,
+                () ->
+                        ApkSigningBlockUtils.findSignature(
+                                signedApk,
+                                zipSections,
+                                ApkSigningBlockUtils.VERSION_SOURCE_STAMP,
+                                result));
+    }
+
+    @Test
+    public void testSignApk_stampBlock_whenV2SignaturePresent() throws Exception {
+        List<ApkSigner.SignerConfig> signersList =
+                Collections.singletonList(
+                        getDefaultSignerConfigFromResources(FIRST_RSA_2048_SIGNER_RESOURCE_NAME));
+        ApkSigner.SignerConfig sourceStampSigner =
+                getDefaultSignerConfigFromResources(SECOND_RSA_2048_SIGNER_RESOURCE_NAME);
+
+        DataSource signedApk =
+                sign(
+                        "original.apk",
+                        new ApkSigner.Builder(signersList)
+                                .setV1SigningEnabled(true)
+                                .setV2SigningEnabled(true)
+                                .setSourceStampSignerConfig(sourceStampSigner));
+
+        SignatureInfo signatureInfo =
+                getSignatureInfoFromApk(
+                        signedApk,
+                        ApkSigningBlockUtils.VERSION_SOURCE_STAMP,
+                        SourceStampSigner.SOURCE_STAMP_BLOCK_ID);
+        assertNotNull(signatureInfo.signatureBlock);
+    }
+
+    @Test
+    public void testSignApk_stampBlock_whenV3SignaturePresent() throws Exception {
+        List<ApkSigner.SignerConfig> signersList =
+                Collections.singletonList(
+                        getDefaultSignerConfigFromResources(FIRST_RSA_2048_SIGNER_RESOURCE_NAME));
+        ApkSigner.SignerConfig sourceStampSigner =
+                getDefaultSignerConfigFromResources(SECOND_RSA_2048_SIGNER_RESOURCE_NAME);
+
+        DataSource signedApk =
+                sign(
+                        "original.apk",
+                        new ApkSigner.Builder(signersList)
+                                .setV1SigningEnabled(true)
+                                .setV3SigningEnabled(true)
+                                .setSourceStampSignerConfig(sourceStampSigner));
+
+        SignatureInfo signatureInfo =
+                getSignatureInfoFromApk(
+                        signedApk,
+                        ApkSigningBlockUtils.VERSION_SOURCE_STAMP,
+                        SourceStampSigner.SOURCE_STAMP_BLOCK_ID);
+        assertNotNull(signatureInfo.signatureBlock);
+    }
+
     private RSAPublicKey getRSAPublicKeyFromSigningBlock(DataSource apk, int signatureVersionId)
             throws Exception {
         int signatureVersionBlockId;
@@ -992,11 +1097,8 @@
                 throw new Exception(
                         "Invalid signature version ID specified: " + signatureVersionId);
         }
-        ApkUtils.ZipSections zipSections = ApkUtils.findZipSections(apk);
-        ApkSigningBlockUtils.Result result = new ApkSigningBlockUtils.Result(signatureVersionId);
         SignatureInfo signatureInfo =
-                ApkSigningBlockUtils.findSignature(
-                        apk, zipSections, signatureVersionBlockId, result);
+                getSignatureInfoFromApk(apk, signatureVersionId, signatureVersionBlockId);
         // FORMAT:
         // * length prefixed sequence of length prefixed signers
         //   * length-prefixed signed data
@@ -1024,8 +1126,17 @@
         // byte indicating the number of padding bits in the public key. Read this first byte to
         // allow parsing the rest of the RSAPublicKey as a sequence.
         subjectPublicKeyBuffer.get();
-        RSAPublicKey rsaPublicKey = Asn1BerParser.parse(subjectPublicKeyBuffer, RSAPublicKey.class);
-        return rsaPublicKey;
+        return Asn1BerParser.parse(subjectPublicKeyBuffer, RSAPublicKey.class);
+    }
+
+    private SignatureInfo getSignatureInfoFromApk(
+            DataSource apk, int signatureVersionId, int signatureVersionBlockId)
+            throws IOException, ZipFormatException,
+                    ApkSigningBlockUtils.SignatureNotFoundException {
+        ApkUtils.ZipSections zipSections = ApkUtils.findZipSections(apk);
+        ApkSigningBlockUtils.Result result = new ApkSigningBlockUtils.Result(signatureVersionId);
+        return ApkSigningBlockUtils.findSignature(
+                apk, zipSections, signatureVersionBlockId, result);
     }
 
     /**