fix: added branchName to state for redirection (#28822)

app/client/src/sagas/DatasourcesSagas.ts

@@ -161,6 +161,7 @@ import {
   getCurrentEnvironmentDetails,
 } from "@appsmith/selectors/environmentSelectors";
 import { waitForFetchEnvironments } from "@appsmith/sagas/EnvironmentSagas";
+import { getCurrentGitBranch } from "selectors/gitSyncSelectors";
 
 function* fetchDatasourcesSaga(
   action: ReduxAction<{ workspaceId?: string } | undefined>,
@@ -619,12 +620,17 @@ function* redirectAuthorizationCodeSaga(
 ) {
   const { datasourceId, pageId, pluginType } = actionPayload.payload;
   const isImport: string = yield select(getWorkspaceIdForImport);
+  const branchName: string | undefined = yield select(getCurrentGitBranch);
 
   if (pluginType === PluginType.API) {
     const currentEnvironment: string = yield select(
       getCurrentEditingEnvironmentId,
     );
-    window.location.href = `/api/v1/datasources/${datasourceId}/pages/${pageId}/code?environmentId=${currentEnvironment}`;
+    let windowLocation = `/api/v1/datasources/${datasourceId}/pages/${pageId}/code?environmentId=${currentEnvironment}`;
+    if (!!branchName) {
+      windowLocation = windowLocation + `&branchName=` + branchName;
+    }
+    window.location.href = windowLocation;
   } else {
     try {
       // Get an "appsmith token" from the server

app/server/appsmith-server/src/main/java/com/appsmith/server/controllers/ce/DatasourceControllerCE.java

@@ -158,13 +158,14 @@ public class DatasourceControllerCE {
     public Mono<Void> getTokenRequestUrl(
             @PathVariable String datasourceId,
             @PathVariable String pageId,
-            ServerWebExchange serverWebExchange,
+            @RequestParam String environmentId,
-            @RequestParam String environmentId) {
+            @RequestParam(name = FieldName.BRANCH_NAME, required = false) String branchName,
+            ServerWebExchange serverWebExchange) {
         log.debug(
                 "Going to retrieve token request URL for datasource with id: {} and page id: {}", datasourceId, pageId);
         return authenticationService
                 .getAuthorizationCodeURLForGenericOAuth2(
-                        datasourceId, environmentId, pageId, serverWebExchange.getRequest())
+                        datasourceId, environmentId, pageId, branchName, serverWebExchange.getRequest())
                 .flatMap(url -> {
                     serverWebExchange.getResponse().setStatusCode(HttpStatus.FOUND);
                     serverWebExchange.getResponse().getHeaders().setLocation(URI.create(url));

app/server/appsmith-server/src/main/java/com/appsmith/server/solutions/ce/AuthenticationServiceCE.java

@@ -19,7 +19,7 @@ public interface AuthenticationServiceCE {
      * @return a url String to continue the authorization flow
      */
     Mono<String> getAuthorizationCodeURLForGenericOAuth2(
-            String datasourceId, String environmentId, String pageId, ServerHttpRequest httpRequest);
+            String datasourceId, String environmentId, String pageId, String branchName, ServerHttpRequest httpRequest);
 
     /**
      * This is the method that handles callback for generic OAuth2. We will be retrieving and storing token information here

app/server/appsmith-server/src/main/java/com/appsmith/server/solutions/ce/AuthenticationServiceCEImpl.java

@@ -101,7 +101,11 @@ public class AuthenticationServiceCEImpl implements AuthenticationServiceCE {
      * @return a url String to continue the authorization flow
      */
     public Mono<String> getAuthorizationCodeURLForGenericOAuth2(
-            String datasourceId, String environmentId, String pageId, ServerHttpRequest httpRequest) {
+            String datasourceId,
+            String environmentId,
+            String pageId,
+            String branchName,
+            ServerHttpRequest httpRequest) {
         // This is the only database access that is controlled by ACL
         // The rest of the queries in this flow will not have context information
 
@@ -130,6 +134,9 @@ public class AuthenticationServiceCEImpl implements AuthenticationServiceCE {
                     OAuth2 oAuth2 = (OAuth2)
                             datasourceStorage.getDatasourceConfiguration().getAuthentication();
                     final String redirectUri = redirectHelper.getRedirectDomain(httpRequest.getHeaders());
+                    final String state = StringUtils.hasText(branchName)
+                            ? String.join(",", pageId, datasourceId, trueEnvironmentId, redirectUri, branchName)
+                            : String.join(",", pageId, datasourceId, trueEnvironmentId, redirectUri);
                     // Adding basic uri components
                     UriComponentsBuilder uriComponentsBuilder = UriComponentsBuilder.fromUriString(
                                     oAuth2.getAuthorizationUrl())
@@ -138,7 +145,7 @@ public class AuthenticationServiceCEImpl implements AuthenticationServiceCE {
                             .queryParam(REDIRECT_URI, redirectUri + Url.DATASOURCE_URL + "/authorize")
                             // The state is used internally to calculate the redirect url when returning control to the
                             // client
-                            .queryParam(STATE, String.join(",", pageId, datasourceId, trueEnvironmentId, redirectUri));
+                            .queryParam(STATE, state);
                     // Adding optional scope parameter
                     if (oAuth2.getScope() != null && !oAuth2.getScope().isEmpty()) {
                         uriComponentsBuilder.queryParam(
@@ -196,7 +203,7 @@ public class AuthenticationServiceCEImpl implements AuthenticationServiceCE {
                 .switchIfEmpty(Mono.error(new AppsmithException(AppsmithError.UNAUTHORIZED_ACCESS)))
                 .flatMap(localState -> {
                     String[] splitStates = localState.split(",");
-                    if (splitStates.length != 4) {
+                    if (splitStates.length < 4) {
                         return Mono.error(new AppsmithException(AppsmithError.UNAUTHORIZED_ACCESS));
                     } else
                         return datasourceService
@@ -324,13 +331,14 @@ public class AuthenticationServiceCEImpl implements AuthenticationServiceCE {
         final String datasourceId = splitState[1];
         final String environmentId = splitState[2];
         final String redirectOrigin = splitState[3];
+        final String branchName = splitState.length == 5 ? splitState[4] : null;
         String response = SUCCESS;
         if (error != null) {
             response = error;
         }
         final String responseStatus = response;
         return newPageService
-                .getById(pageId)
+                .findByIdAndBranchName(pageId, branchName)
                 .map(newPage -> redirectOrigin + Entity.SLASH + Entity.APPLICATIONS
                         + Entity.SLASH + newPage.getApplicationId()
                         + Entity.SLASH + Entity.PAGES
@@ -339,7 +347,9 @@ public class AuthenticationServiceCEImpl implements AuthenticationServiceCE {
                         + Entity.SLASH + Entity.DATASOURCE
                         + Entity.SLASH + datasourceId
                         + "?response_status="
-                        + responseStatus + "&view_mode=true")
+                        + responseStatus
+                        + "&view_mode=true"
+                        + (StringUtils.hasText(branchName) ? "&branch=" + branchName : ""))
                 .onErrorResume(e -> Mono.just(redirectOrigin + Entity.SLASH + Entity.APPLICATIONS
                         + "?response_status="
                         + responseStatus + "&view_mode=true"));

app/server/appsmith-server/src/test/java/com/appsmith/server/solutions/AuthenticationServiceTest.java

@@ -109,7 +109,7 @@ public class AuthenticationServiceTest {
     @WithUserDetails(value = "api_user")
     public void testGetAuthorizationCodeURL_missingDatasource() {
         Mono<String> authorizationCodeUrlMono = authenticationService.getAuthorizationCodeURLForGenericOAuth2(
-                "invalidId", FieldName.UNUSED_ENVIRONMENT_ID, "irrelevantPageId", null);
+                "invalidId", FieldName.UNUSED_ENVIRONMENT_ID, "irrelevantPageId", null, null);
 
         StepVerifier.create(authorizationCodeUrlMono)
                 .expectErrorMatches(throwable -> throwable instanceof AppsmithException
@@ -150,7 +150,7 @@ public class AuthenticationServiceTest {
         Mono<String> authorizationCodeUrlMono = datasourceMono
                 .map(BaseDomain::getId)
                 .flatMap(datasourceId -> authenticationService.getAuthorizationCodeURLForGenericOAuth2(
-                        datasourceId, defaultEnvironmentId, "irrelevantPageId", null));
+                        datasourceId, defaultEnvironmentId, "irrelevantPageId", null, null));
 
         StepVerifier.create(authorizationCodeUrlMono)
                 .expectErrorMatches(throwable -> throwable instanceof AppsmithException
@@ -228,7 +228,7 @@ public class AuthenticationServiceTest {
         final String datasourceId1 = datasourceMono.map(BaseDomain::getId).block();
 
         Mono<String> authorizationCodeUrlMono = authenticationService.getAuthorizationCodeURLForGenericOAuth2(
-                datasourceId1, defaultEnvironmentId, pageDto.getId(), httpRequest);
+                datasourceId1, defaultEnvironmentId, pageDto.getId(), null, httpRequest);
 
         StepVerifier.create(authorizationCodeUrlMono)
                 .assertNext(url -> {
@@ -248,4 +248,98 @@ public class AuthenticationServiceTest {
                 })
                 .verifyComplete();
     }
+
+    @Test
+    @WithUserDetails(value = "api_user")
+    public void testGetAuthorizationCodeURL_validDatasourceAndBranchName() {
+        LinkedMultiValueMap<String, String> mockHeaders = new LinkedMultiValueMap<>(1);
+        mockHeaders.add(HttpHeaders.REFERER, "https://mock.origin.com/source/uri");
+
+        MockServerHttpRequest httpRequest =
+                MockServerHttpRequest.get("").headers(mockHeaders).build();
+
+        Workspace testWorkspace = new Workspace();
+        testWorkspace.setName("Test-Workspace-OAuth2-git-redirection");
+        testWorkspace = workspaceService.create(testWorkspace).block();
+        String workspaceId = testWorkspace == null ? "" : testWorkspace.getId();
+        Mockito.when(pluginExecutorHelper.getPluginExecutor(Mockito.any()))
+                .thenReturn(Mono.just(new MockPluginExecutor()));
+
+        String defaultEnvironmentId = workspaceService
+                .getDefaultEnvironmentId(workspaceId, environmentPermission.getExecutePermission())
+                .block();
+
+        PageDTO testPage = new PageDTO();
+        testPage.setName("Test-Page-oauth2-git-redirection");
+
+        Application newApp = new Application();
+        newApp.setName(UUID.randomUUID().toString());
+        Application application =
+                applicationPageService.createApplication(newApp, workspaceId).block();
+
+        testPage.setApplicationId(application.getId());
+
+        PageDTO pageDTO = applicationPageService.createPage(testPage).block();
+
+        Mono<Plugin> pluginMono = pluginService.findByName("Installed Plugin Name");
+        // install plugin
+        pluginMono
+                .flatMap(plugin -> pluginService.installPlugin(PluginWorkspaceDTO.builder()
+                        .pluginId(plugin.getId())
+                        .workspaceId(workspaceId)
+                        .status(WorkspacePluginStatus.FREE)
+                        .build()))
+                .block();
+        Datasource datasource = new Datasource();
+        datasource.setName("Valid datasource for OAuth2");
+        datasource.setWorkspaceId(workspaceId);
+        DatasourceConfiguration datasourceConfiguration = new DatasourceConfiguration();
+        datasourceConfiguration.setUrl("http://test.com");
+        OAuth2 authenticationDTO = new OAuth2();
+        authenticationDTO.setClientId("ClientId");
+        authenticationDTO.setClientSecret("ClientSecret");
+        authenticationDTO.setAuthorizationUrl("AuthorizationURL");
+        authenticationDTO.setAccessTokenUrl("AccessTokenURL");
+        authenticationDTO.setScope(Set.of("Scope1", "Scope2"));
+        authenticationDTO.setCustomAuthenticationParameters(Set.of(new Property("key", "value")));
+        datasourceConfiguration.setAuthentication(authenticationDTO);
+        datasource.setDatasourceConfiguration(datasourceConfiguration);
+
+        HashMap<String, DatasourceStorageDTO> storages = new HashMap<>();
+        datasource.setDatasourceStorages(storages);
+        storages.put(
+                defaultEnvironmentId, new DatasourceStorageDTO(null, defaultEnvironmentId, datasourceConfiguration));
+
+        Mono<Datasource> datasourceMono = pluginMono
+                .map(plugin -> {
+                    datasource.setPluginId(plugin.getId());
+                    return datasource;
+                })
+                .flatMap(datasourceService::create)
+                .cache();
+
+        final String datasourceId = datasourceMono.map(BaseDomain::getId).block();
+
+        Mono<String> authorizationCodeUrlMono = authenticationService.getAuthorizationCodeURLForGenericOAuth2(
+                datasourceId, null, pageDTO.getId(), "testBranch", httpRequest);
+
+        StepVerifier.create(authorizationCodeUrlMono)
+                .assertNext(url -> {
+                    assertThat(url)
+                            .matches(Pattern.compile("AuthorizationURL" + "\\?client_id=ClientId"
+                                    + "&response_type=code"
+                                    + "&redirect_uri=https://mock.origin.com/api/v1/datasources/authorize"
+                                    + "&state="
+                                    + String.join(
+                                            ",",
+                                            pageDTO.getId(),
+                                            datasourceId,
+                                            defaultEnvironmentId,
+                                            "https://mock.origin.com",
+                                            "testBranch")
+                                    + "&scope=Scope\\d%20Scope\\d"
+                                    + "&key=value"));
+                })
+                .verifyComplete();
+    }
 }