content: /*
   *  Licensed to the Apache Software Foundation (ASF) under one
   *  or more contributor license agreements.  See the NOTICE file
   *  distributed with this work for additional information
   *  regarding copyright ownership.  The ASF licenses this file
   *  to you under the Apache License, Version 2.0 (the
   *  "License"); you may not use this file except in compliance
   *  with the License.  You may obtain a copy of the License at
   *
   *        http://www.apache.org/licenses/LICENSE-2.0
   *
   *  Unless required by applicable law or agreed to in writing,
   *  software distributed under the License is distributed on an
   *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   *  KIND, either express or implied.  See the License for the
   *  specific language governing permissions and limitations
   *  under the License.
   */
  package org.apache.causeway.security.keycloak.services;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.LinkedHashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.function.Consumer;
  import java.util.function.Predicate;
  
  import org.jspecify.annotations.Nullable;
  import org.springframework.security.core.GrantedAuthority;
  import org.springframework.security.core.authority.AuthorityUtils;
  import org.springframework.security.core.authority.mapping.GrantedAuthoritiesMapper;
  import org.springframework.security.oauth2.client.oidc.userinfo.OidcUserRequest;
  import org.springframework.security.oauth2.client.oidc.userinfo.OidcUserService;
  import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
  import org.springframework.security.oauth2.core.OAuth2Error;
  import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
  import org.springframework.security.oauth2.core.oidc.user.DefaultOidcUser;
  import org.springframework.security.oauth2.core.oidc.user.OidcUser;
  import org.springframework.security.oauth2.jwt.Jwt;
  import org.springframework.security.oauth2.jwt.JwtDecoder;
  import org.springframework.security.oauth2.jwt.JwtException;
  
  import org.apache.causeway.commons.internal.base._Casts;
  import org.apache.causeway.commons.internal.base._NullSafe;
  import org.apache.causeway.core.config.CausewayConfiguration;
  
  import org.jspecify.annotations.NonNull;
  import lombok.RequiredArgsConstructor;
  
  @RequiredArgsConstructor
  public class KeycloakOauth2UserService extends OidcUserService {
  
      private final static OAuth2Error INVALID_REQUEST = new OAuth2Error(OAuth2ErrorCodes.INVALID_REQUEST);
  
      final JwtDecoder jwtDecoder;
      final GrantedAuthoritiesMapper authoritiesMapper;
      final CausewayConfiguration causewayConfiguration;
  
      /**
       * Augments {@link OidcUserService#loadUser(OidcUserRequest)} to add authorities
       * provided by Keycloak.
       * <p>
       * Needed because {@link OidcUserService#loadUser(OidcUserRequest)} (currently)
       * does not provide a hook for adding custom authorities from a
       * {@link OidcUserRequest}.
       */
      @Override
      public OidcUser loadUser(final OidcUserRequest userRequest) throws OAuth2AuthenticationException {
  
          OidcUser user = super.loadUser(userRequest);
  
          Set<GrantedAuthority> authorities = new LinkedHashSet<>();
          authorities.addAll(user.getAuthorities());
          authorities.addAll(extractKeycloakAuthorities(userRequest));
  
          return new DefaultOidcUser(authorities, userRequest.getIdToken(), user.getUserInfo(), "preferred_username");
      }
  
      /**
       * Extracts {@link GrantedAuthority GrantedAuthorities} from the AccessToken in
       * the {@link OidcUserRequest}.
       *
       * @param userRequest
       * @return
       */
      private Collection<? extends GrantedAuthority> extractKeycloakAuthorities(final OidcUserRequest userRequest) {
  
          Jwt token = parseJwt(userRequest.getAccessToken().getTokenValue());
  
          final List<String> combinedRoles = new ArrayList<>();
  
          if(causewayConfiguration.getSecurity().getKeycloak().isExtractClientRoles()) {
              // attempt to parse out 'resource_access.${client_id}.roles'
              var rolePrefix = Optional.ofNullable(causewayConfiguration.getSecurity().getKeycloak().getClientRolePrefix()).orElse("");
              asNonEmptyMap(token.getClaims().get("resource_access"))
              .ifPresent(resourceMap->{
                  final String clientId = userRequest.getClientRegistration().getClientId();
                  asNonEmptyMap(resourceMap.get(clientId))
                  .flatMap(clientResource->asNonEmptyCollection(clientResource.get("roles")))
                  .ifPresent(clientRoles->
                      forEachNonNullIn(clientRoles, clientRole -> combinedRoles.add(rolePrefix + clientRole))
                  );
              });
          }
  
          if (causewayConfiguration.getSecurity().getKeycloak().isExtractRealmRoles()) {
              // attempt to parse out 'realm_access.roles'
              var rolePrefix = Optional.ofNullable(causewayConfiguration.getSecurity().getKeycloak().getRealmRolePrefix()).orElse("");
              asNonEmptyMap(token.getClaims().get("realm_access"))
              .ifPresent(realmAccessMap->{
                  asNonEmptyCollection(realmAccessMap.get("roles"))
                  .ifPresent(realmRoles->{
                      forEachNonNullIn(realmRoles, realmRole -> combinedRoles.add(rolePrefix + realmRole));
                  });
              });
          }
  
          if (causewayConfiguration.getSecurity().getKeycloak().isExtractRoles()) {
              // attempt to parse out 'roles'
              var rolePrefix = Optional.ofNullable(causewayConfiguration.getSecurity().getKeycloak().getRolePrefix()).orElse("");
              asNonEmptyCollection(token.getClaims().get("roles"))
              .ifPresent(roles->{
                  forEachNonNullIn(roles, role -> combinedRoles.add(rolePrefix + role));
              });
          }
  
          Collection<? extends GrantedAuthority> authorities = AuthorityUtils.createAuthorityList(combinedRoles.toArray(new String[]{}));
          return authoritiesMapper == null
                  ? authorities
                  : authoritiesMapper.mapAuthorities(authorities);
  
      }
  
      // -- HELPER
  
      @SuppressWarnings("rawtypes")
      private Optional<Map> asNonEmptyMap(final @Nullable Object x){
          return _Casts.castTo(Map.class, x)
                  .filter(Predicate.not(_NullSafe::isEmpty));
      }
  
      @SuppressWarnings({ "rawtypes", "unchecked" })
      private void forEachNonNullIn(final @NonNull Collection x, final @NonNull Consumer<Object> _do){
          x.stream()
                  .filter(Objects::nonNull)
                  .forEach(_do);
      }
  
      @SuppressWarnings("rawtypes")
      private Optional<Collection> asNonEmptyCollection(final @Nullable Object x){
          return _Casts.castTo(Collection.class, x)
                  .filter(Predicate.not(_NullSafe::isEmpty));
      }
  
      private Jwt parseJwt(final String accessTokenValue) {
          try {
              // Token is already verified by spring security infrastructure
              return jwtDecoder.decode(accessTokenValue);
          } catch (JwtException e) {
              throw new OAuth2AuthenticationException(INVALID_REQUEST, e);
          }
      }
  }
  