Initial commit: split dlp-io library out of DataPacker

2026-07-31 14:58:37 +08:00
commit 3b0c81fd9d
22 changed files with 2509 additions and 0 deletions
@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from ._errors import DlpConfigurationError
+
+PROTOCOL_VERSION = 1
+CONTROL_LIMIT = 64 * 1024
+CHUNK_SIZE = 1024 * 1024
+HANDSHAKE_PREFIX = "DLP_IO 1 "
+
+
+class HelperErrorCode(str, Enum):
+    INVALID_PATH = "invalid_path"
+    INVALID_OFFSET = "invalid_offset"
+    FILE_ACCESS = "file_access"
+    INVALID_REQUEST = "invalid_request"
+    INVALID_REQUEST_TYPE = "invalid_request_type"
+    UNSUPPORTED_VERSION = "unsupported_version"
+    UNKNOWN_OPERATION = "unknown_operation"
+    INTERNAL = "internal"
+
+
+@dataclass(frozen=True)
+class HelperSourceConfig:
+    authkey: bytes
+    protocol_version: int = PROTOCOL_VERSION
+    control_limit: int = CONTROL_LIMIT
+    chunk_size: int = CHUNK_SIZE
+
+    @classmethod
+    def from_authkey(cls, authkey: bytes) -> HelperSourceConfig:
+        if not isinstance(authkey, bytes) or len(authkey) != 32:
+            raise DlpConfigurationError("authkey must contain exactly 32 bytes")
+        return cls(authkey=authkey)
+
+
+_HELPER_TEMPLATE = r'''
+import json
+import os
+import sys
+from multiprocessing.connection import AuthenticationError, Listener
+
+VERSION = __PROTOCOL_VERSION__
+CONTROL_LIMIT = __CONTROL_LIMIT__
+CHUNK_SIZE = __CHUNK_SIZE__
+PIPE_PREFIX = r"\\.\pipe\DlpIo_"
+AUTHKEY = bytes.fromhex("__AUTHKEY_HEX__")
+
+
+def send_json(conn, value):
+    conn.send_bytes(json.dumps(value).encode("utf-8"))
+
+
+def error(conn, code, message):
+    send_json(conn, {
+        "version": VERSION,
+        "ok": False,
+        "error_code": code,
+        "error": message,
+    })
+
+
+def handle_read(conn, request):
+    path = request.get("path")
+    offset = request.get("offset", 0)
+    if not isinstance(path, str) or not path:
+        error(conn, "invalid_path", "path must be a non-empty string")
+        return
+    if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0:
+        error(conn, "invalid_offset", "offset must be a non-negative integer")
+        return
+    try:
+        size = os.path.getsize(path)
+        stream = open(path, "rb")
+    except Exception as exc:
+        error(conn, "file_access", "%s: %s" % (type(exc).__name__, exc))
+        return
+    try:
+        stream.seek(offset)
+        send_json(conn, {"version": VERSION, "ok": True, "size": size})
+        while True:
+            chunk = stream.read(CHUNK_SIZE)
+            if not chunk:
+                conn.send_bytes(b"")
+                return
+            conn.send_bytes(chunk)
+    except (BrokenPipeError, EOFError, OSError):
+        return
+    finally:
+        stream.close()
+
+
+def handle(conn):
+    try:
+        request = json.loads(conn.recv_bytes(maxlength=CONTROL_LIMIT).decode("utf-8"))
+    except Exception as exc:
+        error(conn, "invalid_request", "invalid request: %s" % exc)
+        return True
+    if not isinstance(request, dict):
+        error(conn, "invalid_request_type", "request must be a JSON object")
+        return True
+    if request.get("version") != VERSION:
+        error(conn, "unsupported_version", "unsupported protocol version")
+        return True
+    operation = request.get("op")
+    if operation == "ping":
+        send_json(conn, {
+            "version": VERSION,
+            "ok": True,
+            "pid": os.getpid(),
+            "python": sys.version,
+        })
+        return True
+    if operation == "read":
+        handle_read(conn, request)
+        return True
+    if operation == "quit":
+        send_json(conn, {"version": VERSION, "ok": True})
+        return False
+    error(conn, "unknown_operation", "unknown operation: %r" % operation)
+    return True
+
+
+def main():
+    address = PIPE_PREFIX + os.urandom(16).hex()
+    listener = Listener(address=address, family="AF_PIPE", authkey=AUTHKEY)
+    sys.stdout.write("DLP_IO " + str(VERSION) + " " + address + "\n")
+    sys.stdout.flush()
+    running = True
+    while running:
+        try:
+            conn = listener.accept()
+        except AuthenticationError:
+            continue
+        try:
+            running = handle(conn)
+        except Exception as exc:
+            try:
+                error(conn, "internal", "%s: %s" % (type(exc).__name__, exc))
+            except Exception:
+                pass
+        finally:
+            try:
+                conn.close()
+            except Exception:
+                pass
+    listener.close()
+
+
+try:
+    main()
+except Exception as exc:
+    sys.stderr.write("dlp-io helper fatal: %s: %s\n" % (type(exc).__name__, exc))
+    sys.stderr.flush()
+    sys.exit(1)
+'''
+
+
+def render_helper_source(authkey: bytes) -> str:
+    config = HelperSourceConfig.from_authkey(authkey)
+    replacements = {
+        "__AUTHKEY_HEX__": config.authkey.hex(),
+        "__PROTOCOL_VERSION__": str(config.protocol_version),
+        "__CONTROL_LIMIT__": str(config.control_limit),
+        "__CHUNK_SIZE__": str(config.chunk_size),
+    }
+    source = _HELPER_TEMPLATE
+    for marker, value in replacements.items():
+        source = source.replace(marker, value)
+    return source