Fixes AB#3467552 Introduce locks at the cache layer (#2842)

Fixes
[AB#3467552](https://identitydivision.visualstudio.com/fac9d424-53d2-45c0-91b5-ef6ba7a6bf26/_workitems/edit/3467552)
This is a follow up PR to 1. introduce locks closer to the class where
I/O operation is performed. 2. It also introduces a lock at the
BrokerOauth2TokenCache layer to ensure the atomicity of the operations
-> save(), updateMetaDataCache() and loadAggregatedAccountData.

---------

Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: siddhijain <30181294+siddhijain@users.noreply.github.com>

Author: siddhijain

changelog.txt

@@ -1,8 +1,9 @@
 vNext
 ----------
+- [MINOR] Introduce locks at NameValueStorageFileManagerSimpleCacheImpl layer (#2842)
 
 Version 23.2.0
-----------
+
 - [MINOR] Add additional allowed origins for PasskeyWebListener (#2839)
 - [PATCH] Add JavascriptInterface rules to consumer proguard rules (#2837)
 - [MINOR] Add optimized saveAndLoadAggregatedAccountData method in BrokerOAuth2TokenCache (#2832)

common4j/src/main/com/microsoft/identity/common/java/cache/BrokerOAuth2TokenCache.java

@@ -107,6 +107,23 @@ public class BrokerOAuth2TokenCache
     private final MicrosoftFamilyOAuth2TokenCache mFociCache;
     private final int mUid;
     private ProcessUidCacheFactory mDelegate = null;
+
+    /**
+     * Static map of locks for ensuring atomicity of compound save/update/load operations.
+     * Since each thread creates a new BrokerOAuth2TokenCache instance, we need locks
+     * shared across all instances to prevent race conditions when multiple instances
+     * access the same underlying storage (FOCI cache or UID-specific cache).
+     * <p>
+     * Keys are cache identifiers:
+     * <ul>
+     *   <li>"FOCI" - for Family of Client IDs cache (shared by all FOCI apps)</li>
+     *   <li>"UID_<uid>" - for process UID-specific caches (one per UID)</li>
+     * </ul>
+     * This ensures that operations on the same logical cache are serialized, even
+     * when invoked through different BrokerOAuth2TokenCache instances.
+     */
+    private static final ConcurrentHashMap<String, Object> CACHE_OPERATION_LOCKS = new ConcurrentHashMap<>();
+
     /**
      * Shared, process-wide registry of in-memory augmented account/credential caches keyed by
      * storage (SharedPreferences) name.
@@ -1745,10 +1762,14 @@ private List<ICacheRecord> saveAndLoadAggregatedAccountDataOptimized(
                 "refreshTokenRecord, familyId, authScheme)";
 
         final ICacheRecord cacheRecord;
-        synchronized (this) {
-            final long saveAndLoadStartTime = System.currentTimeMillis();
+        final long saveAndLoadStartTime = System.currentTimeMillis();
+
+        final boolean isFoci = !StringUtil.isNullOrEmpty(familyId);
+        // Determine lock key based on which cache we're operating on
+        final String lockKey = isFoci ? "FOCI" : "UID_" + mUid;
+        final Object lock = CACHE_OPERATION_LOCKS.computeIfAbsent(lockKey, k -> new Object());
 
-            final boolean isFoci = !StringUtil.isNullOrEmpty(familyId);
+        synchronized (lock) {
             final MsalOAuth2TokenCache targetCache;
 
             if (isFoci) {
@@ -1808,6 +1829,5 @@ private List<ICacheRecord> saveAndLoadAggregatedAccountDataOptimized(
                     saveAndLoadStartTime);
             return cacheRecordList;
         }
-
     }
 }

common4j/src/main/com/microsoft/identity/common/java/cache/NameValueStorageFileManagerSimpleCacheImpl.java

@@ -25,16 +25,21 @@
 import lombok.NonNull;
 
 import com.google.gson.Gson;
+import com.microsoft.identity.common.java.flighting.CommonFlight;
+import com.microsoft.identity.common.java.flighting.CommonFlightsManager;
 import com.microsoft.identity.common.java.interfaces.IPlatformComponents;
 import com.microsoft.identity.common.java.interfaces.INameValueStorage;
 import com.microsoft.identity.common.java.logging.Logger;
 import com.microsoft.identity.common.java.util.StringUtil;
 
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 /**
  * A simple metadata store definition that uses INameValueStorage to persist, read,
@@ -59,6 +64,12 @@ public abstract class NameValueStorageFileManagerSimpleCacheImpl<T> implements I
     private final boolean mForceReinsertionOfDuplicates;
     private final Gson mGson = new Gson();
 
+    // Lock per storage name to prevent corruption when multiple instances access the same storage file
+    private static final Map<String, ReentrantReadWriteLock> STORAGE_LOCKS = new ConcurrentHashMap<>();
+    private final ReentrantReadWriteLock metadataCacheLock;
+
+    private final boolean useLocks;
+
     /**
      * Constructs a new NameValueStorageFileManagerSimpleCacheImpl. Convenience class for persisting
      * lists of arbitrarily-typed data. Duplicate reinsertion is disabled (backcompat) by default.
@@ -93,6 +104,9 @@ public NameValueStorageFileManagerSimpleCacheImpl(@NonNull final IPlatformCompon
         mStorage = components.getStorageSupplier().getUnencryptedNameValueStore(name, String.class);
         mKeySingleEntry = singleKey;
         mForceReinsertionOfDuplicates = forceReinsertionOfDuplicates;
+        // Get or create a lock for this specific storage file
+        metadataCacheLock = STORAGE_LOCKS.computeIfAbsent(name, k -> new ReentrantReadWriteLock());
+        useLocks = CommonFlightsManager.INSTANCE.getFlightsProvider().isFlightEnabled(CommonFlight.USE_LOCKS_IN_NAME_VALUE_STORAGE);
     }
 
     private interface NamedRunnable<V> extends Callable<V> {
@@ -121,6 +135,9 @@ private <V> V execWithTiming(@NonNull final NamedRunnable<V> runnable) {
 
     @Override
     public boolean insert(final T t) {
+        if (useLocks) {
+            return insertWithLocks(t);
+        }
         return execWithTiming(new NamedRunnable<Boolean>() {
             @Override
             public String getName() {
@@ -129,7 +146,7 @@ public String getName() {
 
             @Override
             public Boolean call() {
-                final Set<T> allMetadata = new HashSet<>(getAll());
+                final Set<T> allMetadata = new HashSet<>(getAllInternal());
 
                 if (mForceReinsertionOfDuplicates) {
                     // This is a bit of workaround for Set's default behavior
@@ -147,27 +164,90 @@ public Boolean call() {
         });
     }
 
+    private boolean insertWithLocks(final T t) {
+        metadataCacheLock.writeLock().lock();
+        try {
+            return execWithTiming(new NamedRunnable<Boolean>() {
+                @Override
+                public String getName() {
+                    return "insert";
+                }
+
+                @Override
+                public Boolean call() {
+                    final Set<T> allMetadata = new HashSet<>(getAllInternal());
+
+                    if (mForceReinsertionOfDuplicates) {
+                        // This is a bit of workaround for Set's default behavior
+                        // where items already within the Set are not replaced if they are
+                        // inserted but already exist.
+                        // This makes it behave more like a Map.
+                        allMetadata.remove(t);
+                    }
+
+                    allMetadata.add(t);
+                    final String json = mGson.toJson(allMetadata);
+                    mStorage.put(mKeySingleEntry, json);
+                    return true;
+                }
+            });
+        } finally {
+            metadataCacheLock.writeLock().unlock();
+        }
+    }
+
     @Override
     public boolean remove(final T t) {
-        return execWithTiming(new NamedRunnable<Boolean>() {
-            @Override
-            public String getName() {
-                return "remove";
-            }
+        metadataCacheLock.writeLock().lock();
+        try {
+            return execWithTiming(new NamedRunnable<Boolean>() {
+                @Override
+                public String getName() {
+                    return "remove";
+                }
+
+                @Override
+                public Boolean call() {
+                    final Set<T> allMetadata = new HashSet<>(getAllInternal());
+                    allMetadata.remove(t);
+                    final String json = mGson.toJson(allMetadata);
+                    mStorage.put(mKeySingleEntry, json);
+                    return true;
+                }
+            });
+        } finally {
+            metadataCacheLock.writeLock().unlock();
+        }
 
-            @Override
-            public Boolean call() {
-                final Set<T> allMetadata = new HashSet<>(getAll());
-                allMetadata.remove(t);
-                final String json = mGson.toJson(allMetadata);
-                mStorage.put(mKeySingleEntry, json);
-                return true;
-            }
-        });
     }
 
     @Override
     public List<T> getAll() {
+        metadataCacheLock.readLock().lock();
+        try {
+            return getAllInternal();
+        } finally {
+            metadataCacheLock.readLock().unlock();
+        }
+    }
+
+    private List<T> getAllWithLocks() {
+        metadataCacheLock.readLock().lock();
+        try {
+            return getAllInternal();
+        } finally {
+            metadataCacheLock.readLock().unlock();
+        }
+    }
+
+    /**
+     * Internal helper method to retrieve all items from storage without acquiring a lock.
+     * This method should only be called when a lock (read or write) is already held by the caller.
+     * For external access, use {@link #getAll()} which properly acquires a read lock.
+     *
+     * @return List of all items in the cache
+     */
+    private List<T> getAllInternal() {
         return execWithTiming(new NamedRunnable<List<T>>() {
             @Override
             public String getName() {
@@ -191,6 +271,9 @@ public List<T> call() {
 
     @Override
     public boolean clear() {
+        if (useLocks) {
+            return clearWithLocks();
+        }
         return execWithTiming(new NamedRunnable<Boolean>() {
             @Override
             public String getName() {
@@ -204,4 +287,24 @@ public Boolean call() {
             }
         });
     }
+
+    private boolean clearWithLocks() {
+        metadataCacheLock.writeLock().lock();
+        try {
+            return execWithTiming(new NamedRunnable<Boolean>() {
+                @Override
+                public String getName() {
+                    return "clear";
+                }
+
+                @Override
+                public Boolean call() {
+                    mStorage.clear();
+                    return true;
+                }
+            });
+        } finally {
+            metadataCacheLock.writeLock().unlock();
+        }
+    }
 }

common4j/src/main/com/microsoft/identity/common/java/flighting/CommonFlight.java

@@ -206,7 +206,12 @@ public enum CommonFlight implements IFlightConfig {
      * Flight to re-enable validating signing certificate chain for broker validation
      * We want to disable the check by default but have the ability to bring it back just in case.
      */
-    RE_ENABLE_VALIDATE_SIGNING_CERT_CHAIN_BROKER_APPS("ReEnableValidateSigningCertChainBrokerApps", false);
+    RE_ENABLE_VALIDATE_SIGNING_CERT_CHAIN_BROKER_APPS("ReEnableValidateSigningCertChainBrokerApps", false),
+
+    /**
+     * Flight to enable the use of locks in name value storage to prevent concurrent access issues.
+     */
+    USE_LOCKS_IN_NAME_VALUE_STORAGE("UseLocksInNameValueStorage", false);
 
     private String key;
     private Object defaultValue;