|
| 1 | +"""Mount strategy for Fly sandboxes.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +from pathlib import Path |
| 6 | +from typing import Literal |
| 7 | + |
| 8 | +from ....sandbox.entries.mounts.base import InContainerMountStrategy, Mount, MountStrategyBase |
| 9 | +from ....sandbox.entries.mounts.patterns import RcloneMountPattern |
| 10 | +from ....sandbox.errors import MountConfigError |
| 11 | +from ....sandbox.materialization import MaterializedFile |
| 12 | +from ....sandbox.session.base_sandbox_session import BaseSandboxSession |
| 13 | + |
| 14 | +# Sprite VMs run as the unprivileged ``sprite`` user with passwordless sudo. |
| 15 | +# ``FlySandboxSession.exec`` rejects ``user=`` kwargs, so we prefix privileged |
| 16 | +# commands with ``sudo -n`` instead of escalating through the framework. |
| 17 | +_SUDO = "sudo -n" |
| 18 | +_APT = ( |
| 19 | + f"{_SUDO} env DEBIAN_FRONTEND=noninteractive DEBCONF_NOWARNINGS=yes apt-get -o Dpkg::Use-Pty=0" |
| 20 | +) |
| 21 | +_RCLONE_CHECK = "command -v rclone >/dev/null 2>&1 || test -x /usr/local/bin/rclone" |
| 22 | +_FUSERMOUNT_CHECK = ( |
| 23 | + "command -v fusermount3 >/dev/null 2>&1 || command -v fusermount >/dev/null 2>&1" |
| 24 | +) |
| 25 | +_INSTALL_RCLONE_COMMANDS = ( |
| 26 | + f"{_APT} update -qq", |
| 27 | + f"{_APT} install -y -qq curl unzip ca-certificates fuse", |
| 28 | + f"curl -fsSL https://rclone.org/install.sh | {_SUDO} bash", |
| 29 | +) |
| 30 | +# fuse package brings ``fusermount`` along — install it together with rclone |
| 31 | +# so the FUSE-mode mount path works out-of-the-box on stock sprite images. |
| 32 | +_INSTALL_FUSE_COMMANDS = ( |
| 33 | + f"{_APT} update -qq", |
| 34 | + f"{_APT} install -y -qq fuse", |
| 35 | +) |
| 36 | +_FUSE_ALLOW_OTHER = ( |
| 37 | + f"{_SUDO} chmod a+rw /dev/fuse && " |
| 38 | + f"{_SUDO} touch /etc/fuse.conf && " |
| 39 | + "(grep -qxF user_allow_other /etc/fuse.conf || " |
| 40 | + f"printf '\\nuser_allow_other\\n' | {_SUDO} tee -a /etc/fuse.conf >/dev/null)" |
| 41 | +) |
| 42 | + |
| 43 | + |
| 44 | +async def _ensure_fuse_support(session: BaseSandboxSession) -> None: |
| 45 | + kernel = await session.exec( |
| 46 | + "sh", |
| 47 | + "-lc", |
| 48 | + "test -c /dev/fuse && grep -qw fuse /proc/filesystems", |
| 49 | + shell=False, |
| 50 | + ) |
| 51 | + if not kernel.ok(): |
| 52 | + raise MountConfigError( |
| 53 | + message="Fly cloud bucket mounts require FUSE support in the kernel", |
| 54 | + context={"missing": "fuse"}, |
| 55 | + ) |
| 56 | + |
| 57 | + fusermount = await session.exec("sh", "-lc", _FUSERMOUNT_CHECK, shell=False) |
| 58 | + if not fusermount.ok(): |
| 59 | + apt = await session.exec("sh", "-lc", "command -v apt-get >/dev/null 2>&1", shell=False) |
| 60 | + if not apt.ok(): |
| 61 | + raise MountConfigError( |
| 62 | + message="fusermount is not installed and apt-get is unavailable; " |
| 63 | + "preinstall the fuse package", |
| 64 | + context={"package": "fuse"}, |
| 65 | + ) |
| 66 | + for command in _INSTALL_FUSE_COMMANDS: |
| 67 | + install = await session.exec("sh", "-lc", command, shell=False, timeout=300) |
| 68 | + if not install.ok(): |
| 69 | + raise MountConfigError( |
| 70 | + message="failed to install fuse", |
| 71 | + context={"package": "fuse", "exit_code": install.exit_code}, |
| 72 | + ) |
| 73 | + recheck = await session.exec("sh", "-lc", _FUSERMOUNT_CHECK, shell=False) |
| 74 | + if not recheck.ok(): |
| 75 | + raise MountConfigError( |
| 76 | + message="fuse was installed but fusermount is still not on PATH", |
| 77 | + context={"package": "fuse"}, |
| 78 | + ) |
| 79 | + |
| 80 | + chmod_result = await session.exec("sh", "-lc", _FUSE_ALLOW_OTHER, shell=False, timeout=30) |
| 81 | + if not chmod_result.ok(): |
| 82 | + raise MountConfigError( |
| 83 | + message="failed to make /dev/fuse accessible", |
| 84 | + context={"exit_code": chmod_result.exit_code}, |
| 85 | + ) |
| 86 | + |
| 87 | + |
| 88 | +async def _ensure_rclone(session: BaseSandboxSession) -> None: |
| 89 | + rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False) |
| 90 | + if rclone.ok(): |
| 91 | + return |
| 92 | + |
| 93 | + apt = await session.exec("sh", "-lc", "command -v apt-get >/dev/null 2>&1", shell=False) |
| 94 | + if not apt.ok(): |
| 95 | + raise MountConfigError( |
| 96 | + message="rclone is not installed and apt-get is unavailable; preinstall rclone", |
| 97 | + context={"package": "rclone"}, |
| 98 | + ) |
| 99 | + |
| 100 | + for command in _INSTALL_RCLONE_COMMANDS: |
| 101 | + install = await session.exec("sh", "-lc", command, shell=False, timeout=300) |
| 102 | + if not install.ok(): |
| 103 | + raise MountConfigError( |
| 104 | + message="failed to install rclone", |
| 105 | + context={"package": "rclone", "exit_code": install.exit_code}, |
| 106 | + ) |
| 107 | + |
| 108 | + rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False) |
| 109 | + if not rclone.ok(): |
| 110 | + raise MountConfigError( |
| 111 | + message="rclone was installed but is still not available on PATH", |
| 112 | + context={"package": "rclone"}, |
| 113 | + ) |
| 114 | + |
| 115 | + |
| 116 | +async def _default_user_ids(session: BaseSandboxSession) -> tuple[str, str] | None: |
| 117 | + result = await session.exec("sh", "-lc", "id -u; id -g", shell=False, timeout=30) |
| 118 | + if not result.ok(): |
| 119 | + return None |
| 120 | + |
| 121 | + lines = result.stdout.decode("utf-8", errors="replace").splitlines() |
| 122 | + if len(lines) < 2 or not lines[0].isdigit() or not lines[1].isdigit(): |
| 123 | + return None |
| 124 | + return lines[0], lines[1] |
| 125 | + |
| 126 | + |
| 127 | +def _append_option(args: list[str], option: str, *values: str) -> None: |
| 128 | + if option not in args: |
| 129 | + args.extend([option, *values]) |
| 130 | + |
| 131 | + |
| 132 | +async def _rclone_pattern_for_session( |
| 133 | + session: BaseSandboxSession, |
| 134 | + pattern: RcloneMountPattern, |
| 135 | +) -> RcloneMountPattern: |
| 136 | + if pattern.mode != "fuse": |
| 137 | + return pattern |
| 138 | + |
| 139 | + extra_args = list(pattern.extra_args) |
| 140 | + _append_option(extra_args, "--allow-other") |
| 141 | + user_ids = await _default_user_ids(session) |
| 142 | + if user_ids is not None: |
| 143 | + uid, gid = user_ids |
| 144 | + _append_option(extra_args, "--uid", uid) |
| 145 | + _append_option(extra_args, "--gid", gid) |
| 146 | + |
| 147 | + return pattern.model_copy(update={"extra_args": extra_args}) |
| 148 | + |
| 149 | + |
| 150 | +def _assert_fly_session(session: BaseSandboxSession) -> None: |
| 151 | + if type(session).__name__ != "FlySandboxSession": |
| 152 | + raise MountConfigError( |
| 153 | + message="fly cloud bucket mounts require a FlySandboxSession", |
| 154 | + context={"session_type": type(session).__name__}, |
| 155 | + ) |
| 156 | + |
| 157 | + |
| 158 | +class FlyCloudBucketMountStrategy(MountStrategyBase): |
| 159 | + """Mount rclone-backed cloud storage in Fly sandboxes.""" |
| 160 | + |
| 161 | + type: Literal["fly_cloud_bucket"] = "fly_cloud_bucket" |
| 162 | + pattern: RcloneMountPattern = RcloneMountPattern(mode="fuse") |
| 163 | + |
| 164 | + def _delegate(self) -> InContainerMountStrategy: |
| 165 | + return InContainerMountStrategy(pattern=self.pattern) |
| 166 | + |
| 167 | + async def _delegate_for_session(self, session: BaseSandboxSession) -> InContainerMountStrategy: |
| 168 | + return InContainerMountStrategy( |
| 169 | + pattern=await _rclone_pattern_for_session(session, self.pattern) |
| 170 | + ) |
| 171 | + |
| 172 | + def validate_mount(self, mount: Mount) -> None: |
| 173 | + self._delegate().validate_mount(mount) |
| 174 | + |
| 175 | + async def activate( |
| 176 | + self, |
| 177 | + mount: Mount, |
| 178 | + session: BaseSandboxSession, |
| 179 | + dest: Path, |
| 180 | + base_dir: Path, |
| 181 | + ) -> list[MaterializedFile]: |
| 182 | + _assert_fly_session(session) |
| 183 | + if self.pattern.mode == "fuse": |
| 184 | + await _ensure_fuse_support(session) |
| 185 | + await _ensure_rclone(session) |
| 186 | + delegate = await self._delegate_for_session(session) |
| 187 | + return await delegate.activate(mount, session, dest, base_dir) |
| 188 | + |
| 189 | + async def deactivate( |
| 190 | + self, |
| 191 | + mount: Mount, |
| 192 | + session: BaseSandboxSession, |
| 193 | + dest: Path, |
| 194 | + base_dir: Path, |
| 195 | + ) -> None: |
| 196 | + _assert_fly_session(session) |
| 197 | + await self._delegate().deactivate(mount, session, dest, base_dir) |
| 198 | + |
| 199 | + async def teardown_for_snapshot( |
| 200 | + self, |
| 201 | + mount: Mount, |
| 202 | + session: BaseSandboxSession, |
| 203 | + path: Path, |
| 204 | + ) -> None: |
| 205 | + _assert_fly_session(session) |
| 206 | + await self._delegate().teardown_for_snapshot(mount, session, path) |
| 207 | + |
| 208 | + async def restore_after_snapshot( |
| 209 | + self, |
| 210 | + mount: Mount, |
| 211 | + session: BaseSandboxSession, |
| 212 | + path: Path, |
| 213 | + ) -> None: |
| 214 | + _assert_fly_session(session) |
| 215 | + if self.pattern.mode == "fuse": |
| 216 | + await _ensure_fuse_support(session) |
| 217 | + await _ensure_rclone(session) |
| 218 | + delegate = await self._delegate_for_session(session) |
| 219 | + await delegate.restore_after_snapshot(mount, session, path) |
| 220 | + |
| 221 | + def build_docker_volume_driver_config( |
| 222 | + self, |
| 223 | + mount: Mount, |
| 224 | + ) -> tuple[str, dict[str, str], bool] | None: |
| 225 | + return None |
| 226 | + |
| 227 | + |
| 228 | +__all__ = [ |
| 229 | + "FlyCloudBucketMountStrategy", |
| 230 | +] |
0 commit comments