Browse Source

chore(virtual-printer): log K-profile, filament and load commands from the slicer in full

The publish log line stops at 100 characters, which cut every K-profile
command off before its tray_id, the field that says how BambuStudio numbers
an A2L or AMS-HT tray. Those commands carry no credentials and are now
logged whole, so a support bundle settles it.

--------------

Security hardening (maziggy/bambuddy-security #42)
maziggy 4 ngày trước cách đây
mục cha
commit
5fec296d72

+ 74 - 12
backend/app/main.py

@@ -497,6 +497,12 @@ _printer_kprofiles_primed_since_connect: dict[int, bool] = {}
 # {(printer_id, filename): archive_id}
 _expected_prints: dict[tuple[int, str], int] = {}
 
+# Archives of the running prints Bambuddy dispatched itself: adopted from an
+# expected print, or found with Bambuddy's dispatch marker. _active_prints holds
+# every print on_print_start archived, those started from Bambu Studio, Handy or
+# the printer's screen included, so it is no proof of ownership on its own.
+_owned_print_archive_ids: set[int] = set()
+
 # Track AMS mapping for prints: {archive_id: [global_tray_id_per_slot]}
 # Used by usage tracker to map 3MF slots to physical AMS trays
 _print_ams_mappings: dict[int, list[int]] = {}
@@ -863,7 +869,18 @@ def _is_bambuddy_authorized_print_in_memory(printer_id: int, state: PrinterState
     if printer_manager.get_current_print_user(printer_id):
         return True
 
-    return any(key in _expected_prints or key in _active_prints for key in _build_status_print_keys(printer_id, state))
+    return any(
+        key in _expected_prints or _active_prints.get(key) in _owned_print_archive_ids
+        for key in _build_status_print_keys(printer_id, state)
+    )
+
+
+def _mark_owned_if_dispatched(archive) -> None:
+    """Remember a running print's archive as Bambuddy's when it carries a
+    dispatch marker Bambuddy writes itself: ``billing_run_id`` (minted per
+    dispatch in the scheduler) or ``created_by_id`` (from the queue item)."""
+    if archive.billing_run_id is not None or archive.created_by_id is not None:
+        _owned_print_archive_ids.add(archive.id)
 
 
 async def _is_printer_kill_switch_enabled_cached() -> bool:
@@ -897,6 +914,14 @@ async def _is_bambuddy_authorized_print(printer_id: int, state: PrinterState, db
     if _is_bambuddy_authorized_print_in_memory(printer_id, state):
         return True
 
+    # The printer's built-in calibration runs (Bambuddy's Calibrate, the
+    # printer's screen) play gcode from its read-only system partition, where no
+    # user file can be. Only that path counts: a calibration *name* proves
+    # nothing, since Bambu Studio's pressure-advance print reports a user-style
+    # path, and any print can be given that name.
+    if str(getattr(state, "gcode_file", "") or "").startswith("/usr/"):
+        return True
+
     possible_keys = _build_status_print_keys(printer_id, state)
 
     # In-memory ownership is lost on every Bambuddy restart, so fall back to what
@@ -904,22 +929,47 @@ async def _is_bambuddy_authorized_print(printer_id: int, state: PrinterState, db
     # actually running, rather than to an unrelated one that reuses a filename.
     raw_subtask_id = getattr(state, "subtask_id", None)
     subtask_id = str(raw_subtask_id).strip() if raw_subtask_id is not None else ""
-    if subtask_id in ("", "0"):
+    if subtask_id == "":
+        # Not reported yet: the first frames after a restart can lack it.
         return None
 
     from backend.app.models.archive import PrintArchive
 
-    result = await db.execute(
-        select(PrintArchive)
-        .where(
-            PrintArchive.printer_id == printer_id,
-            PrintArchive.status == "printing",
-            PrintArchive.subtask_id == subtask_id,
+    if subtask_id == "0":
+        # A print started outside the cloud from Bambu Studio, Handy or the
+        # printer's screen reports "0", so there is no id to match: one of
+        # Bambuddy's printing archives on this printer under the name the
+        # printer reports, compared as the printer echoes names (spaces become
+        # underscores). Nothing reported to compare yet: wait.
+        reported = {_normalise_subtask_name(name) for _pid, name in possible_keys} - {""}
+        if not reported:
+            return None
+        dispatched = await db.execute(
+            select(PrintArchive)
+            .where(
+                PrintArchive.printer_id == printer_id,
+                PrintArchive.status == "printing",
+                or_(PrintArchive.billing_run_id.is_not(None), PrintArchive.created_by_id.is_not(None)),
+                # An archive with a real id of its own is a print that reported
+                # one: a "0" print is not it, even under the same name (one whose
+                # completion was missed would otherwise cover it).
+                or_(PrintArchive.subtask_id.is_(None), PrintArchive.subtask_id.in_(("", "0"))),
+            )
+            .order_by(PrintArchive.created_at.desc())
         )
-        .order_by(PrintArchive.created_at.desc())
-        .limit(1)
-    )
-    archive = result.scalar_one_or_none()
+        archive = next((a for a in dispatched.scalars() if _archive_name_forms(a) & reported), None)
+    else:
+        result = await db.execute(
+            select(PrintArchive)
+            .where(
+                PrintArchive.printer_id == printer_id,
+                PrintArchive.status == "printing",
+                PrintArchive.subtask_id == subtask_id,
+            )
+            .order_by(PrintArchive.created_at.desc())
+            .limit(1)
+        )
+        archive = result.scalar_one_or_none()
 
     # An archive row on its own proves nothing: `on_print_start` archives every
     # print it observes, including ones started from Bambu Studio or Handy, and
@@ -931,6 +981,7 @@ async def _is_bambuddy_authorized_print(printer_id: int, state: PrinterState, db
     if archive is not None and (archive.billing_run_id is not None or archive.created_by_id is not None):
         # Rehydrate the fast in-memory path for subsequent status frames. Include
         # both the archive filename and every normalized key reported by MQTT.
+        _owned_print_archive_ids.add(archive.id)
         _active_prints[(printer_id, archive.filename)] = archive.id
         for key in possible_keys:
             _active_prints[key] = archive.id
@@ -1096,6 +1147,7 @@ async def forget_printer(printer_id: int) -> None:
         for key in [key for key in tracked if key[0] == printer_id]:
             archive_ids.add(tracked.pop(key))
     for archive_id in archive_ids:
+        _owned_print_archive_ids.discard(archive_id)
         _print_ams_mappings.pop(archive_id, None)
         _print_cost_center_ids.pop(archive_id, None)
         _print_plate_ids.pop(archive_id, None)
@@ -4648,6 +4700,9 @@ async def on_print_start(printer_id: int, data: dict):
         if expected_archive_id:
             # This is a reprint/scheduled print - use existing archive, don't create new one
             logger.info("Using expected archive %s for print (skipping duplicate)", expected_archive_id)
+            # Bambuddy dispatched it: the kill switch's fast path keeps it as
+            # Bambuddy's once the expected entry is consumed here.
+            _owned_print_archive_ids.add(expected_archive_id)
             from backend.app.models.archive import PrintArchive
 
             result = await db.execute(select(PrintArchive).where(PrintArchive.id == expected_archive_id))
@@ -4988,6 +5043,7 @@ async def on_print_start(printer_id: int, data: dict):
             if subtask_match:
                 logger.info("Resuming archive %s on subtask_id match (%s)", existing_archive.id, subtask_id)
                 _active_prints[(printer_id, existing_archive.filename)] = existing_archive.id
+                _mark_owned_if_dispatched(existing_archive)
                 if existing_archive.energy_start_kwh is None:
                     await _record_energy_start(existing_archive, printer_id, db, context="subtask-resume")
                 if not notification_sent:
@@ -5071,6 +5127,7 @@ async def on_print_start(printer_id: int, data: dict):
                 )
                 # Track this as the active print
                 _active_prints[(printer_id, existing_archive.filename)] = existing_archive.id
+                _mark_owned_if_dispatched(existing_archive)
                 # Attach subtask_id retroactively so future restarts can resume.
                 # Compare for inequality (not "is empty") to also pick up reprint
                 # dispatches that mint a fresh id — see #1807 for the bogus
@@ -7887,6 +7944,7 @@ async def on_print_complete(printer_id: int, data: dict):
     for key in possible_keys:
         archive_id = _active_prints.pop(key, None)
         if archive_id:
+            _owned_print_archive_ids.discard(archive_id)
             logger.info("Found archive %s with key %s", archive_id, key)
             # Also clean up any other keys pointing to this archive
             keys_to_remove = [k for k, v in _active_prints.items() if v == archive_id]
@@ -10175,6 +10233,10 @@ def _evict_stale_expected_prints() -> None:
     """
     # Use monotonic time so the TTL is unaffected by system clock adjustments
     # (e.g. NTP sync, DST changes).
+    # A print that ended without its archive being found at completion leaves
+    # its id here; only running prints are owned.
+    _owned_print_archive_ids.intersection_update(_active_prints.values())
+
     cutoff = time.monotonic() - _EXPECTED_PRINT_TTL_SECONDS
     stale_keys = [k for k, t in _expected_print_registered_at.items() if t < cutoff]
     if not stale_keys:

+ 16 - 0
backend/app/services/virtual_printer/mqtt_server.py

@@ -23,6 +23,12 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+# Slicer commands logged in full, not cut at 100 characters: they say how the
+# slicer addresses an AMS tray (see _handle_publish).
+_COMMANDS_LOGGED_WHOLE = frozenset(
+    {"extrusion_cali_sel", "extrusion_cali_set", "ams_filament_setting", "ams_change_filament"}
+)
+
 # Default MQTT port for Bambu printers (MQTT over TLS)
 MQTT_PORT = 8883
 
@@ -1230,6 +1236,16 @@ class SimpleMQTTServer:
                 )
                 return
 
+            # The line above stops at 100 characters, before the tray_id of a
+            # K-profile command: how BambuStudio numbers an A2L or AMS-HT tray,
+            # which support bundles are the only source for. These carry no
+            # credentials, so they are logged whole.
+            print_command = data.get("print") if isinstance(data, dict) else None
+            if isinstance(print_command, dict) and print_command.get("command") in _COMMANDS_LOGGED_WHOLE:
+                logger.info(
+                    "MQTT %s from slicer: %s", print_command["command"], json.dumps(print_command, sort_keys=True)
+                )
+
             # Env-flagged command trace (#1622): every slicer-originated publish
             # gets a line in vp_wire/<vp>_cmd.jsonl alongside the printer-side
             # responses captured in mqtt_bridge. Off by default.

+ 3 - 1
backend/tests/integration/test_external_print_confirmation_1898.py

@@ -76,7 +76,7 @@ class _StubArchiveService:
         return archive
 
 
-async def _drive_print_start(test_engine, printer, *, download_ok: bool, extra_patches=()) -> dict:
+async def _drive_print_start(test_engine, printer, *, download_ok: bool, extra_patches=(), state_fields=None) -> dict:
     """Run ``on_print_start`` against the test database.
 
     ``download_ok`` picks the branch: False leaves the 3MF unreachable and the
@@ -90,6 +90,8 @@ async def _drive_print_start(test_engine, printer, *, download_ok: bool, extra_p
         sdcard_reported=True,
         subtask_id="",  # a local print; a bare mock attribute would be stored as its repr
     )
+    for name, value in (state_fields or {}).items():
+        setattr(state, name, value)
 
     patches = [
         patch("backend.app.main.async_session", session_maker),

+ 233 - 0
backend/tests/integration/test_kill_switch_owned_prints.py

@@ -0,0 +1,233 @@
+"""The printer kill switch acts on prints Bambuddy did not start.
+
+Two ways it did not:
+
+- The fast check counted every print in ``_active_prints`` as Bambuddy's, and
+  ``on_print_start`` puts every print it archives there, those started from
+  Bambu Studio, Handy or the printer's screen included. A print running when the
+  switch was turned on, or one archived before the switch looked, was never
+  stopped.
+- A LAN print reports ``subtask_id`` "0". The database check answered "not sure
+  yet" for it and waited, every frame, until the print was archived and the
+  first hole let it through.
+
+Bambuddy's own prints are now the ones it dispatched: an expected print, the
+print user it recorded, or an archive it adopted for a dispatch or that carries
+its dispatch marker. A "0" print is matched to Bambuddy's printing archive on
+that printer by file name; if there is none it waits only while a queue job of
+Bambuddy's is printing there.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from backend.app import main as main_module
+from backend.app.main import register_expected_print
+from backend.app.models.archive import PrintArchive
+from backend.tests.integration.test_external_print_confirmation_1898 import (  # noqa: F401
+    SUBTASK,
+    _clear_print_state,
+    _drive_print_start,
+)
+
+pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
+
+
+@pytest.fixture(autouse=True)
+def _clear_owned(monkeypatch):
+    main_module._owned_print_archive_ids.clear()
+    monkeypatch.setattr(main_module.printer_manager, "get_current_print_user", lambda printer_id: None)
+    yield
+    main_module._owned_print_archive_ids.clear()
+
+
+def _status(subtask_id="0"):
+    return SimpleNamespace(
+        current_print=None, subtask_name=SUBTASK, subtask_id=subtask_id, gcode_file=f"{SUBTASK}.gcode.3mf"
+    )
+
+
+async def test_a_print_bambuddy_did_not_start_stays_unauthorized_once_archived(test_engine, printer_factory):
+    """The AP-1 scenario: archived by on_print_start, the print used to pass the
+    fast check from then on."""
+    printer_id = (await printer_factory()).id
+
+    await _drive_print_start(test_engine, SimpleNamespace(id=printer_id), download_ok=True)
+
+    assert any(key[0] == printer_id for key in main_module._active_prints)
+    assert main_module._is_bambuddy_authorized_print_in_memory(printer_id, _status()) is False
+
+
+async def test_a_print_bambuddy_dispatched_is_its_own_once_adopted(test_engine, db_session, printer_factory):
+    printer_id = (await printer_factory()).id
+    archive = PrintArchive(
+        printer_id=printer_id,
+        filename=f"{SUBTASK}.gcode.3mf",
+        file_path=f"archives/test/{SUBTASK}.gcode.3mf",
+        file_size=1,
+        print_name=SUBTASK,
+        status="pending",
+    )
+    db_session.add(archive)
+    await db_session.commit()
+    register_expected_print(printer_id, f"{SUBTASK}.gcode.3mf", archive.id)
+
+    await _drive_print_start(test_engine, SimpleNamespace(id=printer_id), download_ok=False)
+
+    # Without the expected entries (print start consumes the names the printer
+    # reported; the rest are evicted later), the print is still Bambuddy's.
+    main_module._expected_prints.clear()
+    assert archive.id in main_module._owned_print_archive_ids
+    assert main_module._is_bambuddy_authorized_print_in_memory(printer_id, _status()) is True
+
+
+def _db(archive, dispatched_queue_item_id=None):
+    """Both lookups: by subtask id (one row) and, for "0", Bambuddy's printing rows."""
+    rows = [archive] if archive is not None else []
+    return SimpleNamespace(
+        execute=AsyncMock(return_value=SimpleNamespace(scalar_one_or_none=lambda: archive, scalars=lambda: rows)),
+        scalar=AsyncMock(return_value=dispatched_queue_item_id),
+    )
+
+
+def _archive(archive_id, filename, *, marker="run-1", print_name=None):
+    return SimpleNamespace(
+        id=archive_id, filename=filename, print_name=print_name, billing_run_id=marker, created_by_id=None
+    )
+
+
+async def test_a_lan_print_with_no_dispatch_behind_it_is_unauthorized():
+    assert await main_module._is_bambuddy_authorized_print(9, _status("0"), _db(None)) is False
+
+
+async def test_a_lan_print_matching_bambuddys_dispatched_archive_is_its_own():
+    """After a restart: no id to match, so the file name and the dispatch marker."""
+    archive = _archive(31, f"{SUBTASK}.gcode.3mf")
+
+    assert await main_module._is_bambuddy_authorized_print(9, _status("0"), _db(archive)) is True
+    assert 31 in main_module._owned_print_archive_ids
+
+
+async def test_a_lan_print_waits_while_bambuddy_has_a_job_printing_there():
+    """Nothing of Bambuddy's by that name, but a queue job is printing there."""
+    assert await main_module._is_bambuddy_authorized_print(9, _status("0"), _db(None, 401)) is None
+
+
+async def test_a_name_the_printer_echoes_with_underscores_still_matches():
+    """Bambuddy uploads "My Part" and the printer reports "My_Part"."""
+    archive = _archive(33, "My Part.gcode.3mf")
+    state = SimpleNamespace(current_print=None, subtask_name="My_Part", subtask_id="0", gcode_file="My_Part.gcode.3mf")
+
+    assert await main_module._is_bambuddy_authorized_print(9, state, _db(archive)) is True
+
+
+async def test_a_bambuddy_archive_of_another_name_does_not_cover_a_lan_print():
+    other = _archive(34, "Other_job.gcode.3mf")
+
+    assert await main_module._is_bambuddy_authorized_print(9, _status("0"), _db(other)) is False
+
+
+async def test_only_bambuddys_printing_archives_on_that_printer_are_read():
+    from sqlalchemy.dialects import sqlite
+
+    db = _db(None)
+    await main_module._is_bambuddy_authorized_print(9, _status("0"), db)
+
+    query = db.execute.await_args.args[0]
+    sql = str(query.compile(dialect=sqlite.dialect(), compile_kwargs={"literal_binds": True}))
+    assert "print_archives.printer_id = 9" in sql
+    assert "print_archives.status = 'printing'" in sql
+    assert "print_archives.billing_run_id IS NOT NULL" in sql
+
+
+async def test_a_lan_print_that_names_nothing_yet_waits():
+    state = SimpleNamespace(current_print=None, subtask_name="", subtask_id="0", gcode_file="")
+    db = _db(None)
+
+    assert await main_module._is_bambuddy_authorized_print(9, state, db) is None
+    db.execute.assert_not_awaited()
+
+
+async def test_the_printers_built_in_calibration_is_left_alone():
+    """Bambuddy's Calibrate and the printer's screen play gcode from the
+    printer's system partition (measured: /usr/etc/print/<model>/...)."""
+    state = SimpleNamespace(
+        current_print=None,
+        subtask_name="auto_cali_for_user_param.gcode",
+        subtask_id="0",
+        gcode_file="/usr/etc/print/O1C2/auto_cali_for_user_param.gcode",
+    )
+    db = _db(None)
+
+    assert await main_module._is_bambuddy_authorized_print(9, state, db) is True
+    db.execute.assert_not_awaited()
+
+
+async def test_a_print_named_like_a_calibration_is_still_a_print():
+    """Bambu Studio's pressure-advance print reports /data/Metadata/plate_1.gcode
+    and subtask auto_pa_line_calib_mode (measured on an H2C), exactly as a user
+    print given that name would."""
+    state = SimpleNamespace(
+        current_print=None,
+        subtask_name="pa_line_calib_mode",
+        subtask_id="0",
+        gcode_file="/data/Metadata/plate_1.gcode",
+    )
+
+    assert await main_module._is_bambuddy_authorized_print(9, state, _db(None)) is False
+
+
+async def test_a_lan_print_is_not_covered_by_an_archive_that_has_its_own_id():
+    from sqlalchemy.dialects import sqlite
+
+    db = _db(None)
+    await main_module._is_bambuddy_authorized_print(9, _status("0"), db)
+
+    sql = str(db.execute.await_args.args[0].compile(dialect=sqlite.dialect(), compile_kwargs={"literal_binds": True}))
+    assert "print_archives.subtask_id IS NULL" in sql
+
+
+async def test_a_missing_subtask_id_still_waits():
+    """The first frames after a restart can arrive before the field is filled."""
+    db = _db(None)
+
+    assert await main_module._is_bambuddy_authorized_print(9, _status(None), db) is None
+    db.execute.assert_not_awaited()
+
+
+async def test_a_finished_print_is_no_longer_owned():
+    main_module._owned_print_archive_ids.add(41)
+    main_module._active_prints[(9, f"{SUBTASK}.gcode.3mf")] = 41
+
+    main_module._active_prints.clear()
+    main_module._evict_stale_expected_prints()
+
+    assert 41 not in main_module._owned_print_archive_ids
+
+
+@pytest.mark.parametrize(("marker", "owned"), [("run-7", True), (None, False)])
+async def test_a_print_resumed_after_a_restart_is_owned_only_with_bambuddys_marker(
+    test_engine, db_session, printer_factory, marker, owned
+):
+    """on_print_start resumes the printing archive it finds by name; that
+    archive is Bambuddy's only if Bambuddy dispatched it."""
+    printer_id = (await printer_factory()).id
+    archive = PrintArchive(
+        printer_id=printer_id,
+        filename=f"{SUBTASK}.gcode.3mf",
+        file_path=f"archives/test/{SUBTASK}.gcode.3mf",
+        file_size=1,
+        print_name=SUBTASK,
+        status="printing",
+        billing_run_id=marker,
+    )
+    db_session.add(archive)
+    await db_session.commit()
+
+    await _drive_print_start(
+        test_engine, SimpleNamespace(id=printer_id), download_ok=False, state_fields={"progress": 40}
+    )
+
+    assert (archive.id in main_module._owned_print_archive_ids) is owned

+ 1 - 0
backend/tests/unit/test_forget_printer.py

@@ -25,6 +25,7 @@ NOT_FORGOTTEN = {
     "_print_ams_mappings": "keyed by archive (dropped with the printer's expected/active prints)",
     "_print_cost_center_ids": "keyed by archive (dropped with the printer's expected/active prints)",
     "_print_plate_ids": "keyed by archive (dropped with the printer's expected/active prints)",
+    "_owned_print_archive_ids": "keyed by archive (dropped with the printer's expected/active prints)",
     "_ams_assignment_locks": "a lock carries nothing over; popping one a task holds would let a second be made",
     "_fallback_recovery_locks": "a lock carries nothing over (see its comment on popping)",
     "_HMS_FAILURE_REASONS": "constant",

+ 47 - 0
backend/tests/unit/test_vp_mqtt_server.py

@@ -714,3 +714,50 @@ class TestHandleClientTCPKeepaliveTuning:
                 "a macOS sleep leaves the pre-sleep socket in _clients "
                 "for ~2 h until the default SO_KEEPALIVE probes fire."
             )
+
+
+class TestSlicerCommandsTheLogKeepsWhole:
+    """The publish log line stops at 100 characters, which cut every K-profile
+    command off before its tray_id: the field that says how BambuStudio numbers
+    an A2L or AMS-HT tray. Those few commands are logged whole; nothing they
+    carry is a credential."""
+
+    def _publish(self, caplog, command: dict):
+        import logging
+
+        server = _make_server()
+        server._client_serials["c1"] = server.serial
+        writer = MagicMock()
+        writer.write = MagicMock()
+        writer.drain = AsyncMock()
+        payload = _build_publish_payload("device/01P00AABCDEFGHI/request", {"print": command})
+        with caplog.at_level(logging.INFO, logger="backend.app.services.virtual_printer.mqtt_server"):
+            asyncio.run(server._handle_publish(0x30, payload, writer, "c1"))
+        return caplog.text
+
+    @pytest.mark.parametrize(
+        "command",
+        ["extrusion_cali_sel", "extrusion_cali_set", "ams_filament_setting", "ams_change_filament"],
+    )
+    def test_the_whole_command_is_logged(self, caplog, command):
+        body = {
+            "command": command,
+            "ams_id": 16,
+            "slot_id": 2,
+            "nozzle_diameter": "0.4",
+            "filament_id": "GFA00",
+            "setting_id": "GFSA00_25",
+            "tray_id": 26,
+        }
+
+        text = self._publish(caplog, body)
+
+        assert '"tray_id": 26' in text
+
+    def test_other_commands_keep_the_short_line(self, caplog):
+        text = self._publish(
+            caplog,
+            {"command": "project_file", "url": "ftp:///x.3mf", "param": "Metadata/plate_1.gcode" * 5, "tray_id": 26},
+        )
+
+        assert '"tray_id": 26' not in text