WifiConfigStore: Store integrity data in same file

Store the computed integrity data back in the same config store file.
The previous approach of storing the integrity data in a separate file
made config store updates non-atomic (look at associated bug for
details).

New approach:
a) Store the integrity data at the start of each config store XML file
(version + integrity == metadata for each file).
b) Uprev the config store version to 2 to support the new format.
c) Since we need to an-in place integrity check, when we write the file
For writes:
i) We fill up the integrity fields with the zeroes for the expected
number of bytes in the store file contents.
ii) Compute the integrity for the entire file contents.
iii) Rewrite the document metadata (version & integrity data) with the
newly computed integrity data in store file contents.
iv) Write the file contents to disk.
For reads:
i) Parse out the version & integrity contents from the file contents.
ii) Rewrite the document metadata (version & integrity data) with zeroed
integrity data in store file contents.
iii) Compute the integrity data for the modified file contents created
from (ii) and validate the result with the parsed value from (i).
iv) If the integrity check passes, continue with the parsing of the
document, else abort.
d) Since we need fixed size fields in the integrity fields, remove
storage of keystore alias string from |EncryptedData|. This can anyway
be trivially computed from the config store file name.

Bug: 138482990
Test: Verified that the device does not lose any stored networks on
reboot when the config store file is not modified.
Test:  Verified that the device discards all stored networks on reboot
when the config store file is modified.
Test: atest com.android.server.wifi
Test: Will send for full regression test.

Change-Id: I528d3402cb047cca3793be5f1386c4bb60c39a10

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));
     }
 }