Implement domain removal listener to clean up OAuth providers on domain deletion

Author: Damans227

plugins/user-authenticators/oauth2/src/main/java/org/apache/cloudstack/oauth2/OAuth2AuthManagerImpl.java

@@ -19,11 +19,14 @@
 package org.apache.cloudstack.oauth2;
 
 import com.cloud.domain.Domain;
+import com.cloud.domain.DomainVO;
+import com.cloud.user.DomainManager;
 import com.cloud.user.DomainService;
 import com.cloud.user.dao.UserDao;
 import com.cloud.utils.component.Manager;
 import com.cloud.utils.component.ManagerBase;
 import com.cloud.utils.exception.CloudRuntimeException;
+import org.apache.cloudstack.framework.messagebus.MessageBus;
 import org.apache.cloudstack.api.ApiConstants;
 import org.apache.cloudstack.auth.UserOAuth2Authenticator;
 import org.apache.cloudstack.framework.config.ConfigKey;
@@ -58,6 +61,9 @@ public class OAuth2AuthManagerImpl extends ManagerBase implements OAuth2AuthMana
     @Inject
     private DomainService _domainService;
 
+    @Inject
+    private MessageBus _messageBus;
+
     protected static Map<String, UserOAuth2Authenticator> userOAuth2AuthenticationProvidersMap = new HashMap<>();
 
     private List<UserOAuth2Authenticator> userOAuth2AuthenticationProviders;
@@ -74,10 +80,26 @@ public List<Class<?>> getAuthCommands() {
     @Override
     public boolean start() {
         initializeUserOAuth2AuthenticationProvidersMap();
+        addDomainRemovalListener();
         logger.info("OAUTH plugin loaded");
         return true;
     }
 
+    private void addDomainRemovalListener() {
+        _messageBus.subscribe(DomainManager.MESSAGE_PRE_REMOVE_DOMAIN_EVENT, (senderAddress, subject, args) -> {
+            try {
+                long domainId = ((DomainVO) args).getId();
+                List<OauthProviderVO> providers = _oauthProviderDao.listByDomain(domainId);
+                for (OauthProviderVO provider : providers) {
+                    _oauthProviderDao.expunge(provider.getId());
+                    logger.debug("Removed OAuth provider {} for deleted domain {}", provider.getProvider(), domainId);
+                }
+            } catch (Exception e) {
+                logger.error("Failed to remove OAuth providers for deleted domain", e);
+            }
+        });
+    }
+
     protected boolean isOAuthPluginEnabled(Long domainId) {
         return OAuth2AuthManager.isPluginEnabledForDomain(domainId);
     }

plugins/user-authenticators/oauth2/src/main/java/org/apache/cloudstack/oauth2/dao/OauthProviderDao.java

@@ -27,6 +27,8 @@ public interface OauthProviderDao extends GenericDao<OauthProviderVO, Long> {
 
     public List<OauthProviderVO> listByDomainIncludingGlobal(Long domainId);
 
+    public List<OauthProviderVO> listByDomain(Long domainId);
+
     public OauthProviderVO findByProviderAndDomainWithGlobalFallback(String provider, Long domainId);
 
 }

plugins/user-authenticators/oauth2/src/main/java/org/apache/cloudstack/oauth2/dao/OauthProviderDaoImpl.java

@@ -54,6 +54,13 @@ public List<OauthProviderVO> listByDomainIncludingGlobal(Long domainId) {
         return listBy(sc);
     }
 
+    @Override
+    public List<OauthProviderVO> listByDomain(Long domainId) {
+        SearchCriteria<OauthProviderVO> sc = createSearchCriteria();
+        sc.addAnd("domainId", SearchCriteria.Op.EQ, domainId);
+        return listBy(sc);
+    }
+
     @Override
     public OauthProviderVO findByProviderAndDomainWithGlobalFallback(String provider, Long domainId) {
         OauthProviderVO providerVO = null;

plugins/user-authenticators/oauth2/src/test/java/org/apache/cloudstack/oauth2/OAuth2AuthManagerImplTest.java

@@ -20,9 +20,12 @@
 package org.apache.cloudstack.oauth2;
 
 import com.cloud.domain.Domain;
+import com.cloud.domain.DomainVO;
 import com.cloud.user.DomainService;
 import com.cloud.utils.exception.CloudRuntimeException;
 import org.apache.cloudstack.api.ApiConstants;
+import org.apache.cloudstack.framework.messagebus.MessageBus;
+import org.apache.cloudstack.framework.messagebus.MessageSubscriber;
 import org.apache.cloudstack.oauth2.api.command.DeleteOAuthProviderCmd;
 import org.apache.cloudstack.oauth2.api.command.RegisterOAuthProviderCmd;
 import org.apache.cloudstack.oauth2.api.command.UpdateOAuthProviderCmd;
@@ -49,6 +52,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class OAuth2AuthManagerImplTest {
@@ -63,6 +67,9 @@ public class OAuth2AuthManagerImplTest {
     @Mock
     DomainService _domainService;
 
+    @Mock
+    MessageBus _messageBus;
+
     AutoCloseable closeable;
     @Before
     public void setUp() {
@@ -636,4 +643,37 @@ public void testDuplicateGlobalProviderRejected() {
         }
     }
 
+    @Test
+    public void testDomainDeletionCleansUpOAuthProviders() {
+        Long domainId = 42L;
+
+        OauthProviderVO provider1 = new OauthProviderVO();
+        provider1.setProvider("github");
+        provider1.setDomainId(domainId);
+
+        OauthProviderVO provider2 = new OauthProviderVO();
+        provider2.setProvider("google");
+        provider2.setDomainId(domainId);
+
+        when(_oauthProviderDao.listByDomain(domainId)).thenReturn(Arrays.asList(provider1, provider2));
+        when(_oauthProviderDao.expunge(Mockito.anyLong())).thenReturn(true);
+
+        // Capture the subscriber registered during start()
+        doNothing().when(_authManager).initializeUserOAuth2AuthenticationProvidersMap();
+        Mockito.doAnswer(invocation -> {
+            String subject = invocation.getArgument(0);
+            MessageSubscriber subscriber = invocation.getArgument(1);
+            // Simulate domain removal event
+            DomainVO domain = Mockito.mock(DomainVO.class);
+            when(domain.getId()).thenReturn(domainId);
+            subscriber.onPublishMessage("", subject, domain);
+            return null;
+        }).when(_messageBus).subscribe(Mockito.eq(com.cloud.user.DomainManager.MESSAGE_PRE_REMOVE_DOMAIN_EVENT), Mockito.any());
+
+        _authManager.start();
+
+        verify(_oauthProviderDao).listByDomain(domainId);
+        verify(_oauthProviderDao, Mockito.times(2)).expunge(Mockito.anyLong());
+    }
+
 }