Merge pull request #36 from gooddata/LX-573

jeskepetr · web-flow · commit 009e4a43e0e9 · 2024-10-16T10:13:15.000+02:00
feat: Apply custom auth attributes on authentication call
diff --git a/gooddata-server-oauth2-autoconfigure/src/main/kotlin/AuthenticationStoreClient.kt b/gooddata-server-oauth2-autoconfigure/src/main/kotlin/AuthenticationStoreClient.kt
@@ -177,6 +177,8 @@ interface AuthenticationStoreClient {
  * @property oauthSubjectIdClaim name of the claim in ID token that will be used for finding the user in organization.
  * Defaults to `null` and it means that `sub` claim will be used.
  * @property jitEnabled the switch for enabling/disabling of the JIT provisioning
+ * @property oauthCustomAuthAttributes map for additional oauth authentication attributes to be sent
+ * in authentication request
  *
  * @see AuthenticationStoreClient
  */
@@ -189,6 +191,7 @@ data class Organization(
     val oauthIssuerId: String? = null,
     val oauthSubjectIdClaim: String? = null,
     val jitEnabled: Boolean? = null,
+    val oauthCustomAuthAttributes: Map<String, String>? = null,
 )
 
 /**
diff --git a/gooddata-server-oauth2-autoconfigure/src/main/kotlin/ServerOAuth2AutoConfiguration.kt b/gooddata-server-oauth2-autoconfigure/src/main/kotlin/ServerOAuth2AutoConfiguration.kt
@@ -15,7 +15,7 @@
  */
 package com.gooddata.oauth2.server
 
-import com.gooddata.oauth2.server.oauth2.client.FederationAwareOauth2AuthorizationRequestResolver
+import com.gooddata.oauth2.server.oauth2.client.CustomAttrsAwareOauth2AuthorizationRequestResolver
 import java.util.Base64
 import org.springframework.beans.factory.ObjectProvider
 import org.springframework.beans.factory.annotation.Value
@@ -213,7 +213,7 @@ class ServerOAuth2AutoConfiguration {
     fun federationAwareAuthorizationRequestResolver(
         urlSafeStateAuthorizationRequestResolver: ServerOAuth2AuthorizationRequestResolver,
         cookieService: ReactiveCookieService,
-    ) = FederationAwareOauth2AuthorizationRequestResolver(urlSafeStateAuthorizationRequestResolver, cookieService)
+    ) = CustomAttrsAwareOauth2AuthorizationRequestResolver(urlSafeStateAuthorizationRequestResolver, cookieService)
 
     @Bean
     @Suppress("LongParameterList", "LongMethod")
diff --git a/gooddata-server-oauth2-autoconfigure/src/main/kotlin/oauth2/client/CustomAttrsAwareOauth2AuthorizationRequestResolver.kt b/gooddata-server-oauth2-autoconfigure/src/main/kotlin/oauth2/client/CustomAttrsAwareOauth2AuthorizationRequestResolver.kt
@@ -17,15 +17,19 @@ package com.gooddata.oauth2.server.oauth2.client
 
 import com.gooddata.oauth2.server.ReactiveCookieService
 import com.gooddata.oauth2.server.SPRING_EXTERNAL_IDP
+import com.gooddata.oauth2.server.getOrganizationFromContext
 import org.springframework.security.oauth2.client.web.server.ServerOAuth2AuthorizationRequestResolver
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest
 import org.springframework.web.server.ServerWebExchange
 import reactor.core.publisher.Mono
 import reactor.kotlin.core.publisher.switchIfEmpty
 
 /**
- * The implementation of [ServerOAuth2AuthorizationRequestResolver] that is able to append external identity provider
- * (OIDC federation) info to authorization requests based on [SPRING_EXTERNAL_IDP] cookie in the [ServerWebExchange].
+ * The implementation of [ServerOAuth2AuthorizationRequestResolver] that is able to append ad-hoc authentication attrs
+ * to authorization requests.
+ *
+ * Firstly it can add an external identity provider (OIDC federation) info to authorization requests
+ * based on [SPRING_EXTERNAL_IDP] cookie in the [ServerWebExchange].
  *
  * It wraps the default [ServerOAuth2AuthorizationRequestResolver] which is responsible for building of the original
  * authorization request with standard parameters.
@@ -34,31 +38,35 @@ import reactor.kotlin.core.publisher.switchIfEmpty
  *
  * The new query parameter is Cognito-specific for now, so this does not ensure support for other identity providers.
  *
+ * Secondly, it can add additional authentication attributes if present in the organization definition.
+ *
  * @param defaultResolver the default [ServerOAuth2AuthorizationRequestResolver] to be wrapped
  * @param cookieService the [ReactiveCookieService] to be used for cookie handling
  */
-class FederationAwareOauth2AuthorizationRequestResolver(
+class CustomAttrsAwareOauth2AuthorizationRequestResolver(
     private val defaultResolver: ServerOAuth2AuthorizationRequestResolver,
     private val cookieService: ReactiveCookieService,
 ) : ServerOAuth2AuthorizationRequestResolver {
     override fun resolve(exchange: ServerWebExchange?): Mono<OAuth2AuthorizationRequest> =
         defaultResolver.resolve(exchange).flatMap { authorizationRequest ->
-            enhanceRequestByExternalIdentityParam(authorizationRequest, exchange)
+            enhanceRequestByAdditionalParams(authorizationRequest, exchange)
         }
 
     override fun resolve(
         exchange: ServerWebExchange?,
         clientRegistrationId: String?,
     ): Mono<OAuth2AuthorizationRequest> =
         defaultResolver.resolve(exchange, clientRegistrationId).flatMap { authorizationRequest ->
-            enhanceRequestByExternalIdentityParam(authorizationRequest, exchange)
+            enhanceRequestByAdditionalParams(authorizationRequest, exchange)
         }
 
     /**
-     * Enhances the provided [authorizationRequest] with external identity provider info based
-     * on the [SPRING_EXTERNAL_IDP] cookie existence. If the cookie is present, it is cleared.
+     * Enhances the provided [authorizationRequest] with external additional authentication attributes.
+     * Adds additional authentication attributes from the organization definition if they are present.
+     * Adds identity provider info based on the [SPRING_EXTERNAL_IDP] cookie existence. If the cookie is present,
+     * it is cleared.
      */
-    private fun enhanceRequestByExternalIdentityParam(
+    private fun enhanceRequestByAdditionalParams(
         authorizationRequest: OAuth2AuthorizationRequest,
         exchange: ServerWebExchange?,
     ): Mono<OAuth2AuthorizationRequest> = Mono.justOrEmpty(exchange)
@@ -72,11 +80,23 @@ class FederationAwareOauth2AuthorizationRequestResolver(
                         additionalParams[COGNITO_EXTERNAL_PROVIDER_ID_PARAM_NAME] = externalIdp
                     }
                     .build()
+            }.switchIfEmpty {
+                Mono.just(authorizationRequest)
+            }.flatMap { request ->
+                getOrganizationFromContext().flatMap { organization ->
+                    Mono.just(OAuth2AuthorizationRequest.from(request)
+                        .additionalParameters { additionalParams ->
+                            // if organization contains additional authentication attributes, add them to the request
+                            organization.oauthCustomAuthAttributes?.takeIf { it.isNotEmpty() }
+                                ?.forEach { (key, value) ->
+                                    additionalParams[key] = value
+                                }
+                        }
+                        .build()
+                    )
+                }
             }
         }
-        .switchIfEmpty {
-            Mono.just(authorizationRequest)
-        }
 
     companion object {
         /**
diff --git a/gooddata-server-oauth2-autoconfigure/src/test/kotlin/oauth2/client/CustomAttrsAwareOauth2AuthorizationRequestResolverTest.kt b/gooddata-server-oauth2-autoconfigure/src/test/kotlin/oauth2/client/CustomAttrsAwareOauth2AuthorizationRequestResolverTest.kt
@@ -15,6 +15,9 @@
  */
 package com.gooddata.oauth2.server.oauth2.client
 
+import com.gooddata.oauth2.server.CookieSerializerTest.Companion.ORGANIZATION
+import com.gooddata.oauth2.server.Organization
+import com.gooddata.oauth2.server.OrganizationWebFilter.Companion.orgContextWrite
 import com.gooddata.oauth2.server.ReactiveCookieService
 import com.gooddata.oauth2.server.SPRING_EXTERNAL_IDP
 import io.mockk.every
@@ -31,7 +34,7 @@ import reactor.core.publisher.Mono
 import reactor.test.StepVerifier
 
 @Suppress("ReactiveStreamsUnusedPublisher")
-class FederationAwareOauth2AuthorizationRequestResolverTest {
+class CustomAttrsAwareOauth2AuthorizationRequestResolverTest {
 
     private val defaultResolver: ServerOAuth2AuthorizationRequestResolver = mockk {
         every { resolve(any()) } returns Mono.just(DEFAULT_AUTH_REQUEST)
@@ -42,11 +45,14 @@ class FederationAwareOauth2AuthorizationRequestResolverTest {
         every { invalidateCookie(any(), any()) } just runs
     }
 
-    private val resolver = FederationAwareOauth2AuthorizationRequestResolver(defaultResolver, cookieService)
+    private val resolver = CustomAttrsAwareOauth2AuthorizationRequestResolver(defaultResolver, cookieService)
 
     @Test
     fun `resolve without cookie`() {
-        StepVerifier.create(oauthRequestToUri(resolver.resolve(EXCHANGE)))
+        StepVerifier.create(
+            oauthRequestToUri(resolver.resolve(EXCHANGE))
+                .orgContextWrite(ORGANIZATION)
+        )
             .expectNext("https://example.com/oauth2/authorize?response_type=code&client_id=client-id")
             .verifyComplete()
     }
@@ -55,10 +61,16 @@ class FederationAwareOauth2AuthorizationRequestResolverTest {
     fun `resolve with cookie and invalidates it`() {
         every { cookieService.decodeCookie(EXCHANGE, SPRING_EXTERNAL_IDP) } returns Mono.just("external-idp-id")
 
-        StepVerifier.create(oauthRequestToUri(resolver.resolve(EXCHANGE)))
+        StepVerifier.create(
+            oauthRequestToUri(resolver.resolve(EXCHANGE))
+                .orgContextWrite(ORGANIZATION_WITH_CUSTOM_AUTH_ATTRS)
+        )
             .expectNext(
                 "https://example.com/oauth2/authorize" +
-                    "?response_type=code&client_id=client-id&identity_provider=external-idp-id"
+                    "?response_type=code&client_id=client-id" +
+                    "&identity_provider=external-idp-id" +
+                    "&organization=org_vswH67L51ZQW67PS"
+
             )
             .verifyComplete()
 
@@ -67,11 +79,31 @@ class FederationAwareOauth2AuthorizationRequestResolverTest {
 
     @Test
     fun `resolve without cookie and client registration id`() {
-        StepVerifier.create(oauthRequestToUri(resolver.resolve(EXCHANGE, "registration-id")))
+        StepVerifier.create(
+            oauthRequestToUri(
+                resolver.resolve(EXCHANGE, "registration-id")
+            )
+                .orgContextWrite(ORGANIZATION)
+        )
             .expectNext("https://example.com/oauth2/authorize?response_type=code&client_id=client-id")
             .verifyComplete()
     }
 
+    @Test
+    fun `resolve without cookie and client registration id with custom authentication attributes`() {
+        StepVerifier.create(
+            oauthRequestToUri(
+                resolver.resolve(EXCHANGE, "registration-id")
+            )
+                .orgContextWrite(ORGANIZATION_WITH_CUSTOM_AUTH_ATTRS)
+        )
+            .expectNext(
+                "https://example.com/oauth2/authorize" +
+                    "?response_type=code&client_id=client-id&organization=org_vswH67L51ZQW67PS"
+            )
+            .verifyComplete()
+    }
+
     companion object {
         private val DEFAULT_AUTH_REQUEST = OAuth2AuthorizationRequest
             .authorizationCode()
@@ -85,5 +117,8 @@ class FederationAwareOauth2AuthorizationRequestResolverTest {
 
         private fun oauthRequestToUri(request: Mono<OAuth2AuthorizationRequest>): Mono<String> =
             request.map(OAuth2AuthorizationRequest::getAuthorizationRequestUri)
+
+        val ORGANIZATION_WITH_CUSTOM_AUTH_ATTRS =
+            Organization(id = "org", oauthCustomAuthAttributes = mapOf("organization" to "org_vswH67L51ZQW67PS"))
     }
 }
