user-management/src/main/java/com/cleverthis/authservice/service/impl/KeycloakIdentityManagerService.java

package com.cleverthis.authservice.service.impl;

import com.cleverthis.authservice.dto.TokenResponse;
import com.cleverthis.authservice.dto.VerificationResponse;
import com.cleverthis.authservice.exception.AuthenticationException;
import com.cleverthis.authservice.exception.LogoutException;
import com.cleverthis.authservice.exception.TokenVerificationException;
import com.cleverthis.authservice.service.IdentityManagerService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

/**
 * Keycloak-specific implementation of the IdentityManagerService. Uses WebClient to interact with
 * Keycloak's REST API endpoints.
 */
@Service
@Slf4j
public class KeycloakIdentityManagerService implements IdentityManagerService {

    private final WebClient webClient;


    @Value("${keycloak.server-url}")
    private String keycloakServerUrl;
    @Value("${keycloak.realm}")
    private String realm;
    @Value("${keycloak.client-id}")
    private String clientId;
    @Value("${keycloak.client-secret}")
    private String clientSecret;
    @Value("${keycloak.grant-type}")
    private String grantType; // Should be 'password' for ROPC


    private String getTokenEndpoint() {
        return String.format("%s/realms/%s/protocol/openid-connect/token", this.keycloakServerUrl,
              this.realm);
    }

    private String getIntrospectionEndpoint() {
        return String.format("%s/realms/%s/protocol/openid-connect/token/introspect",
              this.keycloakServerUrl, this.realm);
    }

    private String getRevocationEndpoint() {
        return String.format("%s/realms/%s/protocol/openid-connect/revoke", this.keycloakServerUrl,
              this.realm);
    }

    @Autowired
    public KeycloakIdentityManagerService(WebClient keycloakWebClient) {
        this.webClient = keycloakWebClient;
    }

    @Override
    public Mono<TokenResponse> login(String username, String password) {
        log.debug("Attempting login for user: {}", username);
        MultiValueMap<String, String> formData = new LinkedMultiValueMap<>();
        formData.add("client_id", this.clientId);
        formData.add("client_secret", this.clientSecret);
        formData.add("grant_type", this.grantType); // Use configured grant type (password)
        formData.add("username", username);
        formData.add("password", password);
        formData.add("scope", "openid email profile"); // Request standard scopes

        return this.webClient.post().uri(getTokenEndpoint())
              .contentType(MediaType.APPLICATION_FORM_URLENCODED)
              .body(BodyInserters.fromFormData(formData)).retrieve()
              // Handle specific HTTP status codes
              .onStatus(status -> status.is4xxClientError() || status.is5xxServerError(),
                    clientResponse -> clientResponse.bodyToMono(String.class)
                          .flatMap(errorBody -> {
                              log.error("Keycloak login failed with status {}: {}",
                                    clientResponse.statusCode(), errorBody);
                              return Mono.error(new AuthenticationException(
                                    "Login failed: Invalid credentials or Keycloak error. "
                                          + "Status: " + clientResponse.statusCode()));
                          }))
              .bodyToMono(TokenResponse.class)
              .doOnSuccess(response -> log.debug("Login successful for user: {}", username))
              .doOnError(error -> log.error("Error during login for user {}: {}", username,
                    error.getMessage()));
    }

    @Override
    public Mono<VerificationResponse> verifyToken(String accessToken) {
        log.debug("Verifying access token");
        MultiValueMap<String, String> formData = new LinkedMultiValueMap<>();
        formData.add("client_id", this.clientId);
        formData.add("client_secret", this.clientSecret);
        formData.add("token", accessToken);

        return this.webClient.post().uri(getIntrospectionEndpoint())
              .contentType(MediaType.APPLICATION_FORM_URLENCODED)
              .body(BodyInserters.fromFormData(formData)).retrieve()
              .onStatus(status -> status.is4xxClientError() || status.is5xxServerError(),
                    clientResponse -> clientResponse.bodyToMono(String.class)
                          .flatMap(errorBody -> {
                              log.error("Keycloak token introspection failed with status {}: {}",
                                    clientResponse.statusCode(), errorBody);
                              return Mono.error(new TokenVerificationException(
                                    "Token introspection failed. Status: "
                                          + clientResponse.statusCode()));
                          }))
              .bodyToMono(VerificationResponse.class).flatMap(response -> {
                  if (!response.isActive()) {
                      log.warn("Token verification failed: Token is inactive.");
                      return Mono.error(
                            new TokenVerificationException("Token is invalid or expired."));
                  }
                  log.debug("Token verification successful. User: {}", response.getUsername());
                  return Mono.just(response);
              }).doOnError(error -> log.error("Error during token verification: {}",
                    error.getMessage()));
    }

    @Override
    public Mono<Void> logout(String refreshToken) {
        log.debug("Attempting logout (token revocation)");
        MultiValueMap<String, String> formData = new LinkedMultiValueMap<>();
        formData.add("client_id", this.clientId);
        formData.add("client_secret", this.clientSecret);
        formData.add("token", refreshToken);
        formData.add("token_type_hint", "refresh_token"); // Hint for Keycloak

        return this.webClient.post().uri(getRevocationEndpoint())
              .contentType(MediaType.APPLICATION_FORM_URLENCODED)
              .body(BodyInserters.fromFormData(formData)).retrieve()
              .onStatus(status -> status.is4xxClientError() || status.is5xxServerError(),
                    clientResponse -> clientResponse.bodyToMono(String.class)
                          .flatMap(errorBody -> {
                              // Keycloak might return 400 if token is already invalid, which
                              // is okay for logout
                              if (clientResponse.statusCode() == HttpStatus.BAD_REQUEST) {
                                  log.warn(
                                        "Token revocation returned 400 "
                                              + "(possibly already invalid/revoked): {}",
                                        errorBody);
                                  return Mono.empty(); // Treat as success in logout scenario
                              }
                              log.error("Keycloak token revocation failed with status {}: {}",
                                    clientResponse.statusCode(), errorBody);
                              return Mono.error(new LogoutException("Logout failed. Status: "
                                    + clientResponse.statusCode()));
                          }))
              .bodyToMono(Void.class) // Expecting empty body on success (2xx)
              .doOnSuccess(v -> log.debug("Logout successful (token revoked)."))
              .doOnError(error -> log.error("Error during logout: {}", error.getMessage()));
    }
}