diff --git a/go/cmd/exporter/downloader.go b/go/cmd/exporter/downloader.go index 733f506e944..7d17efab746 100644 --- a/go/cmd/exporter/downloader.go +++ b/go/cmd/exporter/downloader.go @@ -7,7 +7,6 @@ import ( "hash/crc32" "io" "log/slog" - "os" "path/filepath" "strings" "sync" @@ -28,9 +27,9 @@ var flateWriterPool = sync.Pool{ // 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 pre-compressed Deflate data 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) { +// JSON, pre-compresses Deflate data in memory, queues individual JSON uploads, +// and sends the metadata and compressed payload to routerCh. +func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, client clients.CloudStorage, inCh <-chan string, routerCh chan<- processedVuln, writeCh chan<- writeMsg, wg *sync.WaitGroup) { defer wg.Done() for path := range inCh { // Process object. @@ -83,18 +82,6 @@ func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, clien compressedBytes := compBuf.Bytes() crc := crc32.ChecksumIEEE(b) - // Cache pre-compressed Deflate payload to local scratch disk. - localPath := filepath.Join(scratchDir, vuln.GetId()+".deflate") - //nolint:gosec // G703: Staging temporary file in scratch directory - if err := os.WriteFile(localPath, compressedBytes, 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() { @@ -136,11 +123,12 @@ func downloadThenProcessor(ctx context.Context, cancel context.CancelFunc, clien select { case routerCh <- processedVuln{ meta: vulnMeta{ - id: vuln.GetId(), - modified: vuln.GetModified().AsTime(), - crc32: crc, - uncompSize: uint64(len(b)), - compSize: uint64(len(compressedBytes)), + id: vuln.GetId(), + modified: vuln.GetModified().AsTime(), + crc32: crc, + uncompSize: uint64(len(b)), + compSize: uint64(len(compressedBytes)), + deflateData: compressedBytes, }, ecosystems: ecoNames, hasVanir: hasVanir, diff --git a/go/cmd/exporter/exporter.go b/go/cmd/exporter/exporter.go index 660fa53db53..b7d8e977cc1 100644 --- a/go/cmd/exporter/exporter.go +++ b/go/cmd/exporter/exporter.go @@ -45,7 +45,7 @@ func main() { 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.") - cleanUpScratchDir := flag.Bool("cleanup-scratch-dir", false, "Whether to delete the temporary scratch directory on exit. Defaults to false.") + cleanUpScratchDir := flag.Bool("cleanup-scratch-dir", true, "Whether to delete the temporary scratch directory on exit. Defaults to true.") flag.Parse() @@ -114,7 +114,7 @@ func main() { var processorWg sync.WaitGroup for range *numWorkers / 2 { processorWg.Add(1) - go downloadThenProcessor(ctx, cancel, vulnClient, scratchDir, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg) + go downloadThenProcessor(ctx, cancel, vulnClient, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg) } var writerWg sync.WaitGroup @@ -190,7 +190,7 @@ func ecosystemRouter(ctx context.Context, inCh <-chan processedVuln, outCh chan< workers := make(map[string]*ecosystemWorker) var workersWg sync.WaitGroup vulnCounter := 0 - var vanirVulnIDs []string + var vanirVulns []vanirVuln allEcosystemWorker := newAllEcosystemWorker(ctx, scratchDir, outCh, &workersWg) @@ -209,7 +209,7 @@ RouterLoop: vulnCounter++ if vuln.hasVanir { - vanirVulnIDs = append(vanirVulnIDs, vuln.meta.id) + vanirVulns = append(vanirVulns, vanirVuln{id: vuln.meta.id, deflateData: vuln.meta.deflateData}) } for _, eco := range vuln.ecosystems { @@ -237,8 +237,8 @@ RouterLoop: allEcosystemWorker.Finish() workersWg.Wait() - if len(vanirVulnIDs) > 0 && ctx.Err() == nil { - writeVanir(ctx, vanirVulnIDs, outCh, scratchDir) + if len(vanirVulns) > 0 && ctx.Err() == nil { + writeVanir(ctx, vanirVulns, outCh) } if ctx.Err() == nil { diff --git a/go/cmd/exporter/exporter_test.go b/go/cmd/exporter/exporter_test.go index 8e2209864f9..85416a10f16 100644 --- a/go/cmd/exporter/exporter_test.go +++ b/go/cmd/exporter/exporter_test.go @@ -137,7 +137,7 @@ func TestExporterPipeline_EndToEnd(t *testing.T) { var processorWg sync.WaitGroup processorWg.Add(1) - go downloadThenProcessor(ctx, cancel, vulnStorage, scratchDir, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg) + go downloadThenProcessor(ctx, cancel, vulnStorage, gcsPathToProcessorCh, processorToRouterCh, writeCh, &processorWg) var writerWg sync.WaitGroup writerWg.Add(1) diff --git a/go/cmd/exporter/worker.go b/go/cmd/exporter/worker.go index f90ea9ce4dd..31632d4a39f 100644 --- a/go/cmd/exporter/worker.go +++ b/go/cmd/exporter/worker.go @@ -32,13 +32,20 @@ const ( ecosystemsFilename = "ecosystems.txt" ) -// vulnMeta holds the ID, modified time, CRC32, and pre-compression sizes for a vulnerability. +// vulnMeta holds the ID, modified time, CRC32, pre-compression sizes, and deflate payload for a vulnerability. type vulnMeta struct { - id string - modified time.Time - crc32 uint32 - uncompSize uint64 - compSize uint64 + id string + modified time.Time + crc32 uint32 + uncompSize uint64 + compSize uint64 + deflateData []byte +} + +// vanirVuln holds the ID and deflate payload for a vulnerability with Vanir signatures. +type vanirVuln struct { + id string + deflateData []byte } // csvEntry holds the modified time and the relative entry path. @@ -234,7 +241,7 @@ func writeModifiedIDCSV(ctx context.Context, path string, csvData []csvEntry, ou write(ctx, path, buf.Bytes(), "text/csv", outCh) } -// writeZIP constructs and writes a zip file by streaming pre-compressed local files via CreateRaw. +// writeZIP constructs and writes a zip file by streaming pre-compressed in-memory payloads via CreateRaw. 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 vulnMeta) int { @@ -262,16 +269,9 @@ func writeZIP(ctx context.Context, path string, allVulns []vulnMeta, outCh chan< logger.ErrorContext(ctx, "failed to create raw vuln in zip file", slog.String("id", vuln.id), slog.Any("err", err)) continue } - localPath := filepath.Join(scratchDir, vuln.id+".deflate") - f, err := os.Open(localPath) - if err != nil { - logger.ErrorContext(ctx, "failed to open local vuln deflate file", slog.String("path", localPath), slog.Any("err", err)) - continue - } - if _, err := io.Copy(w, f); err != nil { + if _, err := w.Write(vuln.deflateData); err != nil { logger.ErrorContext(ctx, "failed to write vuln deflate data to 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)) @@ -283,26 +283,20 @@ func writeZIP(ctx context.Context, path string, allVulns []vulnMeta, outCh chan< } // 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+".deflate") - //nolint:gosec // G703: Reading temporary local file in scratch directory - f, err := os.Open(localPath) - if err != nil { - logger.ErrorContext(ctx, "failed to open local vuln file for vanir", slog.String("id", id), slog.Any("err", err)) - continue - } - fr := flate.NewReader(f) +// by decompressing the in-memory deflate payloads and marshaling the combined JSON array in memory. +func writeVanir(ctx context.Context, vanirVulns []vanirVuln, outCh chan<- writeMsg) { + logger.InfoContext(ctx, "constructing vanir file", slog.Int("count", len(vanirVulns))) + slices.SortFunc(vanirVulns, func(a, b vanirVuln) int { + return cmp.Compare(a.id, b.id) + }) + + vulns := make([]json.RawMessage, 0, len(vanirVulns)) + for _, v := range vanirVulns { + fr := flate.NewReader(bytes.NewReader(v.deflateData)) data, err := io.ReadAll(fr) _ = fr.Close() - _ = f.Close() if err != nil { - logger.ErrorContext(ctx, "failed to decompress local vuln file for vanir", slog.String("id", id), slog.Any("err", err)) + logger.ErrorContext(ctx, "failed to decompress vuln file for vanir", slog.String("id", v.id), slog.Any("err", err)) continue } vulns = append(vulns, data)