docker-java-sdk-demo/src/main/java/space/anyi/docker/OjSandboxAPI.java

package space.anyi.docker;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
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 org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * Docker OJ（Online Judge）代码沙箱 MVP
 *
 * <p>判题流程：一次提交只使用<b>一个</b>容器——
 * 先注入源码并编译，再通过 docker exec 逐用例运行，前一个用例与后一个用例共享同一容器。</p>
 *
 * <pre>
 * 等价命令:
 *   docker run -d --name oj-sandbox -w /code --user nobody --read-only \
 *     --network none --cap-drop ALL -m 64m --cpus 1 --pids-limit 64 \
 *     --tmpfs /code --tmpfs /tmp amazoncorretto:17 sh -c "sleep 300"
 *   docker exec <cid> sh -c "printf %s '<BASE64>' | base64 -d > /code/Main.java"
 *   docker exec <cid> javac -encoding UTF-8 -d /code /code/Main.java
 *   docker exec -i <cid> java -Xmx48m -cp /code Main < case1.in
 *   docker exec -i <cid> java -Xmx48m -cp /code Main < case2.in
 *   docker rm -f <cid>
 * </pre>
 */
public class OjSandboxAPI {
    private static final Logger log = LoggerFactory.getLogger(OjSandboxAPI.class);

    /** 判题镜像：本机已存在的 amazoncorretto JDK 镜像，自带 javac/java */
    private static final String JUDGE_IMAGE = "amazoncorretto:17";
    /** 容器内代码目录（挂在可写 tmpfs 上，与只读根文件系统配合） */
    private static final String CODE_DIR = "/code";
    /** 源码文件名（选手主类必须叫 Main） */
    private static final String SOURCE_FILE = "Main.java";
    /** 编译阶段超时（毫秒），远超单文件 javac 的正常耗时 */
    private static final long COMPILE_TIMEOUT_MS = 20_000;
    /** 容器内常驻主进程，让容器保持运行以便后续 exec */
    private static final String KEEP_ALIVE_CMD = "sleep 300";

    private final DockerClient dockerClient;

    public OjSandboxAPI() {
        this.dockerClient = DockerClientFactory.createDockerClient();
    }

    // ==================== 判题结果模型 ====================

    /** 判题结果分类 */
    public enum JudgeVerdict {
        /** Accepted 通过 */
        AC,
        /** Wrong Answer 答案错误 */
        WA,
        /** Compile Error 编译错误 */
        CE,
        /** Runtime Error 运行时错误（进程非 0 退出） */
        RE,
        /** Time Limit Exceeded 超出时间限制 */
        TLE,
        /** Memory Limit Exceeded 超出内存限制（MVP 保留位） */
        MLE,
        /** Output Limit Exceeded 输出超出限制（MVP 保留位） */
        OLE,
        /** System Error 判题系统自身异常 */
        SE
    }

    /** 判题配置：单次提交的资源限制 */
    public record JudgeConfig(int timeLimitMs, int memoryLimitMb, int outputLimitBytes) {
    }

    /** 单个测试用例：标准输入与期望输出 */
    public record TestCase(String stdin, String expected) {
    }

    /** 单个用例的运行结果（verdict 为空表示进程正常退出，等待输出比对后定 AC/WA） */
    public record RunResult(String stdout, String stderr, int exitCode, long durationMs,
                            JudgeVerdict verdict) {
    }

    /** 单个用例的最终判定（含实际输出，便于展示与比对） */
    public record JudgeResult(int caseNo, TestCase testCase, RunResult run) {
    }

    /** 一次提交的完整判题结果 */
    public record JudgeOutcome(boolean compileOk, String compileOutput, List<JudgeResult> results) {
    }

    // ==================== 对外入口 ====================

    /**
     * 判题入口：一次提交只创建一个容器，完成源码注入 + 编译 + 逐用例运行
     * @param source 选手提交的 Java 源码（主类名必须为 Main）
     * @param config 判题资源配置
     * @param cases 测试用例列表
     * @return 判题结果（编译失败时 results 为空）
     */
    public JudgeOutcome judge(String source, JudgeConfig config, List<TestCase> cases) {
        String containerId = null;
        try {
            containerId = createSandboxContainer(config);
            log.info("沙箱容器已创建: {}", containerId);

            // 1) 注入源码
            injectSource(containerId, source);

            // 2) 编译（等价 docker exec <cid> javac）
            String compileOutput = compile(containerId);
            if (!"".equals(compileOutput)) {
                log.warn("编译失败，判定 CE:\n{}", compileOutput);
                return new JudgeOutcome(false, compileOutput, List.of());
            }
            log.info("编译通过");

            // 3) 逐用例运行（等价 docker exec -i <cid> java ... < case.in）
            List<JudgeResult> results = new ArrayList<>();
            for (int i = 0; i < cases.size(); i++) {
                TestCase tc = cases.get(i);
                RunResult run = runCase(containerId, config, tc);
                JudgeVerdict verdict = finalVerdict(run, tc);
                RunResult judged = new RunResult(run.stdout(), run.stderr(), run.exitCode(),
                        run.durationMs(), verdict);
                results.add(new JudgeResult(i + 1, tc, judged));
                log.info("用例 {}/{} 判定: {} (耗时 {}ms, 退出码 {})",
                        i + 1, cases.size(), verdict, judged.durationMs(), judged.exitCode());
                // TLE 时进程无法在容器内单独终止，整容器已被杀死，后续用例无法继续
                if (verdict == JudgeVerdict.TLE) {
                    log.warn("用例超时已将沙箱容器杀死，其余用例不再执行");
                    break;
                }
            }
            return new JudgeOutcome(true, compileOutput, results);
        } finally {
            if (containerId != null) {
                // 等价命令: docker rm -f <cid>
                dockerClient.removeContainerCmd(containerId).withForce(true).exec();
                log.info("沙箱容器已清理: {}", containerId);
            }
        }
    }

    // ==================== 容器创建 ====================

    /**
     * 创建并启动沙箱容器：只读根 + tmpfs 可写代码区 + 资源/权限全部收紧
     * @return 容器 ID
     */
    private String createSandboxContainer(JudgeConfig config) {
        // 等价命令:
        //   docker run -d -w /code --user nobody --read-only --network none \
        //     --cap-drop ALL -m 64m --cpus 1 --pids-limit 64 \
        //     --tmpfs /code --tmpfs /tmp amazoncorretto:17 sh -c "sleep 300"
        HostConfig hostConfig = buildSandboxHostConfig(config);
        CreateContainerResponse container = dockerClient.createContainerCmd(JUDGE_IMAGE)
                .withName("oj-sandbox-" + System.nanoTime())
                .withHostConfig(hostConfig)
                .withWorkingDir(CODE_DIR)
                .withUser("nobody")
                .withEnv("HOME=/tmp")
                .withEntrypoint("sh", "-c")
                .withCmd(KEEP_ALIVE_CMD)
                .exec();
        dockerClient.startContainerCmd(container.getId()).exec();
        return container.getId();
    }

    /**
     * 沙箱安全基线：
     * 无网络、丢弃所有内核 capability、只读根文件系统（避免篡改）、
     * 内存/CPU/进程数限制，以及非 root 用户运行
     */
    private HostConfig buildSandboxHostConfig(JudgeConfig config) {
        long memoryBytes = config.memoryLimitMb() * 1024L * 1024L;
        // /code 与 /tmp 挂载可写 tmpfs（1777 即任何人可读写），与只读根互补
        Map<String, String> tmpFs = Map.of(
                CODE_DIR, "size=8m,mode=1777",
                "/tmp", "size=8m,mode=1777");
        return HostConfig.newHostConfig()
                .withMemory(memoryBytes)                 // 等价 -m 64m
                .withMemorySwap(memoryBytes * 2)         // 需保持 swap >= memory，否则 409
                .withNanoCPUs(1_000_000_000L)            // 等价 --cpus 1
                .withPidsLimit(64L)                      // 等价 --pids-limit 64（防 fork 炸弹）
                .withNetworkMode("none")                 // 等价 --network none（禁止联网）
                .withCapDrop(Capability.ALL)             // 等价 --cap-drop ALL（丢弃全部特权）
                .withReadonlyRootfs(true)                // 等价 --read-only（防篡改镜像层）
                .withTmpFs(tmpFs);                       // 等价 --tmpfs /code --tmpfs /tmp
    }

    // ==================== 编译阶段 ====================

    /**
     * 把源码写入容器内的 /code（tmpfs，可写）
     * 注意两点：
     * 1) 根文件系统只读时，docker cp（copyArchiveToContainerCmd）会被 daemon 拒绝
     *    （400 container rootfs is marked read-only），故改用 exec 注入；
     * 2) 不通过 exec 的 stdin 传内容——httpclient5 半关闭语义下它不会向容器发送
     *    EOF，程序（如 cat）会一直等待输入。因此把源码 base64 编码后放进命令行参数，
     *    在容器内解码写盘，天然与 stdin 无关。
     */
    private void injectSource(String containerId, String source) {
        String encoded = Base64.getEncoder()
                .encodeToString(source.getBytes(StandardCharsets.UTF_8));
        // 等价命令: docker exec <cid> sh -c "printf %s '<base64>' | base64 -d > /code/Main.java"
        // base64 字符集仅含 [A-Za-z0-9+/=]，不含单引号，可安全内联到 shell 命令
        ExecCreateCmdResponse exec = dockerClient.execCreateCmd(containerId)
                .withCmd("sh", "-c",
                        "printf %s '" + encoded + "' | base64 -d > " + CODE_DIR + "/" + SOURCE_FILE)
                .withAttachStdout(true)
                .withAttachStderr(true)
                .exec();
        try {
            dockerClient.execStartCmd(exec.getId())
                    .withDetach(false)
                    .exec(new ResultCallback.Adapter<>() {
                    }).awaitCompletion(COMPILE_TIMEOUT_MS, TimeUnit.MILLISECONDS);
            log.info("源码 [{}] 已注入沙箱容器 ({} 字节)", SOURCE_FILE, source.length());
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.error("源码注入被中断: {}", e.getMessage());
        }
    }

    /**
     * 在容器内编译源码；成功返回空串，失败返回 javac 诊断信息
     */
    private String compile(String containerId) {
        StringBuilder output = new StringBuilder();
        try {
            // 等价命令: docker exec <cid> javac -encoding UTF-8 -d /code /code/Main.java
            ExecCreateCmdResponse exec = dockerClient.execCreateCmd(containerId)
                    .withCmd("javac", "-encoding", "UTF-8", "-d", CODE_DIR,
                            CODE_DIR + "/" + SOURCE_FILE)
                    .withAttachStdout(true)
                    .withAttachStderr(true)
                    .exec();
            boolean completed = dockerClient.execStartCmd(exec.getId())
                    .withDetach(false)
                    .exec(new ResultCallback.Adapter<>() {
                        @Override
                        public void onNext(Frame frame) {
                            output.append(new String(frame.getPayload(), StandardCharsets.UTF_8));
                        }
                    }).awaitCompletion(COMPILE_TIMEOUT_MS, TimeUnit.MILLISECONDS);
            if (!completed) {
                return "编译超时: " + output;
            }

            int exitCode = dockerClient.inspectExecCmd(exec.getId()).exec().getExitCode();
            return exitCode == 0 ? "" : output.toString();
        } catch (Exception e) {
            log.error("编译执行异常", e);
            return "编译系统异常: " + e.getMessage();
        }
    }

    // ==================== 运行阶段 ====================

    /**
     * 在沙箱容器内通过 exec 运行一个用例：
     * 标准输入经 withStdIn 注入，标准输出/错误输出按帧收集，并施加时间与输出上限
     */
    private RunResult runCase(String containerId, JudgeConfig config, TestCase tc) {
        long startNanos = System.nanoTime();
        StringBuilder stdout = new StringBuilder();
        StringBuilder stderr = new StringBuilder();
        int[] bytes = {0};
        boolean[] truncated = {false};

        // JVM 堆上限要比容器内存限制小，避免 JVM 启动即顶到 cgroup 内存上限
        int heapMb = Math.max(8, config.memoryLimitMb() - 16);
        // 用例输入与真实 OJ 一致以换行结尾：
        // 例如 Scanner 读取最后一个 token 仍需分隔符才能结束，而 httpclient5 无法向容器传达 stdin EOF
        String input = tc.stdin().endsWith("\n") ? tc.stdin() : tc.stdin() + "\n";
        try {
            // 等价命令: docker exec -i <cid> java -Xmx48m -cp /code Main < case1.in
            ExecCreateCmdResponse exec = dockerClient.execCreateCmd(containerId)
                    .withAttachStdin(true)
                    .withAttachStdout(true)
                    .withAttachStderr(true)
                    .withCmd("java", "-Xmx" + heapMb + "m", "-cp", CODE_DIR, "Main")
                    .exec();

            boolean timedOut = false;
            ResultCallback.Adapter<Frame> callback = dockerClient.execStartCmd(exec.getId())
                    .withDetach(false)
                    .withStdIn(new ByteArrayInputStream(input.getBytes(StandardCharsets.UTF_8)))
                    .exec(new ResultCallback.Adapter<>() {
                        @Override
                        public void onNext(Frame frame) {
                            byte[] payload = frame.getPayload();
                            bytes[0] += payload.length;
                            // 超过输出上限后不再累积，避免内存被输出撑爆
                            if (bytes[0] > config.outputLimitBytes()) {
                                truncated[0] = true;
                                return;
                            }
                            String text = new String(payload, StandardCharsets.UTF_8);
                            if (frame.getStreamType() == StreamType.STDOUT ||
                                    frame.getStreamType() == StreamType.RAW) {
                                stdout.append(text);
                            } else if (frame.getStreamType() == StreamType.STDERR) {
                                stderr.append(text);
                            }
                        }
                    });

            boolean completed = callback.awaitCompletion(config.timeLimitMs(), TimeUnit.MILLISECONDS);
            if (!completed) {
                // exec 进程无法单独终止，只能整容器 kill（等价 docker kill <cid>）
                log.warn("用例运行超时 ({}ms)，杀死沙箱容器", config.timeLimitMs());
                dockerClient.killContainerCmd(containerId).withSignal("KILL").exec();
                callback.close();
                timedOut = true;
            }

            long durationMs = (System.nanoTime() - startNanos) / 1_000_000;

            if (timedOut) {
                return new RunResult(stdout.toString(), stderr.toString(), -1, durationMs,
                        JudgeVerdict.TLE);
            }
            int exitCode = dockerClient.inspectExecCmd(exec.getId()).exec().getExitCode();
            if (truncated[0]) {
                return new RunResult(stdout.toString(), stderr.toString(), exitCode, durationMs,
                        JudgeVerdict.OLE);
            }
            if (exitCode != 0) {
                return new RunResult(stdout.toString(), stderr.toString(), exitCode, durationMs,
                        JudgeVerdict.RE);
            }
            // 进程正常退出，输出比对由 finalVerdict 决定 AC/WA
            return new RunResult(stdout.toString(), stderr.toString(), exitCode, durationMs, null);
        } catch (Exception e) {
            log.error("用例运行系统异常", e);
            long durationMs = (System.nanoTime() - startNanos) / 1_000_000;
            return new RunResult(stdout.toString(), stderr.toString(), -1, durationMs,
                    JudgeVerdict.SE);
        }
    }

    /**
     * 对运行结果做最终判定：非空 verdict 直接返回，为空则比对输出
     */
    private JudgeVerdict finalVerdict(RunResult run, TestCase tc) {
        if (run.verdict() != null) {
            return run.verdict();
        }
        // 输出比对：忽略行尾空白、结尾空行、回车符差异
        boolean matched = normalizeText(run.stdout()).equals(normalizeText(tc.expected()));
        return matched ? JudgeVerdict.AC : JudgeVerdict.WA;
    }

    /**
     * 基准输出规整：\r\n 归一为 \n、去除每行行尾空白、丢弃末尾空行
     */
    private static String normalizeText(String text) {
        String normalized = (text == null ? "" : text)
                .replace("\r\n", "\n").replace("\r", "\n");
        List<String> lines = new ArrayList<>(normalized.lines().map(String::stripTrailing).toList());
        while (!lines.isEmpty() && lines.get(lines.size() - 1).isEmpty()) {
            lines.remove(lines.size() - 1);
        }
        return String.join("\n", lines);
    }
}