# syntax=docker/dockerfile:1.7-labs
ARG BASE_STAGE=dev
# Built from docker/Dockerfile.microvm; CI checks this matches hack/microvm-image-tag.sh.
ARG MICROVM_IMAGE=public.ecr.aws/n4e0e1y0/beta9-microvm:d22ea0c29c1c

FROM golang:1.25-bookworm AS golang

RUN apt-get update && apt-get install -y curl git



FROM debian:bookworm as dropbear
ENV DROPBEAR_VERSION=2025.88
ENV DROPBEAR_SHA256=783f50ea27b17c16da89578fafdb6decfa44bb8f6590e5698a4e4d3672dc53d4

RUN apt update && apt install curl bzip2 zlib1g-dev make gcc -y  
RUN mkdir /dropbear && \
    curl -fsSL https://dropbear.nl/mirror/releases/dropbear-${DROPBEAR_VERSION}.tar.bz2 -o /dropbear/dropbear.tar.bz2 && \
    echo "${DROPBEAR_SHA256}  /dropbear/dropbear.tar.bz2" | sha256sum -c -
RUN cd /dropbear && tar -xvf dropbear.tar.bz2 && cd dropbear-${DROPBEAR_VERSION} && ./configure --enable-static --enable-pam=no CPPFLAGS="-DDISABLE_PAM" LDFLAGS="-static" && make PROGRAMS="dropbear" && mv dropbear /usr/local/bin/dropbear 

# runc
# ========================
# FIXME: once we upgrade to 1.23, we can just use "FROM golang"
FROM golang:1.23-bookworm AS runc

WORKDIR /workspace

COPY docker/patches/runc/runc-criu-production-log-level.patch /tmp/runc-criu-production-log-level.patch
COPY docker/patches/runc/runc-criu-cpu-compatibility.patch /tmp/runc-criu-cpu-compatibility.patch
COPY docker/patches/buildah/runtime-helper-reaping.patch /tmp/buildah-runtime-helper-reaping.patch

RUN apt-get update && apt-get install -y \
    make \
    btrfs-progs \
    git \
    go-md2man \
    iptables \
    libapparmor-dev \
    libglib2.0-dev \
    libgpgme11-dev \
    libseccomp-dev \
    libselinux1-dev \
    libbtrfs-dev \
    pkg-config

RUN <<EOT
set -eux
git clone -b dlm/force-compat-mode-on-restore https://github.com/beam-cloud/runc.git /workspace/runc \
  && cd /workspace/runc \
  && git reset --hard d85a78033ba2905f5c57afdc81b73a12be72c03c \
  && git apply --check /tmp/runc-criu-production-log-level.patch \
  && git apply /tmp/runc-criu-production-log-level.patch \
  && git apply --check /tmp/runc-criu-cpu-compatibility.patch \
  && git apply /tmp/runc-criu-cpu-compatibility.patch

cd /workspace/runc
make
make install
/usr/local/sbin/runc --version
cd ..

git clone https://github.com/containers/buildah.git /workspace/buildah
cd /workspace/buildah
git reset --hard 9bd608b3c0eaa1476b62763d0bd225745e49d542
git apply --check /tmp/buildah-runtime-helper-reaping.patch
git apply /tmp/buildah-runtime-helper-reaping.patch
make BUILDTAGS="cni"
make install
EOT

ENV CGO_ENABLED="0"  \
    DISABLE_DOCS="1" \
    GO_DYN_FLAGS=""  \
    BUILDTAGS="containers_image_openpgp"

RUN <<EOT
set -eux
git clone https://github.com/containers/skopeo.git /workspace/skopeo
cd /workspace/skopeo
git reset --hard ac66ca51ec5730b239f757907ccf96bb2d89b83c
make
make install
/usr/local/bin/skopeo --version
EOT


# gVisor (runsc)
# ========================
FROM debian:bookworm AS gvisor

ARG TARGETARCH
ARG GVISOR_VERSION=20260714.0
# beam-cloud/gvisor@205db7fdc1ad (checkpoint version ebbaed3f9dd5)
ARG GVISOR_AMD64_RELEASE=release-20260714.0-beam.16
ARG GVISOR_AMD64_SHA512=823945794ce3b3d26c8dc8899d16d27d198a65f9e0cfc297c7423e68b6b48870107806d20a48ffb608287d3085efddd975dc66357fe71b9655da12e5d2fff5d3

WORKDIR /tmp

RUN apt-get update && apt-get install -y curl

RUN <<EOT
set -eux

# Determine the correct architecture for gVisor
if [ "$TARGETARCH" = "amd64" ]; then
    GVISOR_ARCH="x86_64"
    GVISOR_BASE_URL="https://github.com/beam-cloud/gvisor/releases/download/${GVISOR_AMD64_RELEASE}"
    RUNSC_URL="${GVISOR_BASE_URL}/runsc.${GVISOR_ARCH}"
    RUNSC_SHA512="$GVISOR_AMD64_SHA512"
elif [ "$TARGETARCH" = "arm64" ]; then
    GVISOR_ARCH="aarch64"
    GVISOR_BASE_URL="https://storage.googleapis.com/gvisor/releases/release/${GVISOR_VERSION}/${GVISOR_ARCH}"
    RUNSC_URL="${GVISOR_BASE_URL}/runsc"
    RUNSC_SHA512_URL="${RUNSC_URL}.sha512"
else
    echo "Unsupported architecture: $TARGETARCH" && exit 1
fi

# Download runsc and its checksum to /tmp
curl -fsSL "$RUNSC_URL" -o /tmp/runsc
if [ -n "${RUNSC_SHA512:-}" ]; then
    printf '%s  runsc\n' "$RUNSC_SHA512" > /tmp/runsc.sha512
else
    curl -fsSL "$RUNSC_SHA512_URL" -o /tmp/runsc.sha512
fi

# Verify checksum (cd to /tmp where both files are)
cd /tmp && sha512sum -c runsc.sha512

# Move to final location and make executable
mv /tmp/runsc /usr/local/bin/runsc
chmod +x /usr/local/bin/runsc

# Verify installation and check available platforms
/usr/local/bin/runsc --version
/usr/local/bin/runsc help 2>&1 | grep -i platform || { echo "Platform options unavailable" >&2; exit 1; }
EOT

# qemu-storage-daemon
# ========================
# jammy's qemu-utils (6.2) predates qemu-storage-daemon being packaged, so
# build the tools from source against the same ubuntu:22.04 userland as the
# final image. Tools-only build: no system emulation, no user emulation.
FROM ubuntu:22.04 AS qemu
ARG QEMU_VERSION=8.2.10

RUN apt-get update && apt-get install -y \
    curl xz-utils python3 python3-pip pkg-config ninja-build \
    libglib2.0-dev zlib1g-dev flex bison && \
    pip3 install meson==1.5.2

RUN <<EOT
set -eux
curl -fsSL https://download.qemu.org/qemu-${QEMU_VERSION}.tar.xz | tar -xJ -C /tmp
cd /tmp/qemu-${QEMU_VERSION}
./configure --disable-system --disable-user --disable-docs --disable-guest-agent --enable-tools
ninja -C build storage-daemon/qemu-storage-daemon
install -m 0755 build/storage-daemon/qemu-storage-daemon /usr/local/bin/qemu-storage-daemon
/usr/local/bin/qemu-storage-daemon --version
EOT

# nvidia-container-toolkit
# ========================
FROM golang:1.23-bookworm AS nvidia-container-toolkit

ARG NVIDIA_CONTAINER_TOOLKIT_COMMIT=df7cc6c99980c04837a5ddb8ae34ce2211fc9126

WORKDIR /workspace

RUN apt-get update && apt-get install -y build-essential

RUN <<EOT
set -eux
git clone https://github.com/beam-cloud/nvidia-container-toolkit.git .
git reset --hard ${NVIDIA_CONTAINER_TOOLKIT_COMMIT}
make build
make binaries
EOT

# Cloud Hypervisor, virtiofsd and the guest kernel (microvm sandboxes)
# ========================
FROM ${MICROVM_IMAGE} AS microvm

# beta9 worker
# ========================
FROM golang AS worker

WORKDIR /workspace
ENV GONOSUMDB="github.com/beam-cloud/*"

COPY go.mod go.sum ./
RUN go mod download
COPY --exclude=sdk . .

# gVisor's nosplit CPU probes need optimization to stay within Go's stack limit.
RUN go build -gcflags "all=-N" -gcflags "gvisor.dev/gvisor/pkg/cpuid=" -o /usr/local/bin/worker-debug ./cmd/worker/main.go
RUN go build -o /usr/local/bin/worker ./cmd/worker/main.go
RUN go build -o /usr/local/bin/agent ./cmd/agent
# Guest init for microvm sandboxes: static, PID 1 inside the VM.
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /usr/local/bin/beam-vminit ./cmd/vminit

# uv
# ========================
FROM golang AS uv
RUN curl -LsSf https://astral.sh/uv/0.6.13/install.sh | sh

# final image
# ========================
FROM nvidia/cuda:12.8.0-base-ubuntu22.04 AS release
FROM release AS dev

FROM ${BASE_STAGE} AS final
ARG TARGETARCH
ARG GOPROC_VERSION=0.1.18
ARG ZSTD_VERSION=1.4.8+dfsg-3build1
ENV DEBIAN_FRONTEND="noninteractive"
WORKDIR /workspace

RUN apt-get update && \
    apt-get install -y curl gpg && \
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
        gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg && \
    curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
        sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
        tee /etc/apt/sources.list.d/nvidia-container-toolkit.list && \
    apt-get update && \
    apt-get install -y psmisc tini iproute2 iptables kmod "zstd=${ZSTD_VERSION}" \
        # qcow durable disks: qemu-img + qemu-storage-daemon, kernel NBD client,
        # and ext4 tooling for formatting freshly created volumes. libglib2.0-0
        # is the runtime dependency of the qemu-storage-daemon built above.
        qemu-utils nbd-client e2fsprogs libglib2.0-0

RUN curl -fsSL https://tailscale.com/install.sh | sh
RUN apt-get install -y --no-install-recommends nvidia-container-toolkit-base nvidia-container-toolkit
RUN <<'EOT'
set -eux
mv /usr/bin/nvidia-container-cli /usr/bin/nvidia-container-cli.real
cat >/usr/local/bin/nvidia-container-cli <<'EOF'
#!/usr/bin/env bash
set -euo pipefail

real_cli="${NVIDIA_CONTAINER_CLI_REAL:-/usr/bin/nvidia-container-cli.real}"
has_no_persistenced=0
for arg in "$@"; do
  if [ "$arg" = "--no-persistenced" ]; then
    has_no_persistenced=1
    break
  fi
done

if [ "$has_no_persistenced" -eq 0 ]; then
  args=()
  inserted=0
  for arg in "$@"; do
    args+=("$arg")
    if [ "$arg" = "configure" ]; then
      args+=("--no-persistenced")
      inserted=1
    fi
  done
  if [ "$inserted" -eq 1 ]; then
    exec "$real_cli" "${args[@]}"
  fi
fi

exec "$real_cli" "$@"
EOF
chmod +x /usr/local/bin/nvidia-container-cli
ln -s /usr/local/bin/nvidia-container-cli /usr/bin/nvidia-container-cli
EOT
RUN apt-get update && apt-get install -y fuse3 libfuse2 libfuse3-dev libfuse-dev bash-completion

# buildah run --network=private builds its bridge from the CNI plugins.
RUN apt-get install -y libbtrfs-dev libgpgme11-dev containernetworking-plugins
RUN mkdir -p /etc/containers/registries.conf.d
RUN curl -L https://raw.githubusercontent.com/containers/shortnames/refs/heads/main/shortnames.conf \
    -o /etc/containers/registries.conf.d/shortnames.conf
RUN echo 'unqualified-search-registries = ["docker.io"]' > /etc/containers/registries.conf

ARG CRIU_COMMIT=b47c692bb3a73ae45ee1c4ab215a8f6869d40efd
COPY docker/patches/criu/criu-network-remap.patch /tmp/criu-network-remap.patch
COPY docker/patches/criu/criu-skip-non-cuda-resume.patch /tmp/criu-skip-non-cuda-resume.patch
COPY docker/patches/criu/criu-skip-fresh-zero-pages.patch /tmp/criu-skip-fresh-zero-pages.patch
COPY docker/patches/criu/criu-ignore-erms-compat.patch /tmp/criu-ignore-erms-compat.patch

RUN <<EOT
set -eux
if [ "$(uname -m)" = "x86_64" ] || [ "$(uname -m)" = "aarch64" ]; then
    apt-get update && apt-get install -y \
        gcc \
        make \
        python3-protobuf \
        libbsd-dev \
        libdrm-dev \
        libnet1-dev \
        libnl-3-dev \
        libcap-dev \
        uuid-dev \
        gnutls-dev \
        libnftables-dev \
        liblz4-dev \
        pkg-config \
        protobuf-c-compiler \
        protobuf-compiler \
        libprotobuf-c-dev \
        git
    git clone https://github.com/checkpoint-restore/criu.git
    cd criu
    git fetch origin ${CRIU_COMMIT}
    git reset --hard ${CRIU_COMMIT}
    git apply --check --unidiff-zero /tmp/criu-network-remap.patch
    git apply --unidiff-zero /tmp/criu-network-remap.patch
    git apply --check --unidiff-zero /tmp/criu-skip-non-cuda-resume.patch
    git apply --unidiff-zero /tmp/criu-skip-non-cuda-resume.patch
    git apply --check /tmp/criu-skip-fresh-zero-pages.patch
    git apply /tmp/criu-skip-fresh-zero-pages.patch
    git apply --check /tmp/criu-ignore-erms-compat.patch
    git apply /tmp/criu-ignore-erms-compat.patch

    if [ "$(uname -m)" = "x86_64" ]; then
        make install-lib install-criu install-compel install-cuda_plugin PREFIX=/usr SKIP_PIP_INSTALL=1
        stat /usr/lib/criu/cuda_plugin.so
    else
        make install-lib install-criu install-compel PREFIX=/usr SKIP_PIP_INSTALL=1
    fi
    criu --version
    ldd "$(command -v criu)" | grep -q liblz4
    cd .. && rm -rf criu
fi
EOT

RUN mkdir -p /etc/criu && printf '%s\n' \
    'compress-region 262144' \
    'decompress-threads 32' \
    > /etc/criu/runc.conf

# Hard-disable any AMDGPU plugin that might exist from base layers or future package updates
RUN mkdir -p /usr/lib/criu/plugins.disabled && \
    if [ -e /usr/lib/criu/amdgpu_plugin.so ]; then \
        mv /usr/lib/criu/amdgpu_plugin.so /usr/lib/criu/plugins.disabled/ ; \
    fi && \
    true

ENV MOUNT_S3_URL_ARM64="https://s3.amazonaws.com/mountpoint-s3-release/1.8.0/arm64/mount-s3-1.8.0-arm64.tar.gz"
ENV MOUNT_S3_URL_AMD64="https://s3.amazonaws.com/mountpoint-s3-release/1.8.0/x86_64/mount-s3-1.8.0-x86_64.tar.gz"

RUN MOUNT_S3_URL=$(if [ "$TARGETARCH" = "arm64" ]; then echo $MOUNT_S3_URL_ARM64; else echo $MOUNT_S3_URL_AMD64; fi) && \
    curl -O $MOUNT_S3_URL && \
    TARBALL=$(basename $MOUNT_S3_URL) && \
    mkdir mount-s3 && \
    tar -xzf $TARBALL -C mount-s3 && \
    mv mount-s3/bin/mount-s3 /usr/local/bin/ms3 && chmod +x /usr/local/bin/ms3 && \
    rm -rf mount-s3 $TARBALL

ARG CUDA_CHECKPOINT_COMMIT=00d5cce84c628088d6caa203fc4af40c1538b6f7
ARG CUDA_CHECKPOINT_SHA256=707fa7f54136824d6c1d6dd724b9b1717610f831033c00d06da474de363a06db

RUN <<EOT
set -eux
case "$TARGETARCH" in
    amd64)
        curl -fsSL https://raw.githubusercontent.com/NVIDIA/cuda-checkpoint/${CUDA_CHECKPOINT_COMMIT}/bin/x86_64_Linux/cuda-checkpoint -o /usr/local/bin/cuda-checkpoint
        echo "${CUDA_CHECKPOINT_SHA256}  /usr/local/bin/cuda-checkpoint" | sha256sum -c -
        chmod +x /usr/local/bin/cuda-checkpoint
        ;;
    arm64)
        echo "cuda-checkpoint is not available for linux/arm64; skipping optional helper"
        ;;
    *) echo "Unsupported architecture: $TARGETARCH" && exit 1 ;;
esac
EOT

RUN apt-get install -y wget git busybox-static

COPY docker/goproc.sha256 /tmp/goproc.sha256
RUN cd /tmp && \
    wget "https://github.com/beam-cloud/goproc/releases/download/v${GOPROC_VERSION}/goproc-${GOPROC_VERSION}-linux-${TARGETARCH}" && \
    sha256sum --check --ignore-missing goproc.sha256 && \
    install -m 0755 "goproc-${GOPROC_VERSION}-linux-${TARGETARCH}" /usr/local/bin/goproc && \
    rm goproc.sha256 "goproc-${GOPROC_VERSION}-linux-${TARGETARCH}"

RUN apt-get remove -y curl && \
    apt-get clean && apt-get autoremove -y && apt-get autopurge -y && \
    rm -rf /var/lib/apt/lists/* /var/log/*


ENV BUILDAH_RUNTIME=/usr/local/bin/runsc

RUN mkdir -p /workspace/etc
RUN printf 'nameserver 1.1.1.1\nnameserver 8.8.8.8\nnameserver 9.9.9.9\noptions timeout:2 attempts:2\n' > /workspace/etc/resolv.conf
RUN echo '127.0.0.1 localhost' > /workspace/etc/hosts
RUN echo 'hosts: files dns' > /workspace/etc/nsswitch.conf
COPY ./bin/beta9-cli /usr/local/bin/beta9
RUN chmod +x /usr/local/bin/beta9
COPY ./bin/snapshot /usr/local/bin/snapshot
RUN chmod +x /usr/local/bin/snapshot

COPY --from=runc /usr/local/sbin/runc /usr/local/sbin/runc
COPY --from=runc /usr/local/bin/buildah /usr/local/bin/buildah
COPY --from=runc /usr/local/bin/skopeo /usr/local/bin/skopeo
COPY --from=runc /workspace/skopeo/default-policy.json /etc/containers/policy.json
COPY --from=gvisor /usr/local/bin/runsc /usr/local/bin/runsc
COPY --from=qemu /usr/local/bin/qemu-storage-daemon /usr/bin/qemu-storage-daemon
COPY --from=nvidia-container-toolkit /workspace/nvidia-container-runtime* /usr/bin/

# LD_PRELOAD libraries
COPY ./bin/volume_cache_x86.so /usr/local/lib/volume_cache_x86.so
COPY ./bin/volume_cache_arm.so /usr/local/lib/volume_cache_arm.so

RUN if [ "$(uname -m)" = "x86_64" ]; then \
    cp /usr/local/lib/volume_cache_x86.so /usr/local/lib/volume_cache.so; \
else \
    cp /usr/local/lib/volume_cache_arm.so /usr/local/lib/volume_cache.so; \
fi

COPY --from=worker /usr/local/bin/worker /usr/local/bin/worker
COPY --from=worker /usr/local/bin/worker-debug /usr/local/bin/worker-debug
COPY --from=worker /usr/local/bin/agent /usr/local/bin/agent
COPY --from=worker /usr/local/bin/agent /usr/local/bin/beam-agent
COPY --from=worker /usr/local/bin/beam-vminit /usr/local/bin/beam-vminit
COPY --from=microvm /microvm/bin/ /usr/local/bin/
COPY --from=microvm /microvm/share/ /usr/local/share/beam/microvm/
COPY ./sdk/src/beta9 /workspace/sdk
COPY --from=uv /root/.local/bin/uv /usr/local/bin/uv
COPY --from=dropbear /usr/local/bin/dropbear /usr/local/bin/dropbear

RUN mkdir -p /etc/criu
RUN touch /etc/criu/default.conf
# add the following options to the default.conf file on new lines
RUN echo "verbosity 4" >> /etc/criu/default.conf
RUN echo "display-stats" >> /etc/criu/default.conf
RUN echo "timeout 60" >> /etc/criu/default.conf

# Create gVisor runtime directories
RUN mkdir -p /run/gvisor

ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics,ngx,video
ENV _BUILDAH_STARTED_IN_USERNS="" \
    BUILDAH_ISOLATION=oci

# Configure buildah/containers for better concurrency
# image_parallel_copies controls how many layers can be pushed/pulled simultaneously
RUN mkdir -p /etc/containers && \
    echo '[engine]' > /etc/containers/containers.conf && \
    echo 'image_parallel_copies = 8' >> /etc/containers/containers.conf