Skip to content

Commit

Permalink
feat: make custom heracles client to support the user fetching use cases
Browse files Browse the repository at this point in the history
  • Loading branch information
sumandas0 committed Dec 5, 2023
1 parent 4fddf2a commit 269b34c
Show file tree
Hide file tree
Showing 27 changed files with 544 additions and 190 deletions.
2 changes: 1 addition & 1 deletion auth-agents-common/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@

<dependency>
<groupId>org.apache.atlas</groupId>
<artifactId>client-keycloak</artifactId>
<artifactId>client-auth</artifactId>
<version>${project.version}</version>
</dependency>

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@
import org.keycloak.representations.idm.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.commons.configuration.Configuration;

import java.util.*;
import java.util.concurrent.Callable;
Expand All @@ -40,7 +39,8 @@
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static org.apache.atlas.keycloak.client.AtlasKeycloakClient.getKeycloakClient;
import static org.apache.atlas.auth.client.keycloak.AtlasKeycloakClient.getKeycloakClient;
import static org.apache.atlas.auth.client.heracles.AtlasHeraclesClient.getHeraclesClient;
import static org.apache.atlas.repository.Constants.*;
import static org.apache.atlas.repository.util.AccessControlUtils.ARGO_SERVICE_USER_NAME;
import static org.apache.atlas.repository.util.AccessControlUtils.BACKEND_SERVICE_USER_NAME;
Expand Down Expand Up @@ -277,10 +277,8 @@ public RangerUserStore loadUserStoreIfUpdated(long lastUpdatedTime) throws Atlas
}

Map<String, Set<String>> userGroupMapping = new HashMap<>();

List<UserRepresentation> kUsers = getKeycloakClient().getAllUsers();
List<UserRepresentation> kUsers = getHeraclesClient().getAllUsers();
LOG.info("Found {} keycloak users", kUsers.size());

List<Callable<Object>> callables = new ArrayList<>();
kUsers.forEach(x -> callables.add(new UserGroupsFetcher(x, userGroupMapping)));

Expand Down Expand Up @@ -417,13 +415,13 @@ public RangerRole call() throws Exception {
//get all users for Roles
Thread usersFetcher = new Thread(() -> {
int start = 0;
int size = 1500;
int size = 100;
boolean found = true;
Set<UserRepresentation> ret = new HashSet<>();

do {
try {
Set<UserRepresentation> userRepresentations = getKeycloakClient().getRoleUserMembers(kRole.getName(), start, size);
Set<UserRepresentation> userRepresentations = getHeraclesClient().getRoleUserMembers(kRole.getName(), start, size);
if (CollectionUtils.isNotEmpty(userRepresentations)) {
ret.addAll(userRepresentations);
start += size;
Expand Down
2 changes: 1 addition & 1 deletion client-keycloak/pom.xml → client-auth/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
</parent>
<modelVersion>4.0.0</modelVersion>

<artifactId>client-keycloak</artifactId>
<artifactId>client-auth</artifactId>

<properties>
<maven.compiler.source>8</maven.compiler.source>
Expand Down
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
package org.apache.atlas.keycloak.client;
package org.apache.atlas.auth.client.auth;

import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.micrometer.core.instrument.Timer;
import org.apache.atlas.auth.client.config.AuthConfig;
import org.apache.atlas.auth.client.heracles.RetrofitHeraclesClient;
import org.apache.atlas.auth.client.keycloak.RetrofitKeycloakClient;
import okhttp3.*;
import okhttp3.logging.HttpLoggingInterceptor;
import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.keycloak.client.config.KeycloakConfig;
import org.apache.atlas.keycloak.client.service.AtlasKeycloakAuthService;
import org.apache.atlas.service.metrics.MetricUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -26,9 +28,9 @@
import static org.apache.atlas.AtlasErrorCode.BAD_REQUEST;
import static org.apache.atlas.AtlasErrorCode.RESOURCE_NOT_FOUND;

abstract class AbstractKeycloakClient {
public class AbstractAuthClient {

private final static Logger LOG = LoggerFactory.getLogger(AbstractKeycloakClient.class);
private final static Logger LOG = LoggerFactory.getLogger(AbstractAuthClient.class);
private static final Map<Integer, AtlasErrorCode> ERROR_CODE_MAP = new HashMap<>();

private static final int DEFAULT_KEYCLOAK_RETRY = 3;
Expand All @@ -38,19 +40,20 @@ abstract class AbstractKeycloakClient {
private static final String INTEGRATION = "integration";
private static final String KEYCLOAK = "keycloak";

protected final KeycloakConfig keycloakConfig;
protected final RetrofitKeycloakClient retrofit;
protected final AuthConfig authConfig;
protected final RetrofitKeycloakClient retrofitKeycloakClient;
protected final RetrofitHeraclesClient retrofitHeraclesClient;

private final AtlasKeycloakAuthService authService;
private final KeycloakAuthenticationService authService;
private MetricUtils metricUtils = null;

static {
ERROR_CODE_MAP.put(HTTP_NOT_FOUND, RESOURCE_NOT_FOUND);
ERROR_CODE_MAP.put(HTTP_BAD_REQUEST, BAD_REQUEST);
}

public AbstractKeycloakClient(KeycloakConfig keycloakConfig) {
this.keycloakConfig = keycloakConfig;
public AbstractAuthClient(AuthConfig authConfig) {
this.authConfig = authConfig;
this.metricUtils = new MetricUtils();
HttpLoggingInterceptor httpInterceptor = new HttpLoggingInterceptor();
httpInterceptor.setLevel(HttpLoggingInterceptor.Level.BODY);
Expand All @@ -64,11 +67,15 @@ public AbstractKeycloakClient(KeycloakConfig keycloakConfig) {
.writeTimeout(TIMEOUT_IN_SEC, TimeUnit.SECONDS)
.readTimeout(TIMEOUT_IN_SEC, TimeUnit.SECONDS)
.build();
this.retrofit = new Retrofit.Builder().client(okHttpClient)
.baseUrl(this.keycloakConfig.getAuthServerUrl())
this.retrofitKeycloakClient = new Retrofit.Builder().client(okHttpClient)
.baseUrl(this.authConfig.getAuthServerUrl())
.addConverterFactory(JacksonConverterFactory.create(new ObjectMapper())).build()
.create(RetrofitKeycloakClient.class);
authService = new AtlasKeycloakAuthService(keycloakConfig);
this.retrofitHeraclesClient = new Retrofit.Builder().client(okHttpClient)
.baseUrl(this.authConfig.getHeraclesApiServerUrl())
.addConverterFactory(JacksonConverterFactory.create(new ObjectMapper().disable(DeserializationFeature.FAIL_ON_IGNORED_PROPERTIES))).build()
.create(RetrofitHeraclesClient.class);
authService = new KeycloakAuthenticationService(authConfig);
}

/**
Expand Down Expand Up @@ -97,7 +104,6 @@ public Response intercept(@NonNull Chain chain) throws IOException {
return chain.proceed(request);
}
};

/**
* Called only during auth failures.
*/
Expand All @@ -110,8 +116,8 @@ public Request authenticate(Route route, @NonNull Response response) {
}
LOG.info("Keycloak: Current keycloak token status, Expired: {}", authService.isTokenExpired());
return response.request().newBuilder()
.addHeader(AUTHORIZATION, BEARER + authService.getAuthToken())
.build();
.addHeader(AUTHORIZATION, BEARER + authService.getAuthToken())
.build();
}

private int responseCount(Response response) {
Expand Down Expand Up @@ -139,8 +145,9 @@ protected <T> retrofit2.Response<T> processResponse(retrofit2.Call<T> req) throw
throw new AtlasBaseException(ERROR_CODE_MAP.getOrDefault(response.code(), BAD_REQUEST), errMsg);
} catch (Exception e) {
LOG.error("Keycloak: request failed, request: {} {}, Exception: {}", req.request().method(), req.request().url(), e);
throw new AtlasBaseException(BAD_REQUEST, "Keycloak request failed");
throw new AtlasBaseException(BAD_REQUEST, "Auth request failed");
}
}


}
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
package org.apache.atlas.keycloak.client.service;
package org.apache.atlas.auth.client.auth;

import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;
import okhttp3.logging.HttpLoggingInterceptor;
import org.apache.atlas.auth.client.config.AuthConfig;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.keycloak.client.RetrofitKeycloakClient;
import org.apache.atlas.keycloak.client.config.KeycloakConfig;
import org.apache.atlas.auth.client.keycloak.RetrofitKeycloakClient;
import org.jetbrains.annotations.NotNull;
import org.keycloak.representations.AccessTokenResponse;
import org.slf4j.Logger;
Expand All @@ -19,9 +19,9 @@

import static org.apache.atlas.AtlasErrorCode.BAD_REQUEST;

public final class AtlasKeycloakAuthService {
public final class KeycloakAuthenticationService {

public final static Logger LOG = LoggerFactory.getLogger(AtlasKeycloakAuthService.class);
public final static Logger LOG = LoggerFactory.getLogger(KeycloakAuthenticationService.class);

private final static String GRANT_TYPE = "grant_type";
private static final String CLIENT_ID = "client_id";
Expand All @@ -30,14 +30,14 @@ public final class AtlasKeycloakAuthService {
private static final int TIMEOUT_IN_SECS = 60;

private final RetrofitKeycloakClient retrofit;
private final KeycloakConfig keycloakConfig;
private final AuthConfig authConfig;
private AccessTokenResponse currentAccessToken;
private long expirationTime = -1;

public AtlasKeycloakAuthService(KeycloakConfig keycloakConfig) {
this.keycloakConfig = keycloakConfig;
public KeycloakAuthenticationService(AuthConfig authConfig) {
this.authConfig = authConfig;
this.retrofit = new Retrofit.Builder().client(getOkHttpClient())
.baseUrl(this.keycloakConfig.getAuthServerUrl())
.baseUrl(this.authConfig.getAuthServerUrl())
.addConverterFactory(JacksonConverterFactory.create(new ObjectMapper())).build()
.create(RetrofitKeycloakClient.class);
}
Expand Down Expand Up @@ -70,7 +70,7 @@ public String getAuthToken() {
synchronized (this) {
if (isTokenExpired()) {
try {
retrofit2.Response<AccessTokenResponse> resp = this.retrofit.grantToken(this.keycloakConfig.getRealmId(), getTokenRequest()).execute();
retrofit2.Response<AccessTokenResponse> resp = this.retrofit.grantToken(this.authConfig.getRealmId(), getTokenRequest()).execute();
if (resp.isSuccessful()) {
currentAccessToken = resp.body();
expirationTime = currentTime() + currentAccessToken.getExpiresIn() - EXPIRY_OFFSET_SEC;
Expand All @@ -97,7 +97,7 @@ public boolean isTokenExpired() {
}

private RequestBody getTokenRequest() {
return new FormBody.Builder().addEncoded(CLIENT_ID, this.keycloakConfig.getClientId()).addEncoded(CLIENT_SECRET, this.keycloakConfig.getClientSecret()).addEncoded(GRANT_TYPE, this.keycloakConfig.getGrantType()).build();
return new FormBody.Builder().addEncoded(CLIENT_ID, this.authConfig.getClientId()).addEncoded(CLIENT_SECRET, this.authConfig.getClientSecret()).addEncoded(GRANT_TYPE, this.authConfig.getGrantType()).build();
}

private long currentTime() {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
package org.apache.atlas.auth.client.config;

import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.commons.lang.StringUtils;
import org.codehaus.jettison.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.Optional;

import static org.apache.atlas.ApplicationProperties.ATLAS_CONFIGURATION_DIRECTORY_PROPERTY;

public class AuthConfig {
private static final Logger LOG = LoggerFactory.getLogger(AuthConfig.class);

public String authServerUrl;
public String realmId;
public String clientId;
public String clientSecret;
public String grantType;
public String heraclesApiServerUrl;

private static final String KEYCLOAK_PROPERTIES = "keycloak.json";
private static final String DEFAULT_GRANT_TYPE = "client_credentials";
private static final String KEY_REALM_ID = "realm";
private static final String KEY_AUTH_SERVER_URL = "auth-server-url";
private static final String KEY_CLIENT_ID = "resource";
private static final String KEY_CREDENTIALS = "credentials";
private static final String KEY_SECRET = "secret";

public String getAuthServerUrl() {
return authServerUrl;
}

public String getRealmId() {
return realmId;
}

public String getClientId() {
return clientId;
}

public String getClientSecret() {
return clientSecret;
}

public String getGrantType() {
return grantType;
}

public String getHeraclesApiServerUrl() {
return heraclesApiServerUrl;
}

public static AuthConfig getConfig() throws AtlasBaseException {
String confLocation = System.getProperty(ATLAS_CONFIGURATION_DIRECTORY_PROPERTY);
Optional<File> confFile = getConfigurationFile(confLocation);

if (confFile.isPresent()) {
try {
JSONObject object = new JSONObject(readFileToString(confFile.get()));
return buildAuthConfigFromJson(object);
} catch (Exception e) {
LOG.error("Error parsing Keycloak configuration: ", e);
throw new AtlasBaseException(AtlasErrorCode.KEYCLOAK_INIT_FAILED, "Error parsing Keycloak configuration");
}
} else {
throw new AtlasBaseException(AtlasErrorCode.KEYCLOAK_INIT_FAILED, "Keycloak configuration file not found in location " + confLocation);
}
}

private static Optional<File> getConfigurationFile(String confLocation) {
if (StringUtils.isNotEmpty(confLocation)) {
File confFile = new File(confLocation, KEYCLOAK_PROPERTIES);
if (confFile.exists()) {
return Optional.of(confFile);
}
}
return Optional.empty();
}

private static String readFileToString(File file) throws Exception {
return new String(Files.readAllBytes(file.toPath()), StandardCharsets.UTF_8);
}

private static AuthConfig buildAuthConfigFromJson(JSONObject object) throws Exception {
String realmId = object.getString(KEY_REALM_ID);
String authServerUrl = object.getString(KEY_AUTH_SERVER_URL) + "/";
String clientId = object.getString(KEY_CLIENT_ID);
String grantType = DEFAULT_GRANT_TYPE;
String clientSecret = object.getJSONObject(KEY_CREDENTIALS).getString(KEY_SECRET);

LOG.info("Keycloak configuration: REALM_ID:{}, AUTH_SERVER_URL:{}", realmId, authServerUrl);
return AuthConfigBuilder.builder().realId(realmId).authServerUrl(authServerUrl).clientId(clientId).grantType(grantType).clientSecret(clientSecret).build();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
package org.apache.atlas.auth.client.config;

import org.apache.atlas.AtlasConfiguration;

public class AuthConfigBuilder {

private String authServerUrl;
private String realmId;
private String clientId;
private String clientSecret;
private String grantType = "client_credentials";

private AuthConfigBuilder() {
}

public static AuthConfigBuilder builder() {
return new AuthConfigBuilder();
}

public AuthConfigBuilder authServerUrl(String authServerUrl) {
this.authServerUrl = authServerUrl;
return this;
}

public AuthConfigBuilder realId(String realId) {
this.realmId = realId;
return this;
}

public AuthConfigBuilder clientId(String clientId) {
this.clientId = clientId;
return this;
}

public AuthConfigBuilder clientSecret(String clientSecret) {
this.clientSecret = clientSecret;
return this;
}

public AuthConfigBuilder grantType(String grantType) {
this.grantType = grantType;
return this;
}

public AuthConfig build() {
AuthConfig authConfig = new AuthConfig();
authConfig.authServerUrl = authServerUrl;
authConfig.realmId = realmId;
authConfig.clientId = clientId;
authConfig.clientSecret = clientSecret;
authConfig.grantType = grantType;
authConfig.heraclesApiServerUrl= AtlasConfiguration.HERACLES_API_SERVER_URL.getString()+"/";
return authConfig;
}
}
Loading

0 comments on commit 269b34c

Please sign in to comment.