Skip to content

Commit

Permalink
Merge pull request cdapio#15376 from cdapio/cache-provisioned-credential
Browse files Browse the repository at this point in the history 
[CDAP-20860] Cache provisioned credentials in memory for task workers and preview runners.
  • Loading branch information
@itsankit-google
itsankit-google authored Oct 25, 2023
2 parents f0d6ddc + 4191641 commit d1d9ff4
Show file tree
Hide file tree
Showing 10 changed files with 164 additions and 62 deletions.
1 change: 1 addition & 0 deletions cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/deploy/RemoteConfigurator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ public ListenableFuture<ConfigResponse> config() {
try {
RunnableTaskRequest request = RunnableTaskRequest.getBuilder(ConfiguratorTask.class.getName())
.withParam(GSON.toJson(deploymentInfo))
.withNamespace(deploymentInfo.getNamespaceId().getNamespace())
.build();

byte[] result = remoteTaskExecutor.runTask(request);
Expand Down
61 changes: 38 additions & 23 deletions ...rc/main/java/io/cdap/cdap/internal/app/worker/sidecar/GcpMetadataHttpHandlerInternal.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,13 @@

package io.cdap.cdap.internal.app.worker.sidecar;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonSyntaxException;
import com.google.inject.Singleton;
import io.cdap.cdap.api.common.HttpErrorStatusProvider;
import io.cdap.cdap.common.BadRequestException;
import io.cdap.cdap.common.ForbiddenException;
import io.cdap.cdap.common.conf.CConfiguration;
Expand All @@ -32,7 +34,6 @@
import io.cdap.cdap.internal.namespace.credential.RemoteNamespaceCredentialProvider;
import io.cdap.cdap.proto.BasicThrowable;
import io.cdap.cdap.proto.codec.BasicThrowableCodec;
import io.cdap.cdap.proto.credential.CredentialProvisioningException;
import io.cdap.cdap.proto.credential.NamespaceCredentialProvider;
import io.cdap.cdap.proto.credential.NotFoundException;
import io.cdap.cdap.proto.credential.ProvisionedCredential;
Expand All @@ -49,6 +50,8 @@
import java.net.URL;
import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import javax.ws.rs.DELETE;
import javax.ws.rs.GET;
import javax.ws.rs.PUT;
Expand All @@ -74,6 +77,8 @@ public class GcpMetadataHttpHandlerInternal extends AbstractAppFabricHttpHandler
private final NamespaceCredentialProvider credentialProvider;
private final GcpWorkloadIdentityInternalAuthenticator gcpWorkloadIdentityInternalAuthenticator;
private GcpMetadataTaskContext gcpMetadataTaskContext;
private final LoadingCache<ProvisionedCredentialCacheKey,
ProvisionedCredential> credentialLoadingCache;

/**
* Constructs the {@link GcpMetadataHttpHandlerInternal}.
Expand All @@ -89,6 +94,18 @@ public GcpMetadataHttpHandlerInternal(CConfiguration cConf,
new GcpWorkloadIdentityInternalAuthenticator(gcpMetadataTaskContext);
this.credentialProvider = new RemoteNamespaceCredentialProvider(remoteClientFactory,
this.gcpWorkloadIdentityInternalAuthenticator);
this.credentialLoadingCache = CacheBuilder.newBuilder()
// Provisioned credential expire after 60mins, assuming 20% buffer in cache exp (0.8*60).
.expireAfterWrite(48, TimeUnit.MINUTES)
.build(new CacheLoader<ProvisionedCredentialCacheKey, ProvisionedCredential>() {
@Override
public ProvisionedCredential load(ProvisionedCredentialCacheKey
provisionedCredentialCacheKey) throws Exception {
return fetchTokenFromCredentialProvider(
provisionedCredentialCacheKey.getGcpMetadataTaskContext(),
provisionedCredentialCacheKey.getScopes());
}
});
}

/**
Expand Down Expand Up @@ -139,31 +156,29 @@ public void token(HttpRequest request, HttpResponder responder,
// needed when initializing
// io.cdap.cdap.common.guice.DFSLocationModule$LocationFactoryProvider#get
// in io.cdap.cdap.internal.app.worker.TaskWorkerTwillRunnable.
LOG.warn("The GCP Metadata Task Context has been identified as null.");
GcpTokenResponse gcpTokenResponse = new GcpTokenResponse("Bearer", "invalidToken", 3599);
responder.sendJson(HttpResponseStatus.OK, GSON.toJson(gcpTokenResponse));
return;
}

try {
// fetch token from credential provider
GcpTokenResponse gcpTokenResponse =
Retries.callWithRetries(() -> fetchTokenFromCredentialProvider(scopes),
RetryStrategies.fromConfiguration(cConf, Constants.Service.TASK_WORKER + "."));
ProvisionedCredential provisionedCredential =
credentialLoadingCache.get(
new ProvisionedCredentialCacheKey(this.gcpMetadataTaskContext, scopes));
GcpTokenResponse gcpTokenResponse = new GcpTokenResponse("Bearer",
provisionedCredential.get(),
Duration.between(Instant.now(), provisionedCredential.getExpiration()).getSeconds());
responder.sendJson(HttpResponseStatus.OK, GSON.toJson(gcpTokenResponse));
return;
} catch (NotFoundException e) {
} catch (ExecutionException e) {
if (!(e.getCause() instanceof NotFoundException)) {
LOG.error("Failed to fetch token from credential provider", e.getCause());
throw e;
}
// if credential identity not found,
// fallback to gcp metadata server for backward compatibility.
} catch (Exception ex) {
if (ex instanceof HttpErrorStatusProvider) {
HttpResponseStatus status = HttpResponseStatus.valueOf(
((HttpErrorStatusProvider) ex).getStatusCode());
responder.sendJson(status, exceptionToJson(ex));
} else {
LOG.warn("Failed to fetch token from credential provider", ex);
responder.sendJson(HttpResponseStatus.INTERNAL_SERVER_ERROR, exceptionToJson(ex));
}
return;
}

if (metadataServiceTokenEndpoint == null) {
Expand All @@ -177,17 +192,16 @@ public void token(HttpRequest request, HttpResponder responder,
responder.sendJson(HttpResponseStatus.OK,
fetchTokenFromMetadataServer(scopes).getResponseBodyAsString());
} catch (Exception ex) {
LOG.warn("Failed to fetch token from metadata service", ex);
LOG.error("Failed to fetch token from metadata server", ex);
responder.sendJson(HttpResponseStatus.INTERNAL_SERVER_ERROR, exceptionToJson(ex));
}
}

private GcpTokenResponse fetchTokenFromCredentialProvider(String scopes) throws NotFoundException,
IOException, CredentialProvisioningException {
ProvisionedCredential provisionedCredential =
this.credentialProvider.provision(gcpMetadataTaskContext.getNamespace(), scopes);
return new GcpTokenResponse("Bearer", provisionedCredential.get(),
Duration.between(Instant.now(), provisionedCredential.getExpiration()).getSeconds());
private ProvisionedCredential fetchTokenFromCredentialProvider(
GcpMetadataTaskContext gcpMetadataTaskContext, String scopes) throws Exception {
return Retries.callWithRetries(() ->
this.credentialProvider.provision(gcpMetadataTaskContext.getNamespace(), scopes),
RetryStrategies.fromConfiguration(cConf, Constants.Service.TASK_WORKER + "."));
}

private HttpResponse fetchTokenFromMetadataServer(String scopes) throws IOException {
Expand Down Expand Up @@ -229,6 +243,7 @@ public void setContext(FullHttpRequest request, HttpResponder responder)
public void clearContext(HttpRequest request, HttpResponder responder) {
this.gcpMetadataTaskContext = null;
this.gcpWorkloadIdentityInternalAuthenticator.setGcpMetadataTaskContext(gcpMetadataTaskContext);
this.credentialLoadingCache.invalidateAll();
LOG.trace("Context cleared.");
responder.sendStatus(HttpResponseStatus.OK);
}
Expand Down
72 changes: 72 additions & 0 deletions ...src/main/java/io/cdap/cdap/internal/app/worker/sidecar/ProvisionedCredentialCacheKey.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
/*
* Copyright © 2023 Cask Data, Inc.
*
* Licensed 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 io.cdap.cdap.internal.app.worker.sidecar;

import io.cdap.cdap.proto.security.GcpMetadataTaskContext;
import java.util.Objects;

/**
* Defines the contents of key used for
* caching {@link io.cdap.cdap.proto.credential.ProvisionedCredential}.
*/
public final class ProvisionedCredentialCacheKey {
private final GcpMetadataTaskContext gcpMetadataTaskContext;
private final String scopes;
private transient Integer hashCode;

public ProvisionedCredentialCacheKey(GcpMetadataTaskContext gcpMetadataTaskContext,
String scopes) {
this.gcpMetadataTaskContext = gcpMetadataTaskContext;
this.scopes = scopes;
}

public GcpMetadataTaskContext getGcpMetadataTaskContext() {
return gcpMetadataTaskContext;
}

public String getScopes() {
return scopes;
}

@Override
public boolean equals(Object o) {
if (!(o instanceof ProvisionedCredentialCacheKey)) {
return false;
}
ProvisionedCredentialCacheKey that = (ProvisionedCredentialCacheKey) o;
return Objects.equals(gcpMetadataTaskContext.getNamespace(),
that.gcpMetadataTaskContext.getNamespace())
&& Objects.equals(gcpMetadataTaskContext.getUserCredential().toString(),
that.gcpMetadataTaskContext.getUserCredential().toString())
&& Objects.equals(gcpMetadataTaskContext.getUserId(),
that.gcpMetadataTaskContext.getUserId())
&& Objects.equals(gcpMetadataTaskContext.getUserIp(),
that.gcpMetadataTaskContext.getUserIp())
&& Objects.equals(scopes, that.scopes);
}

@Override
public int hashCode() {
Integer hashCode = this.hashCode;
if (hashCode == null) {
this.hashCode = hashCode = Objects.hash(gcpMetadataTaskContext.getNamespace(),
gcpMetadataTaskContext.getUserCredential().toString(),
gcpMetadataTaskContext.getUserId(), gcpMetadataTaskContext.getUserIp(), scopes);
}
return hashCode;
}
}
3 changes: 1 addition & 2 deletions cdap-app-fabric/src/test/java/io/cdap/cdap/internal/app/worker/TaskWorkerMetricsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,8 +105,7 @@ public void afterTest() {
public void testSimpleRequest() throws IOException {
String taskClassName = TaskWorkerServiceTest.TestRunnableClass.class.getName();
RunnableTaskRequest req = RunnableTaskRequest.getBuilder(taskClassName)
.withParam("100")
.build();
.withParam("100").withNamespace("testNamespace").build();
String reqBody = GSON.toJson(req);
HttpResponse response = HttpRequests.execute(
HttpRequest.post(uri.resolve("/v3Internal/worker/run").toURL())
Expand Down
19 changes: 10 additions & 9 deletions cdap-app-fabric/src/test/java/io/cdap/cdap/internal/app/worker/TaskWorkerServiceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,8 +145,8 @@ public void testPeriodicRestartWithInflightRequest() throws IOException {

// Post valid request
String want = "5000";
RunnableTaskRequest req = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName()).withParam(want)
.build();
RunnableTaskRequest req = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName())
.withParam(want).withNamespace("testNamespace").build();
String reqBody = GSON.toJson(req);
HttpResponse response = HttpRequests.execute(
HttpRequest.post(uri.resolve("/v3Internal/worker/run").toURL())
Expand Down Expand Up @@ -224,8 +224,8 @@ public void testRestartAfterMultipleExecutions() throws IOException {

// Post valid request
String want = "100";
RunnableTaskRequest req = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName()).withParam(want)
.build();
RunnableTaskRequest req = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName())
.withParam(want).withNamespace("testNamespace").build();
String reqBody = GSON.toJson(req);
HttpResponse response = HttpRequests.execute(
HttpRequest.post(uri.resolve("/v3Internal/worker/run").toURL())
Expand All @@ -248,8 +248,8 @@ public void testStartAndStopWithValidRequest() throws IOException {

// Post valid request
String want = "100";
RunnableTaskRequest req = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName()).withParam(want)
.build();
RunnableTaskRequest req = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName())
.withParam(want).withNamespace("testNamespace").build();
String reqBody = GSON.toJson(req);
HttpResponse response = HttpRequests.execute(
HttpRequest.post(uri.resolve("/v3Internal/worker/run").toURL())
Expand All @@ -267,7 +267,8 @@ public void testStartAndStopWithInvalidRequest() throws Exception {
URI uri = URI.create(String.format("http://%s:%s", addr.getHostName(), addr.getPort()));

// Post invalid request
RunnableTaskRequest noClassReq = RunnableTaskRequest.getBuilder("NoClass").build();
RunnableTaskRequest noClassReq = RunnableTaskRequest.getBuilder("NoClass")
.withNamespace("testNamespace").withParam("100").build();
String reqBody = GSON.toJson(noClassReq);
HttpResponse response = HttpRequests.execute(
HttpRequest.post(uri.resolve("/v3Internal/worker/run").toURL())
Expand All @@ -287,8 +288,8 @@ public void testConcurrentRequests() throws Exception {
InetSocketAddress addr = taskWorkerService.getBindAddress();
URI uri = URI.create(String.format("http://%s:%s", addr.getHostName(), addr.getPort()));

RunnableTaskRequest request = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName()).
withParam("1000").build();
RunnableTaskRequest request = RunnableTaskRequest.getBuilder(TestRunnableClass.class.getName())
.withParam("1000").withNamespace("testNamespace").build();

String reqBody = GSON.toJson(request);
List<Callable<HttpResponse>> calls = new ArrayList<>();
Expand Down
13 changes: 8 additions & 5 deletions ...mmon/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,13 +207,16 @@ public void run(FullHttpRequest request, HttpResponder responder) {
RunnableTaskContext runnableTaskContext = new RunnableTaskContext(
runnableTaskRequest);
try {
if (runnableTaskRequest.getParam().getEmbeddedTaskRequest() != null
&& runnableTaskRequest.getParam().getEmbeddedTaskRequest().getNamespace() != null) {
// set the GcpMetadataTaskContext before running the task.
NamespaceId namespaceId = new NamespaceId(
NamespaceId namespaceId;
if (runnableTaskRequest.getParam().getEmbeddedTaskRequest() != null) {
// For system app tasks
namespaceId = new NamespaceId(
runnableTaskRequest.getParam().getEmbeddedTaskRequest().getNamespace());
GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, cConf);
} else {
namespaceId = new NamespaceId(runnableTaskRequest.getNamespace());
}
// set the GcpMetadataTaskContext before running the task.
GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, cConf);
runnableTaskLauncher.launchRunnableTask(runnableTaskContext);
TaskDetails taskDetails = new TaskDetails(metricsCollectionService,
startTime,
Expand Down
6 changes: 3 additions & 3 deletions cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ public void testFailedMetrics() throws Exception {
RemoteTaskExecutor remoteTaskExecutor = new RemoteTaskExecutor(cConf, mockMetricsCollector, remoteClientFactory,
RemoteTaskExecutor.Type.TASK_WORKER);
RunnableTaskRequest runnableTaskRequest = RunnableTaskRequest.getBuilder(InValidRunnableClass.class.getName()).
withParam("param").build();
withParam("param").withNamespace("testNamespace").build();
try {
remoteTaskExecutor.runTask(runnableTaskRequest);
} catch (RemoteExecutionException e) {
Expand All @@ -203,7 +203,7 @@ public void testSuccessMetrics() throws Exception {
RemoteTaskExecutor remoteTaskExecutor = new RemoteTaskExecutor(cConf, mockMetricsCollector, remoteClientFactory,
RemoteTaskExecutor.Type.TASK_WORKER);
RunnableTaskRequest runnableTaskRequest = RunnableTaskRequest.getBuilder(ValidRunnableClass.class.getName()).
withParam("param").build();
withParam("param").withNamespace("testNamespace").build();
remoteTaskExecutor.runTask(runnableTaskRequest);
mockMetricsCollector.stopAndWait();
Assert.assertSame(1, metricCollectors.size());
Expand All @@ -224,7 +224,7 @@ public void testRetryMetrics() throws Exception {
RemoteTaskExecutor remoteTaskExecutor = new RemoteTaskExecutor(cConf, mockMetricsCollector, remoteClientFactory,
RemoteTaskExecutor.Type.TASK_WORKER);
RunnableTaskRequest runnableTaskRequest = RunnableTaskRequest.getBuilder(ValidRunnableClass.class.getName()).
withParam("param").build();
withParam("param").withNamespace("testNamespace").build();
try {
remoteTaskExecutor.runTask(runnableTaskRequest);
} catch (Exception e) {
Expand Down
1 change: 1 addition & 0 deletions ...main/java/io/cdap/cdap/security/spi/credential/GcpWorkloadIdentityCredentialProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,7 @@ private String exchangeTokenViaSts(String token, String scopes, String audience)
throws IOException {

// replace comma with space, see:
// https://cloud.google.com/functions/docs/securing/function-identity#access_tokens
// https://cloud.google.com/iam/docs/reference/sts/rest/v1/TopLevel/token#request-body
scopes = Arrays.stream(scopes.split(",")).map(String::trim)
.filter(s -> !s.isEmpty()).distinct().collect(Collectors.joining(" "));
Expand Down
38 changes: 21 additions & 17 deletions cdap-kubernetes/src/main/java/io/cdap/cdap/master/environment/k8s/KubeMasterEnvironment.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@
import java.io.FileWriter;
import java.io.IOException;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
Expand Down Expand Up @@ -549,32 +550,35 @@ public void onNamespaceCreation(String cdapNamespace, Map<String, String> proper

@Override
public void createIdentity(String k8sNamespace, String identity) throws ApiException {
if (identity.equals("default")) {
// skip creating default service account as it already exists.
return;
}

KubeUtil.validateRFC1123LabelName(identity);
LOG.info("Creating credential identity: {}", identity);
V1ObjectMeta serviceAccountMetadata = new V1ObjectMeta();
serviceAccountMetadata.setName(identity);
V1ServiceAccount serviceAccount = new V1ServiceAccount();
serviceAccount.setMetadata(serviceAccountMetadata);
try {
coreV1Api.createNamespacedServiceAccount(k8sNamespace, serviceAccount,
null, null, null, null);
LOG.info("Creating credential identity: {}", identity);
createK8sSaIfNotExists(k8sNamespace, identity);
} catch (ApiException e) {
if (e.getCode() == 409) {
// ignore, the SA already exists.
return;
}
LOG.error(
String.format("Unable to create the service account %s with status %s and body: %s",
serviceAccount.getMetadata().getName(), e.getCode(), e.getResponseBody()), e);
identity, e.getCode(), e.getResponseBody()), e);
throw e;
}
}

private void createK8sSaIfNotExists(String k8sNamespace, String serviceAccountName)
throws ApiException {
try {
coreV1Api.readNamespacedServiceAccount(serviceAccountName, k8sNamespace, null);
} catch (ApiException e) {
if (e.getCode() != HttpURLConnection.HTTP_NOT_FOUND) {
throw e;
}
V1ObjectMeta serviceAccountMetadata = new V1ObjectMeta();
serviceAccountMetadata.setName(serviceAccountName);
V1ServiceAccount serviceAccount = new V1ServiceAccount();
serviceAccount.setMetadata(serviceAccountMetadata);
coreV1Api.createNamespacedServiceAccount(k8sNamespace, serviceAccount,
null, null, null, null);
}
}

@Override
public void deleteIdentity(String k8sNamespace, @Nullable String identity) throws ApiException {
if (identity == null || identity.equals("default")) {
Expand Down
Loading

0 comments on commit d1d9ff4

Please sign in to comment.