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

"""宿主侧 vsock notify 最小实现（P1，feat-chv-muse-p1）。

方向只一个：客机 → 宿主。客机连宿主 CID 2 的
VSOCK_NOTIFY_PORT 端口发 ready/心跳（客内脚本
cloud/chv/guest/vsock_notify.py）；宿主在
``<vsock socket>_<port>`` 这个 Unix socket 上监听接收。

形态依据（CHV 官方 docs/vsock.md，upstream 原文）：
- ``--vsock cid=3,socket=/tmp/ch.vsock``：cid 是客机 CID，
  socket 是宿主侧 Unix socket 路径；
- 客 → 宿主：宿主须先在 ``<socket>_<port>`` 上监听
  （如 ``/tmp/ch.vsock_1234``），客机连众所周知的 CID 2
  （Host）；
- 只支持 stream socket。

只做就绪辅助：wait_ready() 收到 ready 即返回，供
_run_with_vm/常驻路径在 SSH 轮询之外多一个早期信号；
SSH 仍是唯一命令通道与就绪判定——vsock 超时、无消息、
监听失败一律返回 None，绝不因此判任务失败、绝不抛到
调用方。消息是 JSON 行：

    {"type": "ready"|"heartbeat", "vm": "<vm 名>", "ts": <unix 秒>}

不做宿主 → 客机方向（客机不监听，host.py 协议仍走 SSH）。
"""

import json
import logging
import os
import socket
import time

_LOGGER = logging.getLogger("cloud.cloud_hypervisor")

# 宿主 CID（vsock 众所周知值，CHV 官方 docs/vsock.md）。
HOST_CID = 2
# 客机 CID 默认值（与 CHV_VSOCK_CID 默认值一致；0/1/2 保留）。
DEFAULT_GUEST_CID = 3
# notify 端口：客机连 (CID 2, 本端口) 发 ready/心跳；
# 宿主监听 <socket>_<本端口>。与客内脚本的端口须一致。
VSOCK_NOTIFY_PORT = 1024
# 单行上限：客机发来的行超长即丢弃该连接（防脏数据撑内存）。
MAX_LINE_BYTES = 65536
# 合法消息类型。
VALID_TYPES = ("ready", "heartbeat")


def vsock_socket_path(state_dir: str, vm_name: str) -> str:
    """该 VM 的 vsock socket 路径（放进 VmConfig 的值）。

    与 build_vm_config 的派生一致：任务盘同目录的
    ``state_dir/<vm>/vsock.sock``。vm_name 须与
    api_socket/任务盘用的名字相同（executor.vm_name(task)
    的返回值，含 chv- 前缀），这里不二次加前缀。"""
    return os.path.join(state_dir, vm_name, "vsock.sock")


def notify_socket_path(socket_path: str,
                       port: int = VSOCK_NOTIFY_PORT) -> str:
    """宿主侧监听路径：``<socket>_<port>``（CHV 官方约定）。"""
    return "%s_%d" % (socket_path, port)


def encode_notify(message_type: str, vm: str,
                  ts: float | None = None) -> bytes:
    """编码一条 notify 消息为 JSON 行（含换行）。

    type 非法或 vm 为空即 ValueError（调用方是可信的客内
    脚本/宿主代码，编码错误是编程错误，早炸）。"""
    if message_type not in VALID_TYPES:
        raise ValueError("vsock notify 类型非法：%r" % (message_type,))
    if not vm:
        raise ValueError("vsock notify vm 不能为空")
    if ts is None:
        ts = time.time()
    line = json.dumps({"type": message_type, "vm": vm, "ts": ts},
                      separators=(",", ":"))
    return (line + "\n").encode("utf-8")


def decode_notify(line: bytes | str) -> dict | None:
    """解码一行 notify 消息：合法返回 dict，否则返回 None。

    客机字节不可信，任何畸形（非 JSON、非对象、类型不在
    ready/heartbeat 内、vm 空/非字符串、ts 非数字）一律
    None，绝不抛异常。"""
    try:
        if isinstance(line, bytes):
            line = line.decode("utf-8")
        obj = json.loads(line)
    except (UnicodeDecodeError, ValueError):
        return None
    if not isinstance(obj, dict):
        return None
    if obj.get("type") not in VALID_TYPES:
        return None
    if not isinstance(obj.get("vm"), str) or not obj.get("vm"):
        return None
    if not isinstance(obj.get("ts"), (int, float)) \
            or isinstance(obj.get("ts"), bool):
        return None
    return obj


def _bind_notify_listener(listener, notify_path: str) -> bool:
    """bind notify socket：残留文件先删一次再重试一次；
    两次都失败即 False（调用方记 warning 后返回 None，
    与原来直抛到外层 except 的行为一致）。"""
    try:
        listener.bind(notify_path)
        return True
    except OSError:
        pass
    try:
        os.unlink(notify_path)
    except OSError:
        pass
    try:
        listener.bind(notify_path)
        return True
    except OSError:
        return False


def listen_notify(notify_path: str, timeout: float,
                  want_type: str | None = None) -> dict | None:
    """在 notify_path（Unix socket）上监听，等一条 notify。

    收到首条合法消息即返回（want_type 给定时跳过其它类型
    继续等，直到超时）；超时、客机断开未发合法消息、监听
    失败一律返回 None。监听用的 socket 文件用完即删（尽力
    而为）；bind 前已存在的同名残留文件先删一次再 bind。
    timeout <= 0 即不等，直接返回 None。"""
    if timeout is None or timeout <= 0:
        return None
    deadline = time.monotonic() + timeout
    listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        if not _bind_notify_listener(listener, notify_path):
            _LOGGER.warning("vsock notify 监听失败：%s", notify_path)
            return None
        listener.listen(1)
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return None
            listener.settimeout(remaining)
            try:
                conn, _ = listener.accept()
            except socket.timeout:
                return None
            except OSError:
                _LOGGER.warning("vsock notify accept 失败：%s",
                                notify_path)
                return None
            try:
                return _read_notify(conn, deadline, want_type)
            finally:
                try:
                    conn.close()
                except OSError:
                    pass
    except OSError:
        _LOGGER.warning("vsock notify 监听失败：%s", notify_path)
        return None
    finally:
        try:
            listener.close()
        except OSError:
            pass
        try:
            os.unlink(notify_path)
        except OSError:
            pass


def _read_notify(conn, deadline: float,
                 want_type: str | None) -> dict | None:
    """从一条已接受的连接上读 JSON 行，直到合法消息或超时/
    断开。want_type 给定时跳过其它合法类型继续读。"""
    try:
        stream = conn.makefile("rb")
    except OSError:
        return None
    try:
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return None
            try:
                conn.settimeout(remaining)
                line = stream.readline(MAX_LINE_BYTES + 1)
            except (OSError, ValueError):
                return None
            if not line:
                return None
            if len(line) > MAX_LINE_BYTES:
                return None
            msg = decode_notify(line)
            if msg is None:
                continue
            if want_type is not None and msg["type"] != want_type:
                continue
            return msg
    finally:
        try:
            stream.close()
        except OSError:
            pass


def wait_ready(state_dir: str, vm_name: str,
               timeout: float = 30,
               port: int = VSOCK_NOTIFY_PORT) -> dict | None:
    """等一台 VM 的 vsock ready（就绪辅助，非 SSH 判定）。

    在 ``state_dir/<vm>/vsock.sock_<port>`` 上监听，直到
    收到 {"type": "ready"} 或超时。返回 ready 消息 dict，
    超时/失败返回 None。调用方（_run_with_vm/常驻路径）
    拿到 None 时继续走 SSH 就绪轮询，绝不因此失败。"""
    path = notify_socket_path(
        vsock_socket_path(state_dir, vm_name), port)
    return listen_notify(path, timeout, want_type="ready")