Use the authentication response cookie fields to search for flow cookies

Author: itaihanski

descopesdk/src/main/java/com/descope/android/DescopeFlowCoordinator.kt

@@ -454,9 +454,10 @@ document.head.appendChild(element)
         try {
             val jwtServerResponse = JwtServerResponse.fromJson(data, emptyList())
             // take tokens from cookies if missing
-            val cookieString = CookieManager.getInstance().getCookie(url)
-            jwtServerResponse.sessionJwt = jwtServerResponse.sessionJwt ?: findJwtInCookies(cookieString, name = SESSION_COOKIE_NAME)
-            jwtServerResponse.refreshJwt = jwtServerResponse.refreshJwt ?: findJwtInCookies(cookieString, name = REFRESH_COOKIE_NAME)
+            val respCookieString = CookieManager.getInstance().getCookie("https://${jwtServerResponse.cookieDomain}${jwtServerResponse.cookiePath}")
+            val urlCookieString = CookieManager.getInstance().getCookie(url)
+            jwtServerResponse.sessionJwt = jwtServerResponse.sessionJwt ?: findJwtInCookies(SESSION_COOKIE_NAME, respCookieString, urlCookieString)
+            jwtServerResponse.refreshJwt = jwtServerResponse.refreshJwt ?: findJwtInCookies(REFRESH_COOKIE_NAME, respCookieString, urlCookieString)
             val authResponse = jwtServerResponse.convert()
             logger.debug("Flow received an authentication response", data)
             handleSuccess(authResponse)
@@ -654,12 +655,13 @@ private fun String.escapeForBackticks() = replace("\\", "\\\\")
 
 // Cookies
 
-internal fun findJwtInCookies(cookieString: String?, name: String): String? {
-    // split and aggregate all cookies 
-    val cookies = mutableListOf<HttpCookie>().apply {
+internal fun findJwtInCookies(name: String, vararg cookieStrings: String?): String? {
+    val cookies = mutableListOf<HttpCookie>()
+    cookieStrings.forEach { cookieString ->
+        // split and aggregate all cookies 
         cookieString?.split("; ")?.forEach {
             try {
-                addAll(HttpCookie.parse(it))
+                cookies.addAll(HttpCookie.parse(it))
             } catch (_: Exception) {
             }
         }

descopesdk/src/main/java/com/descope/internal/http/Responses.kt

@@ -17,6 +17,8 @@ internal data class JwtServerResponse(
     var refreshJwt: String?,
     val user: UserResponse?,
     val firstSeen: Boolean,
+    val cookieDomain: String,
+    val cookiePath: String,
 ) {
     companion object {
         fun fromJson(json: String, cookies: List<HttpCookie>) = JSONObject(json).run {
@@ -36,6 +38,8 @@ internal data class JwtServerResponse(
                 refreshJwt = stringOrEmptyAsNull("refreshJwt") ?: refreshJwt,
                 user = optJSONObject("user")?.run { UserResponse.fromJson(this) },
                 firstSeen = optBoolean("firstSeen"),
+                cookieDomain = stringOrEmptyAsNull("cookieDomain") ?: "",
+                cookiePath = stringOrEmptyAsNull("cookiePath") ?: "",
             )
         }
     }

descopesdk/src/test/java/com/descope/internal/routes/TestUtils.kt

@@ -106,7 +106,9 @@ internal val mockJwtResponse = JwtServerResponse(
         totp = true,
         saml = true,
         scim = true,
-    )
+    ),
+    cookieDomain = "domain",
+    cookiePath = "/path",
 )
 
 private object MockInfo: SystemInfo {

descopesdk/src/test/java/com/descope/session/TokenTest.kt

@@ -45,7 +45,7 @@ class TokenTest {
             cookies.add(HttpCookie("name$i", "value$i"))
         }
         cookies.add(HttpCookie(REFRESH_COOKIE_NAME, jwtForP123))
-        val refreshJwt = findJwtInCookies(cookies.joinToString(separator = "; "), name = REFRESH_COOKIE_NAME)
+        val refreshJwt = findJwtInCookies(name = REFRESH_COOKIE_NAME, cookies.joinToString(separator = "; "))
         assertEquals(jwtForP123, refreshJwt)
     }
 
@@ -58,7 +58,7 @@ class TokenTest {
         }
         cookies.add(HttpCookie(REFRESH_COOKIE_NAME, jwtForP123))
 
-        val refreshJwt = findJwtInCookies(cookies.joinToString(separator = "; "), name = REFRESH_COOKIE_NAME)
+        val refreshJwt = findJwtInCookies(name = REFRESH_COOKIE_NAME, cookies.joinToString(separator = "; "))
         assertEquals(jwtForP456, refreshJwt)
     }
 
@@ -71,7 +71,7 @@ class TokenTest {
         }
         cookies.add(HttpCookie(SESSION_COOKIE_NAME, jwtForP123))
 
-        var refreshJwt = findJwtInCookies(cookies.joinToString(separator = "; "), name = SESSION_COOKIE_NAME)
+        var refreshJwt = findJwtInCookies(name = SESSION_COOKIE_NAME, cookies.joinToString(separator = "; "))
         assertEquals(laterJwtForP123, refreshJwt)
 
         // try again with a different order
@@ -82,8 +82,47 @@ class TokenTest {
         }
         cookies.add(HttpCookie(SESSION_COOKIE_NAME, laterJwtForP123))
 
-        refreshJwt = findJwtInCookies(cookies.joinToString(separator = "; "), name = SESSION_COOKIE_NAME)
+        refreshJwt = findJwtInCookies(name = SESSION_COOKIE_NAME, cookies.joinToString(separator = "; "))
         assertEquals(laterJwtForP123, refreshJwt)
     }
 
+    @Test
+    fun cookies_oneDsr_multipleStrings() {
+        val cookies = mutableListOf<HttpCookie>()
+        for (i in 0 until 10) {
+            cookies.add(HttpCookie("name$i", "value$i"))
+        }
+        cookies.add(HttpCookie(REFRESH_COOKIE_NAME, jwtForP123))
+        val moreCookies = mutableListOf<HttpCookie>()
+        for (i in 10 until 20) {
+            moreCookies.add(HttpCookie("name$i", "value$i"))
+        }
+        val refreshJwt = findJwtInCookies(name = REFRESH_COOKIE_NAME, moreCookies.joinToString(separator = "; "), cookies.joinToString(separator = "; "))
+        assertEquals(jwtForP123, refreshJwt)
+    }
+
+    @Test
+    fun cookies_multipleDs_multipleStrings() {
+        val cookies = mutableListOf<HttpCookie>()
+        cookies.add(HttpCookie(SESSION_COOKIE_NAME, laterJwtForP123))
+        for (i in 0 until 10) {
+            cookies.add(HttpCookie("name$i", "value$i"))
+        }
+
+        val moreCookies = mutableListOf<HttpCookie>()
+        moreCookies.add(HttpCookie(SESSION_COOKIE_NAME, laterJwtForP123))
+        for (i in 0 until 10) {
+            moreCookies.add(HttpCookie("name$i", "value$i"))
+        }
+        moreCookies.add(HttpCookie(SESSION_COOKIE_NAME, jwtForP123))
+        
+        
+
+        var refreshJwt = findJwtInCookies(name = SESSION_COOKIE_NAME, cookies.joinToString(separator = "; "), moreCookies.joinToString(separator = "; "))
+        assertEquals(laterJwtForP123, refreshJwt)
+
+        // try again with a different order
+        refreshJwt = findJwtInCookies(name = SESSION_COOKIE_NAME, moreCookies.joinToString(separator = "; "), cookies.joinToString(separator = "; "))
+        assertEquals(laterJwtForP123, refreshJwt)
+    }
 }