Source code for swh.vulns.grpc_server

# Copyright (C) 2023-2026  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

"""
A simple tool to start the swh-vulns gRPC server in Rust.
"""

import logging
import multiprocessing
import os
from pathlib import Path
import shlex
import shutil
import socket
import subprocess
from typing import Tuple, Union

from .naive_server import run as run_naive_server

logger = logging.getLogger(__name__)




class ExecutableNotFound(EnvironmentError):
    pass



def _find_project_root():
    return Path(__file__).resolve().parent.parent.parent


def _find_binary(name, *base_dirs):
    for base_dir in base_dirs:
        for profile in ("debug", "release"):
            candidate = Path(base_dir) / profile / name
            if candidate.exists():
                return candidate
    return None




def find_free_port():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]





def build_rust_grpc_server_cmdline(**config):
    logger.debug("Building gRPC server command line")
    port = config.get("port")
    if port is None:
        port = find_free_port()
        logger.debug("Port not configured, using random port %s", port)

    project_root = _find_project_root()
    print("project_root", project_root)
    binary = _find_binary("swh-vulns-grpc-serve", project_root / "target")
    if not binary:
        binary = shutil.which("swh-vulns-grpc-serve")
    if not binary:
        raise ExecutableNotFound("swh-vulns-grpc-serve executable not found")

    cmd = [str(binary)]
    cmd += ["--db", str(config["db"])]
    cmd += ["--commit2vuln", str(config["commit2vuln"])]
    cmd += ["--subgraphwccs", str(config["subgraphwccs"])]
    cmd += ["--graph", str(config["graph"])]
    cmd += ["--bind", f"127.0.0.1:{port}"]
    if config.get("statsd_host"):
        cmd += ["--statsd-host", str(config["statsd_host"])]
    logger.debug("Configuration: %r", config)
    logger.info("Starting gRPC server: %s", " ".join(shlex.quote(x) for x in cmd))
    return cmd, port





def spawn_rust_grpc_server(**config):
    """Returns the process and the port it is listening on"""
    cmd, port = build_rust_grpc_server_cmdline(**config)
    env = dict(os.environ)
    if config.get("profile") == "debug":
        env.setdefault("RUST_LOG", "debug,h2=info")
    server = subprocess.Popen(cmd, env=env)
    return server, port





def spawn_naive_grpc_server(**config) -> Tuple[multiprocessing.Process, int]:
    """Returns the process and the port it is listening on"""
    port = find_free_port()
    server = multiprocessing.Process(
        target=run_naive_server, args=(f"[::]:{port}", config["vulnerabilities"])
    )
    server.start()
    return server, port





def stop_grpc_server(
    server: Union[subprocess.Popen, multiprocessing.Process], timeout: int = 15
):
    server.terminate()
    try:
        if isinstance(server, subprocess.Popen):
            server.wait(timeout=timeout)
        else:
            server.join(timeout=timeout)
    except subprocess.TimeoutExpired:
        logger.warning("Server did not terminate, sending kill signal...")
        server.kill()