diff --git a/AGENTS.md b/AGENTS.md index e21df58abbb..eaa8ed5214b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -220,9 +220,14 @@ All Go microservices are compiled using a single, unified multi-target Dockerfil - Git client daemon/utility to precompute and cache git operations required by other services. - Performs intensive Git tasks like computing commit graphs and generating patch IDs. +8. **`recoverer`**: + - Daemon that subscribes to failed task recovery Pub/Sub messages. + - Repairs and retries failed GCS writes, reimports missing vulnerability records from sources (via Gitter, GCS bucket, or REST), and handles GCS generation mismatches. + ### Internal Shared Libraries (`go/internal/`) - **`api/`**: Shared package containing the core gRPC public server implementation of the OSV API. - **`worker/`**: Core engine and subscriber logic for the Go worker. +- **`recoverer/`**: Core engine and handlers for the Go recoverer. - **`database/`**: Shared Datastore client and repository models (specifically [`go/internal/database/datastore/`](go/internal/database/datastore/)). - *Design Pattern*: Models here **mirror** the Datastore models defined in the Python library ([`osv/models.py`](osv/models.py)). - *Consistency Testing*: To prevent synchronization drift between Go and Python database models, a database validation test is maintained under [`go/internal/database/datastore/internal/validate/`](go/internal/database/datastore/internal/validate/) (run via `run_validate.sh`). @@ -258,14 +263,14 @@ Contains deployment setups, workers running in GKE, Cloud Functions, and the use - **Deployment Target**: **Google Cloud Run** (managed via Cloud Deploy pipeline `osv-website`). ### 3. Workers (`gcp/workers/`) -- **`worker` (`gcp/workers/worker/`)**: **Base Environment**. Retains shared Poetry dependencies and base Dockerfile for Python workers (`recoverer`, `vanir_signatures`); legacy worker daemon replaced by Go worker under `go/cmd/worker/`. +- **`worker` (`gcp/workers/worker/`)**: **Base Environment**. Retains shared Poetry dependencies and base Dockerfile for Python workers (`vanir_signatures`); legacy worker daemon replaced by Go worker under `go/cmd/worker/`. - **ClusterFuzz Worker (`gcp/workers/oss_fuzz_worker/`, `gcp/workers/oss_fuzz_importer/`)**: **Barely Maintained**. Siloed workloads for OSS-Fuzz integration. - **Deployment Target**: **GKE** (managed via Cloud Deploy pipeline `oss-fuzz-workers`). - **`vanir_signatures`**: **Active (Python)**. Used for signature generation/verification. -- **`recoverer`**: **Active (Python)**. Used to recover/repair states; scheduled for migration to Go in the future. ### 4. Indexer (`gcp/indexer/`) - **Status**: **Active (Go)**. - Handles indexing, but is not under active development. - **Deployment Target**: **GKE** (managed via Cloud Deploy pipeline `gke-indexer`). + diff --git a/Makefile b/Makefile index ec5afd5c3c4..747c828dd3b 100644 --- a/Makefile +++ b/Makefile @@ -18,9 +18,6 @@ run-cmd := poetry run lib-tests: ## Run core Python library tests ./run_tests.sh -recoverer-tests: ## Run Python recoverer tests - cd gcp/workers/recoverer && ./run_tests.sh - vanir-signatures-tests: ## Run Vanir signatures tests cd gcp/workers/vanir_signatures && ./run_tests.sh @@ -123,7 +120,7 @@ run-api-server-test: @cd go && go build -o ./api-devserver ./cmd/api-devserver && (GOOGLE_CLOUD_PROJECT=oss-vdb-test OSV_VULNERABILITIES_BUCKET=osv-test-vulnerabilities ./api-devserver $(ARGS); EXIT_CODE=$$?; rm -f ./api-devserver; exit $$EXIT_CODE) # TODO: API integration tests. -all-tests: lib-tests recoverer-tests website-tests vulnfeed-tests bindings-tests go-tests ## Run all tests +all-tests: lib-tests website-tests vulnfeed-tests bindings-tests go-tests ## Run all tests reimport-tui: ## Run the reimport TUI tool test -f $(HOME)/.config/gcloud/application_default_credentials.json || (echo "GCP Application Default Credentials not set, try 'gcloud auth application-default login'"; exit 1) diff --git a/README.md b/README.md index 08ca31b15e1..2eeb692fc95 100644 --- a/README.md +++ b/README.md @@ -44,8 +44,8 @@ consists of: | `gcp/functions` | The Cloud Function for publishing PyPI vulnerabilities (maintained, but not developed) | | `gcp/indexer` | The determine version `indexer` | | `gcp/website` | The backend of the osv.dev web interface, with the frontend in `frontend3`
Blog posts (in `blog`) | -| `gcp/workers/` | Python workers (`recoverer`, `vanir_signatures`, and `oss_fuzz_worker`) | -| `go/` | Go module for shared libraries and commands (`cmd/exporter`, `cmd/recordchecker`) | +| `gcp/workers/` | Python workers (`vanir_signatures` and `oss_fuzz_worker`) | +| `go/` | Go module for shared libraries and commands (`cmd/api`, `cmd/importer`, `cmd/worker`, `cmd/exporter`, `cmd/recoverer`, `cmd/relations`, etc.) | | `osv/` | The core OSV Python library, used in basically all Python services
OSV ecosystem package versioning helpers in `ecosystems/`
Datastore model definitions in `models.py` | | `tools/` | Misc scripts/tools, mostly intended for development (datastore stuff, linting)
The `indexer-api-caller` for indexer calling | | `vulnfeeds/` | Go module for (mostly) the NVD CVE conversion
The Alpine feed converter (`cmd/alpine`)
The Debian feed converter (`tools/debian`, which is written in Python) | diff --git a/deployment/build-and-stage.yaml b/deployment/build-and-stage.yaml index 522a39ac740..529916e5dd4 100644 --- a/deployment/build-and-stage.yaml +++ b/deployment/build-and-stage.yaml @@ -62,7 +62,7 @@ steps: args: ['push', '--all-tags', 'gcr.io/oss-vdb/worker-base'] waitFor: ['build-worker-base', 'cloud-build-queue'] -# Build/push core worker/recoverer images. +# Build/push core worker images. - name: gcr.io/cloud-builders/docker args: ['build', '-t', 'gcr.io/oss-vdb/worker:latest', '-t', 'gcr.io/oss-vdb/worker:$COMMIT_SHA', '-f', 'gcp/workers/worker/Dockerfile', '.'] id: 'build-worker' @@ -71,15 +71,6 @@ steps: args: ['push', '--all-tags', 'gcr.io/oss-vdb/worker'] waitFor: ['build-worker', 'cloud-build-queue'] -- name: gcr.io/cloud-builders/docker - args: ['build', '-t', 'gcr.io/oss-vdb/recoverer:latest', '-t', 'gcr.io/oss-vdb/recoverer:$COMMIT_SHA', '.'] - dir: 'gcp/workers/recoverer' - id: 'build-recoverer' - waitFor: ['build-worker'] -- name: gcr.io/cloud-builders/docker - args: ['push', '--all-tags', 'gcr.io/oss-vdb/recoverer'] - waitFor: ['build-recoverer', 'cloud-build-queue'] - - name: gcr.io/cloud-builders/docker args: ['build', '-t', 'gcr.io/oss-vdb/oss-fuzz-worker:latest', '-t', 'gcr.io/oss-vdb/oss-fuzz-worker:$COMMIT_SHA', '-f', 'gcp/workers/oss_fuzz_worker/Dockerfile', '.'] id: 'build-oss-fuzz-worker' @@ -224,6 +215,20 @@ steps: args: ['push', '--all-tags', 'gcr.io/oss-vdb/osv-server'] waitFor: ['build-osv-server', 'cloud-build-queue'] +- name: 'gcr.io/cloud-builders/docker' + entrypoint: 'bash' + args: ['-c', 'docker pull gcr.io/oss-vdb/recoverer:latest || exit 0'] + id: 'pull-recoverer' + waitFor: ['setup'] +- name: gcr.io/cloud-builders/docker + args: ['buildx', 'build', '-t', 'gcr.io/oss-vdb/recoverer:latest', '-t', 'gcr.io/oss-vdb/recoverer:$COMMIT_SHA', '--target', 'recoverer', '--build-context', 'bindings=../bindings', '-f', 'Dockerfile', '--cache-from', 'gcr.io/oss-vdb/recoverer:latest', '--pull', '.'] + dir: 'go' + id: 'build-recoverer' + waitFor: ['pull-recoverer', 'build-osv-server'] +- name: gcr.io/cloud-builders/docker + args: ['push', '--all-tags', 'gcr.io/oss-vdb/recoverer'] + waitFor: ['build-recoverer', 'cloud-build-queue'] + # Build/push staging-api-test images to gcr.io/oss-vdb-test. - 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', '.'] diff --git a/deployment/clouddeploy/gke-workers/base/core/recoverer.yaml b/deployment/clouddeploy/gke-workers/base/core/recoverer.yaml index 54c498a36a9..e207a1e6d34 100644 --- a/deployment/clouddeploy/gke-workers/base/core/recoverer.yaml +++ b/deployment/clouddeploy/gke-workers/base/core/recoverer.yaml @@ -30,11 +30,15 @@ spec: - name: recoverer image: recoverer imagePullPolicy: Always - resources: - requests: - cpu: "10m" - memory: "256Mi" - limits: - cpu: "200m" - memory: "512Mi" + env: + - name: GITTER_HOST + value: http://gitter-service:8888 + resources: + requests: + cpu: "10m" + memory: "256Mi" + limits: + cpu: "200m" + memory: "512Mi" + diff --git a/gcp/workers/cloudbuild.yaml b/gcp/workers/cloudbuild.yaml index 3076be4ecc1..2d580a21819 100644 --- a/gcp/workers/cloudbuild.yaml +++ b/gcp/workers/cloudbuild.yaml @@ -33,14 +33,6 @@ steps: args: ['poetry', 'sync'] waitFor: ['-'] -- name: 'gcr.io/oss-vdb/ci' - id: 'recoverer-tests' - dir: gcp/workers/recoverer - args: ['bash', '-ex', 'run_tests.sh'] - env: - - DATASTORE_EMULATOR_PORT=8005 - - GITTER_PORT=8891 - waitFor: ['init', 'sync'] - name: 'gcr.io/oss-vdb/ci' id: 'vanir-signatures-tests' diff --git a/gcp/workers/recoverer/Dockerfile b/gcp/workers/recoverer/Dockerfile deleted file mode 100644 index 46b69d97793..00000000000 --- a/gcp/workers/recoverer/Dockerfile +++ /dev/null @@ -1,19 +0,0 @@ -# Copyright 2025 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 - -COPY recoverer.py /usr/local/bin -RUN chmod 755 /usr/local/bin/recoverer.py -ENTRYPOINT ["recoverer.py"] \ No newline at end of file diff --git a/gcp/workers/recoverer/recoverer.py b/gcp/workers/recoverer/recoverer.py deleted file mode 100644 index 7a29a153b02..00000000000 --- a/gcp/workers/recoverer/recoverer.py +++ /dev/null @@ -1,364 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2025 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. -"""OSV failed task recoverer.""" - -import base64 -import datetime -import functools -import logging -import os -import requests -import sys -import time - -from google.cloud import ndb -from google.cloud import pubsub_v1 -from google.cloud import storage - -import osv -import osv.models -import osv.sources -from osv.logs import setup_gcp_logging - -_FAILED_TASKS_SUBSCRIPTION = os.getenv('FAILED_TASKS_SUBSCRIPTION', 'recovery') -_TASKS_TOPIC = os.getenv('WORKER_TASK_TOPIC', 'tasks') - -_ndb_client = None -_storage_client = None - - -def ndb_client(): - """Get the ndb client. - Lazily initialized to allow testing with datastore emulator.""" - global _ndb_client - database_id = os.getenv('DATASTORE_DATABASE_ID') - if not database_id: - database_id = None - if _ndb_client is None: - _ndb_client = ndb.Client(database=database_id) - return _ndb_client - - -def storage_client(): - """Get the storage client. - Lazily initialized.""" - global _storage_client - if _storage_client is None: - _storage_client = storage.Client() - return _storage_client - - -def handle_gcs_retry(message: pubsub_v1.types.PubsubMessage) -> bool: - """Handle a failed GCS write.""" - try: - vuln = osv.vulnerability_pb2.Vulnerability.FromString(message.data) - except Exception: - logging.error( - 'gcs_retry: failed to decode protobuf. Ignoring message.', - # chuck the data into the GCP log fields in case it's useful. - extra={ - 'json_fields': { - 'data': base64.encodebytes(message.data).decode() - } - }) - return True - logging.info('gcs_retry: vulnerability: %s', vuln.id) - modified = vuln.modified.ToDatetime(datetime.UTC) - bucket = osv.gcs.get_osv_bucket() - path = os.path.join(osv.gcs.VULN_PB_PATH, vuln.id + '.pb') - pb_blob = bucket.get_blob(path) - # Check that the record hasn't been written/updated in the meantime. - if pb_blob and pb_blob.custom_time and pb_blob.custom_time >= modified: - logging.warning( - 'gcs_retry: %s was modified before message was processed: ' - 'message: %s, blob: %s', vuln.id, modified, pb_blob.custom_time) - # TODO(michaelkedar): trigger a reimport of the record. - return True - - pb_blob = bucket.blob(path) - pb_blob.custom_time = modified - try: - pb_blob.upload_from_string( - message.data, content_type='application/octet-stream') - return True - except Exception: - logging.exception('gcs_retry: failed to upload %s protobuf to GCS', vuln.id) - return False - - -def handle_gcs_missing(message: pubsub_v1.types.PubsubMessage) -> bool: - """Handle a failed GCS read.""" - vuln_id = message.attributes.get('id') - logging.info('gcs_missing: vulnerability: %s', vuln_id) - if not vuln_id: - logging.error('gcs_missing: message missing id attribute: %s', message) - return True - - with ndb_client().context(): - vuln = osv.Vulnerability.get_by_id(vuln_id) - if not vuln: - logging.error('gcs_missing: Vulnerability entity not found for %s', - vuln_id) - return True - - try: - source, path = osv.sources.parse_source_id(vuln.source_id) - except ValueError: - logging.error('gcs_missing: invalid source_id for %s: %s', vuln_id, - vuln.source_id) - return True - - logging.info('gcs_missing: triggering re-import for %s (%s)', vuln_id, - vuln.source_id) - publisher = pubsub_v1.PublisherClient() - project = os.environ['GOOGLE_CLOUD_PROJECT'] - topic_path = publisher.topic_path(project, _TASKS_TOPIC) - # TODO(michaelkedar): duplicating the download logic here is annoying - # (especially for git repos) and not very robust. - # We should instead force the importer to process these again. - data = download_vuln_data(path, source) - publisher.publish( - topic_path, - data=data, - content_encoding='', - type='update', - source=source, - path=path, - original_sha256='', - deleted='false', - skip_hash_check='true', - req_timestamp=str(int(time.time())), - work_pool='reimport', - ) - - return True - - -@functools.lru_cache -def _get_default_branch(org: str, repo: str) -> str: - """Get the default branch for a GitHub repository.""" - try: - response = requests.get( - f'https://api.github.com/repos/{org}/{repo}', timeout=60) - if response.status_code == 200: - return response.json().get('default_branch', 'main') - logging.warning( - 'Failed to get default branch for %s/%s (HTTP %d). Defaulting to main.', - org, repo, response.status_code) - except requests.RequestException as e: - logging.warning( - 'Request failed when getting default branch for %s/%s: %s. ' - 'Defaulting to main.', org, repo, e) - return 'main' - - -def download_vuln_data(vuln_path: str, source: str) -> bytes: - """Download vulnerability data from the source repository.""" - source_repo: osv.SourceRepository = osv.SourceRepository.get_by_id(source) - if not source_repo: - raise ValueError(f'Source repository {source} not found in Datastore') - - raw_data = None - if source_repo.type == osv.SourceRepositoryType.REST_ENDPOINT: - url = f'{source_repo.link.rstrip("/")}/{vuln_path.lstrip("/")}' - response = requests.get(url, timeout=60) - response.raise_for_status() - raw_data = response.content - elif source_repo.type == osv.SourceRepositoryType.BUCKET: - bucket = storage_client().bucket(source_repo.bucket) - blob = bucket.blob(vuln_path) - raw_data = blob.download_as_bytes() - elif source_repo.type == osv.SourceRepositoryType.GIT: - # Cheeky workaround: all of our GIT repos are on GitHub - # So download the files directly from the GitHub API - # instead of cloning the repos. - org, _, repo = source_repo.repo_url.removeprefix( - 'https://github.com/').removesuffix('.git').partition('/') - if source_repo.name == 'oss-fuzz': - # oss-fuzz uses ssh - org = 'google' - repo = 'oss-fuzz-vulns' - branch = source_repo.repo_branch - if not branch: - branch = _get_default_branch(org, repo) - url = f'https://raw.githubusercontent.com/{org}/{repo}/{branch}/{vuln_path}' - response = requests.get(url, timeout=60) - response.raise_for_status() - raw_data = response.content - else: - raise ValueError(f'unhandled source type {source_repo.type}') - - if raw_data is None: - raise ValueError( - f'Failed to download vulnerability data for {vuln_path} from {source}') - - vuln = osv.parse_vulnerabilities_from_data( - raw_data, source_repo.extension, key_path=source_repo.key_path)[0] - return vuln.SerializeToString(deterministic=True) - - -def handle_gcs_gen_mismatch(message: pubsub_v1.types.PubsubMessage) -> bool: - """Handle a generation mismatch when attempting update a part of a record. - e.g. If a record was reimported while its aliases were being updated. - """ - vuln_id = message.attributes.get('id') - field = message.attributes.get('field') - logging.info('gcs_gen_mismatch: vulnerability: %s, field: %s', vuln_id, field) - if not vuln_id or not field: - logging.error('gcs_gen_mismatch: message missing id or field attribute: %s', - message) - return True - - with ndb_client().context(): - result = osv.gcs.get_by_id_with_generation(vuln_id) - if result is None: - logging.error('gcs_gen_mismatch: vulnerability not in GCS - %s', vuln_id) - logging.info('trying with gcs_missing') - return handle_gcs_missing(message) - vuln_proto, generation = result - - def transaction(): - vuln: osv.Vulnerability = osv.Vulnerability.get_by_id(vuln_id) - if vuln is None: - logging.error('vulnerability not in Datastore - %s', vuln_id) - # TODO(michaelkedar): What to do in this case? - return - modified = vuln.modified - - fields = field.split(',') - for f in fields: - if f == 'aliases': - alias_group = osv.AliasGroup.query( - osv.AliasGroup.bug_ids == vuln_id).get() - if alias_group is None: - aliases = [] - aliases_modified = datetime.datetime.now(datetime.UTC) - else: - aliases = sorted(set(alias_group.bug_ids) - {vuln_id}) - aliases_modified = alias_group.last_modified - # Only update the modified time if it's actually being modified - if vuln_proto.aliases != aliases: - vuln_proto.aliases[:] = aliases - if aliases_modified > modified: - modified = aliases_modified - else: - modified = datetime.datetime.now(datetime.UTC) - - elif f == 'upstream': - upstream_group = osv.UpstreamGroup.query( - osv.UpstreamGroup.db_id == vuln_id).get() - if upstream_group is None: - upstream = [] - upstream_modified = datetime.datetime.now(datetime.UTC) - else: - upstream = upstream_group.upstream_ids - upstream_modified = upstream_group.last_modified - # Only update the modified time if it's actually being modified - if vuln_proto.upstream != upstream: - vuln_proto.upstream[:] = upstream - if upstream_modified > modified: - modified = upstream_modified - else: - modified = datetime.datetime.now(datetime.UTC) - - elif f == 'related': - related_group = osv.RelatedGroup.get_by_id(vuln_id) - if related_group is None: - related = [] - related_modified = datetime.datetime.now(datetime.UTC) - else: - related = related_group.related_ids - related_modified = related_group.last_modified - # Only update the modified time if it's actually being modified - if vuln_proto.related != related: - vuln_proto.related[:] = related - if related_modified > modified: - modified = related_modified - else: - modified = datetime.datetime.now(datetime.UTC) - - vuln_proto.modified.FromDatetime(modified) - osv.ListedVulnerability.from_vulnerability(vuln_proto).put() - vuln.modified = modified - vuln.put() - - try: - ndb.transaction(transaction) - except Exception: - logging.exception( - 'gcs_gen_mismatch: Datastore transaction failed for %s %s', vuln_id, - field) - return False - try: - osv.gcs.upload_vulnerability(vuln_proto, generation) - return True - except Exception: - logging.exception('gcs_gen_mismatch: Writing to bucket failed for %s %s', - vuln_id, field) - return False - - -def handle_generic(message: pubsub_v1.types.PubsubMessage) -> bool: - """Generic message handler.""" - task_type = message.attributes.get('type', 'unknown') - logging.error('`%s` task could not be processed: %s', task_type, message) - # TODO(michaelkedar): We should store these somewhere. - return True - - -HANDLERS = { - 'gcs_retry': handle_gcs_retry, - 'gcs_missing': handle_gcs_missing, - 'gcs_gen_mismatch': handle_gcs_gen_mismatch, -} - - -def handle_task(message: pubsub_v1.types.PubsubMessage) -> bool: - """Handle a 'failed-tasks' message.""" - task_type = message.attributes.get('type') - handler = HANDLERS.get(task_type, handle_generic) - return handler(message) - - -def main(): - project = osv.utils.get_google_cloud_project() - if not project: - logging.error('GOOGLE_CLOUD_PROJECT not set') - sys.exit(1) - - with pubsub_v1.SubscriberClient() as subscriber: - subscription = subscriber.subscription_path(project, - _FAILED_TASKS_SUBSCRIPTION) - - while True: - response = subscriber.pull(subscription=subscription, max_messages=1) - if not response.received_messages: - continue - - message = response.received_messages[0].message - ack_id = response.received_messages[0].ack_id - # Try handle the task - # If successful (returned True), acknowledge it. - # Otherwise, nack the task to trigger it to be redelivered. - if handle_task(message): - subscriber.acknowledge(subscription=subscription, ack_ids=[ack_id]) - else: - subscriber.modify_ack_deadline( - subscription=subscription, ack_ids=[ack_id], ack_deadline_seconds=0) - - -if __name__ == '__main__': - setup_gcp_logging('recoverer') - main() diff --git a/gcp/workers/recoverer/recoverer_test.py b/gcp/workers/recoverer/recoverer_test.py deleted file mode 100644 index 69c26e63891..00000000000 --- a/gcp/workers/recoverer/recoverer_test.py +++ /dev/null @@ -1,319 +0,0 @@ -# Copyright 2025 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. -"""Recoverer tests.""" -import datetime -import os -import unittest -import unittest.mock - -from google.cloud import ndb -from google.cloud import pubsub_v1 - -import osv -from osv import tests - -import recoverer - - -class RecovererTest(unittest.TestCase): - """Recoverer tests.""" - - def setUp(self): - with ndb.Client().context(): - osv.SourceRepository( - id='test', - name='test', - db_prefix=['TEST-'], - ).put() - osv.AliasGroup( - bug_ids=['CVE-456', 'OSV-123', 'TEST-123'], - last_modified=datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - ).put() - osv.UpstreamGroup( - db_id='TEST-123', - upstream_ids=['TEST-1', 'TEST-12'], - last_modified=datetime.datetime(2025, 3, 3, tzinfo=datetime.UTC), - ).put() - osv.Bug( - id='TEST-123', - db_id='TEST-123', - status=1, - source='test', - public=True, - import_last_modified=datetime.datetime( - 2025, 1, 1, tzinfo=datetime.UTC), - ).put() - osv.Vulnerability( - id='TEST-123', - source_id='test:TEST-123.yaml', - modified=datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC), - ).put() - return super().setUp() - - def test_handle_gcs_retry(self): - """Test standard handle_gcs_retry.""" - vuln = osv.vulnerability_pb2.Vulnerability() - vuln.id = 'TEST-555' - modified = datetime.datetime(2025, 5, 5, tzinfo=datetime.UTC) - vuln.modified.FromDatetime(modified) - vuln_bytes = vuln.SerializeToString(deterministic=True) - message = pubsub_v1.types.PubsubMessage(data=vuln_bytes) - self.assertTrue(recoverer.handle_gcs_retry(message)) - - # check this was written - bucket = osv.gcs.get_osv_bucket() - blob = bucket.get_blob(os.path.join(osv.gcs.VULN_PB_PATH, 'TEST-555.pb')) - self.assertIsNotNone(blob) - self.assertEqual(blob.custom_time, modified) - - def test_handle_gcs_retry_overwritten(self): - """Test handle_gcs_retry when vuln was written after pubsub message.""" - original_result = osv.gcs.get_by_id_with_generation('TEST-123') - self.assertIsNotNone(original_result) - - old = osv.vulnerability_pb2.Vulnerability() - old.id = 'TEST-123' - modified = datetime.datetime(2020, 1, 1, tzinfo=datetime.UTC) - old.modified.FromDatetime(modified) - old_bytes = old.SerializeToString(deterministic=True) - message = pubsub_v1.types.PubsubMessage(data=old_bytes) - with self.assertLogs(level='WARNING') as cm: - self.assertTrue(recoverer.handle_gcs_retry(message)) - self.assertEqual(1, len(cm.output)) - self.assertIn('TEST-123 was modified before message was processed', - cm.output[0]) - # make sure it wasn't written - new_result = osv.gcs.get_by_id_with_generation('TEST-123') - self.assertIsNotNone(new_result) - self.assertEqual(original_result, new_result) - - def test_handle_gcs_retry_invalid_data(self): - """Test handle_gcs_retry when data is invalid.""" - message = pubsub_v1.types.PubsubMessage(data=b'invalid') - with self.assertLogs(level='ERROR') as cm: - self.assertTrue(recoverer.handle_gcs_retry(message)) - self.assertEqual(1, len(cm.output)) - self.assertIn('failed to decode protobuf', cm.output[0]) - - @unittest.mock.patch('recoverer.download_vuln_data') - @unittest.mock.patch('google.cloud.pubsub_v1.PublisherClient') - def test_handle_gcs_missing(self, mock_publisher, mock_download): - """Test standard handle_gcs_missing""" - mock_download.return_value = b'dummy_data' - message = pubsub_v1.types.PubsubMessage(attributes={'id': 'TEST-123'}) - self.assertTrue(recoverer.handle_gcs_missing(message)) - - mock_download.assert_called_once_with('TEST-123.yaml', 'test') - # Check that the update message was published - mock_publisher.return_value.publish.assert_called_once() - call_args = mock_publisher.return_value.publish.call_args - self.assertEqual(call_args.kwargs['type'], 'update') - self.assertEqual(call_args.kwargs['source'], 'test') - self.assertEqual(call_args.kwargs['path'], 'TEST-123.yaml') - self.assertEqual(call_args.kwargs['skip_hash_check'], 'true') - self.assertEqual(call_args.kwargs['data'], b'dummy_data') - self.assertEqual(call_args.kwargs['content_encoding'], '') - - def test_handle_gcs_gen_mismatch_aliases(self): - """Test handle_gcs_gen_mismatch for aliases.""" - # Set up records - with ndb.Client().context(): - osv.AliasGroup( - bug_ids=['CVE-111', 'OSV-111', 'TEST-111'], - last_modified=datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - ).put() - osv.Bug( - id='TEST-111', - db_id='TEST-111', - status=1, - source='test', - public=True, - import_last_modified=datetime.datetime( - 2025, 1, 1, tzinfo=datetime.UTC), - last_modified=datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC), - ).put() - g = osv.AliasGroup( - bug_ids=['CVE-222', 'TEST-222'], - last_modified=datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - ).put() - osv.Bug( - id='TEST-222', - db_id='TEST-222', - status=1, - source='test', - public=True, - import_last_modified=datetime.datetime( - 2025, 1, 1, tzinfo=datetime.UTC), - last_modified=datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC), - ).put() - g.delete() - osv.AliasGroup( - bug_ids=['CVE-222', 'OSV-222', 'TEST-222'], - last_modified=datetime.datetime(2025, 3, 3, tzinfo=datetime.UTC), - ).put() - g = osv.AliasGroup( - bug_ids=['CVE-333', 'TEST-333'], - last_modified=datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - ).put() - osv.Bug( - id='TEST-333', - db_id='TEST-333', - status=1, - source='test', - public=True, - import_last_modified=datetime.datetime( - 2025, 1, 1, tzinfo=datetime.UTC), - last_modified=datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC), - ).put() - g.delete() - - message = pubsub_v1.types.PubsubMessage(attributes={ - 'id': 'TEST-111', - 'field': 'aliases' - }) - self.assertTrue(recoverer.handle_gcs_gen_mismatch(message)) - vuln = osv.gcs.get_by_id('TEST-111') - self.assertEqual(['CVE-111', 'OSV-111'], vuln.aliases) - self.assertEqual( - datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - vuln.modified.ToDatetime(datetime.UTC)) - - message = pubsub_v1.types.PubsubMessage(attributes={ - 'id': 'TEST-222', - 'field': 'aliases' - }) - self.assertTrue(recoverer.handle_gcs_gen_mismatch(message)) - vuln = osv.gcs.get_by_id('TEST-222') - self.assertEqual(['CVE-222', 'OSV-222'], vuln.aliases) - self.assertEqual( - datetime.datetime(2025, 3, 3, tzinfo=datetime.UTC), - vuln.modified.ToDatetime(datetime.UTC)) - - message = pubsub_v1.types.PubsubMessage(attributes={ - 'id': 'TEST-333', - 'field': 'aliases' - }) - was_now = datetime.datetime.now(datetime.UTC) - self.assertTrue(recoverer.handle_gcs_gen_mismatch(message)) - vuln = osv.gcs.get_by_id('TEST-333') - self.assertEqual([], vuln.aliases) - # check that the time was updated to "now" - self.assertLessEqual(was_now, vuln.modified.ToDatetime(datetime.UTC)) - - def test_handle_gcs_gen_mismatch_upstream(self): - """Test handle_gcs_gen_mismatch for upstream.""" - # Set up records - with ndb.Client().context(): - osv.UpstreamGroup( - db_id='TEST-111', - upstream_ids=['UPSTREAM-1'], - last_modified=datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - ).put() - osv.Bug( - id='TEST-111', - db_id='TEST-111', - status=1, - source='test', - public=True, - import_last_modified=datetime.datetime( - 2025, 1, 1, tzinfo=datetime.UTC), - last_modified=datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC), - ).put() - g = osv.UpstreamGroup( - db_id='TEST-222', - upstream_ids=['UPSTREAM-2'], - last_modified=datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - ).put() - osv.Bug( - id='TEST-222', - db_id='TEST-222', - status=1, - source='test', - public=True, - import_last_modified=datetime.datetime( - 2025, 1, 1, tzinfo=datetime.UTC), - last_modified=datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC), - ).put() - g.delete() - osv.UpstreamGroup( - db_id='TEST-222', - upstream_ids=['UPSTREAM-2', 'UPSTREAM-22'], - last_modified=datetime.datetime(2025, 3, 3, tzinfo=datetime.UTC), - ).put() - g = osv.UpstreamGroup( - db_id='TEST-333', - upstream_ids=['UPSTREAM-3'], - last_modified=datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - ).put() - osv.Bug( - id='TEST-333', - db_id='TEST-333', - status=1, - source='test', - public=True, - import_last_modified=datetime.datetime( - 2025, 1, 1, tzinfo=datetime.UTC), - last_modified=datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC), - ).put() - g.delete() - - message = pubsub_v1.types.PubsubMessage(attributes={ - 'id': 'TEST-111', - 'field': 'upstream' - }) - self.assertTrue(recoverer.handle_gcs_gen_mismatch(message)) - vuln = osv.gcs.get_by_id('TEST-111') - self.assertEqual(['UPSTREAM-1'], vuln.upstream) - self.assertEqual( - datetime.datetime(2025, 2, 2, tzinfo=datetime.UTC), - vuln.modified.ToDatetime(datetime.UTC)) - - message = pubsub_v1.types.PubsubMessage(attributes={ - 'id': 'TEST-222', - 'field': 'upstream' - }) - self.assertTrue(recoverer.handle_gcs_gen_mismatch(message)) - vuln = osv.gcs.get_by_id('TEST-222') - self.assertEqual(['UPSTREAM-2', 'UPSTREAM-22'], vuln.upstream) - self.assertEqual( - datetime.datetime(2025, 3, 3, tzinfo=datetime.UTC), - vuln.modified.ToDatetime(datetime.UTC)) - - message = pubsub_v1.types.PubsubMessage(attributes={ - 'id': 'TEST-333', - 'field': 'upstream' - }) - was_now = datetime.datetime.now(datetime.UTC) - self.assertTrue(recoverer.handle_gcs_gen_mismatch(message)) - vuln = osv.gcs.get_by_id('TEST-333') - self.assertEqual([], vuln.upstream) - # check that the time was updated to "now" - self.assertLessEqual(was_now, vuln.modified.ToDatetime(datetime.UTC)) - - def test_handle_generic(self): - """Test handle_generic.""" - message = pubsub_v1.types.PubsubMessage(attributes={'type': 'test'}) - with self.assertLogs(level='ERROR') as cm: - self.assertTrue(recoverer.handle_generic(message)) - self.assertEqual(1, len(cm.output)) - self.assertIn('`test` task could not be processed', cm.output[0]) - - -def setUpModule(): - """Set up the test module.""" - unittest.enterModuleContext(tests.datastore_emulator()) - - -if __name__ == '__main__': - unittest.main() diff --git a/gcp/workers/recoverer/run_tests.sh b/gcp/workers/recoverer/run_tests.sh deleted file mode 100755 index a4fe3d75a56..00000000000 --- a/gcp/workers/recoverer/run_tests.sh +++ /dev/null @@ -1,22 +0,0 @@ -#!/bin/bash -ex -# Copyright 2025 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. - -cd ../worker - -# Install dependencies only if not running in Cloud Build -if [ -z "$CLOUDBUILD" ]; then - poetry sync -fi -poetry run python ../recoverer/recoverer_test.py diff --git a/go/Dockerfile b/go/Dockerfile index 27094066d46..f1a74d820ea 100644 --- a/go/Dockerfile +++ b/go/Dockerfile @@ -24,7 +24,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). +# relations, recordchecker, generatesitemap, custommetrics, gitter, first_package_finder, api, recoverer). # ==================================================================================== # ======================================================== @@ -148,3 +148,14 @@ RUN CGO_ENABLED=0 go build -o /app/api ./cmd/api/ FROM gcr.io/distroless/static-debian12@sha256:a9fcaedd4c9b59e12dd65d954f0b5044f19b0647a8a3712e77205df9e7b102cd AS api COPY --from=api-build /app/api / ENTRYPOINT ["/api"] + +# ======================================================== +# Target: Recoverer +# ======================================================== +FROM builder AS recoverer-build +RUN CGO_ENABLED=0 go build -o /app/recoverer ./cmd/recoverer/ + +FROM gcr.io/distroless/static-debian12@sha256:a9fcaedd4c9b59e12dd65d954f0b5044f19b0647a8a3712e77205df9e7b102cd AS recoverer +COPY --from=recoverer-build /app/recoverer / +ENTRYPOINT ["/recoverer"] + diff --git a/go/cmd/recoverer/main.go b/go/cmd/recoverer/main.go new file mode 100644 index 00000000000..2022c05d9ef --- /dev/null +++ b/go/cmd/recoverer/main.go @@ -0,0 +1,170 @@ +// 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 implements the OSV failed task recoverer. +package main + +import ( + "context" + "errors" + "flag" + "log/slog" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "cloud.google.com/go/datastore" + "cloud.google.com/go/pubsub/v2" + "cloud.google.com/go/storage" + osvdatastore "github.com/google/osv.dev/go/internal/database/datastore" + "github.com/google/osv.dev/go/internal/gitter" + "github.com/google/osv.dev/go/internal/recoverer" + "github.com/google/osv.dev/go/logger" + "github.com/google/osv.dev/go/osv/clients" + "github.com/hashicorp/go-retryablehttp" + "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" +) + +type retryableHTTPLogger struct{} + +func (r retryableHTTPLogger) Error(msg string, keysAndValues ...any) { + logger.Error(msg, keysAndValues...) +} + +func (r retryableHTTPLogger) Info(msg string, keysAndValues ...any) { + logger.Info(msg, keysAndValues...) +} + +func (r retryableHTTPLogger) Debug(msg string, keysAndValues ...any) { + logger.Debug(msg, keysAndValues...) +} + +func (r retryableHTTPLogger) Warn(msg string, keysAndValues ...any) { + logger.Warn(msg, keysAndValues...) +} + +func main() { + if err := run(); err != nil { + os.Exit(1) + } +} + +func run() error { + logger.InitGlobalLogger() + defer logger.Close() + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + logger.InfoContext(ctx, "recoverer starting") + + project := os.Getenv("GOOGLE_CLOUD_PROJECT") + if project == "" { + logger.ErrorContext(ctx, "GOOGLE_CLOUD_PROJECT environment variable is not set") + + return errors.New("GOOGLE_CLOUD_PROJECT environment variable is not set") + } + + numWorkers := flag.Int("num-workers", 10, "Number of workers used to process recovery tasks") + flag.Parse() + + pubsubSubscription := envOrDefault("FAILED_TASKS_SUBSCRIPTION", "recovery") + tasksTopic := envOrDefault("WORKER_TASK_TOPIC", "tasks") + datastoreID := envOrDefault("DATASTORE_DATABASE_ID", "") + vulnBucket := envOrDefault("OSV_VULNERABILITIES_BUCKET", "osv-test-vulnerabilities") + defaultTaskPool := envOrDefault("DEFAULT_TASK_POOL", "default") + reimportTaskPool := envOrDefault("REIMPORT_TASK_POOL", "reimport") + gitterHost := os.Getenv("GITTER_HOST") + + dsClient, err := datastore.NewClientWithDatabase(ctx, project, datastoreID) + if err != nil { + logger.ErrorContext(ctx, "Failed to create datastore client", slog.Any("error", err)) + + return err + } + defer dsClient.Close() + + gcsClient, err := storage.NewClient(ctx) + if err != nil { + logger.ErrorContext(ctx, "Failed to create storage client", slog.Any("error", err)) + + return err + } + defer gcsClient.Close() + + psClient, err := pubsub.NewClient(ctx, project) + if err != nil { + logger.ErrorContext(ctx, "Failed to create pubsub client", slog.Any("error", err)) + + return err + } + defer psClient.Close() + + retryClient := retryablehttp.NewClient() + retryClient.RetryMax = 3 + retryClient.RetryWaitMin = 1 * time.Second + retryClient.RetryWaitMax = 4 * time.Second + retryClient.Logger = retryableHTTPLogger{} + retryClient.HTTPClient.Transport = otelhttp.NewTransport(http.DefaultTransport) + httpClient := retryClient.StandardClient() + + var gitterClient gitter.Client + if gitterHost != "" { + var err error + gitterClient, err = gitter.NewClient(gitterHost, httpClient) + if err != nil { + logger.ErrorContext(ctx, "Failed to create gitter client", slog.Any("error", err)) + + return err + } + } else { + logger.WarnContext(ctx, "GITTER_HOST is not set; git source recovery may fail") + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + SourceRepo: osvdatastore.NewSourceRepositoryStore(dsClient), + Relations: osvdatastore.NewRelationsStore(dsClient), + GCS: clients.NewGCSClient(gcsClient, vulnBucket), + GCSProvider: clients.NewGCSStorageProvider(gcsClient), + Publisher: &clients.GCPPublisher{Publisher: psClient.Publisher(tasksTopic)}, + }, + GitterClient: gitterClient, + HTTPClient: httpClient, + DefaultTaskPool: defaultTaskPool, + ReimportTaskPool: reimportTaskPool, + }) + + sub := psClient.Subscriber(pubsubSubscription) + sub.ReceiveSettings.MaxOutstandingMessages = *numWorkers + sub.ReceiveSettings.MaxOutstandingBytes = -1 + sub.ReceiveSettings.MaxExtension = 6 * time.Hour + sub.ReceiveSettings.MaxDurationPerAckExtension = 10 * time.Minute + + logger.InfoContext(ctx, "recoverer listening for messages", + slog.String("subscription", pubsubSubscription), + slog.String("topic", tasksTopic)) + + return rec.Run(ctx, sub) +} + +func envOrDefault(key, defaultValue string) string { + if value, exists := os.LookupEnv(key); exists { + return value + } + + return defaultValue +} diff --git a/go/internal/recoverer/recoverer.go b/go/internal/recoverer/recoverer.go new file mode 100644 index 00000000000..335dc80c9ab --- /dev/null +++ b/go/internal/recoverer/recoverer.go @@ -0,0 +1,552 @@ +// 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 recoverer implements the OSV failed task recoverer. +package recoverer + +import ( + "bytes" + "context" + "encoding/base64" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "slices" + "strconv" + "strings" + "time" + + "cloud.google.com/go/datastore" + "cloud.google.com/go/pubsub/v2" + osvdatastore "github.com/google/osv.dev/go/internal/database/datastore" + "github.com/google/osv.dev/go/internal/gitter" + gitterpb "github.com/google/osv.dev/go/internal/gitter/pb/repository" + "github.com/google/osv.dev/go/internal/models" + "github.com/google/osv.dev/go/logger" + "github.com/google/osv.dev/go/osv/clients" + "github.com/ossf/osv-schema/bindings/go/osvschema" + "github.com/tidwall/gjson" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/propagation" + "go.opentelemetry.io/otel/trace" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/timestamppb" + "k8s.io/apimachinery/pkg/util/yaml" +) + +type Stores struct { + DatastoreClient *datastore.Client + SourceRepo models.SourceRepositoryStore + Relations models.RelationsStore + GCS clients.CloudStorage + GCSProvider clients.CloudStorageProvider + Publisher clients.Publisher +} + +type Config struct { + Stores Stores + GitterClient gitter.Client + HTTPClient *http.Client + DefaultTaskPool string + ReimportTaskPool string +} + +type Recoverer struct { + stores Stores + gitterClient gitter.Client + httpClient *http.Client + defaultTaskPool string + reimportTaskPool string +} + +func New(cfg Config) *Recoverer { + defaultPool := cfg.DefaultTaskPool + if defaultPool == "" { + defaultPool = "default" + } + reimportPool := cfg.ReimportTaskPool + if reimportPool == "" { + reimportPool = "reimport" + } + httpClient := cfg.HTTPClient + if httpClient == nil { + httpClient = http.DefaultClient + } + + return &Recoverer{ + stores: cfg.Stores, + gitterClient: cfg.GitterClient, + httpClient: httpClient, + defaultTaskPool: defaultPool, + reimportTaskPool: reimportPool, + } +} + +// Run starts the subscriber loop listening for recovery tasks. +func (r *Recoverer) Run(ctx context.Context, sub *pubsub.Subscriber) error { + return sub.Receive(ctx, r.HandleMessage) +} + +// HandleMessage handles a single Pub/Sub message from the failed tasks subscription. +func (r *Recoverer) HandleMessage(ctx context.Context, m *pubsub.Message) { + taskType := m.Attributes["type"] + logInfo := []any{ + slog.String("type", taskType), + slog.String("id", m.Attributes["id"]), + } + + taskCtx := otel.GetTextMapPropagator().Extract(ctx, propagation.MapCarrier(m.Attributes)) + taskCtx, span := otel.Tracer("recoverer").Start(taskCtx, "process_task", + trace.WithAttributes(attribute.String("task_type", taskType))) + defer span.End() + + var err error + switch taskType { + case "gcs_retry": + err = r.HandleGCSRetry(taskCtx, m) + case "gcs_missing": + err = r.HandleGCSMissing(taskCtx, m) + case "gcs_gen_mismatch": + err = r.HandleGCSGenMismatch(taskCtx, m) + default: + err = r.HandleGeneric(taskCtx, m) + } + + if err != nil { + logger.ErrorContext(taskCtx, "Failed to process recovery task", append(logInfo, slog.Any("error", err))...) + m.Nack() + } else { + m.Ack() + } +} + +// HandleGCSRetry handles a failed GCS write. +func (r *Recoverer) HandleGCSRetry(ctx context.Context, m *pubsub.Message) error { + var vuln osvschema.Vulnerability + if err := proto.Unmarshal(m.Data, &vuln); err != nil { + logger.ErrorContext(ctx, "gcs_retry: failed to decode protobuf. Ignoring message.", + slog.Any("error", err), + slog.String("data_base64", base64.StdEncoding.EncodeToString(m.Data))) + // Non-retryable decoding error, return nil to ack + return nil + } + + vulnID := vuln.GetId() + logger.InfoContext(ctx, "gcs_retry: vulnerability", slog.String("id", vulnID)) + modified := vuln.GetModified().AsTime() + path := fmt.Sprintf("all/pb/%s.pb", vulnID) + + attrs, err := r.stores.GCS.ReadObjectAttrs(ctx, path) + if err == nil { + if !attrs.CustomTime.IsZero() && (attrs.CustomTime.Equal(modified) || attrs.CustomTime.After(modified)) { + logger.WarnContext(ctx, "gcs_retry: was modified before message was processed", + slog.String("id", vulnID), + slog.Time("message_modified", modified), + slog.Time("blob_modified", attrs.CustomTime)) + + return nil + } + } else if !errors.Is(err, clients.ErrNotFound) { + logger.ErrorContext(ctx, "gcs_retry: failed to get object attrs from GCS", slog.String("id", vulnID), slog.Any("error", err)) + + return err + } + + opts := &clients.WriteOptions{ + CustomTime: &modified, + ContentType: "application/octet-stream", + } + if err := r.stores.GCS.WriteObject(ctx, path, m.Data, opts); err != nil { + logger.ErrorContext(ctx, "gcs_retry: failed to upload protobuf to GCS", slog.String("id", vulnID), slog.Any("error", err)) + + return err + } + + return nil +} + +// HandleGCSMissing handles a failed GCS read by triggering a reimport. +func (r *Recoverer) HandleGCSMissing(ctx context.Context, m *pubsub.Message) error { + vulnID := m.Attributes["id"] + logger.InfoContext(ctx, "gcs_missing: vulnerability", slog.String("id", vulnID)) + if vulnID == "" { + logger.ErrorContext(ctx, "gcs_missing: message missing id attribute") + + return nil + } + + var vuln osvdatastore.Vulnerability + key := datastore.NameKey("Vulnerability", vulnID, nil) + if err := r.stores.DatastoreClient.Get(ctx, key, &vuln); err != nil { + if errors.Is(err, datastore.ErrNoSuchEntity) { + logger.ErrorContext(ctx, "gcs_missing: Vulnerability entity not found", slog.String("id", vulnID)) + + return nil + } + logger.ErrorContext(ctx, "gcs_missing: failed to fetch Vulnerability entity", slog.String("id", vulnID), slog.Any("error", err)) + + return err + } + + source, path, ok := strings.Cut(vuln.SourceID, ":") + if !ok || source == "" || path == "" { + logger.ErrorContext(ctx, "gcs_missing: invalid source_id", slog.String("id", vulnID), slog.String("source_id", vuln.SourceID)) + + return nil + } + + logger.InfoContext(ctx, "gcs_missing: triggering re-import", slog.String("id", vulnID), slog.String("source_id", vuln.SourceID)) + + data, err := r.DownloadVulnData(ctx, source, path) + if err != nil { + logger.ErrorContext(ctx, "gcs_missing: failed to download vuln data", slog.String("id", vulnID), slog.String("source_id", vuln.SourceID), slog.Any("error", err)) + + return err + } + + msg := &pubsub.Message{ + Data: data, + Attributes: map[string]string{ + "type": "update", + "source": source, + "path": path, + "original_sha256": "", + "deleted": "false", + "skip_hash_check": "true", + "req_timestamp": strconv.FormatInt(time.Now().Unix(), 10), + "content_encoding": "", + "work_pool": r.reimportTaskPool, + }, + } + otel.GetTextMapPropagator().Inject(ctx, propagation.MapCarrier(msg.Attributes)) + + res := r.stores.Publisher.Publish(ctx, msg) + if _, err := res.Get(ctx); err != nil { + logger.ErrorContext(ctx, "gcs_missing: failed to publish update message", slog.String("id", vulnID), slog.Any("error", err)) + + return err + } + + return nil +} + +// DownloadVulnData downloads and parses vulnerability data from the source repository. +func (r *Recoverer) DownloadVulnData(ctx context.Context, source, vulnPath string) ([]byte, error) { + sourceRepo, err := r.stores.SourceRepo.Get(ctx, source) + if err != nil { + return nil, fmt.Errorf("failed to get source repository %s: %w", source, err) + } + + var rawData []byte + switch sourceRepo.Type { + case models.SourceRepositoryTypeREST: + if sourceRepo.Link == "" { + return nil, fmt.Errorf("REST source repository %s is missing Link", source) + } + reqURL, err := url.JoinPath(sourceRepo.Link, vulnPath) + if err != nil { + return nil, fmt.Errorf("failed to construct REST URL for %s / %s: %w", sourceRepo.Link, vulnPath, err) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP request: %w", err) + } + resp, err := r.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("HTTP request failed for %s: %w", reqURL, err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("HTTP request failed with status %d for %s", resp.StatusCode, reqURL) + } + rawData, err = io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body for %s: %w", reqURL, err) + } + + case models.SourceRepositoryTypeBucket: + if sourceRepo.Bucket == nil || sourceRepo.Bucket.Name == "" { + return nil, fmt.Errorf("bucket source repository %s is missing bucket name", source) + } + if r.stores.GCSProvider == nil { + return nil, errors.New("GCSProvider is not configured") + } + bucketStorage := r.stores.GCSProvider.Bucket(sourceRepo.Bucket.Name) + rawData, err = bucketStorage.ReadObject(ctx, vulnPath) + if err != nil { + return nil, fmt.Errorf("failed to read object %s from bucket %s: %w", vulnPath, sourceRepo.Bucket.Name, err) + } + + case models.SourceRepositoryTypeGit: + if sourceRepo.Git == nil || sourceRepo.Git.URL == "" { + return nil, fmt.Errorf("git source repository %s is missing git url", source) + } + if r.gitterClient == nil { + return nil, errors.New("gitter client is not configured") + } + ref := sourceRepo.Git.LastSyncedCommit + if ref == "" { + ref = sourceRepo.Git.Branch + } + if ref == "" { + ref = "HEAD" + } + req := &gitterpb.FileContentRequest{ + Url: sourceRepo.Git.URL, + Commit: ref, + Path: vulnPath, + } + resp, err := r.gitterClient.GetFileContent(ctx, req) + if err != nil { + return nil, fmt.Errorf("gitter GetFileContent failed for %s (commit %s, path %s): %w", sourceRepo.Git.URL, ref, vulnPath, err) + } + rawData = resp.GetContent() + + default: + return nil, fmt.Errorf("unhandled source repository type: %v", sourceRepo.Type) + } + + vulnProto, err := parseVulnerability(rawData, sourceRepo.Extension, vulnPath, sourceRepo.KeyPath) + if err != nil { + return nil, fmt.Errorf("failed to parse vulnerability data for %s:%s: %w", source, vulnPath, err) + } + + return proto.Marshal(vulnProto) +} + +func parseVulnerability(raw []byte, extension, vulnPath, keyPath string) (*osvschema.Vulnerability, error) { + data := bytes.ToValidUTF8(raw, []byte("\uFFFD")) + ext := strings.ToLower(extension) + if ext == ".yaml" || ext == ".yml" || strings.HasSuffix(strings.ToLower(vulnPath), ".yaml") || strings.HasSuffix(strings.ToLower(vulnPath), ".yml") { + jsonBytes, err := yaml.ToJSON(data) + if err != nil { + return nil, fmt.Errorf("failed to convert YAML to JSON: %w", err) + } + data = jsonBytes + } + + if keyPath != "" { + res := gjson.GetBytes(data, keyPath) + if !res.Exists() { + return nil, fmt.Errorf("key path %q not found in data", keyPath) + } + data = []byte(res.Raw) + } + + var vulnProto osvschema.Vulnerability + unmarshalOptions := protojson.UnmarshalOptions{DiscardUnknown: true} + if err := unmarshalOptions.Unmarshal(data, &vulnProto); err != nil { + return nil, fmt.Errorf("failed to unmarshal OSV proto: %w", err) + } + + return &vulnProto, nil +} + +// HandleGCSGenMismatch handles a generation mismatch when attempting to update a part of a record. +func (r *Recoverer) HandleGCSGenMismatch(ctx context.Context, m *pubsub.Message) error { + vulnID := m.Attributes["id"] + fieldStr := m.Attributes["field"] + logger.InfoContext(ctx, "gcs_gen_mismatch: vulnerability", slog.String("id", vulnID), slog.String("field", fieldStr)) + if vulnID == "" || fieldStr == "" { + logger.ErrorContext(ctx, "gcs_gen_mismatch: message missing id or field attribute") + + return nil + } + + path := fmt.Sprintf("all/pb/%s.pb", vulnID) + attrs, err := r.stores.GCS.ReadObjectAttrs(ctx, path) + if errors.Is(err, clients.ErrNotFound) { + logger.ErrorContext(ctx, "gcs_gen_mismatch: vulnerability not in GCS", slog.String("id", vulnID)) + logger.InfoContext(ctx, "trying with gcs_missing", slog.String("id", vulnID)) + + return r.HandleGCSMissing(ctx, m) + } + if err != nil { + logger.ErrorContext(ctx, "gcs_gen_mismatch: failed to read object attrs from GCS", slog.String("id", vulnID), slog.Any("error", err)) + + return err + } + generation := attrs.Generation + + data, err := r.stores.GCS.ReadObject(ctx, path) + if err != nil { + logger.ErrorContext(ctx, "gcs_gen_mismatch: failed to read object data from GCS", slog.String("id", vulnID), slog.Any("error", err)) + + return err + } + + var baseProto osvschema.Vulnerability + if err := proto.Unmarshal(data, &baseProto); err != nil { + logger.ErrorContext(ctx, "gcs_gen_mismatch: failed to unmarshal proto from GCS", slog.String("id", vulnID), slog.Any("error", err)) + + return nil + } + + var modified time.Time + var finalProto *osvschema.Vulnerability + vulnKey := datastore.NameKey("Vulnerability", vulnID, nil) + dsFound := true + + _, txErr := r.stores.DatastoreClient.RunInTransaction(ctx, func(tx *datastore.Transaction) error { + var dsVuln osvdatastore.Vulnerability + if err := tx.Get(vulnKey, &dsVuln); err != nil { + if errors.Is(err, datastore.ErrNoSuchEntity) { + logger.ErrorContext(ctx, "vulnerability not in Datastore", slog.String("id", vulnID)) + dsFound = false + + return nil + } + + return err + } + modified = dsVuln.Modified + vulnProto := proto.Clone(&baseProto).(*osvschema.Vulnerability) + + fields := strings.Split(fieldStr, ",") + for _, f := range fields { + switch strings.TrimSpace(f) { + case "aliases": + aliasResult, err := r.stores.Relations.GetAliases(ctx, vulnID) + var aliases []string + var aliasesModified time.Time + if errors.Is(err, models.ErrNotFound) { + aliases = []string{} + aliasesModified = time.Now().UTC() + } else if err != nil { + return fmt.Errorf("failed to get aliases for %s: %w", vulnID, err) + } else { + aliases = aliasResult.Aliases + aliasesModified = aliasResult.Modified + } + if !slices.Equal(vulnProto.GetAliases(), aliases) { + vulnProto.Aliases = aliases + if aliasesModified.After(modified) { + modified = aliasesModified + } else { + modified = time.Now().UTC() + } + } + + case "upstream": + upstreamResult, err := r.stores.Relations.GetUpstream(ctx, vulnID) + var upstream []string + var upstreamModified time.Time + if errors.Is(err, models.ErrNotFound) { + upstream = []string{} + upstreamModified = time.Now().UTC() + } else if err != nil { + return fmt.Errorf("failed to get upstream for %s: %w", vulnID, err) + } else { + upstream = upstreamResult.Upstream + upstreamModified = upstreamResult.Modified + } + if !slices.Equal(vulnProto.GetUpstream(), upstream) { + vulnProto.Upstream = upstream + if upstreamModified.After(modified) { + modified = upstreamModified + } else { + modified = time.Now().UTC() + } + } + + case "related": + relatedResult, err := r.stores.Relations.GetRelated(ctx, vulnID) + var related []string + var relatedModified time.Time + if errors.Is(err, models.ErrNotFound) { + related = []string{} + relatedModified = time.Now().UTC() + } else if err != nil { + return fmt.Errorf("failed to get related for %s: %w", vulnID, err) + } else { + related = relatedResult.Related + relatedModified = relatedResult.Modified + } + if !slices.Equal(vulnProto.GetRelated(), related) { + vulnProto.Related = related + if relatedModified.After(modified) { + modified = relatedModified + } else { + modified = time.Now().UTC() + } + } + } + } + + vulnProto.Modified = timestamppb.New(modified) + dsVuln.Modified = modified + listedVuln := osvdatastore.NewListedVulnerabilityFromProto(vulnProto) + listedKey := datastore.NameKey("ListedVulnerability", vulnID, nil) + + if _, err := tx.Put(vulnKey, &dsVuln); err != nil { + return err + } + if _, err := tx.Put(listedKey, listedVuln); err != nil { + return err + } + + finalProto = vulnProto + + return nil + }) + + if txErr != nil { + logger.ErrorContext(ctx, "gcs_gen_mismatch: Datastore transaction failed", + slog.String("id", vulnID), slog.String("field", fieldStr), slog.Any("error", txErr)) + + return txErr + } + + if !dsFound { + return nil + } + + newData, err := proto.Marshal(finalProto) + if err != nil { + logger.ErrorContext(ctx, "gcs_gen_mismatch: failed to marshal proto", slog.String("id", vulnID), slog.Any("error", err)) + + return err + } + opts := &clients.WriteOptions{ + IfGenerationMatches: &generation, + CustomTime: &modified, + ContentType: "application/octet-stream", + } + if err := r.stores.GCS.WriteObject(ctx, path, newData, opts); err != nil { + logger.ErrorContext(ctx, "gcs_gen_mismatch: Writing to bucket failed", + slog.String("id", vulnID), slog.String("field", fieldStr), slog.Any("error", err)) + + return err + } + + return nil +} + +// HandleGeneric handles unhandled task types by logging and acknowledging. +func (r *Recoverer) HandleGeneric(ctx context.Context, m *pubsub.Message) error { + taskType := m.Attributes["type"] + if taskType == "" { + taskType = "unknown" + } + logger.ErrorContext(ctx, fmt.Sprintf("`%s` task could not be processed", taskType), + slog.String("type", taskType), + slog.Any("attributes", m.Attributes)) + + return nil +} diff --git a/go/internal/recoverer/recoverer_test.go b/go/internal/recoverer/recoverer_test.go new file mode 100644 index 00000000000..dde482d5179 --- /dev/null +++ b/go/internal/recoverer/recoverer_test.go @@ -0,0 +1,789 @@ +// 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 recoverer_test + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "cloud.google.com/go/datastore" + "cloud.google.com/go/pubsub/v2" + "github.com/google/go-cmp/cmp" + osvdatastore "github.com/google/osv.dev/go/internal/database/datastore" + "github.com/google/osv.dev/go/internal/gitter" + gitterpb "github.com/google/osv.dev/go/internal/gitter/pb/repository" + "github.com/google/osv.dev/go/internal/models" + "github.com/google/osv.dev/go/internal/recoverer" + "github.com/google/osv.dev/go/osv/clients" + "github.com/google/osv.dev/go/testutils" + "github.com/ossf/osv-schema/bindings/go/osvschema" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/testing/protocmp" + "google.golang.org/protobuf/types/known/timestamppb" +) + +type mockGitterClient struct { + fileContentFunc func(ctx context.Context, req *gitterpb.FileContentRequest) (*gitterpb.FileContentResponse, error) +} + +func (m *mockGitterClient) GetGit(_ context.Context, _ string, _ bool) (io.ReadCloser, error) { + return nil, errors.New("not implemented") +} + +func (m *mockGitterClient) Cache(_ context.Context, _ string) error { + return nil +} + +func (m *mockGitterClient) GetTags(_ context.Context, _ string) (*gitterpb.TagsResponse, error) { + return nil, errors.New("not implemented") +} + +func (m *mockGitterClient) GetAffectedCommits(_ context.Context, _ *gitterpb.AffectedCommitsRequest) (*gitterpb.AffectedCommitsResponse, error) { + return nil, errors.New("not implemented") +} + +func (m *mockGitterClient) GetFileDiffs(_ context.Context, _ *gitterpb.FileDiffsRequest) (*gitterpb.FileDiffsResponse, error) { + return nil, errors.New("not implemented") +} + +func (m *mockGitterClient) GetFileContent(ctx context.Context, req *gitterpb.FileContentRequest) (*gitterpb.FileContentResponse, error) { + if m.fileContentFunc != nil { + return m.fileContentFunc(ctx, req) + } + + return &gitterpb.FileContentResponse{}, nil +} + +var _ gitter.Client = (*mockGitterClient)(nil) + +type mockStorageProvider struct { + storage clients.CloudStorage +} + +func (m *mockStorageProvider) Bucket(_ string) clients.CloudStorage { + return m.storage +} + +func TestHandleGCSRetry(t *testing.T) { + ctx := context.Background() + gcsMock := testutils.NewMockStorage() + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + GCS: gcsMock, + }, + }) + + modified := time.Date(2025, 5, 5, 0, 0, 0, 0, time.UTC) + vuln := &osvschema.Vulnerability{ + Id: "TEST-555", + Modified: timestamppb.New(modified), + } + vulnBytes, err := proto.Marshal(vuln) + if err != nil { + t.Fatalf("Failed to marshal proto: %v", err) + } + + msg := &pubsub.Message{ + Data: vulnBytes, + } + + if err := rec.HandleGCSRetry(ctx, msg); err != nil { + t.Fatalf("HandleGCSRetry failed: %v", err) + } + + // Verify object was written to GCS + attrs, err := gcsMock.ReadObjectAttrs(ctx, "all/pb/TEST-555.pb") + if err != nil { + t.Fatalf("Failed to read object attrs: %v", err) + } + if !attrs.CustomTime.Equal(modified) { + t.Errorf("CustomTime mismatch: got %v, want %v", attrs.CustomTime, modified) + } + + data, err := gcsMock.ReadObject(ctx, "all/pb/TEST-555.pb") + if err != nil { + t.Fatalf("Failed to read object data: %v", err) + } + var storedVuln osvschema.Vulnerability + if err := proto.Unmarshal(data, &storedVuln); err != nil { + t.Fatalf("Failed to unmarshal stored proto: %v", err) + } + if diff := cmp.Diff(vuln, &storedVuln, protocmp.Transform()); diff != "" { + t.Errorf("Stored proto mismatch (-want +got):\n%s", diff) + } +} + +func TestHandleGCSRetry_Overwritten(t *testing.T) { + ctx := context.Background() + gcsMock := testutils.NewMockStorage() + + // Initial newer object in GCS + newerTime := time.Date(2025, 6, 1, 0, 0, 0, 0, time.UTC) + initialVuln := &osvschema.Vulnerability{ + Id: "TEST-123", + Modified: timestamppb.New(newerTime), + } + initialBytes, err := proto.Marshal(initialVuln) + if err != nil { + t.Fatalf("Failed to marshal proto: %v", err) + } + err = gcsMock.WriteObject(ctx, "all/pb/TEST-123.pb", initialBytes, &clients.WriteOptions{ + CustomTime: &newerTime, + }) + if err != nil { + t.Fatalf("Failed to write initial object: %v", err) + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + GCS: gcsMock, + }, + }) + + // Older message + olderTime := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC) + olderVuln := &osvschema.Vulnerability{ + Id: "TEST-123", + Modified: timestamppb.New(olderTime), + } + olderBytes, err := proto.Marshal(olderVuln) + if err != nil { + t.Fatalf("Failed to marshal proto: %v", err) + } + + msg := &pubsub.Message{ + Data: olderBytes, + } + + if err := rec.HandleGCSRetry(ctx, msg); err != nil { + t.Fatalf("HandleGCSRetry returned error: %v", err) + } + + // Verify the object in GCS remains the newer one + attrs, err := gcsMock.ReadObjectAttrs(ctx, "all/pb/TEST-123.pb") + if err != nil { + t.Fatalf("Failed to read object attrs: %v", err) + } + if !attrs.CustomTime.Equal(newerTime) { + t.Errorf("CustomTime was overwritten: got %v, want %v", attrs.CustomTime, newerTime) + } +} + +func TestHandleGCSRetry_InvalidData(t *testing.T) { + ctx := context.Background() + gcsMock := testutils.NewMockStorage() + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + GCS: gcsMock, + }, + }) + + msg := &pubsub.Message{ + Data: []byte("invalid-protobuf-bytes"), + } + + // Should not error (non-retryable invalid data is acknowledged) + if err := rec.HandleGCSRetry(ctx, msg); err != nil { + t.Fatalf("HandleGCSRetry returned error on invalid data: %v", err) + } +} + +func TestHandleGCSMissing_Git(t *testing.T) { + ctx := context.Background() + dsClient := testutils.MustNewDatastoreClientForTesting(t) + gcsMock := testutils.NewMockStorage() + publisherMock := &testutils.MockPublisher{} + + // Setup SourceRepository and Vulnerability in Datastore + sourceRepoStore := osvdatastore.NewSourceRepositoryStore(dsClient) + err := sourceRepoStore.Update(ctx, "test-git", &models.SourceRepository{ + Name: "test-git", + Type: models.SourceRepositoryTypeGit, + Git: &models.SourceRepoGit{ + URL: "https://github.com/google/test-repo.git", + Branch: "main", + LastSyncedCommit: "abcdef123456", + }, + Extension: ".yaml", + }) + if err != nil { + t.Fatalf("Failed to create SourceRepository: %v", err) + } + + vulnKey := datastore.NameKey("Vulnerability", "TEST-GIT-1", nil) + _, err = dsClient.Put(ctx, vulnKey, &osvdatastore.Vulnerability{ + SourceID: "test-git:vulns/TEST-GIT-1.yaml", + Modified: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }) + if err != nil { + t.Fatalf("Failed to put Vulnerability: %v", err) + } + + gitterMock := &mockGitterClient{ + fileContentFunc: func(_ context.Context, req *gitterpb.FileContentRequest) (*gitterpb.FileContentResponse, error) { + if req.GetUrl() != "https://github.com/google/test-repo.git" { + return nil, fmt.Errorf("unexpected URL: %s", req.GetUrl()) + } + if req.GetCommit() != "abcdef123456" { + return nil, fmt.Errorf("unexpected commit: %s", req.GetCommit()) + } + if req.GetPath() != "vulns/TEST-GIT-1.yaml" { + return nil, fmt.Errorf("unexpected path: %s", req.GetPath()) + } + + yamlContent := `id: TEST-GIT-1 +modified: "2025-01-01T00:00:00Z" +summary: "Git test summary" +` + + return &gitterpb.FileContentResponse{Content: []byte(yamlContent)}, nil + }, + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + SourceRepo: sourceRepoStore, + GCS: gcsMock, + Publisher: publisherMock, + }, + GitterClient: gitterMock, + }) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "id": "TEST-GIT-1", + }, + } + + if err := rec.HandleGCSMissing(ctx, msg); err != nil { + t.Fatalf("HandleGCSMissing failed: %v", err) + } + + if len(publisherMock.Messages) != 1 { + t.Fatalf("Expected 1 published message, got %d", len(publisherMock.Messages)) + } + + pubMsg := publisherMock.Messages[0] + if pubMsg.Attributes["type"] != "update" { + t.Errorf("Expected type update, got %s", pubMsg.Attributes["type"]) + } + if pubMsg.Attributes["source"] != "test-git" { + t.Errorf("Expected source test-git, got %s", pubMsg.Attributes["source"]) + } + if pubMsg.Attributes["path"] != "vulns/TEST-GIT-1.yaml" { + t.Errorf("Expected path vulns/TEST-GIT-1.yaml, got %s", pubMsg.Attributes["path"]) + } + if pubMsg.Attributes["skip_hash_check"] != "true" { + t.Errorf("Expected skip_hash_check true, got %s", pubMsg.Attributes["skip_hash_check"]) + } + + var pubVuln osvschema.Vulnerability + if err := proto.Unmarshal(pubMsg.Data, &pubVuln); err != nil { + t.Fatalf("Failed to unmarshal published proto: %v", err) + } + if pubVuln.GetId() != "TEST-GIT-1" { + t.Errorf("Expected vuln ID TEST-GIT-1, got %s", pubVuln.GetId()) + } + if pubVuln.GetSummary() != "Git test summary" { + t.Errorf("Expected summary 'Git test summary', got %s", pubVuln.GetSummary()) + } +} + +func TestHandleGCSMissing_Bucket(t *testing.T) { + ctx := context.Background() + dsClient := testutils.MustNewDatastoreClientForTesting(t) + gcsMock := testutils.NewMockStorage() + bucketMock := testutils.NewMockStorage() + publisherMock := &testutils.MockPublisher{} + + sourceRepoStore := osvdatastore.NewSourceRepositoryStore(dsClient) + err := sourceRepoStore.Update(ctx, "test-bucket", &models.SourceRepository{ + Name: "test-bucket", + Type: models.SourceRepositoryTypeBucket, + Bucket: &models.SourceRepoBucket{ + Name: "test-bucket-name", + }, + Extension: ".json", + }) + if err != nil { + t.Fatalf("Failed to create SourceRepository: %v", err) + } + + vulnKey := datastore.NameKey("Vulnerability", "TEST-BUCKET-1", nil) + _, err = dsClient.Put(ctx, vulnKey, &osvdatastore.Vulnerability{ + SourceID: "test-bucket:vulns/TEST-BUCKET-1.json", + Modified: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }) + if err != nil { + t.Fatalf("Failed to put Vulnerability: %v", err) + } + + jsonContent := `{"id": "TEST-BUCKET-1", "modified": "2025-01-01T00:00:00Z", "summary": "Bucket test summary"}` + err = bucketMock.WriteObject(ctx, "vulns/TEST-BUCKET-1.json", []byte(jsonContent), nil) + if err != nil { + t.Fatalf("Failed to write to bucket mock: %v", err) + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + SourceRepo: sourceRepoStore, + GCS: gcsMock, + GCSProvider: &mockStorageProvider{storage: bucketMock}, + Publisher: publisherMock, + }, + }) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "id": "TEST-BUCKET-1", + }, + } + + if err := rec.HandleGCSMissing(ctx, msg); err != nil { + t.Fatalf("HandleGCSMissing failed: %v", err) + } + + if len(publisherMock.Messages) != 1 { + t.Fatalf("Expected 1 published message, got %d", len(publisherMock.Messages)) + } + pubMsg := publisherMock.Messages[0] + var pubVuln osvschema.Vulnerability + if err := proto.Unmarshal(pubMsg.Data, &pubVuln); err != nil { + t.Fatalf("Failed to unmarshal published proto: %v", err) + } + if pubVuln.GetId() != "TEST-BUCKET-1" { + t.Errorf("Expected vuln ID TEST-BUCKET-1, got %s", pubVuln.GetId()) + } +} + +func TestHandleGCSMissing_REST(t *testing.T) { + ctx := context.Background() + dsClient := testutils.MustNewDatastoreClientForTesting(t) + gcsMock := testutils.NewMockStorage() + publisherMock := &testutils.MockPublisher{} + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/vulns/TEST-REST-1.json" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id": "TEST-REST-1", "modified": "2025-01-01T00:00:00Z", "summary": "REST test summary"}`)) + })) + defer server.Close() + + sourceRepoStore := osvdatastore.NewSourceRepositoryStore(dsClient) + err := sourceRepoStore.Update(ctx, "test-rest", &models.SourceRepository{ + Name: "test-rest", + Type: models.SourceRepositoryTypeREST, + Link: server.URL, + Extension: ".json", + }) + if err != nil { + t.Fatalf("Failed to create SourceRepository: %v", err) + } + + vulnKey := datastore.NameKey("Vulnerability", "TEST-REST-1", nil) + _, err = dsClient.Put(ctx, vulnKey, &osvdatastore.Vulnerability{ + SourceID: "test-rest:vulns/TEST-REST-1.json", + Modified: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }) + if err != nil { + t.Fatalf("Failed to put Vulnerability: %v", err) + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + SourceRepo: sourceRepoStore, + GCS: gcsMock, + Publisher: publisherMock, + }, + HTTPClient: server.Client(), + }) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "id": "TEST-REST-1", + }, + } + + if err := rec.HandleGCSMissing(ctx, msg); err != nil { + t.Fatalf("HandleGCSMissing failed: %v", err) + } + + if len(publisherMock.Messages) != 1 { + t.Fatalf("Expected 1 published message, got %d", len(publisherMock.Messages)) + } + pubMsg := publisherMock.Messages[0] + var pubVuln osvschema.Vulnerability + if err := proto.Unmarshal(pubMsg.Data, &pubVuln); err != nil { + t.Fatalf("Failed to unmarshal published proto: %v", err) + } + if pubVuln.GetId() != "TEST-REST-1" { + t.Errorf("Expected vuln ID TEST-REST-1, got %s", pubVuln.GetId()) + } +} + +func TestHandleGCSGenMismatch_Aliases(t *testing.T) { + ctx := context.Background() + dsClient := testutils.MustNewDatastoreClientForTesting(t) + gcsMock := testutils.NewMockStorage() + relationsStore := osvdatastore.NewRelationsStore(dsClient) + + // Populate Datastore + aliasMod := time.Date(2025, 2, 2, 0, 0, 0, 0, time.UTC) + _, err := dsClient.Put(ctx, datastore.IncompleteKey("AliasGroup", nil), &osvdatastore.AliasGroup{ + VulnIDs: []string{"CVE-111", "OSV-111", "TEST-111"}, + Modified: aliasMod, + }) + if err != nil { + t.Fatalf("Failed to put AliasGroup: %v", err) + } + + vulnMod := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + vulnKey := datastore.NameKey("Vulnerability", "TEST-111", nil) + _, err = dsClient.Put(ctx, vulnKey, &osvdatastore.Vulnerability{ + SourceID: "test:TEST-111.yaml", + Modified: vulnMod, + }) + if err != nil { + t.Fatalf("Failed to put Vulnerability: %v", err) + } + + // GCS has TEST-111 with no aliases + vuln := &osvschema.Vulnerability{ + Id: "TEST-111", + Modified: timestamppb.New(vulnMod), + } + vulnBytes, _ := proto.Marshal(vuln) + err = gcsMock.WriteObject(ctx, "all/pb/TEST-111.pb", vulnBytes, &clients.WriteOptions{ + CustomTime: &vulnMod, + }) + if err != nil { + t.Fatalf("Failed to write to GCS: %v", err) + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + Relations: relationsStore, + GCS: gcsMock, + }, + }) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "id": "TEST-111", + "field": "aliases", + }, + } + + if err := rec.HandleGCSGenMismatch(ctx, msg); err != nil { + t.Fatalf("HandleGCSGenMismatch failed: %v", err) + } + + // Read from GCS and verify aliases were updated + gcsData, err := gcsMock.ReadObject(ctx, "all/pb/TEST-111.pb") + if err != nil { + t.Fatalf("Failed to read GCS object: %v", err) + } + var updatedVuln osvschema.Vulnerability + if err := proto.Unmarshal(gcsData, &updatedVuln); err != nil { + t.Fatalf("Failed to unmarshal GCS proto: %v", err) + } + + wantAliases := []string{"CVE-111", "OSV-111"} + if diff := cmp.Diff(wantAliases, updatedVuln.GetAliases()); diff != "" { + t.Errorf("Aliases mismatch (-want +got):\n%s", diff) + } + if !updatedVuln.GetModified().AsTime().Equal(aliasMod) { + t.Errorf("Modified time mismatch: got %v, want %v", updatedVuln.GetModified().AsTime(), aliasMod) + } + + // Check Datastore Vulnerability modified time + var dv osvdatastore.Vulnerability + if err := dsClient.Get(ctx, vulnKey, &dv); err != nil { + t.Fatalf("Failed to get Vulnerability: %v", err) + } + if !dv.Modified.Equal(aliasMod) { + t.Errorf("Datastore Vulnerability modified mismatch: got %v, want %v", dv.Modified, aliasMod) + } +} + +func TestHandleGCSGenMismatch_Upstream(t *testing.T) { + ctx := context.Background() + dsClient := testutils.MustNewDatastoreClientForTesting(t) + gcsMock := testutils.NewMockStorage() + relationsStore := osvdatastore.NewRelationsStore(dsClient) + + upstreamMod := time.Date(2025, 2, 2, 0, 0, 0, 0, time.UTC) + upstreamKey := datastore.NameKey("UpstreamGroup", "TEST-111", nil) + _, err := dsClient.Put(ctx, upstreamKey, &osvdatastore.UpstreamGroup{ + VulnID: "TEST-111", + UpstreamIDs: []string{"UPSTREAM-1"}, + Modified: upstreamMod, + }) + if err != nil { + t.Fatalf("Failed to put UpstreamGroup: %v", err) + } + + vulnMod := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + vulnKey := datastore.NameKey("Vulnerability", "TEST-111", nil) + _, err = dsClient.Put(ctx, vulnKey, &osvdatastore.Vulnerability{ + SourceID: "test:TEST-111.yaml", + Modified: vulnMod, + }) + if err != nil { + t.Fatalf("Failed to put Vulnerability: %v", err) + } + + vuln := &osvschema.Vulnerability{ + Id: "TEST-111", + Modified: timestamppb.New(vulnMod), + } + vulnBytes, _ := proto.Marshal(vuln) + err = gcsMock.WriteObject(ctx, "all/pb/TEST-111.pb", vulnBytes, &clients.WriteOptions{ + CustomTime: &vulnMod, + }) + if err != nil { + t.Fatalf("Failed to write to GCS: %v", err) + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + Relations: relationsStore, + GCS: gcsMock, + }, + }) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "id": "TEST-111", + "field": "upstream", + }, + } + + if err := rec.HandleGCSGenMismatch(ctx, msg); err != nil { + t.Fatalf("HandleGCSGenMismatch failed: %v", err) + } + + gcsData, err := gcsMock.ReadObject(ctx, "all/pb/TEST-111.pb") + if err != nil { + t.Fatalf("Failed to read GCS object: %v", err) + } + var updatedVuln osvschema.Vulnerability + if err := proto.Unmarshal(gcsData, &updatedVuln); err != nil { + t.Fatalf("Failed to unmarshal GCS proto: %v", err) + } + + wantUpstream := []string{"UPSTREAM-1"} + if diff := cmp.Diff(wantUpstream, updatedVuln.GetUpstream()); diff != "" { + t.Errorf("Upstream mismatch (-want +got):\n%s", diff) + } + if !updatedVuln.GetModified().AsTime().Equal(upstreamMod) { + t.Errorf("Modified time mismatch: got %v, want %v", updatedVuln.GetModified().AsTime(), upstreamMod) + } +} + +func TestHandleGCSGenMismatch_Related(t *testing.T) { + ctx := context.Background() + dsClient := testutils.MustNewDatastoreClientForTesting(t) + gcsMock := testutils.NewMockStorage() + relationsStore := osvdatastore.NewRelationsStore(dsClient) + + relatedMod := time.Date(2025, 2, 2, 0, 0, 0, 0, time.UTC) + relatedKey := datastore.NameKey("RelatedGroup", "TEST-111", nil) + _, err := dsClient.Put(ctx, relatedKey, &osvdatastore.RelatedGroup{ + RelatedIDs: []string{"RELATED-1"}, + Modified: relatedMod, + }) + if err != nil { + t.Fatalf("Failed to put RelatedGroup: %v", err) + } + + vulnMod := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + vulnKey := datastore.NameKey("Vulnerability", "TEST-111", nil) + _, err = dsClient.Put(ctx, vulnKey, &osvdatastore.Vulnerability{ + SourceID: "test:TEST-111.yaml", + Modified: vulnMod, + }) + if err != nil { + t.Fatalf("Failed to put Vulnerability: %v", err) + } + + vuln := &osvschema.Vulnerability{ + Id: "TEST-111", + Modified: timestamppb.New(vulnMod), + } + vulnBytes, _ := proto.Marshal(vuln) + err = gcsMock.WriteObject(ctx, "all/pb/TEST-111.pb", vulnBytes, &clients.WriteOptions{ + CustomTime: &vulnMod, + }) + if err != nil { + t.Fatalf("Failed to write to GCS: %v", err) + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + Relations: relationsStore, + GCS: gcsMock, + }, + }) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "id": "TEST-111", + "field": "related", + }, + } + + if err := rec.HandleGCSGenMismatch(ctx, msg); err != nil { + t.Fatalf("HandleGCSGenMismatch failed: %v", err) + } + + gcsData, err := gcsMock.ReadObject(ctx, "all/pb/TEST-111.pb") + if err != nil { + t.Fatalf("Failed to read GCS object: %v", err) + } + var updatedVuln osvschema.Vulnerability + if err := proto.Unmarshal(gcsData, &updatedVuln); err != nil { + t.Fatalf("Failed to unmarshal GCS proto: %v", err) + } + + wantRelated := []string{"RELATED-1"} + if diff := cmp.Diff(wantRelated, updatedVuln.GetRelated()); diff != "" { + t.Errorf("Related mismatch (-want +got):\n%s", diff) + } + if !updatedVuln.GetModified().AsTime().Equal(relatedMod) { + t.Errorf("Modified time mismatch: got %v, want %v", updatedVuln.GetModified().AsTime(), relatedMod) + } +} + +func TestHandleGCSGenMismatch_MultipleFieldsWithSpaces(t *testing.T) { + ctx := context.Background() + dsClient := testutils.MustNewDatastoreClientForTesting(t) + gcsMock := testutils.NewMockStorage() + relationsStore := osvdatastore.NewRelationsStore(dsClient) + + aliasMod := time.Date(2025, 2, 2, 0, 0, 0, 0, time.UTC) + _, err := dsClient.Put(ctx, datastore.IncompleteKey("AliasGroup", nil), &osvdatastore.AliasGroup{ + VulnIDs: []string{"CVE-222", "TEST-222"}, + Modified: aliasMod, + }) + if err != nil { + t.Fatalf("Failed to put AliasGroup: %v", err) + } + + upstreamMod := time.Date(2025, 3, 3, 0, 0, 0, 0, time.UTC) + upstreamKey := datastore.NameKey("UpstreamGroup", "TEST-222", nil) + _, err = dsClient.Put(ctx, upstreamKey, &osvdatastore.UpstreamGroup{ + VulnID: "TEST-222", + UpstreamIDs: []string{"UPSTREAM-2"}, + Modified: upstreamMod, + }) + if err != nil { + t.Fatalf("Failed to put UpstreamGroup: %v", err) + } + + vulnMod := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + vulnKey := datastore.NameKey("Vulnerability", "TEST-222", nil) + _, err = dsClient.Put(ctx, vulnKey, &osvdatastore.Vulnerability{ + SourceID: "test:TEST-222.yaml", + Modified: vulnMod, + }) + if err != nil { + t.Fatalf("Failed to put Vulnerability: %v", err) + } + + vuln := &osvschema.Vulnerability{ + Id: "TEST-222", + Modified: timestamppb.New(vulnMod), + } + vulnBytes, _ := proto.Marshal(vuln) + err = gcsMock.WriteObject(ctx, "all/pb/TEST-222.pb", vulnBytes, &clients.WriteOptions{ + CustomTime: &vulnMod, + }) + if err != nil { + t.Fatalf("Failed to write to GCS: %v", err) + } + + rec := recoverer.New(recoverer.Config{ + Stores: recoverer.Stores{ + DatastoreClient: dsClient, + Relations: relationsStore, + GCS: gcsMock, + }, + }) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "id": "TEST-222", + "field": "aliases, upstream", + }, + } + + if err := rec.HandleGCSGenMismatch(ctx, msg); err != nil { + t.Fatalf("HandleGCSGenMismatch failed: %v", err) + } + + gcsData, err := gcsMock.ReadObject(ctx, "all/pb/TEST-222.pb") + if err != nil { + t.Fatalf("Failed to read GCS object: %v", err) + } + var updatedVuln osvschema.Vulnerability + if err := proto.Unmarshal(gcsData, &updatedVuln); err != nil { + t.Fatalf("Failed to unmarshal GCS proto: %v", err) + } + + wantAliases := []string{"CVE-222"} + if diff := cmp.Diff(wantAliases, updatedVuln.GetAliases()); diff != "" { + t.Errorf("Aliases mismatch (-want +got):\n%s", diff) + } + wantUpstream := []string{"UPSTREAM-2"} + if diff := cmp.Diff(wantUpstream, updatedVuln.GetUpstream()); diff != "" { + t.Errorf("Upstream mismatch (-want +got):\n%s", diff) + } + if !updatedVuln.GetModified().AsTime().Equal(upstreamMod) { + t.Errorf("Modified time mismatch: got %v, want %v", updatedVuln.GetModified().AsTime(), upstreamMod) + } +} + +func TestHandleGeneric(t *testing.T) { + ctx := context.Background() + rec := recoverer.New(recoverer.Config{}) + + msg := &pubsub.Message{ + Attributes: map[string]string{ + "type": "custom_unknown_task", + }, + } + + if err := rec.HandleGeneric(ctx, msg); err != nil { + t.Fatalf("HandleGeneric returned error: %v", err) + } +}