diff --git a/openviking/storage/viking_fs.py b/openviking/storage/viking_fs.py index 8429f43c5..90e90c2c2 100644 --- a/openviking/storage/viking_fs.py +++ b/openviking/storage/viking_fs.py @@ -30,7 +30,7 @@ from openviking.core.namespace import ( is_accessible as namespace_is_accessible, ) -from openviking.pyagfs.exceptions import AGFSClientError, AGFSDirectoryNotEmptyError, AGFSHTTPError +from openviking.pyagfs.exceptions import AGFSDirectoryNotEmptyError from openviking.resource.watch_storage import is_watch_task_control_uri from openviking.server.error_mapping import is_not_found_error, map_exception from openviking.server.identity import RequestContext, Role @@ -321,49 +321,79 @@ def _ensure_mutable_access(self, uri: str, ctx: Optional[RequestContext]) -> Non resource=normalized_uri, ) - # ========== AGFS Basic Commands ========== + def _agfs_read(self, path: str, offset: int = 0, size: int = -1) -> Any: + """Call AGFS read while tolerating backends that omit offset/size params.""" + if offset == 0 and size == -1: + return self.agfs.read(path) + try: + return self.agfs.read(path, offset, size) + except TypeError: + return self.agfs.read(path) - async def read( + async def _read_path_bytes( self, - uri: str, + path: str, + *, + ctx: Optional[RequestContext] = None, offset: int = 0, size: int = -1, - ctx: Optional[RequestContext] = None, + decrypt: bool = False, + require_exists: bool = False, ) -> bytes: - """Read file""" - self._ensure_access(uri, ctx) - path = self._uri_to_path(uri, ctx=ctx) + """Read bytes from AGFS with optional full-file decrypt and plaintext slicing.""" + if require_exists: + self.agfs.stat(path) - if self._encryptor: - # When encryption is enabled: must read entire file for decryption - result = self.agfs.read(path, 0, -1) - if isinstance(result, bytes): - raw = result - elif result is not None and hasattr(result, "content"): - raw = result.content - else: - raw = b"" + read_offset = 0 if decrypt and self._encryptor else offset + read_size = -1 if decrypt and self._encryptor else size + raw = self._handle_agfs_read(self._agfs_read(path, read_offset, read_size)) + if decrypt: raw = await self._decrypt_content(raw, ctx=ctx) - - # Apply slicing on decrypted plaintext if offset > 0 or size != -1: if size != -1: raw = raw[offset : offset + size] else: raw = raw[offset:] - else: - # When not encrypted: normal read - result = self.agfs.read(path, offset, size) - if isinstance(result, bytes): - raw = result - elif result is not None and hasattr(result, "content"): - raw = result.content - else: - raw = b"" return raw + async def _write_path_bytes( + self, + path: str, + content: bytes, + *, + ctx: Optional[RequestContext] = None, + encrypt: bool = False, + ensure_parent: bool = False, + ) -> Any: + """Write bytes to AGFS with optional parent creation and encryption.""" + if ensure_parent: + await self._ensure_parent_dirs(path) + if encrypt: + content = await self._encrypt_content(content, ctx=ctx) + return self.agfs.write(path, content) + + # ========== AGFS Basic Commands ========== + + async def read( + self, + uri: str, + offset: int = 0, + size: int = -1, + ctx: Optional[RequestContext] = None, + ) -> bytes: + """Read file""" + self._ensure_access(uri, ctx) + path = self._uri_to_path(uri, ctx=ctx) + return await self._read_path_bytes( + path, + ctx=ctx, + offset=offset, + size=size, + decrypt=True, + ) + async def write( self, uri: str, @@ -375,9 +405,7 @@ async def write( path = self._uri_to_path(uri, ctx=ctx) if isinstance(data, str): data = data.encode("utf-8") - - data = await self._encrypt_content(data, ctx=ctx) - return self.agfs.write(path, data) + return await self._write_path_bytes(path, data, ctx=ctx, encrypt=True) async def mkdir( self, @@ -922,7 +950,7 @@ async def abstract( ) file_path = f"{path}/.abstract.md" try: - content_bytes = self._handle_agfs_read(self.agfs.read(file_path)) + content_bytes = await self._read_path_bytes(file_path, ctx=ctx, decrypt=True) except Exception as exc: if not is_not_found_error(exc): mapped = map_exception(exc, resource=uri) @@ -932,10 +960,6 @@ async def abstract( # Fallback to default if .abstract.md doesn't exist return f"# {uri}\n\n[Directory abstract is not ready]" - if self._encryptor: - real_ctx = self._ctx_or_default(ctx) - content_bytes = await self._encryptor.decrypt(real_ctx.account_id, content_bytes) - return self._decode_bytes(content_bytes) async def overview( @@ -960,7 +984,7 @@ async def overview( ) file_path = f"{path}/.overview.md" try: - content_bytes = self._handle_agfs_read(self.agfs.read(file_path)) + content_bytes = await self._read_path_bytes(file_path, ctx=ctx, decrypt=True) except Exception as exc: if not is_not_found_error(exc): mapped = map_exception(exc, resource=uri) @@ -970,10 +994,6 @@ async def overview( # Fallback to default if .overview.md doesn't exist return f"# {uri}\n\n[Directory overview is not ready]" - if self._encryptor: - real_ctx = self._ctx_or_default(ctx) - content_bytes = await self._encryptor.decrypt(real_ctx.account_id, content_bytes) - return self._decode_bytes(content_bytes) async def relations( @@ -1713,9 +1733,7 @@ async def _write_relation_table( table_path = f"{dir_path}/.relations.json" if isinstance(content, str): content = content.encode("utf-8") - - content = await self._encrypt_content(content, ctx=ctx) - self.agfs.write(table_path, content) + await self._write_path_bytes(table_path, content, ctx=ctx, encrypt=True) # ========== Batch Read (backward compatible) ========== @@ -1747,13 +1765,10 @@ async def write_file( """Write file directly.""" self._ensure_access(uri, ctx) path = self._uri_to_path(uri, ctx=ctx) - await self._ensure_parent_dirs(path) if isinstance(content, str): content = content.encode("utf-8") - - content = await self._encrypt_content(content, ctx=ctx) - self.agfs.write(path, content) + await self._write_path_bytes(path, content, ctx=ctx, encrypt=True, ensure_parent=True) async def read_file( self, @@ -1774,29 +1789,13 @@ async def read_file( """ self._ensure_access(uri, ctx) path = self._uri_to_path(uri, ctx=ctx) - # Verify the file exists before reading, because AGFS read returns - # empty bytes for non-existent files instead of raising an error. try: - self.agfs.stat(path) - except Exception: - raise NotFoundError(uri, "file") - try: - content = self.agfs.read(path) - if isinstance(content, bytes): - raw = content - elif content is not None and hasattr(content, "content"): - raw = content.content - else: - raw = b"" - - # If encryption is enabled, always decrypt full file first - if self._encryptor: - raw = await self._decrypt_content(raw, ctx=ctx) - - text = self._decode_bytes(raw) - except Exception: + raw = await self._read_path_bytes(path, ctx=ctx, decrypt=True, require_exists=True) + except Exception as exc: + if not is_not_found_error(exc): + raise raise NotFoundError(uri, "file") - + text = self._decode_bytes(raw) if offset == 0 and limit == -1: return text lines = text.splitlines(keepends=True) @@ -1812,10 +1811,10 @@ async def read_file_bytes( self._ensure_access(uri, ctx) path = self._uri_to_path(uri, ctx=ctx) try: - raw = self._handle_agfs_read(self.agfs.read(path)) - raw = await self._decrypt_content(raw, ctx=ctx) - return raw - except Exception: + return await self._read_path_bytes(path, ctx=ctx, decrypt=True, require_exists=True) + except Exception as exc: + if not is_not_found_error(exc): + raise raise NotFoundError(uri, "file") async def write_file_bytes( @@ -1827,10 +1826,7 @@ async def write_file_bytes( """Write single binary file.""" self._ensure_access(uri, ctx) path = self._uri_to_path(uri, ctx=ctx) - await self._ensure_parent_dirs(path) - - content = await self._encrypt_content(content, ctx=ctx) - self.agfs.write(path, content) + await self._write_path_bytes(path, content, ctx=ctx, encrypt=True, ensure_parent=True) async def append_file( self, @@ -1845,21 +1841,19 @@ async def append_file( try: existing = "" try: - existing_bytes = self._handle_agfs_read(self.agfs.read(path)) - existing_bytes = await self._decrypt_content(existing_bytes, ctx=ctx) + existing_bytes = await self._read_path_bytes( + path, ctx=ctx, decrypt=True, require_exists=True + ) existing = self._decode_bytes(existing_bytes) - except FileNotFoundError: - pass - except AGFSHTTPError as e: - if e.status_code != 404: + except Exception as exc: + if not is_not_found_error(exc): raise - except AGFSClientError: - raise + pass - await self._ensure_parent_dirs(path) final_content = (existing + content).encode("utf-8") - final_content = await self._encrypt_content(final_content, ctx=ctx) - self.agfs.write(path, final_content) + await self._write_path_bytes( + path, final_content, ctx=ctx, encrypt=True, ensure_parent=True + ) except Exception as e: logger.error(f"[VikingFS] Failed to append to file {uri}: {e}") diff --git a/tests/storage/test_vikingfs_encryption_helpers.py b/tests/storage/test_vikingfs_encryption_helpers.py new file mode 100644 index 000000000..ca39a2adf --- /dev/null +++ b/tests/storage/test_vikingfs_encryption_helpers.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Focused tests for VikingFS encrypted helper paths.""" + +import asyncio +import os +import secrets +from pathlib import Path + +import pytest + +from openviking.crypto.encryptor import FileEncryptor +from openviking.crypto.providers import LocalFileProvider +from openviking.server.identity import RequestContext, Role +from openviking.storage.viking_fs import VikingFS +from openviking_cli.session.user_id import UserIdentifier + + +class _FakeAGFS: + def __init__(self): + self.dirs = {"/", "/local"} + self.files = {} + + def mkdir(self, path): + if path in self.files: + raise FileExistsError(path) + if path in self.dirs: + raise FileExistsError(path) + parent = path.rsplit("/", 1)[0] or "/" + if parent not in self.dirs: + raise FileNotFoundError(parent) + self.dirs.add(path) + return path + + def write(self, path, data): + parent = path.rsplit("/", 1)[0] or "/" + if parent not in self.dirs: + raise FileNotFoundError(parent) + self.files[path] = bytes(data) + return path + + def read(self, path, offset=0, size=-1): + if path not in self.files: + raise FileNotFoundError(path) + data = self.files[path] + return data[offset:] if size == -1 else data[offset : offset + size] + + def stat(self, path): + if path in self.dirs: + return {"isDir": True, "size": 0} + if path in self.files: + return {"isDir": False, "size": len(self.files[path])} + raise FileNotFoundError(path) + + +def _ctx(account_id: str, user_id: str, role: Role = Role.USER) -> RequestContext: + return RequestContext(user=UserIdentifier(account_id, user_id, user_id), role=role) + + +def _run(coro): + return asyncio.run(coro) + + +@pytest.fixture +def encryptor(tmp_path: Path) -> FileEncryptor: + key_file = tmp_path / "master.key" + key_file.write_text(secrets.token_bytes(32).hex()) + os.chmod(key_file, 0o600) + return FileEncryptor(LocalFileProvider(key_file=str(key_file))) + + +@pytest.fixture +def encrypted_fs(encryptor: FileEncryptor) -> VikingFS: + return VikingFS(agfs=_FakeAGFS(), encryptor=encryptor) + + +def test_bound_context_helper_paths_encrypt_and_decrypt_consistently(encrypted_fs: VikingFS): + ctx = _ctx("acct-a", "alice") + file_uri = "viking://resources/docs/guide.md" + dir_uri = "viking://resources/docs" + + with encrypted_fs.bind_request_context(ctx): + _run(encrypted_fs.write_file(file_uri, "hello")) + _run(encrypted_fs.append_file(file_uri, "\nworld")) + _run(encrypted_fs.write_file(f"{dir_uri}/.abstract.md", "abstract text")) + _run(encrypted_fs.write_file(f"{dir_uri}/.overview.md", "overview text")) + _run(encrypted_fs.link(dir_uri, file_uri, reason="primary document")) + + assert _run(encrypted_fs.read_file(file_uri)) == "hello\nworld" + assert _run(encrypted_fs.read_file_bytes(file_uri)) == b"hello\nworld" + assert _run(encrypted_fs.abstract(dir_uri)) == "abstract text" + assert _run(encrypted_fs.overview(dir_uri)) == "overview text" + assert _run(encrypted_fs.relations(dir_uri)) == [ + {"uri": file_uri, "reason": "primary document"} + ] + + raw_file = encrypted_fs.agfs.read("/local/acct-a/resources/docs/guide.md") + raw_abstract = encrypted_fs.agfs.read("/local/acct-a/resources/docs/.abstract.md") + raw_relations = encrypted_fs.agfs.read("/local/acct-a/resources/docs/.relations.json") + assert raw_file.startswith(b"OVE1") + assert raw_abstract.startswith(b"OVE1") + assert raw_relations.startswith(b"OVE1") + +def test_read_and_read_file_return_plaintext_while_storage_stays_ciphertext(encrypted_fs: VikingFS): + ctx = _ctx("acct-a", "alice") + uri = "viking://resources/docs/plaintext-proof.txt" + plaintext = "line-1\nline-2\nline-3" + + _run(encrypted_fs.write_file(uri, plaintext, ctx=ctx)) + + assert _run(encrypted_fs.read_file(uri, ctx=ctx)) == plaintext + assert _run(encrypted_fs.read(uri, ctx=ctx)) == plaintext.encode("utf-8") + assert _run(encrypted_fs.read(uri, offset=7, size=6, ctx=ctx)) == b"line-2" + + stored = encrypted_fs.agfs.read("/local/acct-a/resources/docs/plaintext-proof.txt") + assert stored.startswith(b"OVE1") + assert plaintext.encode("utf-8") not in stored