Fixing stale cache issue post snapshot restore

Signed-off-by: Nagaraj G <[email protected]>

Author: Nagaraj G

src/main/java/org/opensearch/security/OpenSearchSecurityPlugin.java

@@ -870,6 +870,8 @@ public void onQueryPhase(SearchContext searchContext, long tookInNanos) {
                     }
                 }
             }.toListener());
+
+            indexModule.addIndexEventListener(cr);
         }
     }
 

src/main/java/org/opensearch/security/configuration/ConfigurationRepository.java

@@ -75,9 +75,13 @@
 import org.opensearch.common.util.concurrent.ThreadContext.StoredContext;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.Strings;
+import org.opensearch.core.index.Index;
+import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.env.Environment;
+import org.opensearch.index.shard.IndexEventListener;
+import org.opensearch.index.shard.IndexShard;
 import org.opensearch.security.auditlog.AuditLog;
 import org.opensearch.security.auditlog.config.AuditConfig;
 import org.opensearch.security.securityconf.DynamicConfigFactory;
@@ -93,8 +97,9 @@
 import org.opensearch.transport.client.Client;
 
 import static org.opensearch.security.support.ConfigConstants.SECURITY_ALLOW_DEFAULT_INIT_USE_CLUSTER_STATE;
+import static org.opensearch.security.support.SnapshotRestoreHelper.isSecurityIndexRestoredFromSnapshot;
 
-public class ConfigurationRepository implements ClusterStateListener {
+public class ConfigurationRepository implements ClusterStateListener, IndexEventListener {
     private static final Logger LOGGER = LogManager.getLogger(ConfigurationRepository.class);
 
     private final String securityIndex;
@@ -668,4 +673,21 @@ public Void run() {
             });
         }
     }
+
+    @Override
+    public void afterIndexShardStarted(IndexShard indexShard) {
+        final ShardId shardId = indexShard.shardId();
+        final Index index = shardId.getIndex();
+
+        // Check if this is a security index shard
+        if (securityIndex.equals(index.getName())) {
+            // Only trigger on primary shard to avoid multiple reloads
+            if (indexShard.routingEntry() != null && indexShard.routingEntry().primary()) {
+                if (isSecurityIndexRestoredFromSnapshot(clusterService, index, securityIndex)) {
+                    LOGGER.info("Security index primary shard {} started - " + "config reloading for snapshot restore", shardId);
+                    reloadConfiguration(CType.values());
+                }
+            }
+        }
+    }
 }

src/main/java/org/opensearch/security/support/SnapshotRestoreHelper.java

@@ -30,14 +30,18 @@
 import java.security.PrivilegedAction;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 import org.opensearch.SpecialPermission;
 import org.opensearch.action.admin.cluster.snapshots.restore.RestoreSnapshotRequest;
 import org.opensearch.action.support.PlainActionFuture;
+import org.opensearch.cluster.RestoreInProgress;
+import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.util.IndexUtils;
+import org.opensearch.core.index.Index;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.repositories.Repository;
 import org.opensearch.security.OpenSearchSecurityPlugin;
@@ -89,6 +93,24 @@ public static SnapshotInfo getSnapshotInfo(RestoreSnapshotRequest restoreRequest
         return snapshotInfo;
     }
 
+    public static boolean isSecurityIndexRestoredFromSnapshot(ClusterService clusterService, Index index, String securityIndex) {
+        final String threadName = Thread.currentThread().getName();
+        boolean isSecurityIndexRestoredFromSnapshot = false;
+
+        try {
+            setCurrentThreadName("[" + ThreadPool.Names.GENERIC + "]");
+            isSecurityIndexRestoredFromSnapshot = Optional.ofNullable(
+                clusterService.state().<RestoreInProgress>custom(RestoreInProgress.TYPE)
+            ).map(restore -> restore.iterator().hasNext() && restore.iterator().next().indices().contains(securityIndex)).orElse(false);
+        } catch (Exception e) {
+            log.warn("Could not determine if index {} was restored from snapshot, assuming new index", index.getName(), e);
+        } finally {
+            setCurrentThreadName(threadName);
+        }
+
+        return isSecurityIndexRestoredFromSnapshot;
+    }
+
     @SuppressWarnings("removal")
     private static void setCurrentThreadName(final String name) {
         final SecurityManager sm = System.getSecurityManager();

src/test/java/org/opensearch/security/configuration/ConfigurationRepositoryTest.java

@@ -15,6 +15,7 @@
 import java.io.IOException;
 import java.nio.file.Path;
 import java.time.Instant;
+import java.util.Collections;
 import java.util.Set;
 import java.util.concurrent.TimeoutException;
 
@@ -34,17 +35,22 @@
 import org.opensearch.cluster.ClusterChangedEvent;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.ClusterStateUpdateTask;
+import org.opensearch.cluster.RestoreInProgress;
 import org.opensearch.cluster.block.ClusterBlocks;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.metadata.MappingMetadata;
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.cluster.node.DiscoveryNodes;
+import org.opensearch.cluster.routing.ShardRouting;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Priority;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.index.Index;
+import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.index.shard.IndexShard;
 import org.opensearch.security.DefaultObjectMapper;
 import org.opensearch.security.auditlog.AuditLog;
 import org.opensearch.security.securityconf.DynamicConfigFactory;
@@ -63,6 +69,7 @@
 
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnitRunner;
 import org.mockito.stubbing.OngoingStubbing;
 
@@ -82,9 +89,11 @@
 import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doCallRealMethod;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -569,6 +578,47 @@ public void getConfigurationsFromIndex_SecurityIndexNotInitiallyReady() throws I
         assertThat(result.size(), is(CType.values().size()));
     }
 
+    @Test
+    public void afterIndexShardStarted_whenSecurityIndexUpdated() throws InterruptedException, TimeoutException {
+        Settings settings = Settings.builder().build();
+        ConfigurationRepository configurationRepository = spy(createConfigurationRepository(settings));
+        IndexShard indexShard = mock(IndexShard.class);
+        ShardRouting shardRouting = mock(ShardRouting.class);
+        ShardId shardId = mock(ShardId.class);
+        Index index = mock(Index.class);
+        ClusterState mockClusterState = mock(ClusterState.class);
+        RestoreInProgress mockRestore = mock(RestoreInProgress.class);
+        RestoreInProgress.Entry mockEntry = mock(RestoreInProgress.Entry.class);
+
+        // Setup mock behavior
+        when(indexShard.shardId()).thenReturn(shardId);
+        when(shardId.getIndex()).thenReturn(index);
+        when(index.getName()).thenReturn(ConfigConstants.OPENDISTRO_SECURITY_DEFAULT_CONFIG_INDEX);
+        when(indexShard.routingEntry()).thenReturn(shardRouting);
+        when(clusterService.state()).thenReturn(mockClusterState);
+        when(mockClusterState.custom(RestoreInProgress.TYPE)).thenReturn(mockRestore);
+
+        // when replica shard updated
+        when(shardRouting.primary()).thenReturn(false);
+        configurationRepository.afterIndexShardStarted(indexShard);
+        verify(configurationRepository, never()).reloadConfiguration(any());
+
+        // when primary shard updated
+        doReturn(true).when(configurationRepository).reloadConfiguration(any());
+        when(shardRouting.primary()).thenReturn(true);
+        when(mockRestore.iterator()).thenReturn(Collections.singletonList(mockEntry).iterator());
+        when(mockEntry.indices()).thenReturn(Collections.singletonList(ConfigConstants.OPENDISTRO_SECURITY_DEFAULT_CONFIG_INDEX));
+        configurationRepository.afterIndexShardStarted(indexShard);
+        verify(configurationRepository).reloadConfiguration(CType.values());
+
+        // When there is error in checking if restored from snapshot
+        Mockito.reset(configurationRepository);
+        when(clusterService.state()).thenThrow(new RuntimeException("ClusterState exception"));
+        when(shardRouting.primary()).thenReturn(true);
+        configurationRepository.afterIndexShardStarted(indexShard);
+        verify(configurationRepository, never()).reloadConfiguration(any());
+    }
+
     void assertClusterState(final ArgumentCaptor<ClusterStateUpdateTask> clusterStateUpdateTaskCaptor) throws Exception {
         final var initializedStateUpdate = clusterStateUpdateTaskCaptor.getValue();
         assertThat(initializedStateUpdate.priority(), is(Priority.IMMEDIATE));