001/*
002 * Copyright © 2025 CUI-OpenSource-Software (info@cuioss.de)
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package de.cuioss.benchmarking.common.repository;
017
018import com.google.gson.JsonObject;
019import de.cuioss.benchmarking.common.http.HttpClientFactory;
020import de.cuioss.benchmarking.common.token.TokenProvider;
021import de.cuioss.benchmarking.common.util.JsonSerializationHelper;
022import de.cuioss.tools.logging.CuiLogger;
023
024import java.io.IOException;
025import java.net.URI;
026import java.net.URLEncoder;
027import java.net.http.HttpClient;
028import java.net.http.HttpRequest;
029import java.net.http.HttpResponse;
030import java.nio.charset.StandardCharsets;
031import java.time.Duration;
032import java.util.ArrayList;
033import java.util.List;
034import java.util.concurrent.atomic.AtomicInteger;
035
036import static de.cuioss.benchmarking.common.util.BenchmarkingLogMessages.ERROR.FAILED_FETCH_TOKEN;
037import static de.cuioss.benchmarking.common.util.BenchmarkingLogMessages.WARN.TOKEN_POOL_EMPTY;
038
039/**
040 * Keycloak-based token repository for fetching real JWT tokens from a Keycloak server.
041 * <p>
042 * This implementation fetches actual tokens from a Keycloak authentication server,
043 * making it suitable for integration benchmarks that need to test against real
044 * authentication infrastructure.
045 * </p>
046 * <p>
047 * Features:
048 * <ul>
049 *   <li>Fetches tokens from Keycloak using password grant</li>
050 *   <li>Maintains a pool of tokens for rotation</li>
051 *   <li>Supports SSL verification configuration</li>
052 *   <li>Provides round-robin token distribution</li>
053 * </ul>
054 *
055 * @author Oliver Wolff
056 */
057public class KeycloakTokenRepository implements TokenProvider {
058
059    private static final CuiLogger LOGGER = new CuiLogger(KeycloakTokenRepository.class);
060    private static final int HTTP_OK = 200;
061    private static final int REQUEST_TIMEOUT_SECONDS = 10;
062
063    private final TokenRepositoryConfig config;
064    private final List<TokenInfo> tokenPool;
065    private final AtomicInteger tokenIndex;
066    private final HttpClient httpClient;
067
068    /**
069     * Creates a new KeycloakTokenRepository with the given configuration.
070     *
071     * @param config the configuration for connecting to Keycloak
072     */
073    public KeycloakTokenRepository(TokenRepositoryConfig config) {
074        this.config = config;
075        this.tokenPool = new ArrayList<>(config.getTokenPoolSize());
076        this.tokenIndex = new AtomicInteger(0);
077
078        this.httpClient = HttpClientFactory.getInsecureClient();
079
080        // Initialize token pool
081        initializeTokenPool();
082    }
083
084    /**
085     * {@inheritDoc}
086     * <p>
087     * Gets the next token from the pool using round-robin rotation.
088     * This ensures even distribution and simulates cache miss scenarios.
089     * If the pool is empty, fetches a single token directly from Keycloak.
090     * </p>
091     */
092    @Override
093    public String getNextToken() {
094        if (tokenPool.isEmpty()) {
095            LOGGER.warn(TOKEN_POOL_EMPTY);
096            return fetchSingleToken();
097        }
098
099        int index = tokenIndex.getAndIncrement() % tokenPool.size();
100        return tokenPool.get(index).accessToken();
101    }
102
103    /**
104     * {@inheritDoc}
105     * <p>
106     * Returns the current size of the token pool.
107     * </p>
108     */
109    @Override
110    public int getTokenPoolSize() {
111        return tokenPool.size();
112    }
113
114    private void initializeTokenPool() {
115        LOGGER.debug("Initializing token pool with %s tokens", config.getTokenPoolSize());
116
117        for (int i = 0; i < config.getTokenPoolSize(); i++) {
118            String token = fetchSingleToken();
119            tokenPool.add(new TokenInfo(token));
120        }
121
122        LOGGER.debug("Token pool initialized with %s tokens", tokenPool.size());
123    }
124
125    private String fetchSingleToken() {
126        String tokenEndpoint = "%s/realms/%s/protocol/openid-connect/token".formatted(
127                config.getKeycloakBaseUrl(), config.getRealm());
128
129        try {
130            // Build form data as URL encoded string
131            String formData = "grant_type=" + URLEncoder.encode("password", StandardCharsets.UTF_8) +
132                    "&client_id=" + URLEncoder.encode(config.getClientId(), StandardCharsets.UTF_8) +
133                    "&client_secret=" + URLEncoder.encode(config.getClientSecret(), StandardCharsets.UTF_8) +
134                    "&username=" + URLEncoder.encode(config.getUsername(), StandardCharsets.UTF_8) +
135                    "&password=" + URLEncoder.encode(config.getPassword(), StandardCharsets.UTF_8);
136
137            HttpRequest request = HttpRequest.newBuilder()
138                    .uri(URI.create(tokenEndpoint))
139                    .timeout(Duration.ofSeconds(REQUEST_TIMEOUT_SECONDS))
140                    .header("Content-Type", "application/x-www-form-urlencoded")
141                    .POST(HttpRequest.BodyPublishers.ofString(formData))
142                    .build();
143
144            HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
145
146            if (response.statusCode() == HTTP_OK) {
147                return extractAccessToken(response);
148            } else {
149                handleTokenFetchError(response);
150                throw new TokenFetchException("Unexpected error - handleTokenFetchError should have thrown");
151            }
152        } catch (IOException e) {
153            throw new TokenFetchException("Error fetching token from Keycloak", e);
154        } catch (InterruptedException e) {
155            Thread.currentThread().interrupt();
156            throw new TokenFetchException("Token fetch interrupted", e);
157        }
158    }
159
160
161    /**
162     * Internal class to hold token information.
163     */
164    private record TokenInfo(String accessToken) {
165
166    }
167
168    private String extractAccessToken(HttpResponse<String> response) {
169        String responseBody = response.body();
170        if (responseBody == null || responseBody.isEmpty()) {
171            throw new TokenFetchException("Empty response body from token endpoint");
172        }
173
174        JsonObject jsonResponse = JsonSerializationHelper.fromJson(responseBody, JsonObject.class);
175        if (jsonResponse == null || !jsonResponse.has("access_token")) {
176            throw new TokenFetchException("No access_token field in response");
177        }
178
179        String token = jsonResponse.get("access_token").getAsString();
180        if (token == null || token.isEmpty()) {
181            throw new TokenFetchException("Access token is null or empty");
182        }
183
184        return token;
185    }
186
187    private void handleTokenFetchError(HttpResponse<String> response) {
188        String errorBody = response.body() != null ? response.body() : "<no body>";
189
190        LOGGER.error(FAILED_FETCH_TOKEN, response.statusCode(), errorBody);
191
192        throw new TokenFetchException(
193                "Failed to fetch token from Keycloak. Status: %d, Body: %s".formatted(
194                        response.statusCode(), errorBody)
195        );
196    }
197
198    /**
199     * Custom exception for token fetch errors.
200     */
201    public static class TokenFetchException extends RuntimeException {
202        public TokenFetchException(String message) {
203            super(message);
204        }
205
206        public TokenFetchException(String message, Throwable cause) {
207            super(message, cause);
208        }
209    }
210
211    /**
212     * {@inheritDoc}
213     * <p>
214     * Refreshes the token pool by fetching new tokens from Keycloak.
215     * This replaces all existing tokens in the pool with fresh ones.
216     * </p>
217     *
218     * @throws TokenFetchException if unable to fetch new tokens from Keycloak
219     */
220    @Override
221    public void refreshTokens() {
222        LOGGER.debug("Refreshing token pool with %s tokens", config.getTokenPoolSize());
223
224        tokenPool.clear();
225
226        for (int i = 0; i < config.getTokenPoolSize(); i++) {
227            String token = fetchSingleToken();
228            tokenPool.add(new TokenInfo(token));
229        }
230
231        // Reset the index to start from the beginning
232        tokenIndex.set(0);
233
234        LOGGER.debug("Token pool refreshed with %s tokens", tokenPool.size());
235    }
236}