modulus/modulus/clangd.py

import fcntl
import hashlib
import json
import os
import pathlib
import shutil
import signal
import subprocess
import sys
import tarfile
import tempfile
import uuid

import modulus.config
import modulus.container

cache_version = 1

def settings(context):
  return {**modulus.config.clangd_defaults, **context.lsp_settings.get("clangd", {})}

def cache_directory(context, image):
  cache_home = pathlib.Path(os.environ.get("XDG_CACHE_HOME", pathlib.Path.home() / ".cache")).resolve()
  config = settings(context)
  key = {name: config[name] for name in ("executable", "query_drivers", "header_roots")}
  key["version"] = cache_version
  digest = hashlib.sha256(json.dumps(key, sort_keys=True).encode()).hexdigest()[:20]
  return cache_home / "modulus" / "clangd" / image.removeprefix("sha256:") / digest

def read_manifest(cache):
  data = json.loads((cache / "manifest.json").read_text(encoding="utf-8"))
  roots = data["roots"]
  if data["version"] != cache_version or not isinstance(roots, list) or not roots:
    raise ValueError("invalid header cache manifest")
  for root in roots:
    path = pathlib.PurePosixPath(root)
    if not path.is_absolute() or root == "/" or ".." in path.parts:
      raise ValueError("invalid header root")
    if not (cache / "rootfs" / root.lstrip("/")).is_dir():
      raise ValueError("incomplete header cache")
  return roots

def extract_headers(stream, destination, roots):
  # Only dereferenced regular files/directories below the discovered header roots.
  allowed = [pathlib.PurePosixPath(root.lstrip("/")) for root in roots]
  with tarfile.open(fileobj=stream, mode="r|") as archive:
    for member in archive:
      path = pathlib.PurePosixPath(member.name)
      if path.is_absolute() or ".." in path.parts or not any(path.is_relative_to(root) for root in allowed):
        raise ValueError(f"unsafe header archive path: {member.name}")
      target = destination / path
      if member.isdir():
        target.mkdir(parents=True, exist_ok=True)
      elif member.isfile():
        target.parent.mkdir(parents=True, exist_ok=True)
        with archive.extractfile(member) as source, target.open("wb") as output:
          shutil.copyfileobj(source, output)
        target.chmod(0o444)
      else:
        raise ValueError(f"non-regular header archive entry: {member.name}")

def prepare(context):
  try:
    image = modulus.container.current_image(context, modulus.container.docker_platform_arch())
    if image is None:
      raise RuntimeError("build image is missing or stale; run configure first")
    cache = cache_directory(context, image)
    cache.parent.mkdir(parents=True, exist_ok=True)
    with (cache.parent / ".lock").open("a") as lock:
      fcntl.flock(lock, fcntl.LOCK_EX)
      if cache.exists():
        read_manifest(cache)
        return 0
      probe = context.tools_dir / "clangd_headers.py"
      result = subprocess.run(
        [*modulus.container.isolated_run(image), "--volume", f"{probe}:/modulus_headers.py:ro", image,
          "python3", "/modulus_headers.py", json.dumps(settings(context))],
        stdout=subprocess.PIPE, text=True, check=True,
      )
      roots = json.loads(result.stdout)
      with tempfile.TemporaryDirectory(prefix=".headers_", dir=cache.parent) as temporary:
        stage = pathlib.Path(temporary) / "cache"
        stage.mkdir()
        command = [*modulus.container.isolated_run(image), image, "tar", "--dereference", "--hard-dereference",
          "--exclude=*/.git", "--exclude=*/__pycache__",
          "-C", "/", "-cf", "-", *[root.lstrip("/") for root in roots]]
        with subprocess.Popen(command, stdout=subprocess.PIPE) as process:
          try:
            extract_headers(process.stdout, stage / "rootfs", roots)
          finally:
            process.stdout.close()
          if process.wait() != 0:
            raise RuntimeError("could not copy container headers")
        (stage / "manifest.json").write_text(json.dumps({"version": cache_version, "roots": roots}) + "\n")
        read_manifest(stage)
        stage.rename(cache)
      print(f"Prepared clangd headers: {cache}", file=sys.stderr)
    return 0
  except (OSError, ValueError, KeyError, RuntimeError, subprocess.SubprocessError, tarfile.TarError) as error:
    print(f"Modulus clangd preparation: {error}", file=sys.stderr)
    return 1

def path_mappings(cache, roots):
  mappings = []
  for root in sorted(roots, key=len, reverse=True):
    client = str(cache / "rootfs" / root.lstrip("/"))
    if any(c in client + root for c in ",=\n\r"):
      raise ValueError("header paths cannot contain clangd mapping delimiters")
    mappings.append(f"{client}={root}")
  return "--path-mappings=" + ",".join(mappings)

def server_arguments(context, argv):
  config = settings(context)
  return [config["executable"], *config["arguments"], *argv,
    "--query-driver=" + ",".join(config["query_drivers"])]

def launch_command(context, argv, image, cache, roots, name):
  root = context.repo_root
  return [
    *modulus.container.isolated_run(image), "--interactive", "--name", name,
    "--volume", f"{root}:{root}:ro",
    "--volume", f"{context.build_root}:{context.build_root}:rw",
    "--env", f"XDG_CACHE_HOME={context.common_dir / 'clangd'}",
    "--workdir", str(root), image, *server_arguments(context, argv), path_mappings(cache, roots),
  ]

def launch(context, argv):
  try:
    if context.docker_marker.exists():
      os.execv(settings(context)["executable"], server_arguments(context, argv))
    image = modulus.container.current_image(context, modulus.container.docker_platform_arch())
    if image is None:
      raise RuntimeError("missing or stale build image; run ./modulus.py configure (Modulus: ./tools/configure.py)")
    databases = [context.default_target.build_dir, *[
      context.repo_root / database["build_dir"] for database in context.lsp_settings.get("compile_databases", [])
    ]]
    for database in databases:
      if not (database / "compile_commands.json").is_file():
        raise RuntimeError(f"missing {database / 'compile_commands.json'}; run the lsp command")
    cache = cache_directory(context, image)
    try:
      roots = read_manifest(cache)
    except (OSError, ValueError, KeyError) as error:
      raise RuntimeError("header cache is missing or invalid; run the lsp command") from error
    name = "modulus_clangd_" + uuid.uuid4().hex
    command = launch_command(context, argv, image, cache, roots, name)
    return run_server(command, name)
  except (OSError, ValueError, RuntimeError) as error:
    print(f"Modulus clangd: {error}", file=sys.stderr)
    return 1

def run_server(command, name):
  # Inherit stdio verbatim; never send build logs or a terminal over the LSP pipe.

  def interrupted(signum, frame):
    raise SystemExit(128 + signum)

  previous = {sig: signal.signal(sig, interrupted) for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP)}
  process = None
  try:
    process = subprocess.Popen(command)
    return process.wait()
  finally:
    for sig, handler in previous.items():
      signal.signal(sig, handler)
    if process is not None:
      try:
        subprocess.run(["docker", "rm", "--force", name], stdout=subprocess.DEVNULL,
          stderr=subprocess.DEVNULL, check=False, timeout=15)
      except (OSError, subprocess.TimeoutExpired) as error:
        print(f"Modulus clangd cleanup: {error}", file=sys.stderr)
      try:
        process.wait(timeout=5)
      except subprocess.TimeoutExpired:
        process.kill()
        process.wait()