"""DatabasePool — LRU leases on Code Mode databases for a project. Code Mode may bind a lease to an existing IDA GUI or to a shared managed idalib worker. The pool therefore owns *client interest*, never an IDA process. Releasing an LRU entry persists managed IDBs but does not implicitly save a GUI, then closes only this TUI's lease; other clients and GUI sessions remain alive. Managed workers exit themselves after their final lease. The historical memory budget remains useful for managed idalib instances, while GUI process memory is only advisory. The active and pinned databases are never released to satisfy it. """ from __future__ import annotations from .project import BinaryRef, Project #: Fallback budget if /proc/meminfo can't be read (MB). _FALLBACK_BUDGET_MB = 2048 def _total_ram_mb() -> int: try: with open("/proc/meminfo") as f: for line in f: if line.startswith("MemTotal:"): return int(line.split()[1]) // 1024 except (OSError, ValueError, IndexError): pass return 0 def _pss_mb(pid: int | None) -> int: """Proportional set size of the leased instance process, in MB. PSS is useful for managed idalib workers. For GUI/shared processes it is only advisory because the TUI neither owns all that memory nor controls process exit. """ if not pid: return 0 try: with open(f"/proc/{pid}/smaps_rollup") as f: for line in f: if line.startswith("Pss:"): return int(line.split()[1]) // 1024 except (OSError, ValueError, IndexError): pass return 0 def _default_spawn(ref: BinaryRef, ttl: int, *, new_database: bool = False): # pragma: no cover - needs IDA from .codemode_client import CodeModeClient return CodeModeClient( ref.staged, ttl=ttl, load_args=ref.load_args, output_database=ref.db, new_database=new_database, ) class DatabasePool: """Live Code Mode database leases, keyed by project label.""" def __init__(self, project: Project, *, budget_mb: int | None = None, ttl: int = 1800, spawn=None, mem_fn=None) -> None: self.project = project self._ttl = ttl self._spawn = spawn or _default_spawn self._mem = mem_fn or (lambda c: _pss_mb(getattr(c, "pid", None))) self._clients: dict[str, object] = {} self._lru: list[str] = [] # least-recently-used first self._pinned: set[str] = set() self._recreate: set[str] = set() # Ctrl+L: next attachment creates a fresh IDB self.active: str | None = None # never evicted if budget_mb is None: ram = _total_ram_mb() budget_mb = (ram * project.memory_pct // 100) if ram else _FALLBACK_BUDGET_MB self.budget_mb = max(budget_mb, 256) self.evicted: list[str] = [] # labels evicted, most recent last # -- residency --------------------------------------------------------- # def resident(self) -> list[str]: """Labels with a live worker, least-recently-used first.""" return list(self._lru) def is_resident(self, label: str) -> bool: return label in self._clients def memory_mb(self) -> int: return sum(self._mem(c) for c in self._clients.values()) def pin(self, label: str, on: bool = True) -> None: """Keep ``label`` resident regardless of the budget.""" self._pinned.add(label) if on else self._pinned.discard(label) def is_pinned(self, label: str) -> bool: return label in self._pinned # -- acquire ----------------------------------------------------------- # def get(self, label: str, progress=None): """A live client for ``label``, attaching or spawning as needed. Do not sweep IDA scratch files here: a registered GUI or another Code Mode client may own the database. Code Mode's registry locks and health probes are the authority for safe discovery and stale-record cleanup. """ client = self._clients.get(label) if client is not None: self._touch(label) return client ref = self.project.by_label(label) if ref is None: raise KeyError(f"no such binary in the project: {label}") def note(msg: str) -> None: if progress is not None: progress(msg) note(f"staging {ref.label}\u2026") self.project.stage(ref) note(f"opening {ref.label}\u2026") fresh = label in self._recreate client = (_default_spawn(ref, self._ttl, new_database=fresh) if self._spawn is _default_spawn else self._spawn(ref, self._ttl)) connect = getattr(client, "connect", None) if connect is not None: connect(progress=progress) if progress is not None else connect() self._clients[label] = client self._recreate.discard(label) self._lru.append(label) self._enforce_budget(protect=label) return client def prewarm(self, label: str, progress=None) -> bool: """Attach a database for ``label`` only if it fits the current budget. Pre-warming must never cost residency: evicting a binary the user actually visited to speculatively load one they haven't is a straight downgrade, and the eviction would also throw away that binary's caches. So this refuses rather than making room, and returns False. The cost of a database not attached yet can only be estimated; the largest resident instance is the best evidence available. With nothing resident we have no evidence at all, so we allow one — that is the case where the budget is certainly free. """ if label in self._clients: return False if self.project.by_label(label) is None: return False used = self.memory_mb() est = max((self._mem(c) for c in self._clients.values()), default=0) if used + est > self.budget_mb: return False self.get(label, progress=progress) # get() enforces the budget protecting the NEW label; if that had to # evict, our estimate was wrong and the speculative lease is the one # that should go — never a binary the user chose. if self.memory_mb() > self.budget_mb and label != self.active: self.evict(label) return False return True def recreate_on_next_open(self, label: str) -> None: """Request a fresh IDB after the current lease has been released.""" if self.project.by_label(label) is None: raise KeyError(f"no such binary in the project: {label}") self._recreate.add(label) def _touch(self, label: str) -> None: if label in self._lru: self._lru.remove(label) self._lru.append(label) def set_active(self, label: str | None) -> None: self.active = label if label: self._touch(label) # -- release ----------------------------------------------------------- # def evict(self, label: str, save: bool = True, save_gui: bool = False) -> bool: """Release a resident lease, persisting a managed database first. A budget-driven eviction must not save somebody's GUI implicitly. GUI saves are reserved for an explicit/defensive ``close_all(save=True)``. """ client = self._clients.pop(label, None) if client is None: return False if label in self._lru: self._lru.remove(label) if save and (save_gui or getattr(client, "backend", None) != "gui"): try: # persist analysis + edits so the next open is a load client.save_database() except Exception: # noqa: BLE001 -- evict regardless pass try: client.close() except Exception: # noqa: BLE001 pass self.evicted.append(label) return True def _evictable(self, protect: str | None) -> str | None: for label in self._lru: # least-recently-used first if label == protect or label == self.active or label in self._pinned: continue return label return None def _enforce_budget(self, protect: str | None = None) -> int: """Release LRU leases until the pool fits its budget. Returns how many.""" n = 0 while self.memory_mb() > self.budget_mb: victim = self._evictable(protect) if victim is None: # everything left is active/pinned/protected break self.evict(victim) n += 1 return n def close_all(self, save: bool = True) -> None: for label in list(self._clients): self.evict(label, save=save, save_gui=save) self.active = None # -- introspection ------------------------------------------------------ # def status(self) -> list[dict]: """Per-binary residency for the switcher UI.""" out = [] for ref in self.project.refs: client = self._clients.get(ref.label) out.append({ "label": ref.label, "source": ref.source, "resident": client is not None, "pinned": ref.label in self._pinned, "active": ref.label == self.active, "analysed": self.project.has_db(ref), "memory_mb": self._mem(client) if client is not None else 0, }) return out def __repr__(self) -> str: # pragma: no cover - debug aid return (f"") # Source compatibility for callers that imported the pre-Code-Mode name. WorkerPool = DatabasePool