Merge changes from topic "configstore_integrity_same_file" into qt-qpr1-dev

* changes:
  WifiConfigStore: Store integrity data in same file
  Revert "WifiConfigStore: Limit integrity checks to single user devices"

5 files changed, 238 insertions, 162 deletions

diff --git a/service/java/com/android/server/wifi/WifiConfigManager.java b/service/java/com/android/server/wifi/WifiConfigManager.java
index 8da2d2cb4..b21893404 100644
--- a/service/java/com/android/server/wifi/WifiConfigManager.java
+++ b/service/java/com/android/server/wifi/WifiConfigManager.java
@@ -3094,7 +3094,7 @@ public class WifiConfigManager {
         if (mDeferredUserUnlockRead) {
             Log.i(TAG, "Handling user unlock before loading from store.");
             List<WifiConfigStore.StoreFile> userStoreFiles =
-                    WifiConfigStore.createUserFiles(mCurrentUserId, UserManager.get(mContext));
+                    WifiConfigStore.createUserFiles(mCurrentUserId);
             if (userStoreFiles == null) {
                 Log.wtf(TAG, "Failed to create user store files");
                 return false;
@@ -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) {
@@ -3133,13 +3133,13 @@ public class WifiConfigManager {
     private boolean loadFromUserStoreAfterUnlockOrSwitch(int userId) {
         try {
             List<WifiConfigStore.StoreFile> userStoreFiles =
-                    WifiConfigStore.createUserFiles(userId, UserManager.get(mContext));
+                    WifiConfigStore.createUserFiles(userId);
             if (userStoreFiles == null) {
                 Log.e(TAG, "Failed to create user store files");
                 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 6989e728f..e189d00e1 100644
--- a/service/java/com/android/server/wifi/WifiConfigStore.java
+++ b/service/java/com/android/server/wifi/WifiConfigStore.java
@@ -27,7 +27,6 @@ import android.os.Environment;
 import android.os.FileUtils;
 import android.os.Handler;
 import android.os.Looper;
-import android.os.UserManager;
 import android.util.Log;
 import android.util.SparseArray;
 import android.util.Xml;
@@ -37,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;
@@ -53,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;
@@ -100,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.
@@ -150,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
      */
@@ -209,7 +222,7 @@ public class WifiConfigStore {
      * @param clock       clock instance to retrieve timestamps for alarms.
      * @param wifiMetrics Metrics instance.
      * @param sharedStore StoreFile instance pointing to the shared store file. This should
-     *                    be retrieved using {@link #createSharedFile(UserManager)} method.
+     *                    be retrieved using {@link #createSharedFile()} method.
      */
     public WifiConfigStore(Context context, Looper looper, Clock clock, WifiMetrics wifiMetrics,
             StoreFile sharedStore) {
@@ -230,8 +243,7 @@ public class WifiConfigStore {
     /**
      * Set the user store files.
      * (Useful for mocking in unit tests).
-     * @param userStores List of {@link StoreFile} created using {@link #createUserFiles(int,
-     * UserManager)}.
+     * @param userStores List of {@link StoreFile} created using {@link #createUserFiles(int)}.
      */
     public void setUserStores(@NonNull List<StoreFile> userStores) {
         Preconditions.checkNotNull(userStores);
@@ -268,11 +280,9 @@ public class WifiConfigStore {
      * @param storeBaseDir Base directory under which the store file is to be stored. The store file
      *                     will be at <storeBaseDir>/wifi/WifiConfigStore.xml.
      * @param fileId Identifier for the file. See {@link StoreFileId}.
-     * @param userManager Instance of UserManager to check if the device is in single user mode.
      * @return new instance of the store file or null if the directory cannot be created.
      */
-    private static @Nullable StoreFile createFile(File storeBaseDir, @StoreFileId int fileId,
-                                                  UserManager userManager) {
+    private static @Nullable StoreFile createFile(File storeBaseDir, @StoreFileId int fileId) {
         File storeDir = new File(storeBaseDir, STORE_DIRECTORY_NAME);
         if (!storeDir.exists()) {
             if (!storeDir.mkdir()) {
@@ -281,23 +291,17 @@ public class WifiConfigStore {
             }
         }
         File file = new File(storeDir, STORE_ID_TO_FILE_NAME.get(fileId));
-        DataIntegrityChecker dataIntegrityChecker = null;
-        // Turn on integrity checking only for single user mode devices.
-        if (userManager.hasUserRestriction(UserManager.DISALLOW_ADD_USER)) {
-            dataIntegrityChecker = new DataIntegrityChecker(file.getAbsolutePath());
-        }
+        DataIntegrityChecker dataIntegrityChecker = new DataIntegrityChecker(file.getName());
         return new StoreFile(file, fileId, dataIntegrityChecker);
     }
 
     /**
      * Create a new instance of the shared store file.
      *
-     * @param userManager Instance of UserManager to check if the device is in single user mode.
      * @return new instance of the store file or null if the directory cannot be created.
      */
-    public static @Nullable StoreFile createSharedFile(UserManager userManager) {
-        return createFile(
-                Environment.getDataMiscDirectory(), STORE_FILE_SHARED_GENERAL, userManager);
+    public static @Nullable StoreFile createSharedFile() {
+        return createFile(Environment.getDataMiscDirectory(), STORE_FILE_SHARED_GENERAL);
     }
 
     /**
@@ -305,16 +309,14 @@ public class WifiConfigStore {
      * The user store file is inside the user's encrypted data directory.
      *
      * @param userId userId corresponding to the currently logged-in user.
-     * @param userManager Instance of UserManager to check if the device is in single user mode.
      * @return List of new instances of the store files created or null if the directory cannot be
      * created.
      */
-    public static @Nullable List<StoreFile> createUserFiles(int userId, UserManager userManager) {
+    public static @Nullable List<StoreFile> createUserFiles(int userId) {
         List<StoreFile> storeFiles = new ArrayList<>();
         for (int fileId : Arrays.asList(
                 STORE_FILE_USER_GENERAL, STORE_FILE_USER_NETWORK_SUGGESTIONS)) {
-            StoreFile storeFile =
-                    createFile(Environment.getDataMiscCeDirectory(userId), fileId, userManager);
+            StoreFile storeFile = createFile(Environment.getDataMiscCeDirectory(userId), fileId);
             if (storeFile == null) {
                 return null;
             }
@@ -409,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
@@ -422,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);
@@ -432,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();
     }
 
     /**
@@ -516,8 +592,7 @@ public class WifiConfigStore {
      * Handles a user switch. This method changes the user specific store files and reads from the
      * new user's store files.
      *
-     * @param userStores List of {@link StoreFile} created using {@link #createUserFiles(int,
-     * UserManager)}.
+     * @param userStores List of {@link StoreFile} created using {@link #createUserFiles(int)}.
      */
     public void switchUserStoresAndRead(@NonNull List<StoreFile> userStores)
             throws XmlPullParserException, IOException {
@@ -564,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.
@@ -584,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<>();
@@ -610,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) {
@@ -628,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) {
@@ -668,21 +782,20 @@ public class WifiConfigStore {
         /**
          * Store the file name for setting the file permissions/logging purposes.
          */
-        private String mFileName;
+        private final String mFileName;
         /**
          * {@link StoreFileId} Type of store file.
          */
-        private @StoreFileId int mFileId;
+        private final @StoreFileId int mFileId;
         /**
-         * The integrity file storing integrity checking data for the store file.
-         * Note: This is only turned on for single user devices.
+         * Integrity checking for the store file.
          */
-        private @Nullable DataIntegrityChecker mDataIntegrityChecker;
+        private final DataIntegrityChecker mDataIntegrityChecker;
 
         public StoreFile(File file, @StoreFileId int fileId,
-                         @Nullable DataIntegrityChecker dataIntegrityChecker) {
+                         @NonNull DataIntegrityChecker dataIntegrityChecker) {
             mAtomicFile = new AtomicFile(file);
-            mFileName = mAtomicFile.getBaseFile().getAbsolutePath();
+            mFileName = file.getAbsolutePath();
             mFileId = fileId;
             mDataIntegrityChecker = dataIntegrityChecker;
         }
@@ -696,7 +809,6 @@ public class WifiConfigStore {
             return mAtomicFile.exists();
         }
 
-
         /**
          * Read the entire raw data from the store file and return in a byte array.
          *
@@ -712,22 +824,6 @@ public class WifiConfigStore {
             } catch (FileNotFoundException e) {
                 return null;
             }
-            if (mDataIntegrityChecker != null) {
-                // Check that the file has not been altered since last writeBufferedRawData()
-                try {
-                    if (!mDataIntegrityChecker.isOk(bytes)) {
-                        Log.wtf(TAG, "Data integrity problem with file: " + mFileName);
-                        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;
         }
 
@@ -763,13 +859,36 @@ public class WifiConfigStore {
                 }
                 throw e;
             }
-            if (mDataIntegrityChecker != null) {
-                // 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/WifiInjector.java b/service/java/com/android/server/wifi/WifiInjector.java
index 762b24b0a..8095bfac6 100644
--- a/service/java/com/android/server/wifi/WifiInjector.java
+++ b/service/java/com/android/server/wifi/WifiInjector.java
@@ -241,7 +241,7 @@ public class WifiInjector {
         mWifiKeyStore = new WifiKeyStore(mKeyStore);
         mWifiConfigStore = new WifiConfigStore(
                 mContext, clientModeImplLooper, mClock, mWifiMetrics,
-                WifiConfigStore.createSharedFile(UserManager.get(mContext)));
+                WifiConfigStore.createSharedFile());
         SubscriptionManager subscriptionManager =
                 mContext.getSystemService(SubscriptionManager.class);
         // Config Manager
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;
-    }
 }