diff --git a/AGENTS.md b/AGENTS.md index eaa8ed5214b..6695b8c682c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -105,7 +105,16 @@ Always format and lint your code before proposing changes. The repository provid - **Rule**: When running Python scripts, always use `poetry run`. ### Go Standards -- Linter: `golangci-lint` (run automatically by the lint script per module). +- Linter: `golangci-lint` +- **Running Go Linters**: Run `golangci-lint` using `go run` directly within the module directory (`go/`, `vulnfeeds/`, or `bindings/go/`): + ```bash + cd go && go run github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.4.0 run ./... + ``` + *(Note: Run outside the sandbox so `go run` can fetch the linter toolchain if not cached).* +- **Formatting Command**: + ```bash + cd go && gofmt -s -w . + ``` - **Rule**: Go code must follow standard Go formatting guidelines. ### Git Commit Guidelines diff --git a/deployment/clouddeploy/gke-workers/base/core/exporter.yaml b/deployment/clouddeploy/gke-workers/base/core/exporter.yaml index b76b3a1cffc..42dbe7c1370 100644 --- a/deployment/clouddeploy/gke-workers/base/core/exporter.yaml +++ b/deployment/clouddeploy/gke-workers/base/core/exporter.yaml @@ -15,24 +15,29 @@ spec: spec: template: spec: - tolerations: - - key: workloadType - operator: Equal - value: highend - nodeSelector: - cloud.google.com/gke-nodepool: highend containers: - name: exporter image: exporter imagePullPolicy: Always + volumeMounts: + - mountPath: "/scratch" + name: "scratch-volume" resources: requests: - cpu: "12" - memory: "100Gi" + cpu: "4" + memory: "4G" + ephemeral-storage: "35G" limits: - cpu: "20" - memory: "150Gi" + cpu: "7" + memory: "8G" + ephemeral-storage: "50G" env: - name: TRACE_SAMPLE_RATE value: "0.0" + - name: SCRATCH_DIR + value: "/scratch" + volumes: + - name: "scratch-volume" + emptyDir: + sizeLimit: "50G" restartPolicy: Never diff --git a/go/cmd/exporter/downloader.go b/go/cmd/exporter/downloader.go index 4858c8480b4..0b66c3d3e24 100644 --- a/go/cmd/exporter/downloader.go +++ b/go/cmd/exporter/downloader.go @@ -3,6 +3,9 @@ package main import ( "context" "log/slog" + "os" + "path/filepath" + "strings" "sync" "github.com/google/osv.dev/go/internal/osvutil" @@ -12,10 +15,11 @@ import ( "google.golang.org/protobuf/proto" ) -// downloader is a worker that receives GCS object handles from inCh, downloads -// the raw protobuf data, unmarshals it into a Vulnerability, and sends the -// result to outCh. -func downloader(ctx context.Context, client clients.CloudStorage, inCh <-chan string, outCh chan<- *osvschema.Vulnerability, wg *sync.WaitGroup) { +// downloadThenProcessor is a worker that receives GCS object handles from inCh, downloads +// the raw protobuf data, unmarshals it into a Vulnerability, marshals it to compact +// JSON, saves it to scratch disk, queues individual JSON uploads, and sends the +// metadata to routerCh. +func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, client clients.CloudStorage, scratchDir string, inCh <-chan string, routerCh chan<- processedVuln, writeCh chan<- writeMsg, wg *sync.WaitGroup) { defer wg.Done() for path := range inCh { // Process object. @@ -38,9 +42,71 @@ func downloader(ctx context.Context, client clients.CloudStorage, inCh <-chan st continue } - // Wait to send the result, or be cancelled. + // Marshal JSON ONCE for this vulnerability. + b, err := marshalToJSON(vuln) + if err != nil { + logger.ErrorContext(ctx, "failed to marshal vulnerability to json", slog.String("id", vuln.GetId()), slog.Any("err", err)) + continue + } + + // Cache to local scratch disk for later ZIP generation. + localPath := filepath.Join(scratchDir, vuln.GetId()+".json") + if err := os.WriteFile(localPath, b, 0600); err != nil { + logger.ErrorContext(ctx, "failed to write cached vulnerability to disk", slog.String("id", vuln.GetId()), slog.Any("err", err)) + // Cancel the exporter context if writing to the scratch disk fails (e.g. disk full) + // to fail fast rather than producing incomplete archives later. + cancel() + + return + } + + hasVanir := false + // Check for Vanir signatures + for _, aff := range vuln.GetAffected() { + spec := aff.GetDatabaseSpecific() + if _, ok := spec.GetFields()["vanir_signatures"]; ok { + hasVanir = true + break + } + } + + ecosystems := make(map[string]struct{}) + for _, aff := range vuln.GetAffected() { + eco := aff.GetPackage().GetEcosystem() + eco, _, _ = strings.Cut(eco, ":") + if eco != "" { + ecosystems[eco] = struct{}{} + } + for _, ref := range aff.GetRanges() { + if ref.GetType() == osvschema.Range_GIT { + ecosystems["GIT"] = struct{}{} + break + } + } + } + if len(ecosystems) == 0 { + ecosystems["[EMPTY]"] = struct{}{} + } + ecoNames := make([]string, 0, len(ecosystems)) + for eco := range ecosystems { + ecoNames = append(ecoNames, eco) + select { + case writeCh <- writeMsg{path: filepath.Join(eco, vuln.GetId()) + ".json", mimeType: "application/json", data: b}: + case <-ctx.Done(): + return + } + } + + // Send processed metadata to router. select { - case outCh <- vuln: + case routerCh <- processedVuln{ + meta: vulnMeta{ + id: vuln.GetId(), + modified: vuln.GetModified().AsTime(), + }, + ecosystems: ecoNames, + hasVanir: hasVanir, + }: case <-ctx.Done(): return } diff --git a/go/cmd/exporter/exporter.go b/go/cmd/exporter/exporter.go index a76c1636a68..000136c22a8 100644 --- a/go/cmd/exporter/exporter.go +++ b/go/cmd/exporter/exporter.go @@ -8,7 +8,7 @@ import ( "log/slog" "os" "os/signal" - "strings" + "path/filepath" "sync" "syscall" @@ -16,7 +16,6 @@ import ( "github.com/google/osv.dev/go/internal/sharding" "github.com/google/osv.dev/go/logger" "github.com/google/osv.dev/go/osv/clients" - "github.com/ossf/osv-schema/bindings/go/osvschema" "go.opentelemetry.io/otel" "google.golang.org/api/option" ) @@ -35,20 +34,37 @@ func main() { ctx, span := otel.Tracer("exporter").Start(ctx, "exporter") defer span.End() + defaultScratchDir := os.Getenv("SCRATCH_DIR") + if defaultScratchDir == "" { + defaultScratchDir = filepath.Join(os.TempDir(), "osv-exporter-scratch") + } + outBucketName := flag.String("bucket", "osv-test-vulnerabilities", "Output bucket or directory name. If -local is true, this is a local path; otherwise, it's a GCS bucket name.") vulnBucketName := flag.String("osv-vulns-bucket", os.Getenv("OSV_VULNERABILITIES_BUCKET"), "GCS bucket to read vulnerability protobufs from. Can also be set with the OSV_VULNERABILITIES_BUCKET environment variable.") uploadToGCS := flag.Bool("upload-to-gcs", false, "If false, writes the output to a local directory specified by -bucket instead of a GCS bucket.") numWorkers := flag.Int("workers", 1000, "The total number of concurrent workers to use for downloading from GCS and writing the output.") breakdownPrefixesStr := flag.String("breakdown-prefixes", "", "Comma-separated list of prefix breakdowns for parallel GCS object listing.") + scratchDirFlag := flag.String("scratch-dir", defaultScratchDir, "Directory to stage temporary JSON and zip files.") flag.Parse() + scratchDir := *scratchDirFlag + if err := os.MkdirAll(scratchDir, 0755); err != nil { + logger.FatalContext(ctx, "failed to create scratch directory", slog.String("dir", scratchDir), slog.Any("err", err)) + } + scratchDir, err := os.MkdirTemp(scratchDir, "osv-exporter-*") + if err != nil { + logger.FatalContext(ctx, "failed to create temp directory in scratch dir", slog.String("dir", scratchDir), slog.Any("err", err)) + } + defer os.RemoveAll(scratchDir) + logger.InfoContext(ctx, "exporter starting", slog.String("bucket", *outBucketName), slog.String("osv-vulns-bucket", *vulnBucketName), slog.Bool("upload-to-gcs", *uploadToGCS), slog.Int("workers", *numWorkers), - slog.String("breakdown-prefixes", *breakdownPrefixesStr)) + slog.String("breakdown-prefixes", *breakdownPrefixesStr), + slog.String("scratch-dir", scratchDir)) if *vulnBucketName == "" { logger.FatalContext(ctx, "OSV_VULNERABILITIES_BUCKET must be set") @@ -74,35 +90,38 @@ func main() { } // The exporter uses a pipeline of channels and worker pools. The data flow is as follows: - // 1. The main goroutine lists GCS objects and sends them to `gcsPathToDownloaderCh`. - // 2. A pool of `downloader` workers receive GCS objects, downloads and unmarshals them into - // OSV vulnerabilities, and send them to `downloaderToRouterCh`. - // 3. The `ecosystemRouter` receives vulnerabilities and dispatches them. It creates a new + // 1. The main goroutine lists GCS objects and sends them to `gcsPathToProcessorCh`. + // 2. A pool of `downloadThenProcessor` workers receive GCS objects, downloads, unmarshals, + // marshals to JSON, saves to scratch disk, queues individual JSON writes to `writeCh`, + // and sends metadata to `processorToRouterCh`. + // 3. The `ecosystemRouter` receives metadata and dispatches it. It creates a new // `ecosystemWorker` for each new ecosystem, and sends all vulnerabilities to a single // `allEcosystemWorker`. - // 4. The `ecosystemWorker`s and the `allEcosystemWorker` process the vulnerabilities and - // generate the final files, sending the data to be written to `routerToWriteCh`. + // 4. The `ecosystemWorker`s and the `allEcosystemWorker` aggregate metadata and generate the + // final zip, csv, and ecosystems.txt files, sending the write requests to `writeCh`. // 5. A pool of `writer` workers receive the file data and write it to the output. - gcsPathToDownloaderCh := make(chan string, 100) - downloaderToRouterCh := make(chan *osvschema.Vulnerability, 100) - routerToWriteCh := make(chan writeMsg, 100) + gcsPathToProcessorCh := make(chan string, 100) + processorToRouterCh := make(chan processedVuln, 100) + writeCh := make(chan writeMsg, 100) breakdownPrefixes := sharding.ExpandBreakdownPrefixes(*breakdownPrefixesStr) - var downloaderWg sync.WaitGroup + var processorWg sync.WaitGroup for range *numWorkers / 2 { - downloaderWg.Add(1) - go downloader(ctx, vulnClient, gcsPathToDownloaderCh, downloaderToRouterCh, &downloaderWg) + processorWg.Add(1) + go downloadThenProcessor(ctx, cancel, vulnClient, scratchDir, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg) } var writerWg sync.WaitGroup for range *numWorkers / 2 { writerWg.Add(1) - go writer(ctx, cancel, routerToWriteCh, outClient, outPrefix, &writerWg) + go writer(ctx, cancel, writeCh, outClient, outPrefix, &writerWg) } var routerWg sync.WaitGroup routerWg.Add(1) - go ecosystemRouter(ctx, downloaderToRouterCh, routerToWriteCh, &routerWg) + go ecosystemRouter(ctx, processorToRouterCh, writeCh, scratchDir, &routerWg) + + logDiskUsage(ctx, scratchDir, "startup") MainLoop: for objName, err := range vulnClient.ObjectsFast(ctx, gcsProtoPrefix, breakdownPrefixes) { @@ -110,17 +129,19 @@ MainLoop: logger.FatalContext(ctx, "failed to list objects", slog.Any("err", err)) } select { - case gcsPathToDownloaderCh <- objName: + case gcsPathToProcessorCh <- objName: case <-ctx.Done(): break MainLoop } } - close(gcsPathToDownloaderCh) - downloaderWg.Wait() - close(downloaderToRouterCh) + close(gcsPathToProcessorCh) + processorWg.Wait() + logDiskUsage(ctx, scratchDir, "downloads_complete") + close(processorToRouterCh) routerWg.Wait() - close(routerToWriteCh) + logDiskUsage(ctx, scratchDir, "export_complete") + close(writeCh) writerWg.Wait() if ctx.Err() != nil { @@ -129,21 +150,48 @@ MainLoop: logger.InfoContext(ctx, "export completed successfully") } -// ecosystemRouter receives vulnerabilities from inCh and fans them out to the +// logDiskUsage logs the scratch filesystem space metrics. +func logDiskUsage(ctx context.Context, dir string, stage string) { + var stat syscall.Statfs_t + if err := syscall.Statfs(dir, &stat); err != nil { + logger.WarnContext(ctx, "failed to get scratch disk usage", slog.String("dir", dir), slog.Any("err", err)) + return + } + if stat.Bsize <= 0 || stat.Blocks == 0 { + return + } + + const gib = 1024 * 1024 * 1024 + blockSize := float64(stat.Bsize) + usedGB := float64(stat.Blocks-stat.Bfree) * blockSize / gib + totalGB := float64(stat.Blocks) * blockSize / gib + freeGB := float64(stat.Bavail) * blockSize / gib + + logger.InfoContext(ctx, "scratch disk usage", + slog.String("stage", stage), + slog.Float64("used_gb", usedGB), + slog.Float64("free_gb", freeGB), + slog.Float64("total_gb", totalGB), + slog.Float64("used_pct", usedGB/totalGB*100), + ) +} + +// ecosystemRouter receives processed vulnerabilities from inCh and fans them out to the // appropriate ecosystemWorker. It creates workers on-demand for each new // ecosystem encountered. It also sends every vulnerability to the allEcosystemWorker. -func ecosystemRouter(ctx context.Context, inCh <-chan *osvschema.Vulnerability, outCh chan<- writeMsg, wg *sync.WaitGroup) { +func ecosystemRouter(ctx context.Context, inCh <-chan processedVuln, outCh chan<- writeMsg, scratchDir string, wg *sync.WaitGroup) { defer wg.Done() logger.InfoContext(ctx, "ecosystem router starting") workers := make(map[string]*ecosystemWorker) var workersWg sync.WaitGroup vulnCounter := 0 + var vanirVulnIDs []string - allEcosystemWorker := newAllEcosystemWorker(ctx, outCh, &workersWg) + allEcosystemWorker := newAllEcosystemWorker(ctx, scratchDir, outCh, &workersWg) RouterLoop: for { - var vuln *osvschema.Vulnerability + var vuln processedVuln var ok bool select { case <-ctx.Done(): @@ -154,38 +202,25 @@ RouterLoop: } } vulnCounter++ - ecosystems := make(map[string]struct{}) - for _, aff := range vuln.GetAffected() { - eco := aff.GetPackage().GetEcosystem() - eco, _, _ = strings.Cut(eco, ":") - if eco != "" { - ecosystems[eco] = struct{}{} - } - for _, ref := range aff.GetRanges() { - if ref.GetType() == osvschema.Range_GIT { - ecosystems["GIT"] = struct{}{} - } - } - } - if len(ecosystems) == 0 { - ecosystems["[EMPTY]"] = struct{}{} + + if vuln.hasVanir { + vanirVulnIDs = append(vanirVulnIDs, vuln.meta.id) } - ecoNames := make([]string, 0, len(ecosystems)) - for eco := range ecosystems { - ecoNames = append(ecoNames, eco) + + for _, eco := range vuln.ecosystems { worker, ok := workers[eco] if !ok { - worker = newEcosystemWorker(ctx, eco, outCh, &workersWg) + worker = newEcosystemWorker(ctx, eco, scratchDir, outCh, &workersWg) workers[eco] = worker } select { - case worker.inCh <- vuln: + case worker.inCh <- vuln.meta: case <-ctx.Done(): break RouterLoop } } select { - case allEcosystemWorker.inCh <- vulnAndEcos{Vulnerability: vuln, ecosystems: ecoNames}: + case allEcosystemWorker.inCh <- vulnAndEcos{meta: vuln.meta, ecosystems: vuln.ecosystems}: case <-ctx.Done(): break RouterLoop } @@ -196,6 +231,11 @@ RouterLoop: } allEcosystemWorker.Finish() workersWg.Wait() + + if len(vanirVulnIDs) > 0 && ctx.Err() == nil { + writeVanir(ctx, vanirVulnIDs, outCh, scratchDir) + } + if ctx.Err() == nil { logger.InfoContext(ctx, "ecosystem router finished, all vulnerabilities dispatched", slog.Int("total_vulnerabilities", vulnCounter)) } else { diff --git a/go/cmd/exporter/exporter_test.go b/go/cmd/exporter/exporter_test.go new file mode 100644 index 00000000000..5d2f233fcfb --- /dev/null +++ b/go/cmd/exporter/exporter_test.go @@ -0,0 +1,233 @@ +package main + +import ( + "archive/zip" + "bytes" + "context" + "encoding/csv" + "encoding/json" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/google/osv.dev/go/testutils" + "github.com/ossf/osv-schema/bindings/go/osvschema" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/structpb" + "google.golang.org/protobuf/types/known/timestamppb" +) + +func TestWriter_GCS_StreamUpload(t *testing.T) { + storage := testutils.NewMockStorage() + tmpDir := t.TempDir() + tmpFile := filepath.Join(tmpDir, "sample.zip") + content := []byte("zip binary contents") + if err := os.WriteFile(tmpFile, content, 0600); err != nil { + t.Fatalf("failed to write temp file: %v", err) + } + + runWriter(t, storage, []writeMsg{ + {path: "all.zip", mimeType: "application/zip", filePath: tmpFile}, + }) + + objPath := filepath.Join("out", "all.zip") + got, err := storage.ReadObject(t.Context(), objPath) + if err != nil { + t.Fatalf("ReadObject(%s) failed: %v", objPath, err) + } + if !bytes.Equal(got, content) { + t.Errorf("expected %q, got %q", content, got) + } +} + +func TestWriter_GCS_StreamSkipsUnchanged(t *testing.T) { + storage := testutils.NewMockStorage() + tmpDir := t.TempDir() + tmpFile := filepath.Join(tmpDir, "sample.zip") + content := []byte("zip binary contents") + if err := os.WriteFile(tmpFile, content, 0600); err != nil { + t.Fatalf("failed to write temp file: %v", err) + } + + objPath := filepath.Join("out", "all.zip") + if err := storage.WriteObject(t.Context(), objPath, content, nil); err != nil { + t.Fatalf("setup WriteObject failed: %v", err) + } + attrsBefore, _ := storage.ReadObjectAttrs(t.Context(), objPath) + + runWriter(t, storage, []writeMsg{ + {path: "all.zip", mimeType: "application/zip", filePath: tmpFile}, + }) + + attrsAfter, err := storage.ReadObjectAttrs(t.Context(), objPath) + if err != nil { + t.Fatalf("ReadObjectAttrs failed: %v", err) + } + if attrsAfter.Generation != attrsBefore.Generation { + t.Errorf("expected generation %d (skipped), got %d", attrsBefore.Generation, attrsAfter.Generation) + } +} + +func TestExporterPipeline_EndToEnd(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + scratchDir := t.TempDir() + vulnStorage := testutils.NewMockStorage() + outStorage := testutils.NewMockStorage() + + time1 := time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC) + time2 := time.Date(2023, 2, 1, 12, 0, 0, 0, time.UTC) + + // Create test vulnerability 1: PyPI + vuln1 := &osvschema.Vulnerability{ + Id: "GHSA-pypi-1", + Modified: timestamppb.New(time1), + Affected: []*osvschema.Affected{ + { + Package: &osvschema.Package{ + Ecosystem: "PyPI", + Name: "requests", + }, + }, + }, + } + pb1, err := proto.Marshal(vuln1) + if err != nil { + t.Fatalf("failed to marshal proto: %v", err) + } + _ = vulnStorage.WriteObject(ctx, "all/pb/vuln1", pb1, nil) + + // Create test vulnerability 2: npm and GIT with Vanir signatures + vanirField, _ := structpb.NewValue("test-signature") + vuln2 := &osvschema.Vulnerability{ + Id: "GHSA-npm-git-2", + Modified: timestamppb.New(time2), + Affected: []*osvschema.Affected{ + { + Package: &osvschema.Package{ + Ecosystem: "npm", + Name: "lodash", + }, + Ranges: []*osvschema.Range{ + { + Type: osvschema.Range_GIT, + }, + }, + DatabaseSpecific: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "vanir_signatures": vanirField, + }, + }, + }, + }, + } + pb2, err := proto.Marshal(vuln2) + if err != nil { + t.Fatalf("failed to marshal proto: %v", err) + } + _ = vulnStorage.WriteObject(ctx, "all/pb/vuln2", pb2, nil) + + gcsPathToProcessorCh := make(chan string, 10) + processorToRouterCh := make(chan processedVuln, 10) + writeCh := make(chan writeMsg, 100) + + var processorWg sync.WaitGroup + processorWg.Add(1) + go downloadThenProcessor(ctx, cancel, vulnStorage, scratchDir, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg) + + var writerWg sync.WaitGroup + writerWg.Add(1) + go writer(ctx, cancel, writeCh, outStorage, "export", &writerWg) + + var routerWg sync.WaitGroup + routerWg.Add(1) + go ecosystemRouter(ctx, processorToRouterCh, writeCh, scratchDir, &routerWg) + + gcsPathToProcessorCh <- "all/pb/vuln1" + gcsPathToProcessorCh <- "all/pb/vuln2" + close(gcsPathToProcessorCh) + + processorWg.Wait() + close(processorToRouterCh) + routerWg.Wait() + close(writeCh) + writerWg.Wait() + + // 1. Verify individual JSON outputs + pypiJSON, err := outStorage.ReadObject(ctx, "export/PyPI/GHSA-pypi-1.json") + if err != nil { + t.Fatalf("expected PyPI/GHSA-pypi-1.json in storage: %v", err) + } + if !bytes.Contains(pypiJSON, []byte(`"id":"GHSA-pypi-1"`)) { + t.Errorf("PyPI JSON content mismatch: %s", string(pypiJSON)) + } + + npmJSON, err := outStorage.ReadObject(ctx, "export/npm/GHSA-npm-git-2.json") + if err != nil { + t.Fatalf("expected npm/GHSA-npm-git-2.json in storage: %v", err) + } + if !bytes.Contains(npmJSON, []byte(`"id":"GHSA-npm-git-2"`)) { + t.Errorf("npm JSON content mismatch: %s", string(npmJSON)) + } + + // 2. Verify all.zip contains all JSON files + allZipBytes, err := outStorage.ReadObject(ctx, "export/all.zip") + if err != nil { + t.Fatalf("expected all.zip: %v", err) + } + zipReader, err := zip.NewReader(bytes.NewReader(allZipBytes), int64(len(allZipBytes))) + if err != nil { + t.Fatalf("failed to open all.zip: %v", err) + } + zipNames := make([]string, 0, len(zipReader.File)) + for _, f := range zipReader.File { + zipNames = append(zipNames, f.Name) + } + if len(zipNames) != 2 { + t.Errorf("expected 2 files in all.zip, got %v", zipNames) + } + + // 3. Verify modified_id.csv ordering (descending by modified time) + csvBytes, err := outStorage.ReadObject(ctx, "export/modified_id.csv") + if err != nil { + t.Fatalf("expected modified_id.csv: %v", err) + } + csvReader := csv.NewReader(bytes.NewReader(csvBytes)) + records, err := csvReader.ReadAll() + if err != nil { + t.Fatalf("failed to parse modified_id.csv: %v", err) + } + // vuln2 has later modified time, so its ecosystems (GIT, npm) should appear before PyPI + if len(records) < 3 { + t.Fatalf("expected at least 3 records in modified_id.csv, got %d", len(records)) + } + if records[0][0] != time2.Format(time.RFC3339Nano) { + t.Errorf("expected first record to be time2, got %s", records[0][0]) + } + + // 4. Verify ecosystems.txt + ecoTxtBytes, err := outStorage.ReadObject(ctx, "export/ecosystems.txt") + if err != nil { + t.Fatalf("expected ecosystems.txt: %v", err) + } + expectedEco := "GIT\nPyPI\nnpm\n" + if string(ecoTxtBytes) != expectedEco { + t.Errorf("expected %q in ecosystems.txt, got %q", expectedEco, string(ecoTxtBytes)) + } + + // 5. Verify Vanir signatures file (GIT/osv_git.json) + vanirBytes, err := outStorage.ReadObject(ctx, "export/GIT/osv_git.json") + if err != nil { + t.Fatalf("expected GIT/osv_git.json: %v", err) + } + var vanirList []map[string]any + if err := json.Unmarshal(vanirBytes, &vanirList); err != nil { + t.Fatalf("failed to unmarshal vanir JSON: %v", err) + } + if len(vanirList) != 1 || vanirList[0]["id"] != "GHSA-npm-git-2" { + t.Errorf("unexpected vanir content: %s", string(vanirBytes)) + } +} diff --git a/go/cmd/exporter/worker.go b/go/cmd/exporter/worker.go index 4029ff2cf02..e5aab4e3243 100644 --- a/go/cmd/exporter/worker.go +++ b/go/cmd/exporter/worker.go @@ -10,6 +10,7 @@ import ( "io" "log/slog" "maps" + "os" "path/filepath" "slices" "strings" @@ -30,18 +31,39 @@ const ( ecosystemsFilename = "ecosystems.txt" ) +// vulnMeta holds the ID and modified time for a vulnerability. +type vulnMeta struct { + id string + modified time.Time +} + +// csvEntry holds the modified time and the relative entry path. +type csvEntry struct { + modified time.Time + path string +} + +// processedVuln holds the metadata, ecosystems, and Vanir flag for a processed vulnerability. +type processedVuln struct { + meta vulnMeta + ecosystems []string + hasVanir bool +} + // ecosystemWorker processes vulnerabilities for a single ecosystem. type ecosystemWorker struct { - ecosystem string - inCh chan *osvschema.Vulnerability + ecosystem string + scratchDir string + inCh chan vulnMeta } // newEcosystemWorker creates and starts a new ecosystemWorker. -func newEcosystemWorker(ctx context.Context, ecosystem string, outCh chan<- writeMsg, wg *sync.WaitGroup) *ecosystemWorker { - ch := make(chan *osvschema.Vulnerability, 100) +func newEcosystemWorker(ctx context.Context, ecosystem string, scratchDir string, outCh chan<- writeMsg, wg *sync.WaitGroup) *ecosystemWorker { + ch := make(chan vulnMeta, 100) worker := &ecosystemWorker{ - ecosystem: ecosystem, - inCh: ch, + ecosystem: ecosystem, + scratchDir: scratchDir, + inCh: ch, } wg.Add(1) go worker.run(ctx, outCh, wg) @@ -49,16 +71,9 @@ func newEcosystemWorker(ctx context.Context, ecosystem string, outCh chan<- writ return worker } -// vulnData holds the ID and marshalled JSON data for a vulnerability. -type vulnData struct { - id string - modified time.Time - data []byte -} - -// run is the main loop for the ecosystemWorker. It receives vulnerabilities, +// run is the main loop for the ecosystemWorker. It receives vulnerability metadata, // aggregates them, and upon completion, writes out the ecosystem-specific -// zip, csv, and (for GIT) vanir files. +// zip and csv files. func (w *ecosystemWorker) run(ctx context.Context, outCh chan<- writeMsg, wg *sync.WaitGroup) { defer wg.Done() ctx, span := otel.Tracer("exporter").Start(ctx, w.ecosystem) @@ -66,39 +81,11 @@ func (w *ecosystemWorker) run(ctx context.Context, outCh chan<- writeMsg, wg *sy logger.InfoContext(ctx, "new ecosystem worker started", slog.String("ecosystem", w.ecosystem)) // 500 size is around the minimum each ecosystem would need, most ecosystems are much bigger. - allVulns := make([]vulnData, 0, 500) - csvData := make([][]string, 0, 500) - var vanirVulns []vulnData + allVulns := make([]vulnMeta, 0, 500) + csvData := make([]csvEntry, 0, 500) for v := range w.inCh { - // Process vulnerability. - b, err := marshalToJSON(v) - if err != nil { - logger.ErrorContext(ctx, "failed to marshal vulnerability to json", slog.String("id", v.GetId()), slog.Any("err", err)) - continue - } - - // Wait to send the result, or be cancelled. - select { - case outCh <- writeMsg{path: filepath.Join(w.ecosystem, v.GetId()) + ".json", mimeType: "application/json", data: b}: - case <-ctx.Done(): - logger.WarnContext(ctx, "ecosystem worker cancelled", slog.String("ecosystem", w.ecosystem), slog.Any("err", ctx.Err())) - return - } - - modified := v.GetModified().AsTime() - allVulns = append(allVulns, vulnData{id: v.GetId(), modified: modified, data: b}) - csvData = append(csvData, []string{modified.Format(time.RFC3339Nano), v.GetId()}) - - // For GIT ecosystem, we want to make a file containing every vulnerability with vanir signatures - if w.ecosystem == gitEcosystem { - for _, aff := range v.GetAffected() { - spec := aff.GetDatabaseSpecific() - if _, ok := spec.GetFields()["vanir_signatures"]; ok { - vanirVulns = append(vanirVulns, vulnData{id: v.GetId(), data: b}) - break - } - } - } + allVulns = append(allVulns, v) + csvData = append(csvData, csvEntry{modified: v.modified, path: v.id}) } if ctx.Err() != nil { @@ -107,10 +94,7 @@ func (w *ecosystemWorker) run(ctx context.Context, outCh chan<- writeMsg, wg *sy logger.InfoContext(ctx, "All vulnerabilities processed", slog.String("ecosystem", w.ecosystem)) writeModifiedIDCSV(ctx, filepath.Join(w.ecosystem, modifiedCSVFilename), csvData, outCh) - writeZIP(ctx, filepath.Join(w.ecosystem, allZipFilename), allVulns, outCh) - if w.ecosystem == gitEcosystem { - writeVanir(ctx, vanirVulns, outCh) - } + writeZIP(ctx, filepath.Join(w.ecosystem, allZipFilename), allVulns, outCh, w.scratchDir) logger.InfoContext(ctx, "ecosystem worker finished processing", slog.String("ecosystem", w.ecosystem)) } @@ -119,24 +103,25 @@ func (w *ecosystemWorker) Finish() { close(w.inCh) } -// vulnAndEcos holds a vulnerability and the list of ecosystems it belongs to. +// vulnAndEcos holds a vulnerability metadata and the list of ecosystems it belongs to. type vulnAndEcos struct { - *osvschema.Vulnerability - + meta vulnMeta ecosystems []string } // allEcosystemWorker processes all vulnerabilities from all ecosystems to create // the global export files. type allEcosystemWorker struct { - inCh chan vulnAndEcos + scratchDir string + inCh chan vulnAndEcos } // newAllEcosystemWorker creates and starts a new allEcosystemWorker. -func newAllEcosystemWorker(ctx context.Context, outCh chan<- writeMsg, wg *sync.WaitGroup) *allEcosystemWorker { +func newAllEcosystemWorker(ctx context.Context, scratchDir string, outCh chan<- writeMsg, wg *sync.WaitGroup) *allEcosystemWorker { ch := make(chan vulnAndEcos, 100) worker := &allEcosystemWorker{ - inCh: ch, + scratchDir: scratchDir, + inCh: ch, } wg.Add(1) go worker.run(ctx, outCh, wg) @@ -153,21 +138,15 @@ func (w *allEcosystemWorker) run(ctx context.Context, outCh chan<- writeMsg, wg logger.InfoContext(ctx, "all-ecosystem worker started") // We have currently about 1.8 million entries, so start at 100k - allVulns := make([]vulnData, 0, 100000) - csvData := make([][]string, 0, 100000) + allVulns := make([]vulnMeta, 0, 100000) + csvData := make([]csvEntry, 0, 100000) ecosystems := make(map[string]struct{}) for v := range w.inCh { - b, err := marshalToJSON(v.Vulnerability) - if err != nil { - logger.ErrorContext(ctx, "failed to marshal vulnerability to json", slog.String("id", v.GetId()), slog.Any("err", err)) - continue - } - modified := v.GetModified().AsTime() - allVulns = append(allVulns, vulnData{id: v.GetId(), modified: modified, data: b}) + allVulns = append(allVulns, v.meta) for _, e := range v.ecosystems { ecosystems[e] = struct{}{} - csvData = append(csvData, []string{modified.Format(time.RFC3339Nano), e + "/" + v.GetId()}) - if len(csvData)%10000 == 0 { + csvData = append(csvData, csvEntry{modified: v.meta.modified, path: e + "/" + v.meta.id}) + if len(csvData)%50000 == 0 { logger.InfoContext(ctx, "processed N vulnerabilities", slog.Int("n", len(csvData))) } } @@ -178,7 +157,7 @@ func (w *allEcosystemWorker) run(ctx context.Context, outCh chan<- writeMsg, wg } writeModifiedIDCSV(ctx, modifiedCSVFilename, csvData, outCh) - writeZIP(ctx, allZipFilename, allVulns, outCh) + writeZIP(ctx, allZipFilename, allVulns, outCh, w.scratchDir) ecos := slices.Collect(maps.Keys(ecosystems)) slices.Sort(ecos) ecoString := strings.Join(ecos, "\n") + "\n" @@ -214,35 +193,58 @@ func write(ctx context.Context, path string, data []byte, mimeType string, outCh } } +// writeStream is a helper to send a streaming file writeMsg to the writer channel. +func writeStream(ctx context.Context, path string, filePath string, mimeType string, outCh chan<- writeMsg) { + select { + case outCh <- writeMsg{path: path, mimeType: mimeType, filePath: filePath}: + case <-ctx.Done(): + } +} + // writeModifiedIDCSV constructs and writes a modified_id.csv file. -func writeModifiedIDCSV(ctx context.Context, path string, csvData [][]string, outCh chan<- writeMsg) { +func writeModifiedIDCSV(ctx context.Context, path string, csvData []csvEntry, outCh chan<- writeMsg) { logger.InfoContext(ctx, "constructing csv file", slog.String("path", path)) - slices.SortFunc(csvData, func(a, b []string) int { + slices.SortFunc(csvData, func(a, b csvEntry) int { return cmp.Or( - -cmp.Compare(a[0], b[0]), // Modified date, descending - cmp.Compare(a[1], b[1]), // path/vuln ID, ascending + -a.modified.Compare(b.modified), // Modified date, descending + cmp.Compare(a.path, b.path), // path/vuln ID, ascending ) }) var buf bytes.Buffer wr := csv.NewWriter(&buf) - if err := wr.WriteAll(csvData); err != nil { - logger.ErrorContext(ctx, "failed writing csv", slog.String("path", path), slog.Any("err", err)) - return + for _, entry := range csvData { + t := entry.modified.UTC().Format(time.RFC3339Nano) + if err := wr.Write([]string{t, entry.path}); err != nil { + logger.ErrorContext(ctx, "failed writing csv line", slog.String("path", path), slog.Any("err", err)) + return + } } wr.Flush() + if err := wr.Error(); err != nil { + logger.ErrorContext(ctx, "failed flushing csv", slog.String("path", path), slog.Any("err", err)) + return + } + logger.InfoContext(ctx, "writing csv file", slog.String("path", path)) write(ctx, path, buf.Bytes(), "text/csv", outCh) } -// writeZIP constructs and writes an all.zip file. -func writeZIP(ctx context.Context, path string, allVulns []vulnData, outCh chan<- writeMsg) { +// writeZIP constructs and writes a zip file by streaming from local files to a temporary zip file. +func writeZIP(ctx context.Context, path string, allVulns []vulnMeta, outCh chan<- writeMsg, scratchDir string) { logger.InfoContext(ctx, "constructing zip file", slog.String("path", path)) - slices.SortFunc(allVulns, func(a, b vulnData) int { + slices.SortFunc(allVulns, func(a, b vulnMeta) int { return cmp.Compare(a.id, b.id) }) - var buf bytes.Buffer - wr := zip.NewWriter(&buf) + + tmpZip, err := os.CreateTemp(scratchDir, "zip-*.tmp") + if err != nil { + logger.ErrorContext(ctx, "failed to create temp zip file", slog.String("path", path), slog.Any("err", err)) + return + } + defer tmpZip.Close() + + wr := zip.NewWriter(tmpZip) for _, vuln := range allVulns { w, err := wr.CreateHeader(&zip.FileHeader{ Name: vuln.id + ".json", @@ -253,29 +255,49 @@ func writeZIP(ctx context.Context, path string, allVulns []vulnData, outCh chan< logger.ErrorContext(ctx, "failed to create vuln json in zip file", slog.String("id", vuln.id), slog.Any("err", err)) continue } - r := bytes.NewReader(vuln.data) - if _, err := io.Copy(w, r); err != nil { + localPath := filepath.Join(scratchDir, vuln.id+".json") + f, err := os.Open(localPath) + if err != nil { + logger.ErrorContext(ctx, "failed to open local vuln json file", slog.String("path", localPath), slog.Any("err", err)) + continue + } + if _, err := io.Copy(w, f); err != nil { logger.ErrorContext(ctx, "failed to write vuln json in zip file", slog.String("id", vuln.id), slog.Any("err", err)) } + f.Close() } if err := wr.Close(); err != nil { logger.ErrorContext(ctx, "failed to close zip writer", slog.String("path", path), slog.Any("err", err)) + return } + logger.InfoContext(ctx, "writing zip file", slog.String("path", path)) - write(ctx, path, buf.Bytes(), "application/zip", outCh) + writeStream(ctx, path, tmpZip.Name(), "application/zip", outCh) } -// writeVanir constructs and writes the osv_git.json file containing vulnerabilities with Vanir signatures. -func writeVanir(ctx context.Context, vanirVulns []vulnData, outCh chan<- writeMsg) { - slices.SortFunc(vanirVulns, func(a, b vulnData) int { return cmp.Compare(a.id, b.id) }) - vulns := make([]json.RawMessage, len(vanirVulns)) - for i, v := range vanirVulns { - vulns[i] = v.data +// writeVanir constructs and writes the osv_git.json file containing vulnerabilities with Vanir signatures +// by reading the cached JSON files from disk and marshaling the combined JSON array in memory. +func writeVanir(ctx context.Context, vanirVulnIDs []string, outCh chan<- writeMsg, scratchDir string) { + logger.InfoContext(ctx, "constructing vanir file", slog.Int("count", len(vanirVulnIDs))) + slices.Sort(vanirVulnIDs) + + vulns := make([]json.RawMessage, 0, len(vanirVulnIDs)) + for _, id := range vanirVulnIDs { + localPath := filepath.Join(scratchDir, id+".json") + data, err := os.ReadFile(localPath) + if err != nil { + logger.ErrorContext(ctx, "failed to read local vuln file for vanir", slog.String("id", id), slog.Any("err", err)) + continue + } + vulns = append(vulns, data) } + finalJSON, err := json.Marshal(vulns) if err != nil { logger.ErrorContext(ctx, "failed to marshal vanir JSON file", slog.Any("err", err)) return } + + logger.InfoContext(ctx, "writing vanir file", slog.String("path", filepath.Join(gitEcosystem, vanirVulnsFilename))) write(ctx, filepath.Join(gitEcosystem, vanirVulnsFilename), finalJSON, "application/json", outCh) } diff --git a/go/cmd/exporter/writer.go b/go/cmd/exporter/writer.go index 29533fa7ba5..d9506b18c66 100644 --- a/go/cmd/exporter/writer.go +++ b/go/cmd/exporter/writer.go @@ -4,6 +4,7 @@ import ( "context" "errors" "hash/crc32" + "io" "log/slog" "os" "path/filepath" @@ -21,6 +22,7 @@ type writeMsg struct { path string mimeType string data []byte + filePath string // If set, stream from local file instead of holding data in memory } // writer is a worker that receives writeMsgs and writes them to either a GCS @@ -29,39 +31,113 @@ func writer(ctx context.Context, cancel context.CancelFunc, inCh <-chan writeMsg defer wg.Done() for msg := range inCh { path := filepath.Join(pathPrefix, msg.path) - if client != nil { - // Skip the upload if the object already has the same content. - if gcsContentUnchanged(ctx, client, path, msg.data) { - continue - } - err := client.WriteObject(ctx, path, msg.data, &clients.WriteOptions{ - ContentType: msg.mimeType, - }) - if err != nil { - logger.Error("failed to write file", slog.String("path", path), slog.Any("err", err)) - cancel() - - break - } + var err error + if msg.filePath != "" { + err = writeFromFile(ctx, client, path, msg.filePath, msg.mimeType) } else { - // Write locally. - dir := filepath.Dir(path) - if err := os.MkdirAll(dir, 0755); err != nil { - logger.Error("failed to create directories", slog.String("dir", dir), slog.Any("err", err)) - cancel() - - break - } - if err := os.WriteFile(path, msg.data, 0600); err != nil { - logger.Error("failed to write file", slog.String("path", path), slog.Any("err", err)) - cancel() - - break - } + err = writeFromMemory(ctx, client, path, msg.data, msg.mimeType) + } + if err != nil { + logger.Error("failed to write output file", slog.String("path", path), slog.Any("err", err)) + cancel() + + break } } } +// writeFromFile handles writing or uploading a file by streaming from a local file path. +func writeFromFile(ctx context.Context, client clients.CloudStorage, path, filePath, mimeType string) error { + if client != nil { + if gcsFileUnchanged(ctx, client, path, filePath) { + return nil + } + f, err := os.Open(filePath) + if err != nil { + return err + } + defer f.Close() + + return client.WriteObjectStream(ctx, path, f, &clients.WriteOptions{ + ContentType: mimeType, + }) + } + + // Write locally: copy from filePath to path + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return err + } + + return copyFile(filePath, path) +} + +// writeFromMemory handles writing or uploading in-memory byte data to GCS or a local file. +func writeFromMemory(ctx context.Context, client clients.CloudStorage, path string, data []byte, mimeType string) error { + if client != nil { + if gcsContentUnchanged(ctx, client, path, data) { + return nil + } + + return client.WriteObject(ctx, path, data, &clients.WriteOptions{ + ContentType: mimeType, + }) + } + + // Write locally + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return err + } + + return os.WriteFile(path, data, 0600) +} + +func copyFile(src, dst string) error { + in, err := os.Open(src) + if err != nil { + return err + } + defer in.Close() + + out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600) + if err != nil { + return err + } + defer out.Close() + + _, err = io.Copy(out, in) + + return err +} + +func gcsFileUnchanged(ctx context.Context, client clients.CloudStorage, path string, filePath string) bool { + attrs, err := client.ReadObjectAttrs(ctx, path) + if err != nil { + if !errors.Is(err, clients.ErrNotFound) { + logger.WarnContext(ctx, "failed to read object attrs, proceeding with upload", slog.String("path", path), slog.Any("err", err)) + } + + return false + } + f, err := os.Open(filePath) + if err != nil { + return false + } + defer f.Close() + h := crc32.New(crc32cTable) + if _, err := io.Copy(h, f); err != nil { + return false + } + if attrs.CRC32C == h.Sum32() { + logger.InfoContext(ctx, "skipping upload, content unchanged", slog.String("path", path)) + + return true + } + + return false +} + // gcsContentUnchanged returns true if the object at path already has the same // CRC32C checksum as data, meaning the upload would be a no-op. Any error // reading the object's attributes (other than ErrNotFound) is logged and diff --git a/go/osv/clients/cloudstorage.go b/go/osv/clients/cloudstorage.go index 4c2f3244b58..012781988cd 100644 --- a/go/osv/clients/cloudstorage.go +++ b/go/osv/clients/cloudstorage.go @@ -18,6 +18,7 @@ package clients import ( "context" "errors" + "io" "iter" "time" ) @@ -73,6 +74,9 @@ type CloudStorage interface { // WriteObject writes a complete byte slice to a storage object. WriteObject(ctx context.Context, path string, data []byte, opts *WriteOptions) error + // WriteObjectStream streams data from an io.Reader to a storage object. + WriteObjectStream(ctx context.Context, path string, r io.Reader, opts *WriteOptions) error + // Objects returns an iterator over objects that match the prefix. Objects(ctx context.Context, prefix string) iter.Seq2[*Object, error] diff --git a/go/osv/clients/gcs_client.go b/go/osv/clients/gcs_client.go index 127f0d9c5ad..2ea6df1aed4 100644 --- a/go/osv/clients/gcs_client.go +++ b/go/osv/clients/gcs_client.go @@ -115,6 +115,32 @@ func (c *GCSClient) WriteObject(ctx context.Context, path string, data []byte, o return err } +func (c *GCSClient) WriteObjectStream(ctx context.Context, path string, r io.Reader, opts *WriteOptions) error { + var err error + for i := range numRetries { + if i > 0 { + // Exponential backoff: 1s, 2s, 4s + time.Sleep(time.Duration(1<<(i-1)) * time.Second) + } + if seeker, ok := r.(io.Seeker); ok && i > 0 { + if _, seekErr := seeker.Seek(0, io.SeekStart); seekErr != nil { + return seekErr + } + } + err = c.writeObjectStreamOnce(ctx, path, r, opts) + if err == nil { + return nil + } + // Check if error is not transient and should not be retried + var apiErr *googleapi.Error + if !errors.As(err, &apiErr) || (apiErr.Code < 500 && apiErr.Code != 429) { + return err + } + } + + return err +} + func (c *GCSClient) writeObjectOnce(ctx context.Context, path string, data []byte, opts *WriteOptions) error { obj := c.bucket.Object(path) @@ -146,6 +172,35 @@ func (c *GCSClient) writeObjectOnce(ctx context.Context, path string, data []byt return writer.Close() } +func (c *GCSClient) writeObjectStreamOnce(ctx context.Context, path string, r io.Reader, opts *WriteOptions) error { + obj := c.bucket.Object(path) + + if opts != nil && opts.IfGenerationMatches != nil { + conds := storage.Conditions{GenerationMatch: *opts.IfGenerationMatches} + if *opts.IfGenerationMatches == 0 { + conds = storage.Conditions{DoesNotExist: true} + } + obj = obj.If(conds) + } + + writer := obj.NewWriter(ctx) + if opts != nil { + if opts.CustomTime != nil { + writer.CustomTime = *opts.CustomTime + } + if opts.ContentType != "" { + writer.ContentType = opts.ContentType + } + } + + if _, err := io.Copy(writer, r); err != nil { + _ = writer.Close() + return err + } + + return writer.Close() +} + func (c *GCSClient) Close() error { return c.client.Close() } diff --git a/go/testutils/gcs.go b/go/testutils/gcs.go index 585adf9a401..2f763253142 100644 --- a/go/testutils/gcs.go +++ b/go/testutils/gcs.go @@ -3,6 +3,7 @@ package testutils import ( "context" "hash/crc32" + "io" "iter" "slices" "strings" @@ -125,6 +126,15 @@ func (c *MockStorage) WriteObject(_ context.Context, path string, data []byte, o return nil } +func (c *MockStorage) WriteObjectStream(ctx context.Context, path string, r io.Reader, opts *clients.WriteOptions) error { + data, err := io.ReadAll(r) + if err != nil { + return err + } + + return c.WriteObject(ctx, path, data, opts) +} + func (c *MockStorage) Objects(_ context.Context, prefix string) iter.Seq2[*clients.Object, error] { // Create a snapshot of the keys to iterate over, so we don't hold the lock. c.mu.RLock()