WifiConfigStore: Store integrity data in same file

am: 952fc8f58d

Change-Id: I67bed9797c92c5935c40612b387846c5e8c663bc

6 files changed, 560 insertions, 162 deletions

diff --git a/service/java/com/android/server/wifi/WifiConfigManager.java b/service/java/com/android/server/wifi/WifiConfigManager.java
index bda1eb7d2..b21893404 100644
--- a/service/java/com/android/server/wifi/WifiConfigManager.java
+++ b/service/java/com/android/server/wifi/WifiConfigManager.java
@@ -3104,7 +3104,7 @@ public class WifiConfigManager {
         }
         try {
             mWifiConfigStore.read();
-        } catch (IOException e) {
+        } catch (IOException | IllegalStateException e) {
             Log.wtf(TAG, "Reading from new store failed. All saved networks are lost!", e);
             return false;
         } catch (XmlPullParserException e) {
@@ -3139,7 +3139,7 @@ public class WifiConfigManager {
                 return false;
             }
             mWifiConfigStore.switchUserStoresAndRead(userStoreFiles);
-        } catch (IOException e) {
+        } catch (IOException | IllegalStateException e) {
             Log.wtf(TAG, "Reading from new store failed. All saved private networks are lost!", e);
             return false;
         } catch (XmlPullParserException e) {
@@ -3214,7 +3214,7 @@ public class WifiConfigManager {
 
         try {
             mWifiConfigStore.write(forceWrite);
-        } catch (IOException e) {
+        } catch (IOException | IllegalStateException e) {
             Log.wtf(TAG, "Writing to store failed. Saved networks maybe lost!", e);
             return false;
         } catch (XmlPullParserException e) {
diff --git a/service/java/com/android/server/wifi/WifiConfigStore.java b/service/java/com/android/server/wifi/WifiConfigStore.java
index a618eb5b5..e189d00e1 100644
--- a/service/java/com/android/server/wifi/WifiConfigStore.java
+++ b/service/java/com/android/server/wifi/WifiConfigStore.java
@@ -36,6 +36,7 @@ import com.android.internal.os.AtomicFile;
 import com.android.internal.util.FastXmlSerializer;
 import com.android.internal.util.Preconditions;
 import com.android.server.wifi.util.DataIntegrityChecker;
+import com.android.server.wifi.util.EncryptedData;
 import com.android.server.wifi.util.XmlUtil;
 
 import org.xmlpull.v1.XmlPullParser;
@@ -52,8 +53,8 @@ import java.io.IOException;
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
-import java.security.DigestException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -99,15 +100,23 @@ public class WifiConfigStore {
 
     private static final String XML_TAG_DOCUMENT_HEADER = "WifiConfigStoreData";
     private static final String XML_TAG_VERSION = "Version";
+    private static final String XML_TAG_HEADER_INTEGRITY = "Integrity";
+    private static final String XML_TAG_INTEGRITY_ENCRYPTED_DATA = "EncryptedData";
+    private static final String XML_TAG_INTEGRITY_IV = "IV";
     /**
      * Current config store data version. This will be incremented for any additions.
      */
-    private static final int CURRENT_CONFIG_STORE_DATA_VERSION = 1;
+    private static final int CURRENT_CONFIG_STORE_DATA_VERSION = 2;
     /** This list of older versions will be used to restore data from older config store. */
     /**
      * First version of the config store data format.
      */
     private static final int INITIAL_CONFIG_STORE_DATA_VERSION = 1;
+    /**
+     * Second version of the config store data format, introduced:
+     *  - Integrity info.
+     */
+    private static final int INTEGRITY_CONFIG_STORE_DATA_VERSION = 2;
 
     /**
      * Alarm tag to use for starting alarms for buffering file writes.
@@ -149,6 +158,11 @@ public class WifiConfigStore {
                 put(STORE_FILE_USER_NETWORK_SUGGESTIONS, STORE_FILE_NAME_USER_NETWORK_SUGGESTIONS);
             }};
 
+    @VisibleForTesting
+    public static final EncryptedData ZEROED_ENCRYPTED_DATA =
+            new EncryptedData(
+                    new byte[EncryptedData.ENCRYPTED_DATA_LENGTH],
+                    new byte[EncryptedData.IV_LENGTH]);
     /**
      * Handler instance to post alarm timeouts to
      */
@@ -276,7 +290,9 @@ public class WifiConfigStore {
                 return null;
             }
         }
-        return new StoreFile(new File(storeDir, STORE_ID_TO_FILE_NAME.get(fileId)), fileId);
+        File file = new File(storeDir, STORE_ID_TO_FILE_NAME.get(fileId));
+        DataIntegrityChecker dataIntegrityChecker = new DataIntegrityChecker(file.getName());
+        return new StoreFile(file, fileId, dataIntegrityChecker);
     }
 
     /**
@@ -395,6 +411,9 @@ public class WifiConfigStore {
      * Serialize all the data from all the {@link StoreData} clients registered for the provided
      * {@link StoreFile}.
      *
+     * This method also computes the integrity of the data being written and serializes the computed
+     * {@link EncryptedData} to the output.
+     *
      * @param storeFile StoreFile that we want to write to.
      * @return byte[] of serialized bytes
      * @throws XmlPullParserException
@@ -408,8 +427,9 @@ public class WifiConfigStore {
         final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
         out.setOutput(outputStream, StandardCharsets.UTF_8.name());
 
-        XmlUtil.writeDocumentStart(out, XML_TAG_DOCUMENT_HEADER);
-        XmlUtil.writeNextValue(out, XML_TAG_VERSION, CURRENT_CONFIG_STORE_DATA_VERSION);
+        // To compute integrity, write zeroes in the integrity fields. Once the integrity is
+        // computed, go back and modfiy the XML fields in place with the computed values.
+        writeDocumentMetadata(out, ZEROED_ENCRYPTED_DATA);
         for (StoreData storeData : storeDataList) {
             String tag = storeData.getName();
             XmlUtil.writeNextSectionStart(out, tag);
@@ -418,7 +438,77 @@ public class WifiConfigStore {
         }
         XmlUtil.writeDocumentEnd(out, XML_TAG_DOCUMENT_HEADER);
 
-        return outputStream.toByteArray();
+        byte[] outBytes = outputStream.toByteArray();
+        EncryptedData encryptedData = storeFile.computeIntegrity(outBytes);
+        if (encryptedData == null) {
+            // should never happen, this is a fatal failure. Abort file write.
+            Log.wtf(TAG, "Failed to compute integrity, failing write");
+            return null;
+        }
+        return rewriteDocumentMetadataRawBytes(outBytes, encryptedData);
+    }
+
+    /**
+     * Helper method to write the metadata at the start of every config store file.
+     * The metadata consists of:
+     * a) Version
+     * b) Integrity data computed on the data contents.
+     */
+    private void writeDocumentMetadata(XmlSerializer out, EncryptedData encryptedData)
+            throws XmlPullParserException, IOException {
+        // First XML header.
+        XmlUtil.writeDocumentStart(out, XML_TAG_DOCUMENT_HEADER);
+        // Next version.
+        XmlUtil.writeNextValue(out, XML_TAG_VERSION, CURRENT_CONFIG_STORE_DATA_VERSION);
+
+        // Next integrity data.
+        XmlUtil.writeNextSectionStart(out, XML_TAG_HEADER_INTEGRITY);
+        XmlUtil.writeNextValue(out, XML_TAG_INTEGRITY_ENCRYPTED_DATA,
+                encryptedData.getEncryptedData());
+        XmlUtil.writeNextValue(out, XML_TAG_INTEGRITY_IV, encryptedData.getIv());
+        XmlUtil.writeNextSectionEnd(out, XML_TAG_HEADER_INTEGRITY);
+    }
+
+    /**
+     * Helper method to generate the raw bytes containing the the metadata at the start of every
+     * config store file.
+     *
+     * NOTE: This does not create a fully formed XML document (the start tag is not closed for
+     * example). This only creates the top portion of the XML which contains the modified
+     * integrity data & version along with the XML prolog (metadata):
+     * <?xml version='1.0' encoding='utf-8' standalone='yes' ?>
+     * <WifiConfigStoreData>
+     * <int name="Version" value="2" />
+     * <Integrity>
+     * <byte-array name="EncryptedData" num="48">!EncryptedData!</byte-array>
+     * <byte-array name="IV" num="12">!IV!</byte-array>
+     * </Integrity>
+     */
+    private byte[] generateDocumentMetadataRawBytes(EncryptedData encryptedData)
+            throws XmlPullParserException, IOException {
+        final XmlSerializer outXml = new FastXmlSerializer();
+        final ByteArrayOutputStream outStream = new ByteArrayOutputStream();
+        outXml.setOutput(outStream, StandardCharsets.UTF_8.name());
+        writeDocumentMetadata(outXml, encryptedData);
+        outXml.endDocument();
+        return outStream.toByteArray();
+    }
+
+    /**
+     * Helper method to rewrite the existing document metadata in the incoming raw bytes in
+     * |inBytes| with the new document metadata created with the provided |encryptedData|.
+     *
+     * NOTE: This assumes that the metadata in existing XML inside |inBytes| aligns exactly
+     * with the new metadata created. So, a simple in place rewrite of the first few bytes (
+     * corresponding to metadata section of the document) from |inBytes| will preserve the overall
+     * document structure.
+     */
+    private byte[] rewriteDocumentMetadataRawBytes(byte[] inBytes, EncryptedData encryptedData)
+            throws XmlPullParserException, IOException {
+        byte[] replaceMetadataBytes = generateDocumentMetadataRawBytes(encryptedData);
+        ByteBuffer outByteBuffer = ByteBuffer.wrap(inBytes);
+        outByteBuffer.put(replaceMetadataBytes);
+        return outByteBuffer.array();
     }
 
     /**
@@ -549,6 +639,10 @@ public class WifiConfigStore {
     /**
      * Deserialize data from a {@link StoreFile} for all {@link StoreData} instances registered.
      *
+     * This method also computes the integrity of the incoming |dataBytes| and compare with
+     * {@link EncryptedData} parsed from |dataBytes|. If the integrity check fails, the data
+     * is discarded.
+     *
      * @param dataBytes The data to parse
      * @param storeFile StoreFile that we read from. Will be used to retrieve the list of clients
      *                  who have data to deserialize from this file.
@@ -569,7 +663,24 @@ public class WifiConfigStore {
 
         // Start parsing the XML stream.
         int rootTagDepth = in.getDepth() + 1;
-        parseDocumentStartAndVersionFromXml(in);
+        XmlUtil.gotoDocumentStart(in, XML_TAG_DOCUMENT_HEADER);
+
+        int version = parseVersionFromXml(in);
+        // Version 2 onwards contains integrity data, so check the integrity of the file.
+        if (version >= INTEGRITY_CONFIG_STORE_DATA_VERSION) {
+            EncryptedData integrityData = parseIntegrityDataFromXml(in, rootTagDepth);
+            // To compute integrity, write zeroes in the integrity fields.
+            dataBytes = rewriteDocumentMetadataRawBytes(dataBytes, ZEROED_ENCRYPTED_DATA);
+            if (!storeFile.checkIntegrity(dataBytes, integrityData)) {
+                Log.wtf(TAG, "Integrity mismatch, discarding data from " + storeFile.mFileName);
+                return;
+            }
+        } else {
+            // When integrity checking is introduced. The existing data will have no related
+            // integrity file for validation. Thus, we will assume the existing data is correct.
+            // Integrity will be computed for the next write.
+            Log.d(TAG, "No integrity data to check; thus vacously true");
+        }
 
         String[] headerName = new String[1];
         Set<StoreData> storeDatasInvoked = new HashSet<>();
@@ -595,15 +706,14 @@ public class WifiConfigStore {
     }
 
     /**
-     * Parse the document start and version from the XML stream.
+     * Parse the version from the XML stream.
      * This is used for both the shared and user config store data.
      *
      * @param in XmlPullParser instance pointing to the XML stream.
      * @return version number retrieved from the Xml stream.
      */
-    private static int parseDocumentStartAndVersionFromXml(XmlPullParser in)
+    private static int parseVersionFromXml(XmlPullParser in)
             throws XmlPullParserException, IOException {
-        XmlUtil.gotoDocumentStart(in, XML_TAG_DOCUMENT_HEADER);
         int version = (int) XmlUtil.readNextValueWithName(in, XML_TAG_VERSION);
         if (version < INITIAL_CONFIG_STORE_DATA_VERSION
                 || version > CURRENT_CONFIG_STORE_DATA_VERSION) {
@@ -613,6 +723,25 @@ public class WifiConfigStore {
     }
 
     /**
+     * Parse the integrity data structure from the XML stream.
+     * This is used for both the shared and user config store data.
+     *
+     * @param in XmlPullParser instance pointing to the XML stream.
+     * @param outerTagDepth Outer tag depth.
+     * @return Instance of {@link EncryptedData} retrieved from the Xml stream.
+     */
+    private static @NonNull EncryptedData parseIntegrityDataFromXml(
+            XmlPullParser in, int outerTagDepth)
+            throws XmlPullParserException, IOException {
+        XmlUtil.gotoNextSectionWithName(in, XML_TAG_HEADER_INTEGRITY, outerTagDepth);
+        byte[] encryptedData =
+                (byte[]) XmlUtil.readNextValueWithName(in, XML_TAG_INTEGRITY_ENCRYPTED_DATA);
+        byte[] iv =
+                (byte[]) XmlUtil.readNextValueWithName(in, XML_TAG_INTEGRITY_IV);
+        return new EncryptedData(encryptedData, iv);
+    }
+
+    /**
      * Dump the local log buffer and other internal state of WifiConfigManager.
      */
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
@@ -653,21 +782,22 @@ public class WifiConfigStore {
         /**
          * Store the file name for setting the file permissions/logging purposes.
          */
-        private String mFileName;
+        private final String mFileName;
         /**
-         * The integrity file storing integrity checking data for the store file.
+         * {@link StoreFileId} Type of store file.
          */
-        private DataIntegrityChecker mDataIntegrityChecker;
+        private final @StoreFileId int mFileId;
         /**
-         * {@link StoreFileId} Type of store file.
+         * Integrity checking for the store file.
          */
-        private @StoreFileId int mFileId;
+        private final DataIntegrityChecker mDataIntegrityChecker;
 
-        public StoreFile(File file, @StoreFileId int fileId) {
+        public StoreFile(File file, @StoreFileId int fileId,
+                         @NonNull DataIntegrityChecker dataIntegrityChecker) {
             mAtomicFile = new AtomicFile(file);
-            mFileName = mAtomicFile.getBaseFile().getAbsolutePath();
-            mDataIntegrityChecker = new DataIntegrityChecker(mFileName);
+            mFileName = file.getAbsolutePath();
             mFileId = fileId;
+            mDataIntegrityChecker = dataIntegrityChecker;
         }
 
         /**
@@ -691,20 +821,8 @@ public class WifiConfigStore {
             byte[] bytes = null;
             try {
                 bytes = mAtomicFile.readFully();
-                // Check that the file has not been altered since last writeBufferedRawData()
-                if (!mDataIntegrityChecker.isOk(bytes)) {
-                    Log.wtf(TAG, "Data integrity problem with file: " + mFileName);
-                    return null;
-                }
             } catch (FileNotFoundException e) {
                 return null;
-            } catch (DigestException e) {
-                // When integrity checking is introduced. The existing data will have no related
-                // integrity file for validation. Thus, we will assume the existing data is correct
-                // and immediately create the integrity file.
-                Log.i(TAG, "isOK() had no integrity data to check; thus vacuously "
-                        + "true. Running update now.");
-                mDataIntegrityChecker.update(bytes);
             }
             return bytes;
         }
@@ -741,11 +859,36 @@ public class WifiConfigStore {
                 }
                 throw e;
             }
-            // There was a legitimate change and update the integrity checker.
-            mDataIntegrityChecker.update(mWriteData);
             // Reset the pending write data after write.
             mWriteData = null;
         }
+
+        /**
+         * Compute integrity of |dataWithZeroedIntegrityFields| to be written to the file.
+         *
+         * @param dataWithZeroedIntegrityFields raw data to be written to the file with the
+         *                                      integrity fields zeroed out for integrity
+         *                                      calculation.
+         * @return Instance of {@link EncryptedData} holding the encrypted integrity data for the
+         * raw data to be written to the file.
+         */
+        public EncryptedData computeIntegrity(byte[] dataWithZeroedIntegrityFields) {
+            return mDataIntegrityChecker.compute(dataWithZeroedIntegrityFields);
+        }
+
+        /**
+         * Check integrity of |dataWithZeroedIntegrityFields| read from the file with the integrity
+         * data parsed from the file.
+         * @param dataWithZeroedIntegrityFields raw data read from the file with the integrity
+         *                                      fields zeroed out for integrity calculation.
+         * @param parsedEncryptedData Instance of {@link EncryptedData} parsed from the integrity
+         *                            fields in the raw data.
+         * @return true if the integrity matches, false otherwise.
+         */
+        public boolean checkIntegrity(byte[] dataWithZeroedIntegrityFields,
+                                      EncryptedData parsedEncryptedData) {
+            return mDataIntegrityChecker.isOk(dataWithZeroedIntegrityFields, parsedEncryptedData);
+        }
     }
 
     /**
diff --git a/service/java/com/android/server/wifi/util/DataIntegrityChecker.java b/service/java/com/android/server/wifi/util/DataIntegrityChecker.java
index 1f3c6b3c1..6f03a4861 100644
--- a/service/java/com/android/server/wifi/util/DataIntegrityChecker.java
+++ b/service/java/com/android/server/wifi/util/DataIntegrityChecker.java
@@ -23,14 +23,7 @@ import android.security.keystore.KeyProperties;
 import android.text.TextUtils;
 import android.util.Log;
 
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
-import java.io.FileOutputStream;
 import java.io.IOException;
-import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
-import java.security.DigestException;
 import java.security.InvalidAlgorithmParameterException;
 import java.security.InvalidKeyException;
 import java.security.KeyStore;
@@ -56,7 +49,6 @@ import javax.crypto.spec.GCMParameterSpec;
 public class DataIntegrityChecker {
     private static final String TAG = "DataIntegrityChecker";
 
-    private static final String FILE_SUFFIX = ".encrypted-checksum";
     private static final String ALIAS_SUFFIX = ".data-integrity-checker-key";
     private static final String CIPHER_ALGORITHM = "AES/GCM/NoPadding";
     private static final String DIGEST_ALGORITHM = "SHA-256";
@@ -65,28 +57,29 @@ public class DataIntegrityChecker {
 
     /**
      * When KEYSTORE_FAILURE_RETURN_VALUE is true, all cryptographic operation failures will not
-     * enforce security and {@link #isOk(byte[])} always return true.
+     * enforce security and {@link #isOk(byte[], EncryptedData)} always return true.
      */
     private static final boolean KEYSTORE_FAILURE_RETURN_VALUE = true;
 
-    private final File mIntegrityFile;
+    private final String mDataFileName;
 
     /**
      * Construct a new integrity checker to update and check if/when a data file was altered
      * outside expected conditions.
      *
-     * @param integrityFilename The {@link File} path prefix for where the integrity data is stored.
-     *                          A file will be created in the name of integrityFile with the suffix
-     *                          {@link DataIntegrityChecker#FILE_SUFFIX} We recommend using the same
-     *                          path as the file for which integrity is performed on.
-     * @throws NullPointerException When integrity file is null or the empty string.
+     * @param dataFileName The full path of the data file for which integrity check is performed.
+     * @throws NullPointerException When data file is empty string.
      */
-    public DataIntegrityChecker(@NonNull String integrityFilename) {
-        if (TextUtils.isEmpty(integrityFilename)) {
-            throw new NullPointerException("integrityFilename must not be null or the empty "
+    public DataIntegrityChecker(@NonNull String dataFileName) {
+        if (TextUtils.isEmpty(dataFileName)) {
+            throw new NullPointerException("dataFileName must not be null or the empty "
                     + "string");
         }
-        mIntegrityFile = new File(integrityFilename + FILE_SUFFIX);
+        mDataFileName = dataFileName;
+    }
+
+    private String getKeyAlias() {
+        return mDataFileName + ALIAS_SUFFIX;
     }
 
     /**
@@ -95,66 +88,55 @@ public class DataIntegrityChecker {
      * Call this method immediately before storing the byte array
      *
      * @param data The data desired to ensure integrity
+     * @return Instance of {@link EncryptedData} containing the encrypted integrity data.
      */
-    public void update(byte[] data) {
+    public EncryptedData compute(byte[] data) {
         if (data == null || data.length < 1) {
-            reportException(new Exception("No data to update"), "No data to update.");
-            return;
+            reportException(new Exception("No data to compute"), "No data to compute.");
+            return null;
         }
         byte[] digest = getDigest(data);
         if (digest == null || digest.length < 1) {
-            return;
+            reportException(new Exception("digest null in compute"),
+                    "digest null in compute");
+            return null;
         }
-        String alias = mIntegrityFile.getName() + ALIAS_SUFFIX;
-        EncryptedData integrityData = encrypt(digest, alias);
-        if (integrityData != null) {
-            writeIntegrityData(integrityData, mIntegrityFile);
-        } else {
-            reportException(new Exception("integrityData null upon update"),
-                    "integrityData null upon update");
+        EncryptedData integrityData = encrypt(digest, getKeyAlias());
+        if (integrityData == null) {
+            reportException(new Exception("integrityData null in compute"),
+                    "integrityData null in compute");
         }
+        return integrityData;
     }
 
+
     /**
      * Check the integrity of a given byte array
      *
      * Call this method immediately before trusting the byte array. This method will return false
-     * when the byte array was altered since the last {@link #update(byte[])}
-     * call, when {@link #update(byte[])} has never been called, or if there is
-     * an underlying issue with the cryptographic functions or the key store.
+     * when the integrity data calculated on the byte array does not match the encrypted integrity
+     * data provided to compare or if there is an underlying issue with the cryptographic functions
+     * or the key store.
      *
-     * @param data The data to check if its been altered
-     * @throws DigestException The integrity mIntegrityFile cannot be read. Ensure
-     *      {@link #isOk(byte[])} is called after {@link #update(byte[])}. Otherwise, consider the
-     *      result vacuously true and immediately call {@link #update(byte[])}.
-     * @return true if the data was not altered since {@link #update(byte[])} was last called
+     * @param data The data to check if its been altered.
+     * @param integrityData Encrypted integrity data to be used for comparison.
+     * @return true if the integrity data computed on |data| matches the provided |integrityData|.
      */
-    public boolean isOk(byte[] data) throws DigestException {
+    public boolean isOk(@NonNull byte[] data, @NonNull EncryptedData integrityData) {
         if (data == null || data.length < 1) {
             return KEYSTORE_FAILURE_RETURN_VALUE;
         }
         byte[] currentDigest = getDigest(data);
         if (currentDigest == null || currentDigest.length < 1) {
+            reportException(new Exception("current digest null"), "current digest null");
             return KEYSTORE_FAILURE_RETURN_VALUE;
         }
-
-        EncryptedData encryptedData = null;
-
-        try {
-            encryptedData = readIntegrityData(mIntegrityFile);
-        } catch (IOException e) {
-            reportException(e, "readIntegrityData had an IO exception");
-            return KEYSTORE_FAILURE_RETURN_VALUE;
-        } catch (ClassNotFoundException e) {
-            reportException(e, "readIntegrityData could not find the class EncryptedData");
+        if (integrityData == null) {
+            reportException(new Exception("integrityData null in isOk"),
+                    "integrityData null in isOk");
             return KEYSTORE_FAILURE_RETURN_VALUE;
         }
-
-        if (encryptedData == null) {
-            // File not found is not considered to be an error.
-            throw new DigestException("No stored digest is available to compare.");
-        }
-        byte[] storedDigest = decrypt(encryptedData);
+        byte[] storedDigest = decrypt(integrityData, getKeyAlias());
         if (storedDigest == null) {
             return KEYSTORE_FAILURE_RETURN_VALUE;
         }
@@ -177,7 +159,7 @@ public class DataIntegrityChecker {
             SecretKey secretKeyReference = getOrCreateSecretKey(keyAlias);
             if (secretKeyReference != null) {
                 cipher.init(Cipher.ENCRYPT_MODE, secretKeyReference);
-                encryptedData = new EncryptedData(cipher.doFinal(data), cipher.getIV(), keyAlias);
+                encryptedData = new EncryptedData(cipher.doFinal(data), cipher.getIV());
             } else {
                 reportException(new Exception("secretKeyReference is null."),
                         "secretKeyReference is null.");
@@ -196,12 +178,12 @@ public class DataIntegrityChecker {
         return encryptedData;
     }
 
-    private byte[] decrypt(EncryptedData encryptedData) {
+    private byte[] decrypt(EncryptedData encryptedData, String keyAlias) {
         byte[] decryptedData = null;
         try {
             Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
             GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, encryptedData.getIv());
-            SecretKey secretKeyReference = getOrCreateSecretKey(encryptedData.getKeyAlias());
+            SecretKey secretKeyReference = getOrCreateSecretKey(keyAlias);
             if (secretKeyReference != null) {
                 cipher.init(Cipher.DECRYPT_MODE, secretKeyReference, spec);
                 decryptedData = cipher.doFinal(encryptedData.getEncryptedData());
@@ -268,30 +250,6 @@ public class DataIntegrityChecker {
         return secretKey;
     }
 
-    private void writeIntegrityData(EncryptedData encryptedData, File file) {
-        try (FileOutputStream fos = new FileOutputStream(file);
-             ObjectOutputStream oos = new ObjectOutputStream(fos)) {
-            oos.writeObject(encryptedData);
-        } catch (FileNotFoundException e) {
-            reportException(e, "writeIntegrityData could not find the integrity file");
-        } catch (IOException e) {
-            reportException(e, "writeIntegrityData had an IO exception");
-        }
-    }
-
-    private EncryptedData readIntegrityData(File file) throws IOException, ClassNotFoundException  {
-        try (FileInputStream fis = new FileInputStream(file);
-             ObjectInputStream ois = new ObjectInputStream(fis)) {
-            return (EncryptedData) ois.readObject();
-        } catch (FileNotFoundException e) {
-            // File not found, this is not considered to be a real error. The file will be created
-            // by the system next time the data file is written. Note that it is not possible for
-            // non system user to delete or modify the file.
-            Log.w(TAG, "readIntegrityData could not find integrity file");
-        }
-        return null;
-    }
-
     private boolean constantTimeEquals(byte[] a, byte[] b) {
         if (a == null && b == null) {
             return true;
@@ -311,7 +269,7 @@ public class DataIntegrityChecker {
     /* TODO(b/128526030): Remove this error reporting code upon resolving the bug. */
     private static final boolean REQUEST_BUG_REPORT = false;
     private void reportException(Exception exception, String error) {
-        Log.wtf(TAG, "An irrecoverable key store error was encountered: " + error);
+        Log.wtf(TAG, "An irrecoverable key store error was encountered: " + error, exception);
         if (REQUEST_BUG_REPORT) {
             SystemProperties.set("dumpstate.options", "bugreportwifi");
             SystemProperties.set("ctl.start", "bugreport");
diff --git a/service/java/com/android/server/wifi/util/EncryptedData.java b/service/java/com/android/server/wifi/util/EncryptedData.java
index 468f28ec0..91342d335 100644
--- a/service/java/com/android/server/wifi/util/EncryptedData.java
+++ b/service/java/com/android/server/wifi/util/EncryptedData.java
@@ -16,22 +16,25 @@
 
 package com.android.server.wifi.util;
 
-import java.io.Serializable;
+import com.android.internal.util.Preconditions;
 
 /**
  * A class to store data created by {@link DataIntegrityChecker}.
  */
-public class EncryptedData implements Serializable {
-    private static final long serialVersionUID = 1337L;
-
-    private byte[] mEncryptedData;
-    private byte[] mIv;
-    private String mKeyAlias;
-
-    public EncryptedData(byte[] encryptedData, byte[] iv, String keyAlias) {
+public class EncryptedData {
+    public static final int ENCRYPTED_DATA_LENGTH = 48;
+    public static final int IV_LENGTH = 12;
+
+    private final byte[] mEncryptedData;
+    private final byte[] mIv;
+
+    public EncryptedData(byte[] encryptedData, byte[] iv) {
+        Preconditions.checkNotNull(encryptedData, iv);
+        Preconditions.checkState(encryptedData.length == ENCRYPTED_DATA_LENGTH,
+                "encryptedData.length=" + encryptedData.length);
+        Preconditions.checkState(iv.length == IV_LENGTH, "iv.length=" + iv.length);
         mEncryptedData = encryptedData;
         mIv = iv;
-        mKeyAlias = keyAlias;
     }
 
     public byte[] getEncryptedData() {
@@ -41,8 +44,4 @@ public class EncryptedData implements Serializable {
     public byte[] getIv() {
         return mIv;
     }
-
-    public String getKeyAlias() {
-        return mKeyAlias;
-    }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiConfigStoreTest.java b/tests/wifitests/src/com/android/server/wifi/WifiConfigStoreTest.java
index 64e762bec..b59e367dd 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiConfigStoreTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiConfigStoreTest.java
@@ -16,9 +16,12 @@
 
 package com.android.server.wifi;
 
+import static com.android.server.wifi.WifiConfigStore.ZEROED_ENCRYPTED_DATA;
+
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
 
+import android.app.test.MockAnswerUtil.AnswerWithArguments;
 import android.app.test.TestAlarmManager;
 import android.content.Context;
 import android.content.pm.PackageManager;
@@ -31,11 +34,16 @@ import androidx.test.filters.SmallTest;
 import com.android.internal.util.ArrayUtils;
 import com.android.server.wifi.WifiConfigStore.StoreData;
 import com.android.server.wifi.WifiConfigStore.StoreFile;
+import com.android.server.wifi.util.DataIntegrityChecker;
+import com.android.server.wifi.util.EncryptedData;
 import com.android.server.wifi.util.XmlUtil;
 
+import libcore.util.HexEncoding;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.xmlpull.v1.XmlPullParser;
@@ -49,6 +57,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 
 /**
  * Unit tests for {@link com.android.server.wifi.WifiConfigStore}.
@@ -64,7 +73,13 @@ public class WifiConfigStoreTest {
     private static final String TEST_DATA_XML_STRING_FORMAT =
             "<?xml version='1.0' encoding='utf-8' standalone='yes' ?>\n"
                     + "<WifiConfigStoreData>\n"
-                    + "<int name=\"Version\" value=\"1\" />\n"
+                    + "<int name=\"Version\" value=\"2\" />\n"
+                    + "<Integrity>\n"
+                    + "<byte-array name=\"EncryptedData\" num=\"48\">000000000000000000000000000000"
+                    + "000000000000000000000000000000000000000000000000000000000000000000"
+                    + "</byte-array>\n"
+                    + "<byte-array name=\"IV\" num=\"12\">000000000000000000000000</byte-array>\n"
+                    + "</Integrity>\n"
                     + "<NetworkList>\n"
                     + "<Network>\n"
                     + "<WifiConfiguration>\n"
@@ -127,19 +142,29 @@ public class WifiConfigStoreTest {
                     + "</DeletedEphemeralSSIDList>\n"
                     + "</WifiConfigStoreData>\n";
 
-    private static final String TEST_DATA_XML_STRING_FORMAT_WITH_ONE_DATA_SOURCE =
+    private static final String TEST_DATA_XML_STRING_FORMAT_V1_WITH_ONE_DATA_SOURCE =
             "<?xml version='1.0' encoding='utf-8' standalone='yes' ?>\n"
                     + "<WifiConfigStoreData>\n"
                     + "<int name=\"Version\" value=\"1\" />\n"
                     + "<%s/>n"
                     + "</WifiConfigStoreData>\n";
-    private static final String TEST_DATA_XML_STRING_FORMAT_WITH_TWO_DATA_SOURCE =
+    private static final String TEST_DATA_XML_STRING_FORMAT_V1_WITH_TWO_DATA_SOURCE =
             "<?xml version='1.0' encoding='utf-8' standalone='yes' ?>\n"
                     + "<WifiConfigStoreData>\n"
                     + "<int name=\"Version\" value=\"1\" />\n"
                     + "<%s/>n"
                     + "<%s/>n"
                     + "</WifiConfigStoreData>\n";
+    private static final String TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE =
+            "<?xml version='1.0' encoding='utf-8' standalone='yes' ?>\n"
+                    + "<WifiConfigStoreData>\n"
+                    + "<int name=\"Version\" value=\"2\" />\n"
+                    + "<Integrity>\n"
+                    + "<byte-array name=\"EncryptedData\" num=\"48\">%s</byte-array>\n"
+                    + "<byte-array name=\"IV\" num=\"12\">%s</byte-array>\n"
+                    + "</Integrity>\n"
+                    + "<%s />\n"
+                    + "</WifiConfigStoreData>\n";
     // Test mocks
     @Mock private Context mContext;
     @Mock private PackageManager mPackageManager;
@@ -147,6 +172,7 @@ public class WifiConfigStoreTest {
     private TestLooper mLooper;
     @Mock private Clock mClock;
     @Mock private WifiMetrics mWifiMetrics;
+    @Mock private DataIntegrityChecker mDataIntegrityChecker;
     private MockStoreFile mSharedStore;
     private MockStoreFile mUserStore;
     private MockStoreFile mUserNetworkSuggestionsStore;
@@ -170,6 +196,10 @@ public class WifiConfigStoreTest {
                 .thenReturn(mAlarmManager.getAlarmManager());
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
         when(mPackageManager.getNameForUid(anyInt())).thenReturn(TEST_CREATOR_NAME);
+        when(mDataIntegrityChecker.compute(any(byte[].class)))
+                .thenReturn(ZEROED_ENCRYPTED_DATA);
+        when(mDataIntegrityChecker.isOk(any(byte[].class), any(EncryptedData.class)))
+                .thenReturn(true);
         mSharedStore = new MockStoreFile(WifiConfigStore.STORE_FILE_SHARED_GENERAL);
         mUserStore = new MockStoreFile(WifiConfigStore.STORE_FILE_USER_GENERAL);
         mUserNetworkSuggestionsStore =
@@ -591,11 +621,11 @@ public class WifiConfigStoreTest {
         assertTrue(mWifiConfigStore.registerStoreData(storeData2));
 
         String fileContentsXmlStringWithOnlyStoreData1 =
-                String.format(TEST_DATA_XML_STRING_FORMAT_WITH_ONE_DATA_SOURCE, storeData1Name);
+                String.format(TEST_DATA_XML_STRING_FORMAT_V1_WITH_ONE_DATA_SOURCE, storeData1Name);
         String fileContentsXmlStringWithOnlyStoreData2 =
-                String.format(TEST_DATA_XML_STRING_FORMAT_WITH_ONE_DATA_SOURCE, storeData2Name);
+                String.format(TEST_DATA_XML_STRING_FORMAT_V1_WITH_ONE_DATA_SOURCE, storeData2Name);
         String fileContentsXmlStringWithStoreData1AndStoreData2 =
-                String.format(TEST_DATA_XML_STRING_FORMAT_WITH_TWO_DATA_SOURCE,
+                String.format(TEST_DATA_XML_STRING_FORMAT_V1_WITH_TWO_DATA_SOURCE,
                         storeData1Name, storeData2Name);
 
         // Scenario 1: StoreData1 in shared store file.
@@ -753,6 +783,293 @@ public class WifiConfigStoreTest {
     }
 
     /**
+     * Tests the read API behaviour when the config store file is version 1.
+     * Expected behaviour: The read should be successful and send the data to the corresponding
+     *                     {@link StoreData} instance.
+     */
+    @Test
+    public void testReadVersion1StoreFile() throws Exception {
+        // Register data container.
+        StoreData sharedStoreData = mock(StoreData.class);
+        when(sharedStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_SHARED_GENERAL);
+        when(sharedStoreData.getName()).thenReturn(TEST_SHARE_DATA);
+        StoreData userStoreData = mock(StoreData.class);
+        when(userStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_USER_GENERAL);
+        when(userStoreData.getName()).thenReturn(TEST_USER_DATA);
+        mWifiConfigStore.registerStoreData(sharedStoreData);
+        mWifiConfigStore.registerStoreData(userStoreData);
+
+        // Read both share and user config store.
+        mWifiConfigStore.setUserStores(mUserStores);
+
+        // Now store some content in the shared and user data files.
+        mUserStore.storeRawDataToWrite(
+                String.format(TEST_DATA_XML_STRING_FORMAT_V1_WITH_ONE_DATA_SOURCE,
+                        TEST_USER_DATA).getBytes());
+        mSharedStore.storeRawDataToWrite(
+                String.format(TEST_DATA_XML_STRING_FORMAT_V1_WITH_ONE_DATA_SOURCE,
+                        TEST_SHARE_DATA).getBytes());
+
+        // Read and verify the data content in the store file (metadata stripped out) has been sent
+        // to the corresponding store data when integrity check passes.
+        mWifiConfigStore.read();
+        verify(sharedStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+        verify(userStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+
+        // We shouldn't perform any data integrity checks on version 1 file.
+        verifyZeroInteractions(mDataIntegrityChecker);
+    }
+
+    /**
+     * Tests the read API behaviour when integrity check fails.
+     * Expected behaviour: The read should return an empty store data.
+     */
+    @Test
+    public void testReadWhenIntegrityCheckFails() throws Exception {
+        // Register data container.
+        StoreData sharedStoreData = mock(StoreData.class);
+        when(sharedStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_SHARED_GENERAL);
+        when(sharedStoreData.getName()).thenReturn(TEST_SHARE_DATA);
+        StoreData userStoreData = mock(StoreData.class);
+        when(userStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_USER_GENERAL);
+        when(userStoreData.getName()).thenReturn(TEST_USER_DATA);
+        mWifiConfigStore.registerStoreData(sharedStoreData);
+        mWifiConfigStore.registerStoreData(userStoreData);
+
+        // Read both share and user config store.
+        mWifiConfigStore.setUserStores(mUserStores);
+
+        // Now store some content in the shared and user data files.
+        mUserStore.storeRawDataToWrite(
+                String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                        HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getEncryptedData()),
+                        HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getIv()),
+                        TEST_USER_DATA).getBytes());
+        mSharedStore.storeRawDataToWrite(
+                String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                        HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getEncryptedData()),
+                        HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getIv()),
+                        TEST_SHARE_DATA).getBytes());
+
+        // Read and verify the data content in the store file (metadata stripped out) has been sent
+        // to the corresponding store data when integrity check passes.
+        mWifiConfigStore.read();
+        verify(sharedStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+        verify(userStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+
+        // Read and verify the data content in the store file (metadata stripped out) has not been
+        // sent to the corresponding store data when integrity check fails.
+        when(mDataIntegrityChecker.isOk(any(byte[].class), any(EncryptedData.class)))
+                .thenReturn(false);
+        mWifiConfigStore.read();
+        verify(sharedStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+        verify(userStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+    }
+
+    /**
+     * Tests the write API behaviour when integrity check fails.
+     * Expected behaviour: The read should return an empty store data.
+     */
+    @Test
+    public void testWriteWhenIntegrityComputeFails() throws Exception {
+        // Register data container.
+        StoreData sharedStoreData = mock(StoreData.class);
+        when(sharedStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_SHARED_GENERAL);
+        when(sharedStoreData.getName()).thenReturn(TEST_SHARE_DATA);
+        when(sharedStoreData.hasNewDataToSerialize()).thenReturn(true);
+        StoreData userStoreData = mock(StoreData.class);
+        when(userStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_USER_GENERAL);
+        when(userStoreData.getName()).thenReturn(TEST_USER_DATA);
+        when(userStoreData.hasNewDataToSerialize()).thenReturn(true);
+        mWifiConfigStore.registerStoreData(sharedStoreData);
+        mWifiConfigStore.registerStoreData(userStoreData);
+
+        // Read both share and user config store.
+        mWifiConfigStore.setUserStores(mUserStores);
+
+        // Reset store file contents & ensure that the user and store data files are empty.
+        mUserStore.storeRawDataToWrite(null);
+        mSharedStore.storeRawDataToWrite(null);
+        assertNull(mUserStore.getStoreBytes());
+        assertNull(mSharedStore.getStoreBytes());
+
+        // Write and verify that the data is written to the config store file when integrity
+        // computation passes.
+        mWifiConfigStore.write(true);
+        assertNotNull(mUserStore.getStoreBytes());
+        assertNotNull(mSharedStore.getStoreBytes());
+        assertTrue(new String(mUserStore.getStoreBytes()).contains(TEST_USER_DATA));
+        assertTrue(new String(mSharedStore.getStoreBytes()).contains(TEST_SHARE_DATA));
+
+        // Reset store file contents & ensure that the user and store data files are empty.
+        mUserStore.storeRawDataToWrite(null);
+        mSharedStore.storeRawDataToWrite(null);
+        assertNull(mUserStore.getStoreBytes());
+        assertNull(mSharedStore.getStoreBytes());
+
+        // Write and verify that the data is not written to the config store file when integrity
+        // computation fails.
+        when(mDataIntegrityChecker.compute(any(byte[].class))).thenReturn(null);
+        mWifiConfigStore.write(true);
+        assertNull(mUserStore.getStoreBytes());
+        assertNull(mSharedStore.getStoreBytes());
+    }
+
+    /**
+     * Tests the write API behaviour to ensure that the integrity data is written to the file.
+     */
+    @Test
+    public void testWriteContainsIntegrityData() throws Exception {
+        byte[] encryptedData = new byte[EncryptedData.ENCRYPTED_DATA_LENGTH];
+        byte[] iv = new byte[EncryptedData.IV_LENGTH];
+        Random random = new Random();
+        random.nextBytes(encryptedData);
+        random.nextBytes(iv);
+        final EncryptedData testEncryptedData = new EncryptedData(encryptedData, iv);
+
+        doAnswer(new AnswerWithArguments() {
+            public EncryptedData answer(byte[] data) {
+                String storeXmlString = new String(data);
+                // Verify that we fill in zeros to the data when we compute integrity.
+                if (storeXmlString.contains(TEST_SHARE_DATA)) {
+                    assertEquals(String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                            HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getEncryptedData()),
+                            HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getIv()),
+                            TEST_SHARE_DATA), storeXmlString);
+                } else if (storeXmlString.contains(TEST_USER_DATA)) {
+                    assertEquals(String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                            HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getEncryptedData()),
+                            HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getIv()),
+                            TEST_USER_DATA), storeXmlString);
+                }
+                return testEncryptedData;
+            }
+        }).when(mDataIntegrityChecker).compute(any(byte[].class));
+
+        // Register data container.
+        StoreData sharedStoreData = mock(StoreData.class);
+        when(sharedStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_SHARED_GENERAL);
+        when(sharedStoreData.getName()).thenReturn(TEST_SHARE_DATA);
+        when(sharedStoreData.hasNewDataToSerialize()).thenReturn(true);
+        StoreData userStoreData = mock(StoreData.class);
+        when(userStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_USER_GENERAL);
+        when(userStoreData.getName()).thenReturn(TEST_USER_DATA);
+        when(userStoreData.hasNewDataToSerialize()).thenReturn(true);
+        mWifiConfigStore.registerStoreData(sharedStoreData);
+        mWifiConfigStore.registerStoreData(userStoreData);
+
+        // Read both share and user config store.
+        mWifiConfigStore.setUserStores(mUserStores);
+
+        // Write and verify that the data is written to the config store file when integrity
+        // computation passes.
+        mWifiConfigStore.write(true);
+
+        // Verify that we fill in zeros to the data when we computed integrity.
+        verify(mDataIntegrityChecker, times(2)).compute(any(byte[].class));
+
+        // Verify the parsed integrity data
+        assertNotNull(mUserStore.getStoreBytes());
+        assertNotNull(mSharedStore.getStoreBytes());
+        String userStoreXmlString = new String(mUserStore.getStoreBytes());
+        String sharedStoreXmlString = new String(mSharedStore.getStoreBytes());
+        assertEquals(String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                HexEncoding.encodeToString(encryptedData).toLowerCase(),
+                HexEncoding.encodeToString(iv).toLowerCase(),
+                TEST_USER_DATA), userStoreXmlString);
+        assertEquals(String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                HexEncoding.encodeToString(encryptedData).toLowerCase(),
+                HexEncoding.encodeToString(iv).toLowerCase(),
+                TEST_SHARE_DATA), sharedStoreXmlString);
+    }
+
+    /**
+     * Tests the read API behaviour to ensure that the integrity data is parsed from the file and
+     * used for checking integrity of the file.
+     */
+    @Test
+    public void testReadParsesIntegrityData() throws Exception {
+        byte[] encryptedData = new byte[EncryptedData.ENCRYPTED_DATA_LENGTH];
+        byte[] iv = new byte[EncryptedData.IV_LENGTH];
+        Random random = new Random();
+        random.nextBytes(encryptedData);
+        random.nextBytes(iv);
+
+        // Register data container.
+        StoreData sharedStoreData = mock(StoreData.class);
+        when(sharedStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_SHARED_GENERAL);
+        when(sharedStoreData.getName()).thenReturn(TEST_SHARE_DATA);
+        when(sharedStoreData.hasNewDataToSerialize()).thenReturn(true);
+        StoreData userStoreData = mock(StoreData.class);
+        when(userStoreData.getStoreFileId())
+                .thenReturn(WifiConfigStore.STORE_FILE_USER_GENERAL);
+        when(userStoreData.getName()).thenReturn(TEST_USER_DATA);
+        when(userStoreData.hasNewDataToSerialize()).thenReturn(true);
+        mWifiConfigStore.registerStoreData(sharedStoreData);
+        mWifiConfigStore.registerStoreData(userStoreData);
+
+        // Read both share and user config store.
+        mWifiConfigStore.setUserStores(mUserStores);
+
+        // Now store some content in the shared and user data files with encrypted data from above.
+        mUserStore.storeRawDataToWrite(
+                String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                        HexEncoding.encodeToString(encryptedData),
+                        HexEncoding.encodeToString(iv),
+                        TEST_USER_DATA).getBytes());
+        mSharedStore.storeRawDataToWrite(
+                String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                        HexEncoding.encodeToString(encryptedData),
+                        HexEncoding.encodeToString(iv),
+                        TEST_SHARE_DATA).getBytes());
+
+        // Read and verify the data content in the store file (metadata stripped out) has been sent
+        // to the corresponding store data when integrity check passes.
+        mWifiConfigStore.read();
+        verify(sharedStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+        verify(userStoreData, times(1)).deserializeData(any(XmlPullParser.class), anyInt());
+
+        // Verify that we parsed the integrity data and used it for checking integrity of the file.
+        ArgumentCaptor<EncryptedData> integrityCaptor =
+                ArgumentCaptor.forClass(EncryptedData.class);
+        ArgumentCaptor<byte[]> dataCaptor = ArgumentCaptor.forClass(byte[].class);
+        // Will be invoked twice for each file - shared & user store file.
+        verify(mDataIntegrityChecker, times(2)).isOk(
+                dataCaptor.capture(), integrityCaptor.capture());
+        // Verify the parsed integrity data
+        assertEquals(2, integrityCaptor.getAllValues().size());
+        EncryptedData parsedEncryptedData1 = integrityCaptor.getAllValues().get(0);
+        assertArrayEquals(encryptedData, parsedEncryptedData1.getEncryptedData());
+        assertArrayEquals(iv, parsedEncryptedData1.getIv());
+        EncryptedData parsedEncryptedData2 = integrityCaptor.getAllValues().get(1);
+        assertArrayEquals(encryptedData, parsedEncryptedData2.getEncryptedData());
+        assertArrayEquals(iv, parsedEncryptedData2.getIv());
+
+        // Verify that we fill in zeros to the data when we performed integrity checked.
+        assertEquals(2, dataCaptor.getAllValues().size());
+        String sharedStoreXmlStringWithZeroedIntegrity =
+                new String(dataCaptor.getAllValues().get(0));
+        assertEquals(String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getEncryptedData()),
+                HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getIv()),
+                TEST_SHARE_DATA), sharedStoreXmlStringWithZeroedIntegrity);
+        String userStoreXmlStringWithZeroedIntegrity = new String(dataCaptor.getAllValues().get(1));
+        assertEquals(String.format(TEST_DATA_XML_STRING_FORMAT_V2_WITH_ONE_DATA_SOURCE,
+                HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getEncryptedData()),
+                HexEncoding.encodeToString(ZEROED_ENCRYPTED_DATA.getIv()),
+                TEST_USER_DATA), userStoreXmlStringWithZeroedIntegrity);
+    }
+
+    /**
      * Mock Store File to redirect all file writes from WifiConfigStore to local buffers.
      * This can be used to examine the data output by WifiConfigStore.
      */
@@ -761,7 +1078,7 @@ public class WifiConfigStoreTest {
         private boolean mStoreWritten;
 
         MockStoreFile(@WifiConfigStore.StoreFileId int fileId) {
-            super(new File("MockStoreFile"), fileId);
+            super(new File("MockStoreFile"), fileId, mDataIntegrityChecker);
         }
 
         @Override
diff --git a/tests/wifitests/src/com/android/server/wifi/util/DataIntegrityCheckerTest.java b/tests/wifitests/src/com/android/server/wifi/util/DataIntegrityCheckerTest.java
index b7076988b..c281b6440 100644
--- a/tests/wifitests/src/com/android/server/wifi/util/DataIntegrityCheckerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/util/DataIntegrityCheckerTest.java
@@ -22,7 +22,6 @@ import org.junit.Ignore;
 import org.junit.Test;
 
 import java.io.File;
-import java.security.DigestException;
 
 /**
  * Unit tests for {@link com.android.server.wifi.util.DataIntegrityChecker}.
@@ -45,8 +44,8 @@ public class DataIntegrityCheckerTest {
                 ".tmp");
         DataIntegrityChecker dataIntegrityChecker = new DataIntegrityChecker(
                 integrityFile.getParent());
-        dataIntegrityChecker.update(sGoodData);
-        assertTrue(dataIntegrityChecker.isOk(sGoodData));
+        EncryptedData encryptedData = dataIntegrityChecker.compute(sGoodData);
+        assertTrue(dataIntegrityChecker.isOk(sGoodData, encryptedData));
     }
 
     /**
@@ -64,25 +63,7 @@ public class DataIntegrityCheckerTest {
                 ".tmp");
         DataIntegrityChecker dataIntegrityChecker = new DataIntegrityChecker(
                 integrityFile.getParent());
-        dataIntegrityChecker.update(sGoodData);
-        assertFalse(dataIntegrityChecker.isOk(sBadData));
-    }
-
-    /**
-     * Verify a corner case where integrity of data that has never been
-     * updated passes and adds the token to the keystore.
-     *
-     * @throws Exception
-     */
-    @Test(expected = DigestException.class)
-    @Ignore
-    public void testIntegrityWithoutUpdate() throws Exception {
-        File tmpFile = File.createTempFile("testIntegrityWithoutUpdate", ".tmp");
-
-        DataIntegrityChecker dataIntegrityChecker = new DataIntegrityChecker(
-                tmpFile.getAbsolutePath());
-
-        // the integrity data is not known, so isOk throws a DigestException
-        assertTrue(dataIntegrityChecker.isOk(sGoodData));
+        EncryptedData encryptedData = dataIntegrityChecker.compute(sGoodData);
+        assertFalse(dataIntegrityChecker.isOk(sBadData, encryptedData));
     }
 }