# 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
import contextlib
import json
import logging
import multiprocessing
from pathlib import Path
import socket
import sqlite3
import threading
import time
from typing import Iterator, List
import grpc
import pytest
from swh.graph.example_dataset import DATASET_DIR
from swh.osv.luigi.commit2vuln import Commit2VulnBv
from swh.osv.luigi.reachability import ComputeConnectedComponents
from swh.osv.luigi.vuln2commit import Vuln2CommitParquet
from .grpc.swhvulns_pb2 import DetectedVulnerability
from .grpc.swhvulns_pb2_grpc import VulnerabilityServiceStub
from .grpc_server import (
ExecutableNotFound,
spawn_naive_grpc_server,
spawn_rust_grpc_server,
)
logger = logging.getLogger(__name__)
# The default on Python < 3.14 on Linux is "fork", which is not safe
# as forking a multithreaded process is problematic
[docs]
class VulnsServerProcess(multiprocessing.context.ForkServerProcess):
def __init__(self, config, *args, **kwargs):
self.config = config
self.q = multiprocessing.get_context("forkserver").Queue()
super().__init__(*args, **kwargs)
[docs]
def run(self):
implementation = self.config["grpc_server"]["implementation"]
try:
if implementation == "rust":
server, port = spawn_rust_grpc_server(**self.config["grpc_server"])
elif implementation == "naive":
server, port = spawn_naive_grpc_server(**self.config["grpc_server"])
else:
raise ValueError(f"Invalid implementation: {implementation}")
self.q.put(
{
"grpc_url": f"localhost:{port}",
"port": port,
"pid": server.pid,
}
)
except Exception as e:
if isinstance(e, ExecutableNotFound):
# hack to add a bit more context and help to the user,
# especially when this is used from another swh package...
# XXX on py>=3.11 we could use e.add_note() instead
e.args = (
*e.args,
"This probably means you need to build the rust grpc server "
"for swh-vulns.",
)
logger.exception(e)
self.q.put(e)
[docs]
def start(self, *args, **kwargs):
super().start()
self.result = self.q.get()
[docs]
class StatsdServer:
def __init__(self):
self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self._sock.bind(("127.0.0.1", 0))
self._sock.settimeout(0.1)
self.host, self.port = self._sock.getsockname()
self._closing = False
self._thread = threading.Thread(target=self._listen)
self._thread.start()
self.datagrams = []
self.new_datagram = threading.Event()
"""Woken up every time a datagram is added to self.datagrams."""
def _listen(self):
while not self._closing:
try:
datagram, addr = self._sock.recvfrom(4096)
except TimeoutError:
continue
self.datagrams.append(datagram)
self.new_datagram.set()
self._sock.close()
[docs]
def close(self):
self._closing = True
[docs]
@pytest.fixture(scope="session")
def vulns_statsd_server():
with contextlib.closing(StatsdServer()) as statsd_server:
yield statsd_server
[docs]
@pytest.fixture(scope="session")
def vulns_test_data_dir(tmp_path_factory) -> Path:
"""Creates an empty directory where to write the fixtures' data"""
return tmp_path_factory.mktemp("vulnerabilities")
[docs]
@pytest.fixture(scope="session")
def vulns_test_db(vulns_test_data_dir: Path) -> Path:
db_path = vulns_test_data_dir / "all.sqlite"
conn = sqlite3.connect(db_path)
conn.execute("""
CREATE TABLE vulnerabilities (
id INTEGER PRIMARY KEY,
filename TEXT NOT NULL,
doc TEXT NOT NULL
)
""")
conn.execute(
"INSERT INTO vulnerabilities (id, filename, doc) VALUES (?, ?, ?)",
(
1,
"TEST-GHSA-0001.json",
json.dumps(
{
"id": "TEST-GHSA-0001",
"modified": "2024-01-01T00:00:00Z",
"affected": [
{
"package": {"ecosystem": "npm", "name": "lodash"},
"ranges": [
{
"type": "GIT",
"events": [
{
"introduced": "0000000000000000000000000000000000000003"
},
{
"fixed": "0000000000000000000000000000000000000018"
},
],
}
],
}
],
"aliases": ["CVE-TEST-0001"],
}
),
),
)
conn.commit()
conn.close()
return db_path
[docs]
@pytest.fixture(scope="session")
def vulns_test_data(vulns_test_data_dir, vulns_test_db) -> Iterator[Path]:
graph_path = DATASET_DIR / "compressed" / "example"
kwargs = dict(
local_graph_path=graph_path.parent,
graph_name=graph_path.name,
vulnerabilities_dir=vulns_test_data_dir,
)
ComputeConnectedComponents(include_vcs_matches=False, **kwargs).run()
Vuln2CommitParquet(**kwargs).run()
Commit2VulnBv(**kwargs).run()
yield vulns_test_data_dir
[docs]
@pytest.fixture(scope="session", params=["naive"])
def vulns_grpc_server_implementation(request):
"""Returns ``"naive"``. This can be overridden in tests that want to test the
Rust implementation by returning ``"rust"``."""
return request.param
[docs]
@pytest.fixture(scope="session")
def naive_vulns_server_data() -> List[DetectedVulnerability]:
"""Data served by the :class:`swh.vulns.naive_server.NaiveVulnerabilityServicer`
spawned by the :func:`spawn_naive_grpc_server` fixture"""
from .grpc.swhvulns_pb2 import (
Vulnerability,
VulnerabilityDatabase,
VulnerabilityDetectionTool,
)
return [
DetectedVulnerability(
swhid=swhid,
vulnerability=Vulnerability(id=["TEST-GHSA-0001", "CVE-TEST-0001"]),
tool=VulnerabilityDetectionTool(name="swh-osv"),
source=VulnerabilityDatabase(name="osv.dev"),
)
for swhid in [
"swh:1:rev:0000000000000000000000000000000000000003", # introduction
"swh:1:rev:0000000000000000000000000000000000000009",
"swh:1:rev:0000000000000000000000000000000000000013",
# swh:1:rev:0000000000000000000000000000000000000018 is the fix
]
]
[docs]
@pytest.fixture(scope="session")
def vulns_grpc_server_config(
vulns_test_db, vulns_statsd_server, vulns_grpc_server_implementation, request
):
config = {
"grpc_server": {
"graph": str(DATASET_DIR / "compressed" / "example"),
"implementation": vulns_grpc_server_implementation,
"db": str(vulns_test_db),
"statsd_host": f"{vulns_statsd_server.host}:{vulns_statsd_server.port}",
}
}
if vulns_grpc_server_implementation == "rust":
vulns_test_data = request.getfixturevalue("vulns_test_data")
config["grpc_server"].update(
{
"commit2vuln": str(vulns_test_data / "commit2vuln_without_cherrypicks"),
"subgraphwccs": str(vulns_test_data / "connected_components.wccs"),
}
)
if vulns_grpc_server_implementation == "naive":
config["grpc_server"]["vulnerabilities"] = request.getfixturevalue(
"naive_vulns_server_data"
)
return config
[docs]
@pytest.fixture(scope="session")
def vulns_grpc_server_process(
vulns_grpc_server_config,
vulns_grpc_server_implementation,
vulns_statsd_server,
request,
):
server = VulnsServerProcess(vulns_grpc_server_config)
yield server
try:
server.kill()
except AttributeError:
# server was never started
pass
[docs]
@pytest.fixture(scope="session")
def vulns_grpc_server_started(vulns_grpc_server_process):
server = vulns_grpc_server_process
server.start()
if isinstance(server.result, Exception):
raise server.result
try:
# wait for the server to be up
for _ in range(100):
try:
socket.create_connection(
("localhost", server.result["port"]), timeout=1.0
)
except ConnectionRefusedError:
time.sleep(0.01)
yield server
finally:
server.kill()
AFFECTED_REVISION = "swh:1:rev:0000000000000000000000000000000000000009"
INTRODUCING_REVISION = "swh:1:rev:0000000000000000000000000000000000000003"
FIX_REVISION = "swh:1:rev:0000000000000000000000000000000000000018"
[docs]
@pytest.fixture(scope="module")
def vulns_grpc_stub(vulns_grpc_server):
with grpc.insecure_channel(vulns_grpc_server) as channel:
stub = VulnerabilityServiceStub(channel)
yield stub
[docs]
@pytest.fixture(scope="module")
def vulns_grpc_server(vulns_grpc_server_started: VulnsServerProcess):
yield vulns_grpc_server_started.result["grpc_url"]
[docs]
@pytest.fixture(scope="module")
def vulns_grpc_server_port(
vulns_grpc_server_started: VulnsServerProcess,
) -> Iterator[int]:
"""Runs a gRPC server with the data configured via :func:`naive_vulns_server_data`.
Returns the port it is listening on"""
yield vulns_grpc_server_started.result["port"]
