fix(authentication) redirection for native login and sso to function within iframes (#13453)

Co-authored-by: Esteban Gutierrez <esteban.gutierrez@acryl.io>

datahub-frontend/app/controllers/AuthenticationController.java

@@ -18,6 +18,7 @@ import auth.sso.SsoManager;
 import client.AuthServiceClient;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.common.urn.CorpuserUrn;
 import com.linkedin.common.urn.Urn;
 import com.typesafe.config.Config;
@@ -25,6 +26,7 @@ import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URLEncoder;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.Base64;
 import java.util.Optional;
 import javax.annotation.Nonnull;
@@ -33,7 +35,6 @@ import org.apache.commons.httpclient.InvalidRedirectLocationException;
 import org.apache.commons.lang3.StringUtils;
 import org.pac4j.core.client.Client;
 import org.pac4j.core.context.CallContext;
-import org.pac4j.core.context.Cookie;
 import org.pac4j.core.context.WebContext;
 import org.pac4j.core.exception.http.FoundAction;
 import org.pac4j.core.exception.http.RedirectionAction;
@@ -70,9 +71,9 @@ public class AuthenticationController extends Controller {
 
   @Inject private org.pac4j.core.config.Config ssoConfig;
 
-  @Inject private PlayCookieSessionStore playCookieSessionStore;
+  @VisibleForTesting @Inject protected PlayCookieSessionStore playCookieSessionStore;
 
-  @Inject private SsoManager ssoManager;
+  @VisibleForTesting @Inject protected SsoManager ssoManager;
 
   @Inject AuthServiceClient authClient;
 
@@ -101,6 +102,10 @@ public class AuthenticationController extends Controller {
     final Optional<String> maybeRedirectPath =
         Optional.ofNullable(request.getQueryString(AUTH_REDIRECT_URI_PARAM));
     String redirectPath = maybeRedirectPath.orElse("/");
+    // If the redirect path is /logOut, we do not want to redirect to the logout page after login.
+    if (redirectPath.equals("/logOut")) {
+      redirectPath = "/";
+    }
     try {
       URI redirectUri = new URI(redirectPath);
       if (redirectUri.getScheme() != null || redirectUri.getAuthority() != null) {
@@ -333,15 +338,37 @@ public class AuthenticationController extends Controller {
     return createSession(userUrnString, accessToken);
   }
 
+  @VisibleForTesting
+  protected Result addRedirectCookie(Result result, CallContext ctx, String redirectPath) {
+    // Set the originally requested path for post-auth redirection. We split off into a separate
+    // cookie from the session
+    // to reduce size of the session cookie
+    FoundAction foundAction = new FoundAction(redirectPath);
+    byte[] javaSerBytes =
+        ((PlayCookieSessionStore) ctx.sessionStore()).getSerializer().serializeToBytes(foundAction);
+    String serialized = Base64.getEncoder().encodeToString(compressBytes(javaSerBytes));
+    Http.CookieBuilder redirectCookieBuilder =
+        Http.Cookie.builder(REDIRECT_URL_COOKIE_NAME, serialized);
+    redirectCookieBuilder.withPath("/");
+    redirectCookieBuilder.withSecure(true);
+    redirectCookieBuilder.withHttpOnly(true);
+    redirectCookieBuilder.withMaxAge(Duration.ofSeconds(86400));
+    redirectCookieBuilder.withSameSite(Http.Cookie.SameSite.NONE);
+
+    return result.withCookies(redirectCookieBuilder.build());
+  }
+
   private Optional<Result> redirectToIdentityProvider(
       Http.RequestHeader request, String redirectPath) {
     CallContext ctx = buildCallContext(request);
 
     final Client client = ssoManager.getSsoProvider().client();
-    configurePac4jSessionStore(ctx, client, redirectPath);
+    configurePac4jSessionStore(ctx, client);
     try {
       final Optional<RedirectionAction> action = client.getRedirectionAction(ctx);
-      return action.map(act -> new PlayHttpActionAdapter().adapt(act, ctx.webContext()));
+      final Optional<Result> maybeResult =
+          action.map(act -> new PlayHttpActionAdapter().adapt(act, ctx.webContext()));
+      return maybeResult.map(result -> addRedirectCookie(result, ctx, redirectPath));
     } catch (Exception e) {
       if (verbose) {
         logger.error(
@@ -370,17 +397,9 @@ public class AuthenticationController extends Controller {
     return new CallContext(webContext, playCookieSessionStore);
   }
 
-  private void configurePac4jSessionStore(CallContext ctx, Client client, String redirectPath) {
+  private void configurePac4jSessionStore(CallContext ctx, Client client) {
     WebContext context = ctx.webContext();
 
-    // Set the originally requested path for post-auth redirection. We split off into a separate
-    // cookie from the session
-    // to reduce size of the session cookie
-    FoundAction foundAction = new FoundAction(redirectPath);
-    byte[] javaSerBytes =
-        ((PlayCookieSessionStore) ctx.sessionStore()).getSerializer().serializeToBytes(foundAction);
-    String serialized = Base64.getEncoder().encodeToString(compressBytes(javaSerBytes));
-    context.addResponseCookie(new Cookie(REDIRECT_URL_COOKIE_NAME, serialized));
     // This is to prevent previous login attempts from being cached.
     // We replicate the logic here, which is buried in the Pac4j client.
     Optional<Object> attempt =

datahub-frontend/test/controllers/AuthenticationControllerTest.java

@@ -0,0 +1,215 @@
+package controllers;
+
+import static auth.AuthUtils.REDIRECT_URL_COOKIE_NAME;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import auth.sso.SsoManager;
+import auth.sso.SsoProvider;
+import client.AuthServiceClient;
+import com.typesafe.config.Config;
+import com.typesafe.config.ConfigFactory;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.commons.compress.utils.Lists;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.pac4j.core.client.Client;
+import org.pac4j.core.context.CallContext;
+import org.pac4j.core.exception.http.FoundAction;
+import org.pac4j.core.exception.http.RedirectionAction;
+import org.pac4j.play.PlayWebContext;
+import org.pac4j.play.store.DataEncrypter;
+import org.pac4j.play.store.PlayCookieSessionStore;
+import play.mvc.Http;
+import play.mvc.Result;
+
+public class AuthenticationControllerTest {
+
+  private AuthenticationController controller;
+  private Config mockConfig;
+  private org.pac4j.core.config.Config ssoConfig;
+  private PlayCookieSessionStore playCookieSessionStore;
+  private SsoManager ssoManager;
+  private AuthServiceClient authClient;
+
+  private class MockSerializer implements org.pac4j.core.util.serializer.Serializer {
+    @Override
+    public String serializeToString(Object var1) {
+      return "";
+    }
+
+    @Override
+    public Object deserializeFromString(String var1) {
+      return new Object();
+    }
+
+    @Override
+    public byte[] serializeToBytes(Object o) {
+      return "serialized-bytes".getBytes();
+    }
+
+    @Override
+    public Object deserializeFromBytes(byte[] bytes) {
+      return new FoundAction("/redirectedPath");
+    }
+  }
+
+  @BeforeEach
+  public void setUp() {
+    // Create mock configurations
+    Map<String, Object> configMap = new HashMap<>();
+    configMap.put("auth.cookie.ttlInHours", 24);
+    configMap.put("auth.cookie.secure", true);
+    configMap.put("auth.cookie.sameSite", "Lax");
+    mockConfig = ConfigFactory.parseMap(configMap);
+
+    // Mock SSO config
+    ssoConfig = new org.pac4j.core.config.Config();
+
+    // Mock session store with custom serializer
+    playCookieSessionStore = new PlayCookieSessionStore(mock(DataEncrypter.class));
+    playCookieSessionStore.setSerializer(new MockSerializer());
+
+    // Mock SSO manager
+    ssoManager = mock(SsoManager.class);
+    SsoProvider mockProvider = mock(SsoProvider.class);
+    when(ssoManager.getSsoProvider()).thenReturn(mockProvider);
+    Client mockClient = mock(Client.class);
+    when(mockProvider.client()).thenReturn(mockClient);
+    when(mockClient.getName()).thenReturn("oidcClient");
+
+    // Mock found action for redirection
+    FoundAction foundAction = new FoundAction("/redirectPath");
+    RedirectionAction redirectAction = foundAction;
+
+    // Setup the client mock to return our redirection action
+    when(mockClient.getRedirectionAction(any(CallContext.class)))
+        .thenReturn(Optional.of(redirectAction));
+
+    // Mock auth service client
+    authClient = mock(AuthServiceClient.class);
+    when(authClient.generateSessionTokenForUser(anyString(), anyString())).thenReturn("mock-token");
+
+    // Create the controller
+    controller = new AuthenticationController(mockConfig);
+    controller.playCookieSessionStore = playCookieSessionStore;
+    controller.ssoManager = ssoManager;
+    controller.authClient = authClient;
+  }
+
+  @Test
+  public void testRedirectCookieCreation() {
+    // Create mock HTTP context
+    Http.RequestBuilder requestBuilder =
+        new Http.RequestBuilder().method("GET").uri("/authenticate?redirect_uri=/dashboard");
+
+    Http.Request request = requestBuilder.build();
+
+    // Create mock context and result
+    CallContext callContext = new CallContext(new PlayWebContext(request), playCookieSessionStore);
+
+    Result initialResult = new Result(302);
+
+    // Call the method to test
+    Result result = controller.addRedirectCookie(initialResult, callContext, "/dashboard");
+
+    // Verify the cookie was added
+    Optional<Http.Cookie> redirectCookie =
+        Lists.newArrayList(result.cookies().iterator()).stream()
+            .filter(cookie -> cookie.name().equals(REDIRECT_URL_COOKIE_NAME))
+            .findFirst();
+
+    assertTrue(redirectCookie.isPresent(), "Redirect cookie should be present");
+    Http.Cookie cookie = redirectCookie.get();
+    assertEquals("/", cookie.path());
+    assertTrue(cookie.secure());
+    assertTrue(cookie.httpOnly());
+    assertEquals(Http.Cookie.SameSite.NONE, cookie.sameSite().get());
+    assertNotNull(cookie.value(), "Cookie value should not be null");
+  }
+
+  @Test
+  public void testSsoRedirectWithCookie() {
+    // Configure SSO to be enabled
+    when(ssoManager.isSsoEnabled()).thenReturn(true);
+
+    // Create mock HTTP context
+    Http.RequestBuilder requestBuilder = new Http.RequestBuilder().method("GET").uri("/sso");
+
+    Http.Request request = requestBuilder.build();
+
+    // Test the SSO method
+    Result result = controller.sso(request);
+
+    // Verify the result code
+    assertEquals(302, result.status());
+
+    // Verify the redirect cookie is added
+    Optional<Http.Cookie> redirectCookie =
+        Lists.newArrayList(result.cookies().iterator()).stream()
+            .filter(cookie -> cookie.name().equals(REDIRECT_URL_COOKIE_NAME))
+            .findFirst();
+
+    assertTrue(redirectCookie.isPresent(), "Redirect cookie should be present for SSO");
+  }
+
+  @Test
+  public void testAuthenticateWithRedirect() {
+    // Configure SSO to be enabled
+    when(ssoManager.isSsoEnabled()).thenReturn(true);
+
+    // Create mock HTTP context
+    Http.RequestBuilder requestBuilder =
+        new Http.RequestBuilder().method("GET").uri("/authenticate?redirect_uri=/dashboard");
+
+    Http.Request request = requestBuilder.build();
+
+    // Test the authenticate method
+    Result result = controller.authenticate(request);
+
+    // Verify the result code
+    assertEquals(302, result.status());
+
+    // Verify the redirect cookie is added
+    Optional<Http.Cookie> redirectCookie =
+        Lists.newArrayList(result.cookies().iterator()).stream()
+            .filter(cookie -> cookie.name().equals(REDIRECT_URL_COOKIE_NAME))
+            .findFirst();
+
+    assertTrue(redirectCookie.isPresent(), "Redirect cookie should be present for authenticate");
+
+    // Verify cookie value contains the encoded path
+    Http.Cookie cookie = redirectCookie.get();
+    assertNotNull(cookie.value(), "Cookie value should not be null");
+  }
+
+  @Test
+  public void testRedirectOnlySetsSecureCookies() {
+    // Create mock HTTP context
+    Http.RequestBuilder requestBuilder =
+        new Http.RequestBuilder().method("GET").uri("/authenticate?redirect_uri=/dashboard");
+
+    Http.Request request = requestBuilder.build();
+
+    // Create mock context and result
+    CallContext callContext = new CallContext(new PlayWebContext(request), playCookieSessionStore);
+
+    Result initialResult = new Result(302);
+
+    // Call the method to test
+    Result result = controller.addRedirectCookie(initialResult, callContext, "/dashboard");
+
+    // Verify all cookies are secure
+    boolean allCookiesSecure =
+        Lists.newArrayList(result.cookies().iterator()).stream().allMatch(Http.Cookie::secure);
+
+    assertTrue(allCookiesSecure, "All cookies should be secure");
+  }
+}

datahub-web-react/src/app/auth/LogIn.tsx

@@ -67,7 +67,7 @@ export type LogInProps = Record<string, never>;
 export const LogIn: React.VFC<LogInProps> = () => {
     const isLoggedIn = useReactiveVar(isLoggedInVar);
     const location = useLocation();
-    const params = QueryString.parse(location.search);
+    const params = QueryString.parse(location.search, { decode: true });
     const maybeRedirectError = params.error_msg;
 
     const themeConfig = useTheme();
@@ -105,7 +105,8 @@ export const LogIn: React.VFC<LogInProps> = () => {
 
     if (isLoggedIn) {
         const maybeRedirectUri = params.redirect_uri;
-        return <Redirect to={(maybeRedirectUri && decodeURIComponent(maybeRedirectUri as string)) || '/'} />;
+        // NOTE we do not decode the redirect_uri because it is already decoded by QueryString.parse
+        return <Redirect to={maybeRedirectUri || '/'} />;
     }
 
     return (