codex-app-server/cloud/chv/disk.py

"""Cloud Hypervisor 每任务磁盘供给（叶子模块）。

单任务单 VM 的文件系统独立性就靠这里：copy 模式每任务
复制一份基础镜像，两个任务绝不写同一文件；reuse 模式
直接返回基础镜像路径（仅单租户实验，生产拦截见执行器
preflight，须同时设 CHV_REUSE_SINGLE_TENANT_EXPERIMENT=1
并审计放行）。

feat-chv-fix（P1-5）配额口径：不再只比基础镜像大小。
check_disk_quota 改为三项：state_dir 已用总量 + 本次
拷贝（基础镜像大小）+ 宿主剩余空间必须容纳本次拷贝；
镜像不可读时不再静默跳过，报 disk_quota_unverified
（preflight 在 wake 路径上会先报 image_missing，这里
是直接调用者的第二道门）。provision_task_disk 对已存在
磁盘校验大小：与源不一致即重拷；源不可读即报失败，
绝不复用脏磁盘。
"""

import os
import shutil

from cloud.chv.errors import ChvFailure


def _dir_total_bytes(path):
    total = 0
    for root, _dirs, files in os.walk(path):
        for name in files:
            try:
                total += os.path.getsize(os.path.join(root, name))
            except OSError:
                pass
    return total


def check_disk_quota(rootfs: str, disk_gb: float, state_dir: str | None = None) -> None:
    """磁盘配额门（feat-chv-fix 口径）：

    - disk_gb 未设时不检查；
    - 镜像不可读时报 disk_quota_unverified，不静默跳过；
    - 镜像大小超过 disk_gb 即 disk_quota_exceeded；
    - state_dir 已给时：已用总量 + 本次拷贝超过 disk_gb
      即 disk_quota_exceeded；宿主剩余空间不足以容纳本次
      拷贝即 disk_quota_exceeded（宿主余量校验）。
    """
    if disk_gb is None or not rootfs:
        return
    try:
        size = os.path.getsize(rootfs)
    except OSError:
        raise ChvFailure(
            "disk_quota_unverified",
            "磁盘镜像不可读，配额无法校验：%s" % rootfs)
    if size > disk_gb * 1024 ** 3:
        raise ChvFailure(
            "disk_quota_exceeded",
            "磁盘镜像 %d 字节超过配额 %.1f GB" % (size, disk_gb))
    if state_dir and os.path.isdir(state_dir):
        used = _dir_total_bytes(state_dir)
        if used + size > disk_gb * 1024 ** 3:
            raise ChvFailure(
                "disk_quota_exceeded",
                "磁盘总量 %d 字节 + 本次 %d 字节超过配额 %.1f GB"
                % (used, size, disk_gb))
        try:
            free = shutil.disk_usage(state_dir).free
        except OSError:
            free = None
        if free is not None and free < size:
            raise ChvFailure(
                "disk_quota_exceeded",
                "宿主剩余空间 %d 字节不足以容纳本次磁盘拷贝 %d 字节"
                % (free, size))


def provision_task_disk(state_dir: str, vm_name: str, rootfs: str, disk_mode: str) -> str:
    """每任务磁盘路径：copy 模式复制到
    state_dir/<vm>/disk.raw；reuse 模式直接返回基础镜像路径。

    feat-chv-fix：已存在磁盘须校验——大小与源不一致即
    重拷（脏磁盘不复用）；源不可读即 ChvFailure，绝不
    静默返回旧磁盘。"""
    if disk_mode == "reuse":
        return rootfs
    task_dir = os.path.join(state_dir, vm_name)
    os.makedirs(task_dir, exist_ok=True)
    task_disk = os.path.join(task_dir, "disk.raw")
    # feat-chv-secfix（L1）：目标是符号链接时先删链接本身，
    # 绝不跟随链接写入目标文件（copyfile 跟随链接会截断并
    # 写进链接目标）。删链接与重建之间的 TOCTOU 窗口已在
    # 文档里明写：state_dir 须 0700 且仅 worker 用户可写，
    # 宿主本地攻击者本就写不进该目录，L1 修的是“链接已
    # 预置”情形下的最后一道门。
    if os.path.islink(task_disk):
        os.remove(task_disk)
    if os.path.exists(task_disk):
        try:
            same = os.path.getsize(task_disk) == os.path.getsize(rootfs)
        except OSError:
            raise ChvFailure(
                "disk_quota_unverified",
                "任务磁盘或源镜像不可读，无法校验复用：%s" % task_disk)
        if same:
            return task_disk
        os.remove(task_disk)
    try:
        shutil.copyfile(rootfs, task_disk)
    except OSError as exc:
        raise ChvFailure(
            "disk_quota_unverified",
            "任务磁盘拷贝失败：%s" % exc)
    return task_disk


def task_disk_path(state_dir: str, vm_name: str) -> str:
    """该 VM 的任务磁盘路径：state_dir/<vm>/disk.raw。"""
    return os.path.join(state_dir, vm_name, "disk.raw")


def delete_task_disk(state_dir: str, vm_name: str) -> bool:
    """删任务磁盘（feat-chv-secfix，H3）：stop/recycle/
    清理路径都调它。返回 True 表示确实删掉了一个文件，
    False 表示本来就没有；删失败抛 OSError（调用方记审计，
    绝不静默——残留盘由 gc_stale_disks 下次再清）。
    符号链接只删链接本身，不跟随。"""
    path = task_disk_path(state_dir, vm_name)
    if os.path.islink(path) or os.path.exists(path):
        os.remove(path)
        return True
    return False


# -- 快照登记与上限（feat-chv-secfix，M3）-------------------------
# 登记文件：state_dir/<vm>/snapshots.json，内容为
# [{url, path（file:// 才有）, size, created_at}]。登记是
# 执行器在宿主侧的账本，不是 VMM 的快照元数据；VMM 不认
# 它，删登记不等于删快照文件（反之亦然），两者在
# prune/record 里成对处理并审计。

def _snapshots_path(state_dir, vm_name):
    return os.path.join(state_dir, vm_name, "snapshots.json")


def _load_snapshots(state_dir, vm_name):
    import json
    path = _snapshots_path(state_dir, vm_name)
    try:
        with open(path, "r") as fh:
            data = json.load(fh)
    except (OSError, ValueError):
        return []
    return data if isinstance(data, list) else []


def _save_snapshots(state_dir, vm_name, entries):
    import json
    path = _snapshots_path(state_dir, vm_name)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp_path = path + ".tmp"
    with open(tmp_path, "w") as fh:
        json.dump(entries, fh)
    os.replace(tmp_path, path)


def prune_snapshots(state_dir: str, vm_name: str,
                    retention_days: float | None,
                    now: float | None = None,
                    delete_files: bool = True) -> tuple[list, list]:
    """按保留期剪枝：删过期条目；delete_files 时同时删
    对应的 file:// 快照文件（只删登记里记的 path，不猜
    文件名）。返回（保留条目，剪掉条目）。"""
    import time
    now = time.time() if now is None else now
    kept, pruned = [], []
    for entry in _load_snapshots(state_dir, vm_name):
        created = entry.get("created_at") or 0
        if retention_days is not None \
                and now - created > retention_days * 86400:
            pruned.append(entry)
            if delete_files and entry.get("path"):
                try:
                    if os.path.exists(entry["path"]):
                        os.remove(entry["path"])
                except OSError:
                    pass
        else:
            kept.append(entry)
    if pruned:
        _save_snapshots(state_dir, vm_name, kept)
    return kept, pruned


def check_snapshot_limits(state_dir: str, vm_name: str, max_count: int,
                          max_total_gb: float,
                          retention_days: float | None) -> list:
    """快照上限预检（snapshot_vm 调 ch-remote 之前调）：
    先按保留期剪枝，再查数量（已有数 + 本次 > 上限即
    snapshot_quota_exceeded）与总量（已有总量已达上限即
    拒绝；本次大小事前未知，总量在 record 时二次收紧，
    见 record_snapshot）。返回剪枝后条目。"""
    kept, _ = prune_snapshots(state_dir, vm_name, retention_days)
    if len(kept) + 1 > max_count:
        raise ChvFailure(
            "snapshot_quota_exceeded",
            "快照数量已达上限 %d（VM %s）" % (max_count, vm_name))
    total = 0
    for entry in kept:
        size = entry.get("size")
        if isinstance(size, (int, float)):
            total += size
        elif entry.get("path") and os.path.exists(entry["path"]):
            try:
                total += os.path.getsize(entry["path"])
            except OSError:
                pass
    if total >= max_total_gb * 1024 ** 3:
        raise ChvFailure(
            "snapshot_quota_exceeded",
            "快照总量 %d 字节已达上限 %.1f GB（VM %s）"
            % (total, max_total_gb, vm_name))
    return kept


def record_snapshot(state_dir: str, vm_name: str, snapshot_url: str, max_total_gb: float) -> tuple[list, list]:
    """快照成功后登记；file:// 目标记实际大小。登记后若
    总量超上限，按最旧优先删登记外的旧快照文件（总量硬
    帽，删动作由调用方审计——本函数返回被删的旧条目）。
    返回（全部条目，被总量帽删掉的旧条目）。"""
    import time
    path = None
    size = None
    if snapshot_url.startswith("file://"):
        path = snapshot_url[len("file://"):]
        try:
            size = os.path.getsize(path)
        except OSError:
            size = None
    entries = _load_snapshots(state_dir, vm_name)
    entries.append({"url": snapshot_url, "path": path,
                    "size": size, "created_at": time.time()})
    _save_snapshots(state_dir, vm_name, entries)
    removed = []
    if max_total_gb is not None:
        cap = max_total_gb * 1024 ** 3
        total = sum(e.get("size") or 0 for e in entries)
        while len(entries) > 1 and total > cap:
            oldest = entries.pop(0)
            if oldest.get("path") and os.path.exists(oldest["path"]):
                try:
                    os.remove(oldest["path"])
                except OSError:
                    pass
            total -= oldest.get("size") or 0
            removed.append(oldest)
        if removed:
            _save_snapshots(state_dir, vm_name, entries)
    return entries, removed