Add SignatureVerifier class to handle public keys from different sources am: 3bd07a5ec6

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Connectivity/+/3375769

Change-Id: I692b47dcb253719d0cdad75c8fc761a006ed0fb7
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyDownloader.java b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyDownloader.java
index d53f007..bd8f7b9 100644
--- a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyDownloader.java
+++ b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyDownloader.java
@@ -15,11 +15,11 @@
  */
 package com.android.server.net.ct;
 
-import android.annotation.NonNull;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
 import android.annotation.RequiresApi;
 import android.app.DownloadManager;
 import android.content.BroadcastReceiver;
-import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
@@ -31,16 +31,12 @@
 
 import com.android.server.net.ct.DownloadHelper.DownloadStatus;
 
+import org.json.JSONException;
+import org.json.JSONObject;
+
 import java.io.IOException;
 import java.io.InputStream;
 import java.security.GeneralSecurityException;
-import java.security.InvalidKeyException;
-import java.security.KeyFactory;
-import java.security.PublicKey;
-import java.security.Signature;
-import java.security.spec.X509EncodedKeySpec;
-import java.util.Base64;
-import java.util.Optional;
 
 /** Helper class to download certificate transparency log files. */
 @RequiresApi(Build.VERSION_CODES.VANILLA_ICE_CREAM)
@@ -51,30 +47,22 @@
     private final Context mContext;
     private final DataStore mDataStore;
     private final DownloadHelper mDownloadHelper;
+    private final SignatureVerifier mSignatureVerifier;
     private final CertificateTransparencyInstaller mInstaller;
 
-    @NonNull private Optional<PublicKey> mPublicKey = Optional.empty();
-
-    @VisibleForTesting
     CertificateTransparencyDownloader(
             Context context,
             DataStore dataStore,
             DownloadHelper downloadHelper,
+            SignatureVerifier signatureVerifier,
             CertificateTransparencyInstaller installer) {
         mContext = context;
+        mSignatureVerifier = signatureVerifier;
         mDataStore = dataStore;
         mDownloadHelper = downloadHelper;
         mInstaller = installer;
     }
 
-    CertificateTransparencyDownloader(Context context, DataStore dataStore) {
-        this(
-                context,
-                dataStore,
-                new DownloadHelper(context),
-                new CertificateTransparencyInstaller());
-    }
-
     void initialize() {
         mInstaller.addCompatibilityVersion(Config.COMPATIBILITY_VERSION);
 
@@ -87,23 +75,14 @@
         }
     }
 
-    void setPublicKey(String publicKey) throws GeneralSecurityException {
-        try {
-            mPublicKey =
-                    Optional.of(
-                            KeyFactory.getInstance("RSA")
-                                    .generatePublic(
-                                            new X509EncodedKeySpec(
-                                                    Base64.getDecoder().decode(publicKey))));
-        } catch (IllegalArgumentException e) {
-            Log.e(TAG, "Invalid public key Base64 encoding", e);
-            mPublicKey = Optional.empty();
+    void startPublicKeyDownload(String publicKeyUrl) {
+        long downloadId = download(publicKeyUrl);
+        if (downloadId == -1) {
+            Log.e(TAG, "Metadata download request failed for " + publicKeyUrl);
+            return;
         }
-    }
-
-    @VisibleForTesting
-    void resetPublicKey() {
-        mPublicKey = Optional.empty();
+        mDataStore.setPropertyLong(Config.PUBLIC_KEY_URL_KEY, downloadId);
+        mDataStore.store();
     }
 
     void startMetadataDownload(String metadataUrl) {
@@ -140,6 +119,11 @@
             return;
         }
 
+        if (isPublicKeyDownloadId(completedId)) {
+            handlePublicKeyDownloadCompleted(completedId);
+            return;
+        }
+
         if (isMetadataDownloadId(completedId)) {
             handleMetadataDownloadCompleted(completedId);
             return;
@@ -150,7 +134,30 @@
             return;
         }
 
-        Log.e(TAG, "Download id " + completedId + " is neither metadata nor content.");
+        Log.i(TAG, "Download id " + completedId + " is not recognized.");
+    }
+
+    private void handlePublicKeyDownloadCompleted(long downloadId) {
+        DownloadStatus status = mDownloadHelper.getDownloadStatus(downloadId);
+        if (!status.isSuccessful()) {
+            handleDownloadFailed(status);
+            return;
+        }
+
+        Uri publicKeyUri = getPublicKeyDownloadUri();
+        if (publicKeyUri == null) {
+            Log.e(TAG, "Invalid public key URI");
+            return;
+        }
+
+        try {
+            mSignatureVerifier.setPublicKeyFrom(publicKeyUri);
+        } catch (GeneralSecurityException | IOException | IllegalArgumentException e) {
+            Log.e(TAG, "Error setting the public Key", e);
+            return;
+        }
+
+        startMetadataDownload(mDataStore.getProperty(Config.METADATA_URL_PENDING));
     }
 
     private void handleMetadataDownloadCompleted(long downloadId) {
@@ -178,7 +185,7 @@
 
         boolean success = false;
         try {
-            success = verify(contentUri, metadataUri);
+            success = mSignatureVerifier.verify(contentUri, metadataUri);
         } catch (IOException | GeneralSecurityException e) {
             Log.e(TAG, "Could not verify new log list", e);
         }
@@ -187,9 +194,16 @@
             return;
         }
 
-        // TODO: validate file content.
+        String version = null;
+        try (InputStream inputStream = mContext.getContentResolver().openInputStream(contentUri)) {
+            version =
+                    new JSONObject(new String(inputStream.readAllBytes(), UTF_8))
+                            .getString("version");
+        } catch (JSONException | IOException e) {
+            Log.e(TAG, "Could not extract version from log list", e);
+            return;
+        }
 
-        String version = mDataStore.getProperty(Config.VERSION_PENDING);
         String contentUrl = mDataStore.getProperty(Config.CONTENT_URL_PENDING);
         String metadataUrl = mDataStore.getProperty(Config.METADATA_URL_PENDING);
         try (InputStream inputStream = mContext.getContentResolver().openInputStream(contentUri)) {
@@ -209,25 +223,10 @@
     }
 
     private void handleDownloadFailed(DownloadStatus status) {
-        Log.e(TAG, "Content download failed with " + status);
+        Log.e(TAG, "Download failed with " + status);
         // TODO(378626065): Report failure via statsd.
     }
 
-    private boolean verify(Uri file, Uri signature) throws IOException, GeneralSecurityException {
-        if (!mPublicKey.isPresent()) {
-            throw new InvalidKeyException("Missing public key for signature verification");
-        }
-        Signature verifier = Signature.getInstance("SHA256withRSA");
-        verifier.initVerify(mPublicKey.get());
-        ContentResolver contentResolver = mContext.getContentResolver();
-
-        try (InputStream fileStream = contentResolver.openInputStream(file);
-                InputStream signatureStream = contentResolver.openInputStream(signature)) {
-            verifier.update(fileStream.readAllBytes());
-            return verifier.verify(signatureStream.readAllBytes());
-        }
-    }
-
     private long download(String url) {
         try {
             return mDownloadHelper.startDownload(url);
@@ -238,6 +237,11 @@
     }
 
     @VisibleForTesting
+    boolean isPublicKeyDownloadId(long downloadId) {
+        return mDataStore.getPropertyLong(Config.PUBLIC_KEY_URL_KEY, -1) == downloadId;
+    }
+
+    @VisibleForTesting
     boolean isMetadataDownloadId(long downloadId) {
         return mDataStore.getPropertyLong(Config.METADATA_URL_KEY, -1) == downloadId;
     }
@@ -247,6 +251,10 @@
         return mDataStore.getPropertyLong(Config.CONTENT_URL_KEY, -1) == downloadId;
     }
 
+    private Uri getPublicKeyDownloadUri() {
+        return mDownloadHelper.getUri(mDataStore.getPropertyLong(Config.PUBLIC_KEY_URL_KEY, -1));
+    }
+
     private Uri getMetadataDownloadUri() {
         return mDownloadHelper.getUri(mDataStore.getPropertyLong(Config.METADATA_URL_KEY, -1));
     }
diff --git a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyFlagsListener.java b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyFlagsListener.java
index 93a7064..f359a2a 100644
--- a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyFlagsListener.java
+++ b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyFlagsListener.java
@@ -32,12 +32,15 @@
     private static final String TAG = "CertificateTransparencyFlagsListener";
 
     private final DataStore mDataStore;
+    private final SignatureVerifier mSignatureVerifier;
     private final CertificateTransparencyDownloader mCertificateTransparencyDownloader;
 
     CertificateTransparencyFlagsListener(
             DataStore dataStore,
+            SignatureVerifier signatureVerifier,
             CertificateTransparencyDownloader certificateTransparencyDownloader) {
         mDataStore = dataStore;
+        mSignatureVerifier = signatureVerifier;
         mCertificateTransparencyDownloader = certificateTransparencyDownloader;
     }
 
@@ -104,8 +107,8 @@
         }
 
         try {
-            mCertificateTransparencyDownloader.setPublicKey(newPublicKey);
-        } catch (GeneralSecurityException e) {
+            mSignatureVerifier.setPublicKey(newPublicKey);
+        } catch (GeneralSecurityException | IllegalArgumentException e) {
             Log.e(TAG, "Error setting the public Key", e);
             return;
         }
diff --git a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyJob.java b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyJob.java
index 6fbf0ba..c5d0413 100644
--- a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyJob.java
+++ b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyJob.java
@@ -24,11 +24,8 @@
 import android.content.IntentFilter;
 import android.os.Build;
 import android.os.SystemClock;
-import android.provider.DeviceConfig;
 import android.util.Log;
 
-import java.util.HashMap;
-
 /** Implementation of the Certificate Transparency job */
 @RequiresApi(Build.VERSION_CODES.VANILLA_ICE_CREAM)
 public class CertificateTransparencyJob extends BroadcastReceiver {
@@ -40,18 +37,14 @@
     private final Context mContext;
     private final DataStore mDataStore;
     private final CertificateTransparencyDownloader mCertificateTransparencyDownloader;
-    // TODO(b/374692404): remove dependency to flags.
-    private final CertificateTransparencyFlagsListener mFlagsListener;
     private final AlarmManager mAlarmManager;
 
     /** Creates a new {@link CertificateTransparencyJob} object. */
     public CertificateTransparencyJob(
             Context context,
             DataStore dataStore,
-            CertificateTransparencyDownloader certificateTransparencyDownloader,
-            CertificateTransparencyFlagsListener flagsListener) {
+            CertificateTransparencyDownloader certificateTransparencyDownloader) {
         mContext = context;
-        mFlagsListener = flagsListener;
         mDataStore = dataStore;
         mCertificateTransparencyDownloader = certificateTransparencyDownloader;
         mAlarmManager = context.getSystemService(AlarmManager.class);
@@ -81,7 +74,15 @@
             Log.w(TAG, "Received unexpected broadcast with action " + intent);
             return;
         }
-        mFlagsListener.onPropertiesChanged(
-                new DeviceConfig.Properties(Config.NAMESPACE_NETWORK_SECURITY, new HashMap<>()));
+        if (Config.DEBUG) {
+            Log.d(TAG, "Starting CT daily job.");
+        }
+
+        mDataStore.setProperty(Config.CONTENT_URL_PENDING, Config.URL_LOG_LIST);
+        mDataStore.setProperty(Config.METADATA_URL_PENDING, Config.URL_SIGNATURE);
+        mDataStore.setProperty(Config.PUBLIC_KEY_URL_PENDING, Config.URL_PUBLIC_KEY);
+        mDataStore.store();
+
+        mCertificateTransparencyDownloader.startPublicKeyDownload(Config.URL_PUBLIC_KEY);
     }
 }
diff --git a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyService.java b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyService.java
index ac55e44..92b2b09 100644
--- a/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyService.java
+++ b/networksecurity/service/src/com/android/server/net/ct/CertificateTransparencyService.java
@@ -28,8 +28,6 @@
 @RequiresApi(Build.VERSION_CODES.VANILLA_ICE_CREAM)
 public class CertificateTransparencyService extends ICertificateTransparencyManager.Stub {
 
-    private final DataStore mDataStore;
-    private final CertificateTransparencyDownloader mCertificateTransparencyDownloader;
     private final CertificateTransparencyFlagsListener mFlagsListener;
     private final CertificateTransparencyJob mCertificateTransparencyJob;
 
@@ -44,15 +42,21 @@
 
     /** Creates a new {@link CertificateTransparencyService} object. */
     public CertificateTransparencyService(Context context) {
-        mDataStore = new DataStore(Config.PREFERENCES_FILE);
-        mCertificateTransparencyDownloader =
-                new CertificateTransparencyDownloader(context, mDataStore);
+        DataStore dataStore = new DataStore(Config.PREFERENCES_FILE);
+        DownloadHelper downloadHelper = new DownloadHelper(context);
+        SignatureVerifier signatureVerifier = new SignatureVerifier(context);
+        CertificateTransparencyDownloader downloader =
+                new CertificateTransparencyDownloader(
+                        context,
+                        dataStore,
+                        downloadHelper,
+                        signatureVerifier,
+                        new CertificateTransparencyInstaller());
+
         mFlagsListener =
-                new CertificateTransparencyFlagsListener(
-                        mDataStore, mCertificateTransparencyDownloader);
+                new CertificateTransparencyFlagsListener(dataStore, signatureVerifier, downloader);
         mCertificateTransparencyJob =
-                new CertificateTransparencyJob(
-                        context, mDataStore, mCertificateTransparencyDownloader, mFlagsListener);
+                new CertificateTransparencyJob(context, dataStore, downloader);
     }
 
     /**
diff --git a/networksecurity/service/src/com/android/server/net/ct/Config.java b/networksecurity/service/src/com/android/server/net/ct/Config.java
index 242f13a..ae30f3a 100644
--- a/networksecurity/service/src/com/android/server/net/ct/Config.java
+++ b/networksecurity/service/src/com/android/server/net/ct/Config.java
@@ -55,4 +55,13 @@
     static final String METADATA_URL_PENDING = "metadata_url_pending";
     static final String METADATA_URL = "metadata_url";
     static final String METADATA_URL_KEY = "metadata_url_key";
+    static final String PUBLIC_KEY_URL_PENDING = "public_key_url_pending";
+    static final String PUBLIC_KEY_URL = "public_key_url";
+    static final String PUBLIC_KEY_URL_KEY = "public_key_url_key";
+
+    // URLs
+    static final String URL_BASE = "https://www.gstatic.com/android/certificate_transparency/";
+    static final String URL_LOG_LIST = URL_BASE + "log_list.json";
+    static final String URL_SIGNATURE = URL_BASE + "log_list.sig";
+    static final String URL_PUBLIC_KEY = URL_BASE + "log_list.pub";
 }
diff --git a/networksecurity/service/src/com/android/server/net/ct/SignatureVerifier.java b/networksecurity/service/src/com/android/server/net/ct/SignatureVerifier.java
new file mode 100644
index 0000000..0b775ca
--- /dev/null
+++ b/networksecurity/service/src/com/android/server/net/ct/SignatureVerifier.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2024 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.server.net.ct;
+
+import android.annotation.NonNull;
+import android.annotation.RequiresApi;
+import android.content.ContentResolver;
+import android.content.Context;
+import android.net.Uri;
+import android.os.Build;
+
+import androidx.annotation.VisibleForTesting;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.security.GeneralSecurityException;
+import java.security.InvalidKeyException;
+import java.security.KeyFactory;
+import java.security.PublicKey;
+import java.security.Signature;
+import java.security.spec.X509EncodedKeySpec;
+import java.util.Base64;
+import java.util.Optional;
+
+/** Verifier of the log list signature. */
+@RequiresApi(Build.VERSION_CODES.VANILLA_ICE_CREAM)
+public class SignatureVerifier {
+
+    private final Context mContext;
+
+    @NonNull private Optional<PublicKey> mPublicKey = Optional.empty();
+
+    public SignatureVerifier(Context context) {
+        mContext = context;
+    }
+
+    @VisibleForTesting
+    Optional<PublicKey> getPublicKey() {
+        return mPublicKey;
+    }
+
+    void resetPublicKey() {
+        mPublicKey = Optional.empty();
+    }
+
+    void setPublicKeyFrom(Uri file) throws GeneralSecurityException, IOException {
+        try (InputStream fileStream = mContext.getContentResolver().openInputStream(file)) {
+            setPublicKey(new String(fileStream.readAllBytes()));
+        }
+    }
+
+    void setPublicKey(String publicKey) throws GeneralSecurityException {
+        setPublicKey(
+                KeyFactory.getInstance("RSA")
+                        .generatePublic(
+                                new X509EncodedKeySpec(Base64.getDecoder().decode(publicKey))));
+    }
+
+    @VisibleForTesting
+    void setPublicKey(PublicKey publicKey) {
+        mPublicKey = Optional.of(publicKey);
+    }
+
+    boolean verify(Uri file, Uri signature) throws GeneralSecurityException, IOException {
+        if (!mPublicKey.isPresent()) {
+            throw new InvalidKeyException("Missing public key for signature verification");
+        }
+        Signature verifier = Signature.getInstance("SHA256withRSA");
+        verifier.initVerify(mPublicKey.get());
+        ContentResolver contentResolver = mContext.getContentResolver();
+
+        try (InputStream fileStream = contentResolver.openInputStream(file);
+                InputStream signatureStream = contentResolver.openInputStream(signature)) {
+            verifier.update(fileStream.readAllBytes());
+            return verifier.verify(signatureStream.readAllBytes());
+        }
+    }
+}
diff --git a/networksecurity/tests/unit/src/com/android/server/net/ct/CertificateTransparencyDownloaderTest.java b/networksecurity/tests/unit/src/com/android/server/net/ct/CertificateTransparencyDownloaderTest.java
index fb55295..87d75e6 100644
--- a/networksecurity/tests/unit/src/com/android/server/net/ct/CertificateTransparencyDownloaderTest.java
+++ b/networksecurity/tests/unit/src/com/android/server/net/ct/CertificateTransparencyDownloaderTest.java
@@ -18,12 +18,14 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
+
 import android.app.DownloadManager;
 import android.content.Context;
 import android.content.Intent;
@@ -33,6 +35,8 @@
 
 import com.android.server.net.ct.DownloadHelper.DownloadStatus;
 
+import org.json.JSONException;
+import org.json.JSONObject;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -67,8 +71,11 @@
     private Context mContext;
     private File mTempFile;
     private DataStore mDataStore;
+    private SignatureVerifier mSignatureVerifier;
     private CertificateTransparencyDownloader mCertificateTransparencyDownloader;
 
+    private long mNextDownloadId = 666;
+
     @Before
     public void setUp() throws IOException, GeneralSecurityException {
         MockitoAnnotations.initMocks(this);
@@ -82,23 +89,37 @@
         mTempFile = File.createTempFile("datastore-test", ".properties");
         mDataStore = new DataStore(mTempFile);
         mDataStore.load();
+        mSignatureVerifier = new SignatureVerifier(mContext);
 
         mCertificateTransparencyDownloader =
                 new CertificateTransparencyDownloader(
-                        mContext, mDataStore, mDownloadHelper, mCertificateTransparencyInstaller);
+                        mContext,
+                        mDataStore,
+                        mDownloadHelper,
+                        mSignatureVerifier,
+                        mCertificateTransparencyInstaller);
     }
 
     @After
     public void tearDown() {
         mTempFile.delete();
-        mCertificateTransparencyDownloader.resetPublicKey();
+        mSignatureVerifier.resetPublicKey();
+    }
+
+    @Test
+    public void testDownloader_startPublicKeyDownload() {
+        String publicKeyUrl = "http://test-public-key.org";
+        long downloadId = preparePublicKeyDownload(publicKeyUrl);
+
+        assertThat(mCertificateTransparencyDownloader.isPublicKeyDownloadId(downloadId)).isFalse();
+        mCertificateTransparencyDownloader.startPublicKeyDownload(publicKeyUrl);
+        assertThat(mCertificateTransparencyDownloader.isPublicKeyDownloadId(downloadId)).isTrue();
     }
 
     @Test
     public void testDownloader_startMetadataDownload() {
         String metadataUrl = "http://test-metadata.org";
-        long downloadId = 666;
-        when(mDownloadHelper.startDownload(metadataUrl)).thenReturn(downloadId);
+        long downloadId = prepareMetadataDownload(metadataUrl);
 
         assertThat(mCertificateTransparencyDownloader.isMetadataDownloadId(downloadId)).isFalse();
         mCertificateTransparencyDownloader.startMetadataDownload(metadataUrl);
@@ -108,8 +129,7 @@
     @Test
     public void testDownloader_startContentDownload() {
         String contentUrl = "http://test-content.org";
-        long downloadId = 666;
-        when(mDownloadHelper.startDownload(contentUrl)).thenReturn(downloadId);
+        long downloadId = prepareContentDownload(contentUrl);
 
         assertThat(mCertificateTransparencyDownloader.isContentDownloadId(downloadId)).isFalse();
         mCertificateTransparencyDownloader.startContentDownload(contentUrl);
@@ -117,16 +137,65 @@
     }
 
     @Test
-    public void testDownloader_metadataDownloadSuccess_startContentDownload() {
-        long metadataId = 123;
-        mDataStore.setPropertyLong(Config.METADATA_URL_KEY, metadataId);
-        when(mDownloadHelper.getDownloadStatus(metadataId))
-                .thenReturn(makeSuccessfulDownloadStatus(metadataId));
-        long contentId = 666;
-        String contentUrl = "http://test-content.org";
-        mDataStore.setProperty(Config.CONTENT_URL_PENDING, contentUrl);
-        when(mDownloadHelper.startDownload(contentUrl)).thenReturn(contentId);
+    public void testDownloader_publicKeyDownloadSuccess_updatePublicKey_startMetadataDownload()
+            throws Exception {
+        long publicKeyId = prepareSuccessfulPublicKeyDownload(writePublicKeyToFile(mPublicKey));
+        long metadataId = prepareMetadataDownload("http://test-metadata.org");
 
+        assertThat(mSignatureVerifier.getPublicKey()).isEmpty();
+        assertThat(mCertificateTransparencyDownloader.isMetadataDownloadId(metadataId)).isFalse();
+        mCertificateTransparencyDownloader.onReceive(
+                mContext, makeDownloadCompleteIntent(publicKeyId));
+
+        assertThat(mSignatureVerifier.getPublicKey()).hasValue(mPublicKey);
+        assertThat(mCertificateTransparencyDownloader.isMetadataDownloadId(metadataId)).isTrue();
+    }
+
+    @Test
+    public void
+            testDownloader_publicKeyDownloadSuccess_updatePublicKeyFail_doNotStartMetadataDownload()
+                    throws Exception {
+        long publicKeyId =
+                prepareSuccessfulPublicKeyDownload(
+                        writeToFile("i_am_not_a_base64_encoded_public_key".getBytes()));
+        long metadataId = prepareMetadataDownload("http://test-metadata.org");
+
+        assertThat(mSignatureVerifier.getPublicKey()).isEmpty();
+        assertThat(mCertificateTransparencyDownloader.isMetadataDownloadId(metadataId)).isFalse();
+        mCertificateTransparencyDownloader.onReceive(
+                mContext, makeDownloadCompleteIntent(publicKeyId));
+
+        assertThat(mSignatureVerifier.getPublicKey()).isEmpty();
+        assertThat(mCertificateTransparencyDownloader.isMetadataDownloadId(metadataId)).isFalse();
+        verify(mDownloadHelper, never()).startDownload(anyString());
+    }
+
+    @Test
+    public void testDownloader_publicKeyDownloadFail_doNotUpdatePublicKey() throws Exception {
+        long publicKeyId =
+                prepareFailedPublicKeyDownload(
+                        // Failure cases where we give up on the download.
+                        DownloadManager.ERROR_INSUFFICIENT_SPACE,
+                        DownloadManager.ERROR_HTTP_DATA_ERROR);
+        Intent downloadCompleteIntent = makeDownloadCompleteIntent(publicKeyId);
+        long metadataId = prepareMetadataDownload("http://test-metadata.org");
+
+        assertThat(mSignatureVerifier.getPublicKey()).isEmpty();
+        assertThat(mCertificateTransparencyDownloader.isMetadataDownloadId(metadataId)).isFalse();
+        mCertificateTransparencyDownloader.onReceive(mContext, downloadCompleteIntent);
+        mCertificateTransparencyDownloader.onReceive(mContext, downloadCompleteIntent);
+
+        assertThat(mSignatureVerifier.getPublicKey()).isEmpty();
+        assertThat(mCertificateTransparencyDownloader.isMetadataDownloadId(metadataId)).isFalse();
+        verify(mDownloadHelper, never()).startDownload(anyString());
+    }
+
+    @Test
+    public void testDownloader_metadataDownloadSuccess_startContentDownload() {
+        long metadataId = prepareSuccessfulMetadataDownload(new File("log_list.sig"));
+        long contentId = prepareContentDownload("http://test-content.org");
+
+        assertThat(mCertificateTransparencyDownloader.isContentDownloadId(contentId)).isFalse();
         mCertificateTransparencyDownloader.onReceive(
                 mContext, makeDownloadCompleteIntent(metadataId));
 
@@ -135,162 +204,270 @@
 
     @Test
     public void testDownloader_metadataDownloadFail_doNotStartContentDownload() {
-        long metadataId = 123;
-        mDataStore.setPropertyLong(Config.METADATA_URL_KEY, metadataId);
-        String contentUrl = "http://test-content.org";
-        mDataStore.setProperty(Config.CONTENT_URL_PENDING, contentUrl);
+        long metadataId =
+                prepareFailedMetadataDownload(
+                        // Failure cases where we give up on the download.
+                        DownloadManager.ERROR_INSUFFICIENT_SPACE,
+                        DownloadManager.ERROR_HTTP_DATA_ERROR);
         Intent downloadCompleteIntent = makeDownloadCompleteIntent(metadataId);
-        // In all these failure cases we give up on the download.
-        when(mDownloadHelper.getDownloadStatus(metadataId))
-                .thenReturn(
-                        makeHttpErrorDownloadStatus(metadataId),
-                        makeStorageErrorDownloadStatus(metadataId));
+        long contentId = prepareContentDownload("http://test-content.org");
 
+        assertThat(mCertificateTransparencyDownloader.isContentDownloadId(contentId)).isFalse();
         mCertificateTransparencyDownloader.onReceive(mContext, downloadCompleteIntent);
         mCertificateTransparencyDownloader.onReceive(mContext, downloadCompleteIntent);
 
-        verify(mDownloadHelper, never()).startDownload(contentUrl);
+        assertThat(mCertificateTransparencyDownloader.isContentDownloadId(contentId)).isFalse();
+        verify(mDownloadHelper, never()).startDownload(anyString());
     }
 
     @Test
     public void testDownloader_contentDownloadSuccess_installSuccess_updateDataStore()
             throws Exception {
-        String version = "456";
-        long contentId = 666;
-        File logListFile = File.createTempFile("log_list", "json");
-        Uri contentUri = Uri.fromFile(logListFile);
-        long metadataId = 123;
+        String newVersion = "456";
+        File logListFile = makeLogListFile(newVersion);
         File metadataFile = sign(logListFile);
-        Uri metadataUri = Uri.fromFile(metadataFile);
-        mCertificateTransparencyDownloader.setPublicKey(
-                Base64.getEncoder().encodeToString(mPublicKey.getEncoded()));
-        setUpContentDownloadCompleteSuccessful(
-                version, metadataId, metadataUri, contentId, contentUri);
+        mSignatureVerifier.setPublicKey(mPublicKey);
+        prepareSuccessfulMetadataDownload(metadataFile);
+        long contentId = prepareSuccessfulContentDownload(logListFile);
         when(mCertificateTransparencyInstaller.install(
-                        eq(Config.COMPATIBILITY_VERSION), any(), eq(version)))
+                        eq(Config.COMPATIBILITY_VERSION), any(), anyString()))
                 .thenReturn(true);
 
-        assertThat(mDataStore.getProperty(Config.VERSION)).isNull();
-        assertThat(mDataStore.getProperty(Config.CONTENT_URL)).isNull();
-        assertThat(mDataStore.getProperty(Config.METADATA_URL)).isNull();
+        assertNoVersionIsInstalled();
         mCertificateTransparencyDownloader.onReceive(
                 mContext, makeDownloadCompleteIntent(contentId));
 
-        verify(mCertificateTransparencyInstaller, times(1))
-                .install(eq(Config.COMPATIBILITY_VERSION), any(), eq(version));
-        assertThat(mDataStore.getProperty(Config.VERSION)).isEqualTo(version);
-        assertThat(mDataStore.getProperty(Config.CONTENT_URL)).isEqualTo(contentUri.toString());
-        assertThat(mDataStore.getProperty(Config.METADATA_URL)).isEqualTo(metadataUri.toString());
+        assertInstallSuccessful(newVersion);
     }
 
     @Test
     public void testDownloader_contentDownloadFail_doNotInstall() throws Exception {
-        mDataStore.setProperty(Config.VERSION_PENDING, "123");
-        long contentId = 666;
+        long contentId =
+                prepareFailedContentDownload(
+                        // Failure cases where we give up on the download.
+                        DownloadManager.ERROR_INSUFFICIENT_SPACE,
+                        DownloadManager.ERROR_HTTP_DATA_ERROR);
         Intent downloadCompleteIntent = makeDownloadCompleteIntent(contentId);
-        // In all these failure cases we give up on the download.
-        when(mDownloadHelper.getDownloadStatus(contentId))
-                .thenReturn(
-                        makeHttpErrorDownloadStatus(contentId),
-                        makeStorageErrorDownloadStatus(contentId));
 
         mCertificateTransparencyDownloader.onReceive(mContext, downloadCompleteIntent);
         mCertificateTransparencyDownloader.onReceive(mContext, downloadCompleteIntent);
-        mCertificateTransparencyDownloader.onReceive(mContext, downloadCompleteIntent);
 
         verify(mCertificateTransparencyInstaller, never()).install(any(), any(), any());
-        assertThat(mDataStore.getProperty(Config.VERSION)).isNull();
-        assertThat(mDataStore.getProperty(Config.CONTENT_URL)).isNull();
-        assertThat(mDataStore.getProperty(Config.METADATA_URL)).isNull();
+        assertNoVersionIsInstalled();
     }
 
     @Test
     public void testDownloader_contentDownloadSuccess_installFail_doNotUpdateDataStore()
             throws Exception {
-        String version = "456";
-        long contentId = 666;
-        File logListFile = File.createTempFile("log_list", "json");
-        Uri contentUri = Uri.fromFile(logListFile);
-        long metadataId = 123;
+        File logListFile = makeLogListFile("456");
         File metadataFile = sign(logListFile);
-        Uri metadataUri = Uri.fromFile(metadataFile);
-        setUpContentDownloadCompleteSuccessful(
-                version, metadataId, metadataUri, contentId, contentUri);
+        mSignatureVerifier.setPublicKey(mPublicKey);
+        prepareSuccessfulMetadataDownload(metadataFile);
+        long contentId = prepareSuccessfulContentDownload(logListFile);
         when(mCertificateTransparencyInstaller.install(
-                        eq(Config.COMPATIBILITY_VERSION), any(), eq(version)))
+                        eq(Config.COMPATIBILITY_VERSION), any(), anyString()))
                 .thenReturn(false);
 
+        assertNoVersionIsInstalled();
         mCertificateTransparencyDownloader.onReceive(
                 mContext, makeDownloadCompleteIntent(contentId));
 
-        assertThat(mDataStore.getProperty(Config.VERSION)).isNull();
-        assertThat(mDataStore.getProperty(Config.CONTENT_URL)).isNull();
-        assertThat(mDataStore.getProperty(Config.METADATA_URL)).isNull();
+        assertNoVersionIsInstalled();
     }
 
     @Test
     public void testDownloader_contentDownloadSuccess_verificationFail_doNotInstall()
-            throws IOException {
-        String version = "456";
-        long contentId = 666;
-        Uri contentUri = Uri.fromFile(File.createTempFile("log_list", "json"));
-        long metadataId = 123;
-        Uri metadataUri = Uri.fromFile(File.createTempFile("log_list-wrong_metadata", "sig"));
-        setUpContentDownloadCompleteSuccessful(
-                version, metadataId, metadataUri, contentId, contentUri);
+            throws Exception {
+        File logListFile = makeLogListFile("456");
+        File metadataFile = File.createTempFile("log_list-wrong_metadata", "sig");
+        mSignatureVerifier.setPublicKey(mPublicKey);
+        prepareSuccessfulMetadataDownload(metadataFile);
+        long contentId = prepareSuccessfulContentDownload(logListFile);
 
+        assertNoVersionIsInstalled();
         mCertificateTransparencyDownloader.onReceive(
                 mContext, makeDownloadCompleteIntent(contentId));
 
         verify(mCertificateTransparencyInstaller, never())
-                .install(eq(Config.COMPATIBILITY_VERSION), any(), eq(version));
-        assertThat(mDataStore.getProperty(Config.VERSION)).isNull();
-        assertThat(mDataStore.getProperty(Config.CONTENT_URL)).isNull();
-        assertThat(mDataStore.getProperty(Config.METADATA_URL)).isNull();
+                .install(eq(Config.COMPATIBILITY_VERSION), any(), anyString());
+        assertNoVersionIsInstalled();
     }
 
     @Test
     public void testDownloader_contentDownloadSuccess_missingVerificationPublicKey_doNotInstall()
             throws Exception {
-        String version = "456";
-        long contentId = 666;
-        File logListFile = File.createTempFile("log_list", "json");
-        Uri contentUri = Uri.fromFile(logListFile);
-        long metadataId = 123;
+        File logListFile = makeLogListFile("456");
         File metadataFile = sign(logListFile);
-        Uri metadataUri = Uri.fromFile(metadataFile);
-        setUpContentDownloadCompleteSuccessful(
-                version, metadataId, metadataUri, contentId, contentUri);
+        mSignatureVerifier.resetPublicKey();
+        prepareSuccessfulMetadataDownload(metadataFile);
+        long contentId = prepareSuccessfulContentDownload(logListFile);
 
+        assertNoVersionIsInstalled();
         mCertificateTransparencyDownloader.onReceive(
                 mContext, makeDownloadCompleteIntent(contentId));
 
         verify(mCertificateTransparencyInstaller, never())
-                .install(eq(Config.COMPATIBILITY_VERSION), any(), eq(version));
+                .install(eq(Config.COMPATIBILITY_VERSION), any(), anyString());
+        assertNoVersionIsInstalled();
+    }
+
+    @Test
+    public void testDownloader_endToEndSuccess_installNewVersion() throws Exception {
+        String newVersion = "456";
+        File logListFile = makeLogListFile(newVersion);
+        File metadataFile = sign(logListFile);
+        File publicKeyFile = writePublicKeyToFile(mPublicKey);
+
+        assertNoVersionIsInstalled();
+
+        // 1. Start download of public key.
+        String publicKeyUrl = "http://test-public-key.org";
+        long publicKeyId = preparePublicKeyDownload(publicKeyUrl);
+
+        mCertificateTransparencyDownloader.startPublicKeyDownload(publicKeyUrl);
+
+        // 2. On successful public key download, set the key and start the metatadata download.
+        setSuccessfulDownload(publicKeyId, publicKeyFile);
+        long metadataId = prepareMetadataDownload("http://test-metadata.org");
+
+        mCertificateTransparencyDownloader.onReceive(
+                mContext, makeDownloadCompleteIntent(publicKeyId));
+
+        // 3. On successful metadata download, start the content download.
+        setSuccessfulDownload(metadataId, metadataFile);
+        long contentId = prepareContentDownload("http://test-content.org");
+
+        mCertificateTransparencyDownloader.onReceive(
+                mContext, makeDownloadCompleteIntent(metadataId));
+
+        // 4. On successful content download, verify the signature and install the new version.
+        setSuccessfulDownload(contentId, logListFile);
+        when(mCertificateTransparencyInstaller.install(
+                        eq(Config.COMPATIBILITY_VERSION), any(), anyString()))
+                .thenReturn(true);
+
+        mCertificateTransparencyDownloader.onReceive(
+                mContext, makeDownloadCompleteIntent(contentId));
+
+        assertInstallSuccessful(newVersion);
+    }
+
+    private void assertNoVersionIsInstalled() {
         assertThat(mDataStore.getProperty(Config.VERSION)).isNull();
         assertThat(mDataStore.getProperty(Config.CONTENT_URL)).isNull();
         assertThat(mDataStore.getProperty(Config.METADATA_URL)).isNull();
     }
 
+    private void assertInstallSuccessful(String version) {
+        assertThat(mDataStore.getProperty(Config.VERSION)).isEqualTo(version);
+        assertThat(mDataStore.getProperty(Config.CONTENT_URL))
+                .isEqualTo(mDataStore.getProperty(Config.CONTENT_URL_PENDING));
+        assertThat(mDataStore.getProperty(Config.METADATA_URL))
+                .isEqualTo(mDataStore.getProperty(Config.METADATA_URL_PENDING));
+    }
+
     private Intent makeDownloadCompleteIntent(long downloadId) {
         return new Intent(DownloadManager.ACTION_DOWNLOAD_COMPLETE)
                 .putExtra(DownloadManager.EXTRA_DOWNLOAD_ID, downloadId);
     }
 
-    private void setUpContentDownloadCompleteSuccessful(
-            String version, long metadataId, Uri metadataUri, long contentId, Uri contentUri)
-            throws IOException {
-        mDataStore.setProperty(Config.VERSION_PENDING, version);
+    private long prepareDownloadId(String url) {
+        long downloadId = mNextDownloadId++;
+        when(mDownloadHelper.startDownload(url)).thenReturn(downloadId);
+        return downloadId;
+    }
 
-        mDataStore.setPropertyLong(Config.METADATA_URL_KEY, metadataId);
-        mDataStore.setProperty(Config.METADATA_URL_PENDING, metadataUri.toString());
-        when(mDownloadHelper.getUri(metadataId)).thenReturn(metadataUri);
+    private long preparePublicKeyDownload(String url) {
+        long downloadId = prepareDownloadId(url);
+        mDataStore.setProperty(Config.PUBLIC_KEY_URL_PENDING, url);
+        return downloadId;
+    }
 
-        mDataStore.setPropertyLong(Config.CONTENT_URL_KEY, contentId);
-        mDataStore.setProperty(Config.CONTENT_URL_PENDING, contentUri.toString());
-        when(mDownloadHelper.getDownloadStatus(contentId))
-                .thenReturn(makeSuccessfulDownloadStatus(contentId));
-        when(mDownloadHelper.getUri(contentId)).thenReturn(contentUri);
+    private long prepareMetadataDownload(String url) {
+        long downloadId = prepareDownloadId(url);
+        mDataStore.setProperty(Config.METADATA_URL_PENDING, url);
+        return downloadId;
+    }
+
+    private long prepareContentDownload(String url) {
+        long downloadId = prepareDownloadId(url);
+        mDataStore.setProperty(Config.CONTENT_URL_PENDING, url);
+        return downloadId;
+    }
+
+    private long prepareSuccessfulDownload(String propertyKey) {
+        long downloadId = mNextDownloadId++;
+        mDataStore.setPropertyLong(propertyKey, downloadId);
+        when(mDownloadHelper.getDownloadStatus(downloadId))
+                .thenReturn(makeSuccessfulDownloadStatus(downloadId));
+        return downloadId;
+    }
+
+    private long prepareSuccessfulDownload(String propertyKey, File file) {
+        long downloadId = prepareSuccessfulDownload(propertyKey);
+        when(mDownloadHelper.getUri(downloadId)).thenReturn(Uri.fromFile(file));
+        return downloadId;
+    }
+
+    private long prepareSuccessfulPublicKeyDownload(File file) {
+        long downloadId = prepareSuccessfulDownload(Config.PUBLIC_KEY_URL_KEY, file);
+        mDataStore.setProperty(
+                Config.METADATA_URL_PENDING, "http://public-key-was-downloaded-here.org");
+        return downloadId;
+    }
+
+    private long prepareSuccessfulMetadataDownload(File file) {
+        long downloadId = prepareSuccessfulDownload(Config.METADATA_URL_KEY, file);
+        mDataStore.setProperty(
+                Config.METADATA_URL_PENDING, "http://metadata-was-downloaded-here.org");
+        return downloadId;
+    }
+
+    private long prepareSuccessfulContentDownload(File file) {
+        long downloadId = prepareSuccessfulDownload(Config.CONTENT_URL_KEY, file);
+        mDataStore.setProperty(
+                Config.CONTENT_URL_PENDING, "http://content-was-downloaded-here.org");
+        return downloadId;
+    }
+
+    private void setSuccessfulDownload(long downloadId, File file) {
+        when(mDownloadHelper.getDownloadStatus(downloadId))
+                .thenReturn(makeSuccessfulDownloadStatus(downloadId));
+        when(mDownloadHelper.getUri(downloadId)).thenReturn(Uri.fromFile(file));
+    }
+
+    private long prepareFailedDownload(String propertyKey, int... downloadManagerErrors) {
+        long downloadId = mNextDownloadId++;
+        mDataStore.setPropertyLong(propertyKey, downloadId);
+        DownloadStatus firstError =
+                DownloadStatus.builder()
+                        .setDownloadId(downloadId)
+                        .setStatus(DownloadManager.STATUS_FAILED)
+                        .setReason(downloadManagerErrors[0])
+                        .build();
+        DownloadStatus[] otherErrors = new DownloadStatus[downloadManagerErrors.length - 1];
+        for (int i = 1; i < downloadManagerErrors.length; i++) {
+            otherErrors[i - 1] =
+                    DownloadStatus.builder()
+                            .setDownloadId(downloadId)
+                            .setStatus(DownloadManager.STATUS_FAILED)
+                            .setReason(downloadManagerErrors[i])
+                            .build();
+        }
+        when(mDownloadHelper.getDownloadStatus(downloadId)).thenReturn(firstError, otherErrors);
+        return downloadId;
+    }
+
+    private long prepareFailedPublicKeyDownload(int... downloadManagerErrors) {
+        return prepareFailedDownload(Config.PUBLIC_KEY_URL_KEY, downloadManagerErrors);
+    }
+
+    private long prepareFailedMetadataDownload(int... downloadManagerErrors) {
+        return prepareFailedDownload(Config.METADATA_URL_KEY, downloadManagerErrors);
+    }
+
+    private long prepareFailedContentDownload(int... downloadManagerErrors) {
+        return prepareFailedDownload(Config.CONTENT_URL_KEY, downloadManagerErrors);
     }
 
     private DownloadStatus makeSuccessfulDownloadStatus(long downloadId) {
@@ -300,20 +477,29 @@
                 .build();
     }
 
-    private DownloadStatus makeStorageErrorDownloadStatus(long downloadId) {
-        return DownloadStatus.builder()
-                .setDownloadId(downloadId)
-                .setStatus(DownloadManager.STATUS_FAILED)
-                .setReason(DownloadManager.ERROR_INSUFFICIENT_SPACE)
-                .build();
+    private File writePublicKeyToFile(PublicKey publicKey)
+            throws IOException, GeneralSecurityException {
+        return writeToFile(Base64.getEncoder().encode(publicKey.getEncoded()));
     }
 
-    private DownloadStatus makeHttpErrorDownloadStatus(long downloadId) {
-        return DownloadStatus.builder()
-                .setDownloadId(downloadId)
-                .setStatus(DownloadManager.STATUS_FAILED)
-                .setReason(DownloadManager.ERROR_HTTP_DATA_ERROR)
-                .build();
+    private File writeToFile(byte[] bytes) throws IOException, GeneralSecurityException {
+        File file = File.createTempFile("temp_file", "tmp");
+
+        try (OutputStream outputStream = new FileOutputStream(file)) {
+            outputStream.write(bytes);
+        }
+
+        return file;
+    }
+
+    private File makeLogListFile(String version) throws IOException, JSONException {
+        File logListFile = File.createTempFile("log_list", "json");
+
+        try (OutputStream outputStream = new FileOutputStream(logListFile)) {
+            outputStream.write(new JSONObject().put("version", version).toString().getBytes(UTF_8));
+        }
+
+        return logListFile;
     }
 
     private File sign(File file) throws IOException, GeneralSecurityException {