Skip to content

fix(codeexecutors): sandbox ContainerCodeExecutor Docker containers #1369

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+363 −81
Open

fix(codeexecutors): sandbox ContainerCodeExecutor Docker containers #1369

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
268 changes: 187 additions & 81 deletions core/src/main/java/com/google/adk/codeexecutors/ContainerCodeExecutor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
/*
* Copyright 2025 Google LLC
* Copyright 2026 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may not in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
Expand All @@ -20,33 +19,69 @@
import static java.util.Objects.requireNonNullElse;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
import com.github.dockerjava.api.model.Container;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.model.Capability;
import com.github.dockerjava.api.model.Frame;
import com.github.dockerjava.api.model.HostConfig;
import com.github.dockerjava.api.model.StreamType;
import com.github.dockerjava.core.DefaultDockerClientConfig;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.command.ExecStartResultCallback;
import com.github.dockerjava.core.command.AttachContainerResultCallback;
import com.google.adk.agents.InvocationContext;
import com.google.adk.codeexecutors.CodeExecutionUtils.CodeExecutionInput;
import com.google.adk.codeexecutors.CodeExecutionUtils.CodeExecutionResult;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Paths;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** A code executor that uses a custom container to execute code. */
public class ContainerCodeExecutor extends BaseCodeExecutor {
/**
* A code executor that runs code in a sandboxed Docker container.
*
* <p>A fresh container is created for every {@link #executeCode} call so code from one session
* cannot observe or affect another session's execution environment. Each container is hardened: no
* network (so the code cannot reach the cloud metadata endpoint), all Linux capabilities dropped,
* no privilege escalation, a read-only root filesystem with a small writable {@code /tmp} tmpfs,
* and memory/PID limits. The code runs as the container's main process wrapped in an in-container
* {@code timeout}, and the container is set to auto-remove on exit; so even if this JVM dies, the
* OS kills the code and Docker removes the container rather than leaving a resource-draining
* "zombie".
*
* <p>This executor holds a {@link DockerClient}; call {@link #close()} (or rely on the registered
* JVM shutdown hook) to release its connections and threads.
*/
public class ContainerCodeExecutor extends BaseCodeExecutor implements AutoCloseable {
private static final Logger logger = LoggerFactory.getLogger(ContainerCodeExecutor.class);
private static final String DEFAULT_IMAGE_TAG = "adk-code-executor:latest";

/** Memory limit for each execution container (512 MiB). */
private static final long MEMORY_LIMIT_BYTES = 512L * 1024 * 1024;

/** Maximum number of processes/threads allowed inside an execution container. */
private static final long PIDS_LIMIT = 128L;

/** Maximum wall-clock time a single execution may run before its container is killed. */
private static final long EXECUTION_TIMEOUT_SECONDS = 60L;

/**
* Extra time to wait for the container to exit after its in-container timeout should have fired.
*/
private static final long COMPLETION_GRACE_SECONDS = 10L;

private final String baseUrl;
private final String image;
private final String dockerPath;
private final DockerClient dockerClient;
private Container container;
private boolean networkEnabled = false;
private long executionTimeoutSeconds = EXECUTION_TIMEOUT_SECONDS;

/**
* Creates a ContainerCodeExecutor from an image.
Expand Down Expand Up @@ -100,17 +135,33 @@ public ContainerCodeExecutor(String baseUrl, String image, String dockerPath) {
this.baseUrl = baseUrl;
this.image = requireNonNullElse(image, DEFAULT_IMAGE_TAG);
this.dockerPath = dockerPath == null ? null : Paths.get(dockerPath).toAbsolutePath().toString();
this.dockerClient = buildDockerClient(baseUrl);
prepareImage();
// Backstop so the client is released even if callers forget to close() this executor.
Runtime.getRuntime().addShutdownHook(new Thread(this::close));
}

if (baseUrl != null) {
var config =
DefaultDockerClientConfig.createDefaultConfigBuilder().withDockerHost(baseUrl).build();
this.dockerClient = DockerClientBuilder.getInstance(config).build();
} else {
this.dockerClient = DockerClientBuilder.getInstance().build();
}
/** Test-only constructor that injects a Docker client and skips image preparation. */
@VisibleForTesting
ContainerCodeExecutor(DockerClient dockerClient, String image) {
this.baseUrl = null;
this.image = requireNonNullElse(image, DEFAULT_IMAGE_TAG);
this.dockerPath = null;
this.dockerClient = dockerClient;
}

/**
* Enables or disables container networking. Networking is disabled by default so executed code
* cannot reach the network, including the cloud metadata endpoint.
*/
public ContainerCodeExecutor setNetworkEnabled(boolean networkEnabled) {
this.networkEnabled = networkEnabled;
return this;
}

initContainer();
Runtime.getRuntime().addShutdownHook(new Thread(this::cleanupContainer));
@VisibleForTesting
void setExecutionTimeoutSeconds(long executionTimeoutSeconds) {
this.executionTimeoutSeconds = executionTimeoutSeconds;
}

@Override
Expand All @@ -129,67 +180,107 @@ public CodeExecutionResult executeCode(
ByteArrayOutputStream stdout = new ByteArrayOutputStream();
ByteArrayOutputStream stderr = new ByteArrayOutputStream();

ExecCreateCmdResponse execCreateCmdResponse =
// Run the code as the container's own main process, wrapped in an in-container `timeout` so the
// OS terminates a runaway even if this JVM dies. A fresh, auto-removing container per execution
// isolates each run and guarantees cleanup.
CreateContainerResponse createContainerResponse =
dockerClient
.execCreateCmd(container.getId())
.withAttachStdout(true)
.withAttachStderr(true)
.withCmd("python3", "-c", codeExecutionInput.code())
.createContainerCmd(image)
.withHostConfig(sandboxHostConfig())
.withCmd(
"timeout",
"-k",
"5",
Long.toString(executionTimeoutSeconds),
"python3",
"-c",
codeExecutionInput.code())
.exec();
String containerId = createContainerResponse.getId();
try {
dockerClient
.execStartCmd(execCreateCmdResponse.getId())
.exec(new ExecStartResultCallback(stdout, stderr))
.awaitCompletion();
// Attach before starting so all output is captured before auto-remove reaps the container.
AttachContainerResultCallback attachCallback =
dockerClient
.attachContainerCmd(containerId)
.withStdOut(true)
.withStdErr(true)
.withFollowStream(true)
.exec(new FrameCapturingCallback(stdout, stderr));
dockerClient.startContainerCmd(containerId).exec();

boolean completed =
attachCallback.awaitCompletion(
executionTimeoutSeconds + COMPLETION_GRACE_SECONDS, TimeUnit.SECONDS);
if (!completed) {
// The in-container timeout should have exited the container already; force-remove it below.
return CodeExecutionResult.builder()
.stderr(
String.format(
"Code execution timed out after %d seconds.", executionTimeoutSeconds))
.build();
}
return CodeExecutionResult.builder()
.stdout(stdout.toString(StandardCharsets.UTF_8))
.stderr(stderr.toString(StandardCharsets.UTF_8))
.build();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Code execution was interrupted.", e);
} finally {
// Backstop: auto-remove usually reaped the container already; force-remove any remnant.
removeContainerQuietly(containerId);
}

return CodeExecutionResult.builder()
.stdout(stdout.toString(StandardCharsets.UTF_8))
.stderr(stderr.toString(StandardCharsets.UTF_8))
.build();
}

private void buildDockerImage() {
if (dockerPath == null) {
throw new IllegalStateException("Docker path is not set.");
}
File dockerfile = new File(dockerPath);
if (!dockerfile.exists()) {
throw new UncheckedIOException(new IOException("Invalid Docker path: " + dockerPath));
/** Builds the hardened {@link HostConfig} applied to every execution container. */
@VisibleForTesting
HostConfig sandboxHostConfig() {
HostConfig hostConfig =
HostConfig.newHostConfig()
.withCapDrop(Capability.ALL)
.withReadonlyRootfs(true)
.withSecurityOpts(ImmutableList.of("no-new-privileges"))
.withMemory(MEMORY_LIMIT_BYTES)
.withPidsLimit(PIDS_LIMIT)
// Remove the container when it exits so an abruptly-killed JVM cannot leak containers.
.withAutoRemove(true)
// A read-only rootfs still needs a small writable scratch space at /tmp.
.withTmpFs(ImmutableMap.of("/tmp", "rw,size=64m"));
if (!networkEnabled) {
hostConfig.withNetworkMode("none");
}
return hostConfig;
}

logger.info("Building Docker image...");
private void removeContainerQuietly(String containerId) {
try {
dockerClient.buildImageCmd(dockerfile).withTag(image).start().awaitCompletion();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Docker image build was interrupted.", e);
dockerClient.removeContainerCmd(containerId).withForce(true).exec();
} catch (RuntimeException e) {
// The container was most likely already auto-removed on exit; this is expected.
logger.debug("Container {} was already removed or could not be removed.", containerId, e);
}
logger.info("Docker image: {} built.", image);
}

private void verifyPythonInstallation() {
ExecCreateCmdResponse execCreateCmdResponse =
dockerClient.execCreateCmd(container.getId()).withCmd("which", "python3").exec();
ByteArrayOutputStream stdout = new ByteArrayOutputStream();
ByteArrayOutputStream stderr = new ByteArrayOutputStream();
try (ExecStartResultCallback callback = new ExecStartResultCallback(stdout, stderr)) {
dockerClient.execStartCmd(execCreateCmdResponse.getId()).exec(callback).awaitCompletion();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Python verification was interrupted.", e);
/** Closes the underlying Docker client, releasing its connections and threads. */
@Override
public void close() {
try {
dockerClient.close();
} catch (IOException e) {
throw new UncheckedIOException(e);
logger.warn("Failed to close docker client", e);
}
}

private void initContainer() {
if (dockerClient == null) {
throw new IllegalStateException("Docker client is not initialized.");
private static DockerClient buildDockerClient(String baseUrl) {
if (baseUrl != null) {
var config =
DefaultDockerClientConfig.createDefaultConfigBuilder().withDockerHost(baseUrl).build();
return DockerClientBuilder.getInstance(config).build();
}
return DockerClientBuilder.getInstance().build();
}

private void prepareImage() {
if (dockerPath != null) {
buildDockerImage();
} else {
Expand All @@ -204,34 +295,49 @@ private void initContainer() {
}
logger.info("Image {} is available.", image);
}
logger.info("Starting container for ContainerCodeExecutor...");
var createContainerResponse =
dockerClient.createContainerCmd(image).withTty(true).withAttachStdin(true).exec();
dockerClient.startContainerCmd(createContainerResponse.getId()).exec();
}

var containers = dockerClient.listContainersCmd().withShowAll(true).exec();
this.container =
containers.stream()
.filter(c -> c.getId().equals(createContainerResponse.getId()))
.findFirst()
.orElseThrow(() -> new IllegalStateException("Failed to find the created container."));
private void buildDockerImage() {
if (dockerPath == null) {
throw new IllegalStateException("Docker path is not set.");
}
File dockerfile = new File(dockerPath);
if (!dockerfile.exists()) {
throw new UncheckedIOException(new IOException("Invalid Docker path: " + dockerPath));
}

logger.info("Container {} started.", container.getId());
verifyPythonInstallation();
logger.info("Building Docker image...");
try {
dockerClient.buildImageCmd(dockerfile).withTag(image).start().awaitCompletion();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Docker image build was interrupted.", e);
}
logger.info("Docker image: {} built.", image);
}

private void cleanupContainer() {
if (container == null) {
return;
/** Routes container output frames into separate stdout and stderr buffers. */
private static final class FrameCapturingCallback extends AttachContainerResultCallback {
private final ByteArrayOutputStream stdout;
private final ByteArrayOutputStream stderr;

FrameCapturingCallback(ByteArrayOutputStream stdout, ByteArrayOutputStream stderr) {
this.stdout = stdout;
this.stderr = stderr;
}
logger.info("[Cleanup] Stopping the container...");
dockerClient.stopContainerCmd(container.getId()).exec();
dockerClient.removeContainerCmd(container.getId()).exec();
logger.info("Container {} stopped and removed.", container.getId());
try {
dockerClient.close();
} catch (IOException e) {
logger.warn("Failed to close docker client", e);

@Override
public void onNext(Frame frame) {
try {
if (frame.getStreamType() == StreamType.STDERR) {
stderr.write(frame.getPayload());
} else {
stdout.write(frame.getPayload());
}
} catch (IOException e) {
throw new UncheckedIOException(e);
}
super.onNext(frame);
}
}
}
Loading
Loading