diff --git a/deployment/build-and-stage.yaml b/deployment/build-and-stage.yaml index fea252bbb90..6b5efd615ed 100644 --- a/deployment/build-and-stage.yaml +++ b/deployment/build-and-stage.yaml @@ -268,11 +268,19 @@ steps: waitFor: ['build-recoverer', 'cloud-build-queue'] # Build/push staging-api-test images to gcr.io/oss-vdb-test. +- name: 'gcr.io/cloud-builders/docker' + entrypoint: 'bash' + args: ['-c', 'docker pull gcr.io/oss-vdb-test/staging-api-test:latest || exit 0'] + id: 'pull-staging-api-test' + waitFor: ['setup'] - name: gcr.io/cloud-builders/docker - args: ['build', '-t', 'gcr.io/oss-vdb-test/staging-api-test:latest', '-t', 'gcr.io/oss-vdb-test/staging-api-test:$COMMIT_SHA', '.'] - dir: 'gcp/workers/staging_api_test' + args: ['buildx', 'build', '--load', '-t', 'gcr.io/oss-vdb-test/staging-api-test:latest', '-t', 'gcr.io/oss-vdb-test/staging-api-test:$COMMIT_SHA', '--target', 'staging-api-test', '--build-context', 'bindings=../bindings', '-f', 'Dockerfile', '--cache-from', 'gcr.io/oss-vdb-test/staging-api-test:latest', '--pull', '.'] + dir: 'go' id: 'build-staging-api-test' - waitFor: ['build-worker'] + waitFor: ['pull-staging-api-test', 'build-recoverer'] + volumes: + - name: 'docker-config' + path: '/root/.docker' - name: gcr.io/cloud-builders/docker args: ['push', '--all-tags', 'gcr.io/oss-vdb-test/staging-api-test'] waitFor: ['build-staging-api-test', 'cloud-build-queue'] diff --git a/gcp/workers/staging_api_test/Dockerfile b/gcp/workers/staging_api_test/Dockerfile deleted file mode 100644 index 8a0070a2301..00000000000 --- a/gcp/workers/staging_api_test/Dockerfile +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -FROM gcr.io/oss-vdb/worker - -WORKDIR /staging_api_test - -COPY retrieve_bugs_from_bucket.py perform_api_calls.py run.sh ./ - -# Add aiohttp lib -RUN cd /env/gcp/workers/worker && POETRY_VIRTUALENVS_CREATE=false poetry add aiohttp - -RUN chmod 755 retrieve_bugs_from_bucket.py perform_api_calls.py run.sh - -ENTRYPOINT ["./run.sh"] diff --git a/gcp/workers/staging_api_test/build.sh b/gcp/workers/staging_api_test/build.sh deleted file mode 100755 index 261f31a0d74..00000000000 --- a/gcp/workers/staging_api_test/build.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/bin/bash -x -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -docker build -t gcr.io/oss-vdb-test/staging-api-test:$1 . && \ -docker build -t gcr.io/oss-vdb-test/staging-api-test:latest . && \ -docker push gcr.io/oss-vdb-test/staging-api-test:$1 && \ -docker push gcr.io/oss-vdb-test/staging-api-test:latest diff --git a/gcp/workers/staging_api_test/perform_api_calls.py b/gcp/workers/staging_api_test/perform_api_calls.py deleted file mode 100755 index 54816774beb..00000000000 --- a/gcp/workers/staging_api_test/perform_api_calls.py +++ /dev/null @@ -1,508 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Mock API queries and send them to the test API endpoint for -performance testing. It is recommended to use two terminals to -run this script concurrently to generate sufficient traffic.""" - -import logging -import aiohttp -import asyncio -import os -import random -import time -import json - -from collections import Counter, defaultdict -from typing import Callable - -import osv -import osv.logs - -BASE_URL = 'https://api.test.osv.dev/v1' -GCP_PROJECT = 'oss-vdb-test' -BUG_DIR = './all_bugs' - -# Total run time in seconds -TOTAL_RUNTIME = 3600 * 5 # 5 hours -# Execute all pending batch size requests within the specified time interval. -FREQUENCY_IN_SECONDS = 1 - -# Number of `vulnerability get` requests to send per second -VULN_QUERY_BATCH_SIZE = 50 -# Number of `version query` requests to send per second -VERSION_QUERY_BATCH_SIZE = 80 -# Number of `package query` requests to send per second -PACKAGE_QUERY_BATCH_SIZE = 20 -# Number of `purl query` requests to send per second -PURL_QUERY_BATCH_SIZE = 30 -# Number of `batch query` requests to send per second -BATCH_QUERY_BATCH_SIZE = 3 -# Number of large `batch query` requests to send per second -LARGE_BATCH_QUERY_BATCH_SIZE = 2 - - -class SimpleBug: - """A simplified bug only contains essential information - for making HTTP requests.""" - - def __init__(self, bug_dict: dict): - self.db_id = bug_dict['db_id'] - # If the package/ecosystem/version value is None, then add a fake value in. - self.package = p if (p := bug_dict.get('project')) else 'foo' - self.ecosystem = e if (e := bug_dict.get('ecosystem')) else 'foo' - self.purl = p if (p := bug_dict.get('purl')) else 'pkg:foo/foo' - - # Use the `affected fuzzy` value as the query version. - # If no 'affected fuzzy' is present, assign a default value. - self.affected_fuzzy = v if (v := - bug_dict.get('affected_fuzzy')) else '1.0.0' - - -def read_from_json(filename: str, ecosystem_map: defaultdict, bug_map: dict, - package_map: defaultdict) -> None: - """Loads bugs from one JSON file into bug dicts. - - Args: - filename: the JSON filename. - - ecosystem_map: - A defaultdict mapping ecosystem names to their bugs. For example: - {'Maven': (CVE-XXXX-XXXX, CVE-XXXX-XXXX), 'PyPI': ()} - - bug_map: - A dict mapping bug ID to its `SimpleBug` object. For example: - {'CVE-XXXX-XXXX,': SimpleBug{}} - - package_map: - A defaultdict mapping package names to their bugs. For example: - {'tensorflow': (CVE-XXXX-XXXX, CVE-XXXX-XXXX), 'curl': ()} - - Returns: - None - """ - with open(filename, "r") as f: - json_file = json.load(f) - for bug_data in json_file: - bug = SimpleBug(bug_data) - ecosystem_map[bug.ecosystem].add(bug.db_id) - package_map[bug.package].add(bug.db_id) - bug_map[bug.db_id] = bug - - -def load_all_bugs() -> tuple[defaultdict, dict, defaultdict]: - """Loads bugs from JSON directory - - Returns: - A defaultdict mapping ecosystem names to their bugs. For example: - {'Maven': (CVE-XXXX-XXXX, CVE-XXXX-XXXX), 'PyPI': ()} - - A dict mapping bug ID to its `SimpleBug` object. For example: - {'CVE-XXXX-XXXX,': SimpleBug{}} - - A defaultdict mapping package names to their bugs. For example: - {'tensorflow': (CVE-XXXX-XXXX, CVE-XXXX-XXXX), 'curl': ()} - """ - - ecosystem_map = defaultdict(set) - bug_map = {} - package_map = defaultdict(set) - for filename in os.listdir(BUG_DIR): - if filename.endswith('.json'): - file_path = os.path.join(BUG_DIR, filename) - read_from_json(file_path, ecosystem_map, bug_map, package_map) - return ecosystem_map, bug_map, package_map - - -async def make_http_request(session: aiohttp.ClientSession, request_url: str, - request_type: str, request_body: dict) -> None: - """Makes one HTTP request - - Args: - session: - The HTTP ClientSession - request_url: - The HTTP request URL - request_type: - The HTTP request type: `GET` or `POST` - request_body: - The HTTP request body in JSON format - """ - try: - timeout = aiohttp.ClientTimeout(sock_connect=300, sock_read=300) - if request_type == 'GET': - async with session.get(request_url) as response: - # Await the response to ensure the server has finished - # and the connection is properly handled. - await response.read() - elif request_type == 'POST': - async with session.post( - request_url, json=request_body, timeout=timeout) as response: - # Await the response to ensure the server has finished - # and the connection is properly handled. - await response.read() - except Exception as e: - # When sending a large number of requests concurrently, - # some may fail due to timeout issues. - # These failures can be ignored as long as the server receives a - # sufficient volume of successful requests. - logging.warning('Error sending request %s with body %s: %s', request_url, - request_body, type(e)) - - -async def make_http_requests_async(request_ids: list, bug_map: dict, url: str, - batch_size: int, - payload_func: Callable) -> None: - """Makes the required number of HTTP requests per second async. - - Args: - request_ids: - A list of bug IDs - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - url: - The request URL - batch_size: - The number of requests to make per second - payload_func: - The payload function, such as `build_batch_payload` - """ - - begin_time = time.monotonic() - logging.info('[%s] Running make request %s for %d seconds', begin_time, - payload_func.__name__, TOTAL_RUNTIME) - - total_run_time = time.monotonic() - begin_time - index = 0 - length = len(request_ids) - async with aiohttp.ClientSession() as session: - while total_run_time < TOTAL_RUNTIME: - start_time = time.monotonic() - - batch_request_ids = request_ids[index:batch_size + index] - if payload_func.__name__ == build_vulnerability_payload.__name__: - for request_id in batch_request_ids: - # OSV getting vulnerability detail is a GET request - asyncio.create_task( - make_http_request(session, f'{url}/{request_id}', 'GET', - payload_func())) - elif payload_func.__name__ == build_batch_payload.__name__: - for _ in range(0, batch_size): - asyncio.create_task( - make_http_request(session, url, 'POST', - payload_func(request_ids, bug_map))) - else: - for request_id in batch_request_ids: - asyncio.create_task( - make_http_request(session, url, 'POST', - payload_func(request_id, bug_map))) - index += batch_size - if index >= length: - index = 0 - - end_time = time.monotonic() - time_elapsed = end_time - start_time - if time_elapsed < FREQUENCY_IN_SECONDS: - await asyncio.sleep(FREQUENCY_IN_SECONDS - time_elapsed) - total_run_time = time.monotonic() - begin_time - - -def build_vulnerability_payload() -> None: - """The vulnerability query doesn't need a request body""" - return None - - -def build_package_payload(request_id: str, bug_map: dict) -> dict[str, any]: - """Builds a package query payload - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - - Returns: - A dict containing package query payload, example: - '"package": {"name": "mruby","ecosystem": "OSS-Fuzz"}}' - """ - - return { - "package": { - "name": bug_map[request_id].package, - "ecosystem": bug_map[request_id].ecosystem - } - } - - -def build_version_payload(request_id: str, bug_map: dict) -> dict: - """Builds a version query payload - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - - Returns: - A dict containing package version query payload, example: - '{"package": { - "name": "mruby","ecosystem": "OSS-Fuzz"}, "version": "2.1.2rc"}' - """ - - return { - "version": bug_map[request_id].affected_fuzzy, - "package": { - "name": bug_map[request_id].package, - "ecosystem": bug_map[request_id].ecosystem - } - } - - -def build_purl_payload(request_id: str, bug_map: dict) -> dict: - """Builds a purl query payload - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - - Returns: - A dict containing package version query payload, example: - '{"package": {"purl": "pkg:golang/github.com/golang-jwt/jwt/v4@4.5.1"}}' - """ - purl = bug_map[request_id].purl - purl_with_version = f'{purl}@{bug_map[request_id].affected_fuzzy}' - - # Use random.choice to select between the two PURL options - chosen_purl = random.choice([purl, purl_with_version]) - - return {"package": {"purl": chosen_purl,}} - - -def build_batch_payload(request_ids: list, - bug_map: dict) -> dict[str, list[dict[str, any]]]: - """Builds a batch query payload - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - - Returns: - A dict containing OSV batch query payload, example: - '{ - "queries": [ - { - "package": { - ... - }, - "version": ... - }, - { - "package": { - ... - }, - "version": ... - }, - ] - }' - """ - size = random.randint(1, 100) - batch_ids = random.sample(request_ids, min(size, len(request_ids))) - queries = [] - for bug_id in batch_ids: - query = {} - query_type = random.choice(['version', 'package', 'purl']) - if query_type == 'version': - query = build_version_payload(bug_id, bug_map) - elif query_type == 'package': - query = build_package_payload(bug_id, bug_map) - elif query_type == 'purl': - query = build_purl_payload(bug_id, bug_map) - queries.append(query) - - return {"queries": [queries]} - - -def get_large_batch_query(package_map: defaultdict) -> list[str]: - """Gets a list of bug IDs for large batch queries. - This list contains bug IDs from the packages with the high - number of vulnerabilities. - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - - Returns: - A dict containing OSV batch query payload, example: - '{ - "queries": [ - { - "package": { - ... - }, - "version": ... - }, - { - "package": { - ... - }, - "version": ... - }, - ] - }' - """ - most_common = 5000 - package_counter = Counter() - for package in package_map: - # filter out invalid package name and Linux Kernel - if package in ('foo', 'Kernel'): - continue - package_counter[package] = len(package_map[package]) - most_vulnerable_packages = package_counter.most_common(most_common) - large_batch_query_ids = [] - for package, package_count in most_vulnerable_packages: - if package_count == 0: - break - large_batch_query_ids.append(package_map[package].pop()) - - random.shuffle(large_batch_query_ids) - return large_batch_query_ids - - -async def send_version_requests(request_ids: list, bug_map: dict) -> None: - """Sends version query requests - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - """ - - url = f'{BASE_URL}/query' - batch_size = VERSION_QUERY_BATCH_SIZE - await make_http_requests_async(request_ids, bug_map, url, batch_size, - build_version_payload) - - -async def send_package_requests(request_ids: list, bug_map: dict) -> None: - """Sends package query requests - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - """ - url = f'{BASE_URL}/query' - batch_size = PACKAGE_QUERY_BATCH_SIZE - await make_http_requests_async(request_ids, bug_map, url, batch_size, - build_package_payload) - - -async def send_purl_requests(request_ids: list, bug_map: dict) -> None: - """Sends purl query requests - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - """ - url = f'{BASE_URL}/query' - batch_size = PURL_QUERY_BATCH_SIZE - await make_http_requests_async(request_ids, bug_map, url, batch_size, - build_purl_payload) - - -async def send_vuln_requests(request_ids: list, bug_map: dict) -> None: - """Sends vulnerability get requests - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - """ - url = f'{BASE_URL}/vulns' - batch_size = VULN_QUERY_BATCH_SIZE - await make_http_requests_async(request_ids, bug_map, url, batch_size, - build_vulnerability_payload) - - -async def send_batch_requests(request_ids: list, bug_map: dict, - batch_size: int) -> None: - """Sends batch query requests - - Args: - request_id: - The bug ID - bug_map: - A dict mapping bug IDs to the corresponding `SimpleBug` objects - batch_size: - The batch query size - """ - url = f'{BASE_URL}/querybatch' - await make_http_requests_async(request_ids, bug_map, url, batch_size, - build_batch_payload) - - -async def main() -> None: - osv.logs.setup_gcp_logging('staging-test') - seed = random.randrange(1000) - logging.info('Random seed %d', seed) - # Log the seed value. This allows us to use the same seed later - # and reproduce this random result for debugging purposes. - random.seed(seed) - - # The `ecosystem_map` can be used to filter our queries for a - # specific ecosystem. - ecosystem_map, bug_map, package_map = load_all_bugs() - vuln_query_ids = list(bug_map.keys()) - package_query_ids = [] - for package in package_map: - # Tests each package once. - package_query_ids.append(package_map[package].pop()) - random.shuffle(package_query_ids) - random.shuffle(vuln_query_ids) - logging.info( - 'It will send vulnerability get requests for %d vulnerabilities.', - len(vuln_query_ids)) - logging.info( - 'It will send package/version/batch query requests for ' - '%d packages within %d ecosystems.', len(package_query_ids), - len(ecosystem_map)) - - # Get all packages with the most frequently occurring number - # of vulnerabilities. - large_batch_query_ids = get_large_batch_query(package_map) - - await asyncio.gather( - send_vuln_requests(vuln_query_ids, bug_map), - send_package_requests(package_query_ids, bug_map), - send_version_requests(package_query_ids, bug_map), - send_purl_requests(package_query_ids, bug_map), - send_batch_requests(package_query_ids, bug_map, BATCH_QUERY_BATCH_SIZE), - send_batch_requests(large_batch_query_ids, bug_map, - LARGE_BATCH_QUERY_BATCH_SIZE), - return_exceptions=True) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/gcp/workers/staging_api_test/retrieve_bugs_from_bucket.py b/gcp/workers/staging_api_test/retrieve_bugs_from_bucket.py deleted file mode 100644 index 63b1e0996df..00000000000 --- a/gcp/workers/staging_api_test/retrieve_bugs_from_bucket.py +++ /dev/null @@ -1,151 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Fetch Bugs from from export bucket""" - -import logging -import os -import random -import json -import sys -import tempfile -import zipfile - -import osv.logs - -from google.cloud import storage - -GCP_PROJECT = 'oss-vdb-test' -BUG_DIR = './all_bugs' -VULN_BUCKET = 'osv-test-vulnerabilities' -ZIP_FILE_PATH = 'all.zip' -ENTRIES_PER_FILE = 10000 # amount of bugs per file - - -def format_bug_for_output(bug: dict[str, any]) -> dict[str, any]: - """Extracts relevant information from a vulnerability record. - - This function processes a vulnerability record and extracts specific fields - needed for further api query usage. - - Args: - bug: a vulnerability record. - - Returns: - A dict storing all the important `Bug` fields that we want to use later - """ - if not bug.get('affected'): - return {'db_id': bug['id']} - - affected_fuzzy = None - affected = random.choice(bug['affected']) - affected_package = affected.get('package') - if not affected_package: - return {'db_id': bug['id']} - - # Store one version for use as the query version later. - if affected.get('versions'): - affected_fuzzy = random.choice(affected['versions']) - - if not affected_fuzzy and affected.get('ranges'): - range_item = random.choice(affected['ranges']) - if range_item and range_item.get('events'): - event = random.choice(range_item['events']) - if event: - affected_fuzzy = random.choice(list(event.values())) - - return { - 'db_id': bug['id'], - 'project': affected_package.get('name', None), - 'ecosystem': affected_package.get('ecosystem', None), - 'purl': affected_package.get('purl', None), - 'affected_fuzzy': affected_fuzzy, - } - - -def download_vuln_zip(tmp_dir: str) -> None: - """Downloads all.zip file from bucket.""" - zip_full_path = os.path.join(VULN_BUCKET, ZIP_FILE_PATH) - logging.info('Start to download %s.', zip_full_path) - storage_client = storage.Client() - bucket = storage_client.get_bucket(VULN_BUCKET) - try: - blob = bucket.blob(ZIP_FILE_PATH) - file_path = os.path.join(tmp_dir, ZIP_FILE_PATH) - blob.download_to_filename(file_path) - except Exception as e: - logging.exception('Failed to download all.zip: %s', e) - sys.exit(1) - logging.info('Downloaded %s.', zip_full_path) - - -def write_to_json(bug_info_list: list[dict[str, any]]) -> None: - """Writes Bugs to JSON files.""" - file_counter = 0 - for i in range(0, len(bug_info_list), ENTRIES_PER_FILE): - try: - file_name = os.path.join(BUG_DIR, f'all_bugs_{file_counter}.json') - with open(file_name, 'w') as f: - # Extract a slice of the list for the current file - end_index = min(i + ENTRIES_PER_FILE, len(bug_info_list)) - json.dump(bug_info_list[i:end_index], f, indent=2) - logging.info('Saved %d entries to %s', ENTRIES_PER_FILE, file_name) - except Exception as e: - logging.exception("Error writing to JSON file %s: %s", file_name, e) - finally: - file_counter += 1 - - -def get_bugs_from_export() -> None: - """Gets all bugs from the exported all.zip and writes to `BUG_DIR`.""" - tmp_dir = os.path.join(BUG_DIR, 'tmp') - os.makedirs(tmp_dir, exist_ok=True) - os.environ['TMPDIR'] = tmp_dir - logging.info('Start to process %s.', ZIP_FILE_PATH) - - with tempfile.TemporaryDirectory() as tmp_dir: - download_vuln_zip(tmp_dir) - all_zip = os.path.join(tmp_dir, ZIP_FILE_PATH) - bug_info_list = [] - with zipfile.ZipFile(all_zip, 'r') as vuln_zip: - for filename in vuln_zip.namelist(): - try: - with vuln_zip.open(filename) as file: - bug = json.load(file) - bug_info_list.append(format_bug_for_output(bug)) - except Exception as e: - logging.warning('Skipping invalid JSON file %s: %s', filename, e) - continue - write_to_json(bug_info_list) - logging.info('All results saved to %s.', BUG_DIR) - - -def main() -> None: - osv.logs.setup_gcp_logging('staging-test') - - if not os.path.exists(BUG_DIR): - seed = random.randrange(1000) - logging.info('Random seed %d', seed) - # Log the seed value. This allows us to use the same seed later - # and reproduce this random result for debugging purposes. - random.seed(seed) - - get_bugs_from_export() - logging.info('Fetching data finished.') - else: - logging.info('%s exists, skipping fetching.', BUG_DIR) - - -if __name__ == '__main__': - main() diff --git a/gcp/workers/staging_api_test/run.sh b/gcp/workers/staging_api_test/run.sh deleted file mode 100755 index 43a0c1392ae..00000000000 --- a/gcp/workers/staging_api_test/run.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash -x -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -if ! python3 ./retrieve_bugs_from_bucket.py; then - echo "Skipping API testing, retrieving bugs failed." - exit 1 -fi - -# `aiohttp` has limits on the number of simultaneous connections. -# Running two instances of the program in parrallel -# can help circumvent this restriction. -python3 ./perform_api_calls.py & -python3 ./perform_api_calls.py & - -# Wait for both background processes to finish -wait \ No newline at end of file diff --git a/go/Dockerfile b/go/Dockerfile index a1cbf86c585..571d2bcd27a 100644 --- a/go/Dockerfile +++ b/go/Dockerfile @@ -23,7 +23,7 @@ # docker build -t osv/importer --target importer --build-context bindings=../bindings -f Dockerfile . # # Select which service to build using the --target flag (e.g. importer, worker, exporter, -# relations, recordchecker, generatesitemap, custommetrics, gitter, first_package_finder, api, recoverer). +# relations, recordchecker, generatesitemap, custommetrics, gitter, first_package_finder, api, recoverer, staging-api-test). # ==================================================================================== # ======================================================== @@ -158,3 +158,13 @@ FROM gcr.io/distroless/static-debian12@sha256:a9fcaedd4c9b59e12dd65d954f0b5044f1 COPY --from=recoverer-build /app/recoverer / ENTRYPOINT ["/recoverer"] +# ======================================================== +# Target: Staging API Test +# ======================================================== +FROM builder AS staging-api-test-build +RUN CGO_ENABLED=0 go build -o /app/staging_api_test ./cmd/staging_api_test/ + +FROM gcr.io/distroless/static-debian12@sha256:a9fcaedd4c9b59e12dd65d954f0b5044f19b0647a8a3712e77205df9e7b102cd AS staging-api-test +COPY --from=staging-api-test-build /app/staging_api_test / +ENTRYPOINT ["/staging_api_test"] + diff --git a/go/cmd/staging_api_test/generator.go b/go/cmd/staging_api_test/generator.go new file mode 100644 index 00000000000..dd03cf00702 --- /dev/null +++ b/go/cmd/staging_api_test/generator.go @@ -0,0 +1,565 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "math/rand/v2" + "net" + "net/http" + "sync" + "sync/atomic" + "time" + + "github.com/google/osv.dev/go/logger" +) + +// GeneratorConfig holds configuration for the API load test generator. +type GeneratorConfig struct { + // BaseURL is the root URL of the OSV API under test (e.g. "https://api.test.osv.dev/v1"). + BaseURL string + + // Duration is the total execution time for the load test. + Duration time.Duration + + // VulnRate is the sustained rate of GET /v1/vulns/{id} requests per second. + // It tests the endpoint for fetching full vulnerability details by OSV ID (e.g. GHSA-..., CVE-...). + VulnRate int + + // VersionRate is the sustained rate of POST /v1/query (package + version) requests per second. + // It simulates standard scanner checks (e.g. osv-scanner) for a specific package version against affected semver/ecosystem ranges. + VersionRate int + + // PackageRate is the sustained rate of POST /v1/query (package name only) requests per second. + // It simulates unversioned queries retrieving all vulnerabilities associated with a given package. + PackageRate int + + // PURLRate is the sustained rate of POST /v1/query (Package URL) requests per second. + // It simulates queries using package URLs, randomly alternating between versioned and unversioned PURLs. + PURLRate int + + // BatchRate is the sustained rate of POST /v1/querybatch requests per second. + // Each batch request bundles up to MaxBatchQuerySize queries across all packages, simulating dependency manifest scans. + BatchRate int + + // LargeBatchRate is the sustained rate of heavy POST /v1/querybatch requests per second. + // Each request contains up to MaxLargeBatchQuerySize queries sampled strictly from the top 5,000 packages with the most vulnerabilities, + // stress-testing backend range matching and result hydration under high match volumes. + LargeBatchRate int + + // MaxBatchQuerySize is the maximum number of queries to bundle in a single standard /v1/querybatch request (1-1000). + MaxBatchQuerySize int + + // MaxLargeBatchQuerySize is the maximum number of queries to bundle in a single heavy /v1/querybatch request (1-1000). + MaxLargeBatchQuerySize int + + // StatsInterval is the interval at which progress and throughput statistics are logged. + StatsInterval time.Duration +} + +// DefaultGeneratorConfig returns standard load testing configuration matching production staging tests. +func DefaultGeneratorConfig() GeneratorConfig { + return GeneratorConfig{ + BaseURL: "https://api.test.osv.dev/v1", + Duration: 5 * time.Hour, + VulnRate: 50, + VersionRate: 80, + PackageRate: 20, + PURLRate: 30, + BatchRate: 3, + LargeBatchRate: 2, + MaxBatchQuerySize: 100, + MaxLargeBatchQuerySize: 100, + StatsInterval: 30 * time.Second, + } +} + +// QueryPayload represents an individual query in OSV API. +type QueryPayload struct { + Version string `json:"version,omitempty"` + Package *PackagePayload `json:"package,omitempty"` +} + +// PackagePayload represents package identifier in OSV API query. +type PackagePayload struct { + Name string `json:"name,omitempty"` + Ecosystem string `json:"ecosystem,omitempty"` + PURL string `json:"purl,omitempty"` +} + +// BatchQueryPayload represents payload for /v1/querybatch. +type BatchQueryPayload struct { + Queries []QueryPayload `json:"queries"` +} + +// GeneratorStats tracks request counts, outcomes, and latencies. +type GeneratorStats struct { + TotalRequests atomic.Uint64 + SuccessCount atomic.Uint64 + ClientErrorCount atomic.Uint64 + ServerErrorCount atomic.Uint64 + NetworkErrorCount atomic.Uint64 + StartTime time.Time +} + +// NewGeneratorStats initializes a new statistics tracker. +func NewGeneratorStats() *GeneratorStats { + return &GeneratorStats{ + StartTime: time.Now(), + } +} + +func (s *GeneratorStats) RecordStatus(code int) { + s.TotalRequests.Add(1) + if code >= 200 && code < 300 { + s.SuccessCount.Add(1) + } else if code >= 400 && code < 500 { + s.ClientErrorCount.Add(1) + } else { + s.ServerErrorCount.Add(1) + } +} + +func (s *GeneratorStats) RecordNetworkError() { + s.TotalRequests.Add(1) + s.NetworkErrorCount.Add(1) +} + +func (s *GeneratorStats) Summary() string { + elapsed := time.Since(s.StartTime) + total := s.TotalRequests.Load() + success := s.SuccessCount.Load() + clientErr := s.ClientErrorCount.Load() + serverErr := s.ServerErrorCount.Load() + netErr := s.NetworkErrorCount.Load() + + rps := 0.0 + if elapsed.Seconds() > 0 { + rps = float64(total) / elapsed.Seconds() + } + + return fmt.Sprintf("Elapsed: %s | Total: %d (%.1f req/s) | 2xx: %d | 4xx: %d | 5xx: %d | NetErr: %d", + elapsed.Truncate(time.Second), total, rps, success, clientErr, serverErr, netErr) +} + +// buildPackagePayload constructs payload for package query. +func buildPackagePayload(id string, vulnMap map[string]*SimpleVuln) ([]byte, error) { + vuln, ok := vulnMap[id] + if !ok { + return nil, fmt.Errorf("vulnerability not found: %s", id) + } + + return json.Marshal(QueryPayload{ + Package: &PackagePayload{ + Name: vuln.Package, + Ecosystem: vuln.Ecosystem, + }, + }) +} + +// buildVersionPayload constructs payload for package version query. +func buildVersionPayload(id string, vulnMap map[string]*SimpleVuln) ([]byte, error) { + vuln, ok := vulnMap[id] + if !ok { + return nil, fmt.Errorf("vulnerability not found: %s", id) + } + + return json.Marshal(QueryPayload{ + Version: vuln.AffectedFuzzy, + Package: &PackagePayload{ + Name: vuln.Package, + Ecosystem: vuln.Ecosystem, + }, + }) +} + +// buildPURLPayload constructs payload for purl query. +func buildPURLPayload(rng *rand.Rand, id string, vulnMap map[string]*SimpleVuln) ([]byte, error) { + vuln, ok := vulnMap[id] + if !ok { + return nil, fmt.Errorf("vulnerability not found: %s", id) + } + + chosenPURL := vuln.PURL + if rng.IntN(2) == 0 { + chosenPURL = fmt.Sprintf("%s@%s", vuln.PURL, vuln.AffectedFuzzy) + } + + return json.Marshal(QueryPayload{ + Package: &PackagePayload{ + PURL: chosenPURL, + }, + }) +} + +// buildBatchPayload constructs payload for batch queries. +func buildBatchPayload(rng *rand.Rand, requestIDs []string, vulnMap map[string]*SimpleVuln, maxBatchQueries int) ([]byte, error) { + if len(requestIDs) == 0 { + return json.Marshal(BatchQueryPayload{Queries: []QueryPayload{}}) + } + + if maxBatchQueries <= 0 { + maxBatchQueries = 100 + } + + sampleSize := rng.IntN(maxBatchQueries) + 1 + if sampleSize > len(requestIDs) { + sampleSize = len(requestIDs) + } + + // Sample random IDs + perm := rng.Perm(len(requestIDs)) + queries := make([]QueryPayload, 0, sampleSize) + + for i := range sampleSize { + vulnID := requestIDs[perm[i]] + vuln, ok := vulnMap[vulnID] + if !ok { + continue + } + + queryType := rng.IntN(3) + switch queryType { + case 0: // version + queries = append(queries, QueryPayload{ + Version: vuln.AffectedFuzzy, + Package: &PackagePayload{ + Name: vuln.Package, + Ecosystem: vuln.Ecosystem, + }, + }) + case 1: // package + queries = append(queries, QueryPayload{ + Package: &PackagePayload{ + Name: vuln.Package, + Ecosystem: vuln.Ecosystem, + }, + }) + case 2: // purl + chosenPURL := vuln.PURL + if rng.IntN(2) == 0 { + chosenPURL = fmt.Sprintf("%s@%s", vuln.PURL, vuln.AffectedFuzzy) + } + queries = append(queries, QueryPayload{ + Package: &PackagePayload{ + PURL: chosenPURL, + }, + }) + } + } + + return json.Marshal(BatchQueryPayload{Queries: queries}) +} + +// newHTTPClient creates an HTTP client optimized for high concurrency load testing. +func newHTTPClient() *http.Client { + transport := &http.Transport{ + Proxy: http.ProxyFromEnvironment, + DialContext: (&net.Dialer{ + Timeout: 30 * time.Second, + KeepAlive: 30 * time.Second, + }).DialContext, + ForceAttemptHTTP2: true, + MaxIdleConns: 2000, + MaxIdleConnsPerHost: 1000, + MaxConnsPerHost: 1000, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } + + return &http.Client{ + Transport: transport, + Timeout: 300 * time.Second, + } +} + +// executeRequest sends an HTTP request and updates stats. +func executeRequest(ctx context.Context, client *http.Client, req *http.Request, stats *GeneratorStats) { + req = req.WithContext(ctx) + resp, err := client.Do(req) + if err != nil { + if ctx.Err() == nil { + logger.Warn("Failed to send HTTP request", + "url", req.URL.String(), + "method", req.Method, + "error", err, + ) + } + stats.RecordNetworkError() + + return + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + bodySnippet, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + logger.Warn("HTTP request returned error status", + "url", req.URL.String(), + "method", req.Method, + "statusCode", resp.StatusCode, + "response", string(bodySnippet), + ) + } + + // Drain response body to enable connection reuse + _, _ = io.Copy(io.Discard, resp.Body) + stats.RecordStatus(resp.StatusCode) +} + +// runVulnWorker dispatches GET /v1/vulns/{id} requests evenly spaced at cfg.VulnRate per second. +func runVulnWorker(ctx context.Context, wg *sync.WaitGroup, client *http.Client, pools *QueryPools, stats *GeneratorStats, cfg GeneratorConfig) { + if len(pools.VulnQueryIDs) == 0 || cfg.VulnRate <= 0 { + return + } + + ticker := time.NewTicker(time.Second / time.Duration(cfg.VulnRate)) + defer ticker.Stop() + + index := 0 + length := len(pools.VulnQueryIDs) + baseURL := cfg.BaseURL + "/vulns" + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + reqID := pools.VulnQueryIDs[index%length] + index++ + url := fmt.Sprintf("%s/%s", baseURL, reqID) + req, err := http.NewRequest(http.MethodGet, url, nil) + if err != nil { + continue + } + wg.Go(func() { + executeRequest(ctx, client, req, stats) + }) + } + } +} + +// runSingleQueryWorker dispatches POST /v1/query requests evenly spaced at the specified rate per second. +func runSingleQueryWorker( + ctx context.Context, + wg *sync.WaitGroup, + client *http.Client, + requestIDs []string, + vulnMap map[string]*SimpleVuln, + stats *GeneratorStats, + rate int, + cfg GeneratorConfig, + payloadBuilder func(string, map[string]*SimpleVuln) ([]byte, error), +) { + if len(requestIDs) == 0 || rate <= 0 { + return + } + + ticker := time.NewTicker(time.Second / time.Duration(rate)) + defer ticker.Stop() + + index := 0 + length := len(requestIDs) + url := cfg.BaseURL + "/query" + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + reqID := requestIDs[index%length] + index++ + payload, err := payloadBuilder(reqID, vulnMap) + if err != nil { + continue + } + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + continue + } + req.Header.Set("Content-Type", "application/json") + wg.Go(func() { + executeRequest(ctx, client, req, stats) + }) + } + } +} + +// runPURLWorker dispatches POST /v1/query (PURL) requests evenly spaced at cfg.PURLRate per second. +func runPURLWorker( + ctx context.Context, + wg *sync.WaitGroup, + client *http.Client, + requestIDs []string, + vulnMap map[string]*SimpleVuln, + stats *GeneratorStats, + cfg GeneratorConfig, + seed uint64, +) { + if len(requestIDs) == 0 || cfg.PURLRate <= 0 { + return + } + + //nolint:gosec // math/rand is sufficient for mock traffic generation + rng := rand.New(rand.NewPCG(seed, seed^0x12345)) + ticker := time.NewTicker(time.Second / time.Duration(cfg.PURLRate)) + defer ticker.Stop() + + index := 0 + length := len(requestIDs) + url := cfg.BaseURL + "/query" + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + reqID := requestIDs[index%length] + index++ + payload, err := buildPURLPayload(rng, reqID, vulnMap) + if err != nil { + continue + } + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + continue + } + req.Header.Set("Content-Type", "application/json") + wg.Go(func() { + executeRequest(ctx, client, req, stats) + }) + } + } +} + +// runBatchWorker dispatches POST /v1/querybatch requests evenly spaced at the specified batchRate per second. +func runBatchWorker( + ctx context.Context, + wg *sync.WaitGroup, + client *http.Client, + requestIDs []string, + vulnMap map[string]*SimpleVuln, + stats *GeneratorStats, + batchRate int, + maxBatchQueries int, + cfg GeneratorConfig, + seed uint64, +) { + if len(requestIDs) == 0 || batchRate <= 0 { + return + } + + //nolint:gosec // math/rand is sufficient for mock traffic generation + rng := rand.New(rand.NewPCG(seed, seed^0x6789A)) + ticker := time.NewTicker(time.Second / time.Duration(batchRate)) + defer ticker.Stop() + + url := cfg.BaseURL + "/querybatch" + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + payload, err := buildBatchPayload(rng, requestIDs, vulnMap, maxBatchQueries) + if err != nil { + continue + } + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + continue + } + req.Header.Set("Content-Type", "application/json") + wg.Go(func() { + executeRequest(ctx, client, req, stats) + }) + } + } +} + +// RunTrafficGenerator orchestrates concurrent traffic generation across all query types. +func RunTrafficGenerator(ctx context.Context, client *http.Client, pools *QueryPools, cfg GeneratorConfig, seed uint64) (*GeneratorStats, error) { + if client == nil { + client = newHTTPClient() + } + + stats := NewGeneratorStats() + + logger.Info("Starting API traffic generator with smooth rate pacing", + "baseURL", cfg.BaseURL, + "duration", cfg.Duration, + "vulnRate", cfg.VulnRate, + "versionRate", cfg.VersionRate, + "packageRate", cfg.PackageRate, + "purlRate", cfg.PURLRate, + "batchRate", cfg.BatchRate, + "largeBatchRate", cfg.LargeBatchRate, + "maxBatchQuerySize", cfg.MaxBatchQuerySize, + "maxLargeBatchQuerySize", cfg.MaxLargeBatchQuerySize, + "totalVulns", len(pools.VulnQueryIDs), + "totalPackages", len(pools.PackageQueryIDs), + "totalLargeBatchPackages", len(pools.LargeBatchQueryIDs), + ) + + genCtx, cancel := context.WithTimeout(ctx, cfg.Duration) + defer cancel() + + var wg sync.WaitGroup + + // Start workers with smooth pacing and isolated rng seeds per worker + wg.Go(func() { + runVulnWorker(genCtx, &wg, client, pools, stats, cfg) + }) + wg.Go(func() { + runSingleQueryWorker(genCtx, &wg, client, pools.PackageQueryIDs, pools.VulnMap, stats, cfg.PackageRate, cfg, buildPackagePayload) + }) + wg.Go(func() { + runSingleQueryWorker(genCtx, &wg, client, pools.PackageQueryIDs, pools.VulnMap, stats, cfg.VersionRate, cfg, buildVersionPayload) + }) + wg.Go(func() { + runPURLWorker(genCtx, &wg, client, pools.PackageQueryIDs, pools.VulnMap, stats, cfg, seed+1) + }) + wg.Go(func() { + runBatchWorker(genCtx, &wg, client, pools.PackageQueryIDs, pools.VulnMap, stats, cfg.BatchRate, cfg.MaxBatchQuerySize, cfg, seed+2) + }) + wg.Go(func() { + runBatchWorker(genCtx, &wg, client, pools.LargeBatchQueryIDs, pools.VulnMap, stats, cfg.LargeBatchRate, cfg.MaxLargeBatchQuerySize, cfg, seed+3) + }) + + // Periodic stats reporter + statsTicker := time.NewTicker(cfg.StatsInterval) + defer statsTicker.Stop() + + doneCh := make(chan struct{}) + go func() { + wg.Wait() + close(doneCh) + }() + + for { + select { + case <-statsTicker.C: + logger.Info("[Stats] " + stats.Summary()) + case <-doneCh: + logger.Info("[Final Stats] " + stats.Summary()) + + return stats, nil + } + } +} diff --git a/go/cmd/staging_api_test/generator_test.go b/go/cmd/staging_api_test/generator_test.go new file mode 100644 index 00000000000..f6dbe125ede --- /dev/null +++ b/go/cmd/staging_api_test/generator_test.go @@ -0,0 +1,174 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "encoding/json" + "math/rand/v2" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestPayloadBuilders(t *testing.T) { + //nolint:gosec // math/rand is sufficient for mock traffic generation + rng := rand.New(rand.NewPCG(42, 42)) + vulnMap := map[string]*SimpleVuln{ + "OSV-001": { + ID: "OSV-001", + Package: "django", + Ecosystem: "PyPI", + PURL: "pkg:pypi/django", + AffectedFuzzy: "3.2.1", + }, + } + + t.Run("PackagePayload", func(t *testing.T) { + payload, err := buildPackagePayload("OSV-001", vulnMap) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + var q QueryPayload + if err := json.Unmarshal(payload, &q); err != nil { + t.Fatalf("failed to unmarshal payload: %v", err) + } + if q.Package == nil || q.Package.Name != "django" || q.Package.Ecosystem != "PyPI" { + t.Errorf("unexpected package payload: %+v", q) + } + }) + + t.Run("VersionPayload", func(t *testing.T) { + payload, err := buildVersionPayload("OSV-001", vulnMap) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + var q QueryPayload + if err := json.Unmarshal(payload, &q); err != nil { + t.Fatalf("failed to unmarshal payload: %v", err) + } + if q.Version != "3.2.1" || q.Package == nil || q.Package.Name != "django" { + t.Errorf("unexpected version payload: %+v", q) + } + }) + + t.Run("PURLPayload", func(t *testing.T) { + payload, err := buildPURLPayload(rng, "OSV-001", vulnMap) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + var q QueryPayload + if err := json.Unmarshal(payload, &q); err != nil { + t.Fatalf("failed to unmarshal payload: %v", err) + } + if q.Package == nil || (!strings.HasPrefix(q.Package.PURL, "pkg:pypi/django")) { + t.Errorf("unexpected purl payload: %+v", q) + } + }) + + t.Run("BatchPayload", func(t *testing.T) { + payload, err := buildBatchPayload(rng, []string{"OSV-001"}, vulnMap, 10) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + var b BatchQueryPayload + if err := json.Unmarshal(payload, &b); err != nil { + t.Fatalf("failed to unmarshal batch payload: %v", err) + } + if len(b.Queries) != 1 { + t.Errorf("expected 1 query in batch, got %d", len(b.Queries)) + } + }) +} + +func TestRunTrafficGenerator(t *testing.T) { + var vulnCalls atomic.Int64 + var queryCalls atomic.Int64 + var batchCalls atomic.Int64 + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + path := r.URL.Path + switch { + case strings.HasPrefix(path, "/vulns/"): + vulnCalls.Add(1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"id":"OSV-001"}`)) + case path == "/query": + queryCalls.Add(1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"vulns":[]}`)) + case path == "/querybatch": + batchCalls.Add(1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"results":[]}`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer ts.Close() + + pools := &QueryPools{ + VulnMap: map[string]*SimpleVuln{ + "OSV-001": { + ID: "OSV-001", + Package: "express", + Ecosystem: "npm", + PURL: "pkg:npm/express", + AffectedFuzzy: "4.17.1", + }, + }, + VulnQueryIDs: []string{"OSV-001"}, + PackageQueryIDs: []string{"OSV-001"}, + LargeBatchQueryIDs: []string{"OSV-001"}, + } + + cfg := GeneratorConfig{ + BaseURL: ts.URL, + Duration: 200 * time.Millisecond, + VulnRate: 20, + VersionRate: 20, + PackageRate: 20, + PURLRate: 20, + BatchRate: 10, + LargeBatchRate: 10, + MaxBatchQuerySize: 10, + MaxLargeBatchQuerySize: 10, + StatsInterval: 100 * time.Millisecond, + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + stats, err := RunTrafficGenerator(ctx, ts.Client(), pools, cfg, 42) + if err != nil { + t.Fatalf("unexpected error running generator: %v", err) + } + + if stats.SuccessCount.Load() == 0 { + t.Errorf("expected at least 1 successful request, got %d", stats.SuccessCount.Load()) + } + if vulnCalls.Load() == 0 { + t.Errorf("expected vulnCalls > 0, got %d", vulnCalls.Load()) + } + if queryCalls.Load() == 0 { + t.Errorf("expected queryCalls > 0, got %d", queryCalls.Load()) + } + if batchCalls.Load() == 0 { + t.Errorf("expected batchCalls > 0, got %d", batchCalls.Load()) + } +} diff --git a/go/cmd/staging_api_test/loader.go b/go/cmd/staging_api_test/loader.go new file mode 100644 index 00000000000..3283a81927d --- /dev/null +++ b/go/cmd/staging_api_test/loader.go @@ -0,0 +1,291 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "archive/zip" + "context" + "encoding/json" + "fmt" + "io" + "math/rand/v2" + "os" + "path/filepath" + "sort" + "strings" + + "cloud.google.com/go/storage" + "github.com/google/osv.dev/go/logger" +) + +// SimpleVuln contains essential information extracted from a vulnerability record for API querying. +type SimpleVuln struct { + ID string `json:"db_id"` + Package string `json:"package"` + Ecosystem string `json:"ecosystem"` + PURL string `json:"purl"` + AffectedFuzzy string `json:"affected_fuzzy"` +} + +// rawVulnerability is a lightweight struct for unmarshaling OSV records from JSON. +type rawVulnerability struct { + ID string `json:"id"` + Affected []rawAffected `json:"affected"` +} + +type rawAffected struct { + Package *rawPackage `json:"package"` + Versions []string `json:"versions"` + Ranges []rawRange `json:"ranges"` +} + +type rawPackage struct { + Name string `json:"name"` + Ecosystem string `json:"ecosystem"` + PURL string `json:"purl"` +} + +type rawRange struct { + Type string `json:"type"` + Events []map[string]string `json:"events"` +} + +// QueryPools holds processed vulnerabilities and partitioned query IDs for API testing. +type QueryPools struct { + VulnMap map[string]*SimpleVuln + EcosystemMap map[string][]string + PackageMap map[string][]string + VulnQueryIDs []string + PackageQueryIDs []string + LargeBatchQueryIDs []string +} + +// formatVuln extracts essential query fields from raw vulnerability JSON. +func formatVuln(rng *rand.Rand, raw rawVulnerability) *SimpleVuln { + vuln := &SimpleVuln{ + ID: raw.ID, + Package: "foo", + Ecosystem: "foo", + PURL: "pkg:foo/foo", + AffectedFuzzy: "1.0.0", + } + + if len(raw.Affected) == 0 { + return vuln + } + + aff := raw.Affected[rng.IntN(len(raw.Affected))] + if aff.Package == nil { + return vuln + } + + if aff.Package.Name != "" { + vuln.Package = aff.Package.Name + } + if aff.Package.Ecosystem != "" { + vuln.Ecosystem = aff.Package.Ecosystem + } + if aff.Package.PURL != "" { + vuln.PURL = aff.Package.PURL + } + + var affectedFuzzy string + if len(aff.Versions) > 0 { + affectedFuzzy = aff.Versions[rng.IntN(len(aff.Versions))] + } + + if affectedFuzzy == "" && len(aff.Ranges) > 0 { + rangeItem := aff.Ranges[rng.IntN(len(aff.Ranges))] + if len(rangeItem.Events) > 0 { + event := rangeItem.Events[rng.IntN(len(rangeItem.Events))] + var values []string + for _, v := range event { + if v != "" { + values = append(values, v) + } + } + if len(values) > 0 { + affectedFuzzy = values[rng.IntN(len(values))] + } + } + } + + if affectedFuzzy != "" { + vuln.AffectedFuzzy = affectedFuzzy + } + + return vuln +} + +// downloadZipFromGCS downloads a zip file from GCS to a temporary local file. +func downloadZipFromGCS(ctx context.Context, client *storage.Client, bucket, objectName string) (string, error) { + tmpFile, err := os.CreateTemp("", "all-vulns-*.zip") + if err != nil { + return "", fmt.Errorf("failed to create temporary file: %w", err) + } + defer tmpFile.Close() + + rc, err := client.Bucket(bucket).Object(objectName).NewReader(ctx) + if err != nil { + _ = os.Remove(tmpFile.Name()) + + return "", fmt.Errorf("failed to read %s/%s from GCS: %w", bucket, objectName, err) + } + defer rc.Close() + + if _, err := io.Copy(tmpFile, rc); err != nil { + _ = os.Remove(tmpFile.Name()) + + return "", fmt.Errorf("failed to download zip from GCS: %w", err) + } + + return tmpFile.Name(), nil +} + +// LoadQueryPoolsFromZip reads all vulnerability JSON files from a zip reader and constructs the query pools. +func LoadQueryPoolsFromZip(ctx context.Context, zipReader *zip.Reader, rng *rand.Rand) (*QueryPools, error) { + vulnMap := make(map[string]*SimpleVuln) + ecosystemMap := make(map[string][]string) + packageMap := make(map[string][]string) + + for _, file := range zipReader.File { + if selectErr := ctx.Err(); selectErr != nil { + return nil, selectErr + } + if !strings.HasSuffix(file.Name, ".json") { + continue + } + + rc, err := file.Open() + if err != nil { + logger.Warn("Failed to open file in zip", "file", file.Name, "error", err) + continue + } + + var raw rawVulnerability + err = json.NewDecoder(rc).Decode(&raw) + _ = rc.Close() + if err != nil || raw.ID == "" { + logger.Warn("Skipping invalid JSON file in zip", "file", file.Name, "error", err) + continue + } + + vuln := formatVuln(rng, raw) + vulnMap[vuln.ID] = vuln + ecosystemMap[vuln.Ecosystem] = append(ecosystemMap[vuln.Ecosystem], vuln.ID) + packageMap[vuln.Package] = append(packageMap[vuln.Package], vuln.ID) + } + + vulnQueryIDs := make([]string, 0, len(vulnMap)) + for id := range vulnMap { + vulnQueryIDs = append(vulnQueryIDs, id) + } + rng.Shuffle(len(vulnQueryIDs), func(i, j int) { + vulnQueryIDs[i], vulnQueryIDs[j] = vulnQueryIDs[j], vulnQueryIDs[i] + }) + + // Make copies of package lists so we can pop from them without destroying the original packageMap + pkgLists := make(map[string][]string, len(packageMap)) + for k, v := range packageMap { + cp := make([]string, len(v)) + copy(cp, v) + pkgLists[k] = cp + } + + var packageQueryIDs []string + for pkg, ids := range pkgLists { + if len(ids) > 0 { + packageQueryIDs = append(packageQueryIDs, ids[len(ids)-1]) + pkgLists[pkg] = ids[:len(ids)-1] + } + } + rng.Shuffle(len(packageQueryIDs), func(i, j int) { + packageQueryIDs[i], packageQueryIDs[j] = packageQueryIDs[j], packageQueryIDs[i] + }) + + // Get large batch query IDs (from top 5000 packages with most vulnerabilities, excluding foo and Kernel) + type pkgCount struct { + pkg string + count int + } + var counts []pkgCount + for pkg, ids := range packageMap { + if pkg == "foo" || pkg == "Kernel" { + continue + } + if len(ids) > 0 { + counts = append(counts, pkgCount{pkg: pkg, count: len(ids)}) + } + } + sort.Slice(counts, func(i, j int) bool { + return counts[i].count > counts[j].count + }) + + const mostCommon = 5000 + limit := mostCommon + if len(counts) < limit { + limit = len(counts) + } + + var largeBatchQueryIDs []string + for i := range limit { + pkg := counts[i].pkg + ids := pkgLists[pkg] + if len(ids) > 0 { + largeBatchQueryIDs = append(largeBatchQueryIDs, ids[len(ids)-1]) + pkgLists[pkg] = ids[:len(ids)-1] + } + } + rng.Shuffle(len(largeBatchQueryIDs), func(i, j int) { + largeBatchQueryIDs[i], largeBatchQueryIDs[j] = largeBatchQueryIDs[j], largeBatchQueryIDs[i] + }) + + return &QueryPools{ + VulnMap: vulnMap, + EcosystemMap: ecosystemMap, + PackageMap: packageMap, + VulnQueryIDs: vulnQueryIDs, + PackageQueryIDs: packageQueryIDs, + LargeBatchQueryIDs: largeBatchQueryIDs, + }, nil +} + +// LoadQueryPools downloads (if needed) and parses vulnerability records from GCS or a local zip file. +func LoadQueryPools(ctx context.Context, gcsClient *storage.Client, bucket, zipPath, localZip string, rng *rand.Rand) (*QueryPools, error) { + zipFilePath := localZip + isTemp := false + if zipFilePath == "" { + logger.Info("Downloading zip file from GCS", "bucket", bucket, "object", zipPath) + var err error + zipFilePath, err = downloadZipFromGCS(ctx, gcsClient, bucket, zipPath) + if err != nil { + return nil, err + } + isTemp = true + } + if isTemp { + defer func() { + _ = os.Remove(zipFilePath) + }() + } + + zipReader, err := zip.OpenReader(zipFilePath) + if err != nil { + return nil, fmt.Errorf("failed to open zip file %s: %w", filepath.Clean(zipFilePath), err) + } + defer zipReader.Close() + + return LoadQueryPoolsFromZip(ctx, &zipReader.Reader, rng) +} diff --git a/go/cmd/staging_api_test/loader_test.go b/go/cmd/staging_api_test/loader_test.go new file mode 100644 index 00000000000..3c49348c697 --- /dev/null +++ b/go/cmd/staging_api_test/loader_test.go @@ -0,0 +1,189 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "archive/zip" + "bytes" + "context" + "math/rand/v2" + "testing" +) + +func TestFormatVuln(t *testing.T) { + //nolint:gosec // math/rand is sufficient for mock traffic generation + rng := rand.New(rand.NewPCG(42, 42)) + + t.Run("NoAffected", func(t *testing.T) { + raw := rawVulnerability{ID: "TEST-001"} + vuln := formatVuln(rng, raw) + if vuln.ID != "TEST-001" { + t.Errorf("expected ID TEST-001, got %s", vuln.ID) + } + if vuln.Package != "foo" || vuln.Ecosystem != "foo" || vuln.PURL != "pkg:foo/foo" || vuln.AffectedFuzzy != "1.0.0" { + t.Errorf("unexpected defaults for vuln: %+v", vuln) + } + }) + + t.Run("WithVersions", func(t *testing.T) { + raw := rawVulnerability{ + ID: "TEST-002", + Affected: []rawAffected{ + { + Package: &rawPackage{ + Name: "mypackage", + Ecosystem: "PyPI", + PURL: "pkg:pypi/mypackage", + }, + Versions: []string{"1.2.3"}, + }, + }, + } + vuln := formatVuln(rng, raw) + if vuln.ID != "TEST-002" { + t.Errorf("expected ID TEST-002, got %s", vuln.ID) + } + if vuln.Package != "mypackage" || vuln.Ecosystem != "PyPI" || vuln.PURL != "pkg:pypi/mypackage" { + t.Errorf("unexpected package details: %+v", vuln) + } + if vuln.AffectedFuzzy != "1.2.3" { + t.Errorf("expected AffectedFuzzy 1.2.3, got %s", vuln.AffectedFuzzy) + } + }) + + t.Run("WithRangesEvents", func(t *testing.T) { + raw := rawVulnerability{ + ID: "TEST-003", + Affected: []rawAffected{ + { + Package: &rawPackage{ + Name: "curl", + Ecosystem: "Debian", + }, + Ranges: []rawRange{ + { + Type: "ECOSYSTEM", + Events: []map[string]string{ + {"introduced": "0", "fixed": "7.88.1"}, + }, + }, + }, + }, + }, + } + vuln := formatVuln(rng, raw) + if vuln.ID != "TEST-003" { + t.Errorf("expected ID TEST-003, got %s", vuln.ID) + } + if vuln.Package != "curl" || vuln.Ecosystem != "Debian" { + t.Errorf("unexpected package details: %+v", vuln) + } + if vuln.AffectedFuzzy != "0" && vuln.AffectedFuzzy != "7.88.1" { + t.Errorf("expected fuzzy version from event, got %s", vuln.AffectedFuzzy) + } + }) +} + +func createTestZip(t *testing.T, files map[string]string) *zip.Reader { + t.Helper() + buf := new(bytes.Buffer) + zw := zip.NewWriter(buf) + + for name, content := range files { + w, err := zw.Create(name) + if err != nil { + t.Fatalf("failed to create zip file entry: %v", err) + } + if _, err := w.Write([]byte(content)); err != nil { + t.Fatalf("failed to write zip file entry content: %v", err) + } + } + + if err := zw.Close(); err != nil { + t.Fatalf("failed to close zip writer: %v", err) + } + + reader, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len())) + if err != nil { + t.Fatalf("failed to create zip reader: %v", err) + } + + return reader +} + +func TestLoadQueryPoolsFromZip(t *testing.T) { + files := map[string]string{ + "vuln1.json": `{ + "id": "OSV-2024-001", + "affected": [{ + "package": {"name": "pkg-a", "ecosystem": "npm", "purl": "pkg:npm/pkg-a"}, + "versions": ["1.0.0", "1.0.1"] + }] + }`, + "vuln2.json": `{ + "id": "OSV-2024-002", + "affected": [{ + "package": {"name": "pkg-a", "ecosystem": "npm"}, + "versions": ["2.0.0"] + }] + }`, + "vuln3.json": `{ + "id": "OSV-2024-003", + "affected": [{ + "package": {"name": "pkg-b", "ecosystem": "PyPI"}, + "versions": ["0.1.0"] + }] + }`, + "vuln4.json": `{ + "id": "OSV-2024-004", + "affected": [{ + "package": {"name": "Kernel", "ecosystem": "Linux"}, + "versions": ["5.10.0"] + }] + }`, + "invalid.json": `not valid json`, + "readme.txt": `hello world`, + } + + zipReader := createTestZip(t, files) + //nolint:gosec // math/rand is sufficient for mock traffic generation + rng := rand.New(rand.NewPCG(123, 123)) + + pools, err := LoadQueryPoolsFromZip(context.Background(), zipReader, rng) + if err != nil { + t.Fatalf("unexpected error loading pools: %v", err) + } + + if len(pools.VulnMap) != 4 { + t.Errorf("expected 4 vulns in VulnMap, got %d", len(pools.VulnMap)) + } + if len(pools.VulnQueryIDs) != 4 { + t.Errorf("expected 4 vulnQueryIDs, got %d", len(pools.VulnQueryIDs)) + } + + // Package query IDs should have 1 entry per distinct package + expectedPackages := map[string]bool{"pkg-a": true, "pkg-b": true, "Kernel": true} + if len(pools.PackageQueryIDs) != len(expectedPackages) { + t.Errorf("expected %d packageQueryIDs, got %d", len(expectedPackages), len(pools.PackageQueryIDs)) + } + + // Large batch queries should exclude "Kernel" and "foo" + for _, id := range pools.LargeBatchQueryIDs { + vuln := pools.VulnMap[id] + if vuln.Package == "Kernel" || vuln.Package == "foo" { + t.Errorf("expected Kernel/foo to be excluded from LargeBatchQueryIDs, found %+v", vuln) + } + } +} diff --git a/go/cmd/staging_api_test/main.go b/go/cmd/staging_api_test/main.go new file mode 100644 index 00000000000..0e2a2d3e44d --- /dev/null +++ b/go/cmd/staging_api_test/main.go @@ -0,0 +1,136 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package main runs staging API performance and load tests by fetching vulnerability +// data from the exported bucket and generating concurrent traffic across OSV endpoints. +package main + +import ( + "context" + "flag" + "fmt" + "math/rand/v2" + "os" + "os/signal" + "syscall" + "time" + + "cloud.google.com/go/storage" + "github.com/google/osv.dev/go/logger" +) + +func envOrDefault(key, fallback string) string { + if val := os.Getenv(key); val != "" { + return val + } + + return fallback +} + +func main() { + if err := run(); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } +} + +func run() error { + logger.InitGlobalLogger() + defer logger.Close() + + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + defaultBucket := envOrDefault("OSV_VULNERABILITIES_BUCKET", "osv-test-vulnerabilities") + defaultAPIURL := envOrDefault("API_BASE_URL", "https://api.test.osv.dev/v1") + defaultZipPath := envOrDefault("ZIP_FILE_PATH", "all.zip") + + bucketFlag := flag.String("bucket", defaultBucket, "GCS bucket to read exported vulnerabilities from") + zipPathFlag := flag.String("zip-path", defaultZipPath, "Path of all.zip inside the GCS bucket") + localZipFlag := flag.String("local-zip", os.Getenv("LOCAL_ZIP"), "Optional local path to all.zip (skips GCS download if specified)") + apiURLFlag := flag.String("api-url", defaultAPIURL, "Base URL for the OSV API under test") + durationFlag := flag.Duration("duration", 5*time.Hour, "Total run duration of the load test (e.g. 5h, 30m)") + seedFlag := flag.Uint64("seed", 0, "Random seed (0 to generate a random seed)") + + vulnRateFlag := flag.Int("vuln-rate", 50, "Number of GET /v1/vulns/{id} requests per second") + versionRateFlag := flag.Int("version-rate", 80, "Number of POST /v1/query (version) requests per second") + packageRateFlag := flag.Int("package-rate", 20, "Number of POST /v1/query (package) requests per second") + purlRateFlag := flag.Int("purl-rate", 30, "Number of POST /v1/query (purl) requests per second") + batchRateFlag := flag.Int("batch-rate", 3, "Number of POST /v1/querybatch (normal) requests per second") + largeBatchRateFlag := flag.Int("large-batch-rate", 2, "Number of POST /v1/querybatch (large) requests per second") + maxBatchSizeFlag := flag.Int("max-batch-size", 100, "Maximum number of queries per standard /v1/querybatch request (1-1000)") + maxLargeBatchSizeFlag := flag.Int("max-large-batch-size", 100, "Maximum number of queries per heavy /v1/querybatch request (1-1000)") + statsIntervalFlag := flag.Duration("stats-interval", 30*time.Second, "Interval between logging summary statistics") + + flag.Parse() + + seed := *seedFlag + if seed == 0 { + seed = rand.Uint64() //nolint:gosec // math/rand is sufficient for mock traffic generation + } + logger.Info("Starting staging API test", "seed", seed) + //nolint:gosec // math/rand is sufficient for mock traffic generation + rng := rand.New(rand.NewPCG(seed, seed^0x5DEECE66D)) + + var gcsClient *storage.Client + if *localZipFlag == "" { + var err error + gcsClient, err = storage.NewClient(ctx) + if err != nil { + logger.Error("Failed to create GCS client", "error", err) + + return fmt.Errorf("failed to create GCS client: %w", err) + } + defer gcsClient.Close() + } + + pools, err := LoadQueryPools(ctx, gcsClient, *bucketFlag, *zipPathFlag, *localZipFlag, rng) + if err != nil { + logger.Error("Failed to load query pools from vulnerability records", "error", err) + + return fmt.Errorf("failed to load query pools: %w", err) + } + + logger.Info("Loaded vulnerability query pools successfully", + "totalVulns", len(pools.VulnQueryIDs), + "totalPackages", len(pools.PackageQueryIDs), + "totalEcosystems", len(pools.EcosystemMap), + "largeBatchPoolSize", len(pools.LargeBatchQueryIDs), + ) + + cfg := GeneratorConfig{ + BaseURL: *apiURLFlag, + Duration: *durationFlag, + VulnRate: *vulnRateFlag, + VersionRate: *versionRateFlag, + PackageRate: *packageRateFlag, + PURLRate: *purlRateFlag, + BatchRate: *batchRateFlag, + LargeBatchRate: *largeBatchRateFlag, + MaxBatchQuerySize: *maxBatchSizeFlag, + MaxLargeBatchQuerySize: *maxLargeBatchSizeFlag, + StatsInterval: *statsIntervalFlag, + } + + _, err = RunTrafficGenerator(ctx, nil, pools, cfg, seed) + if err != nil { + logger.Error("Traffic generator encountered an error", "error", err) + + return err + } + + logger.Info("Staging API test completed successfully") + + return nil +}