diff --git a/osv.dev.code-workspace b/osv.dev.code-workspace new file mode 100644 index 00000000000..2231d56d6c1 --- /dev/null +++ b/osv.dev.code-workspace @@ -0,0 +1,17 @@ +{ + "folders": [ + { + "path": "." + }, + { + "path": "go" + }, + { + "path": "vulnfeeds" + }, + { + "path": "go/cmd/tools/reimport-tui" + } + ], + "settings": {} +} \ No newline at end of file diff --git a/vulnfeeds/cmd/combine-to-osv/main.go b/vulnfeeds/cmd/combine-to-osv/main.go index 60f3ab72ae1..8992181c5bd 100644 --- a/vulnfeeds/cmd/combine-to-osv/main.go +++ b/vulnfeeds/cmd/combine-to-osv/main.go @@ -4,7 +4,6 @@ package main import ( "cmp" "context" - "errors" "flag" "fmt" "io" @@ -32,10 +31,10 @@ import ( ) const ( - defaultOSVOutputPath = "osv-output" - defaultCVE5Path = "cve5" - defaultNVDOSVPath = "nvd" - defaultListFileMaxAge = 2 * 24 * time.Hour + defaultOSVOutputPath = "osv-output" + defaultCVE5Path = "cve5" + defaultNVDOSVPath = "nvd" + parallelStartYear = 2018 ) // CVEWorkItem represents a unit of work for a single CVE. @@ -60,78 +59,26 @@ func cveIDFromPath(p string) models.CVEID { return "" } -// readListFile reads the list file (usually files.txt) from GCS. -// The list file is expected to contain newline-separated GCS object names -// relative to the bucket root (e.g. "prefix/CVE-2023-0001.json"). -// It returns a slice of full GCS paths (e.g. "gs://bucket/prefix/CVE-2023-0001.json"). -func readListFile(ctx context.Context, listObj *storage.ObjectHandle, bucketName string) ([]string, error) { - rc, err := listObj.NewReader(ctx) - if err != nil { - return nil, fmt.Errorf("failed to open reader: %w", err) - } - defer rc.Close() - - content, err := io.ReadAll(rc) - if err != nil { - return nil, fmt.Errorf("failed to read content: %w", err) - } - - lines := strings.Split(string(content), "\n") - fullPaths := make([]string, 0, len(lines)) - for _, line := range lines { - line = strings.TrimSpace(line) - if line == "" { - continue - } - fullPaths = append(fullPaths, fmt.Sprintf("gs://%s/%s", bucketName, line)) - } - - return fullPaths, nil -} - -func listObjects(ctx context.Context, client *storage.Client, pathStr string, maxAge time.Duration) ([]string, error) { +func listObjects(ctx context.Context, client *storage.Client, pathStr string) ([]string, error) { if strings.HasPrefix(pathStr, "gs://") { trimmed := strings.TrimPrefix(pathStr, "gs://") bucketName, prefix, _ := strings.Cut(trimmed, "/") bucket := client.Bucket(bucketName) - listObjectName := gcs.FileListFilename - if prefix != "" { - listObjectName = strings.TrimSuffix(prefix, "/") + "/" + gcs.FileListFilename - } - - listObj := bucket.Object(listObjectName) - attrs, err := listObj.Attrs(ctx) - useListFile := false - if err == nil { - // Check age - age := time.Since(attrs.Updated) - if age <= maxAge { - useListFile = true - } else { - logger.Error("List file is too old (generators might be failing)", slog.String("bucket", bucketName), slog.String("prefix", prefix), slog.Duration("age", age)) - } - } else if !errors.Is(err, storage.ErrObjectNotExist) { - logger.Warn("Failed to get list file attributes", slog.String("bucket", bucketName), slog.String("prefix", prefix), slog.Any("err", err)) - } - - if useListFile { - logger.Info("Using list file", slog.String("bucket", bucketName), slog.String("prefix", prefix)) - fullPaths, err := readListFile(ctx, listObj, bucketName) - if err == nil { - return fullPaths, nil - } - logger.Error("Failed to read list file, falling back to bucket listing", slog.String("bucket", bucketName), slog.String("prefix", prefix), slog.Any("err", err)) + currentYear := time.Now().Year() + var breakdowns []string + for year := parallelStartYear; year <= currentYear; year++ { + breakdowns = append(breakdowns, fmt.Sprintf("CVE-%d-", year)) } - logger.Info("Falling back to bucket listing", slog.String("bucket", bucketName), slog.String("prefix", prefix)) - objs, err := gcs.ListBucketObjects(ctx, bucket, prefix) + objs, err := gcs.ListObjectsFast(ctx, bucket, prefix, breakdowns) if err != nil { return nil, err } + var fullPaths []string for _, obj := range objs { - if strings.HasSuffix(obj, "/") || !strings.HasSuffix(obj, ".json") || strings.HasSuffix(obj, ".metrics.json") { + if !strings.HasSuffix(obj, ".json") || strings.HasSuffix(obj, ".metrics.json") { continue } fullPaths = append(fullPaths, fmt.Sprintf("gs://%s/%s", bucketName, obj)) @@ -256,7 +203,6 @@ func main() { uploadToGCS := flag.Bool("upload-to-gcs", false, "If true, upload to GCS bucket instead of writing to local disk.") numWorkers := flag.Int("workers", 64, "Number of workers to process records") syncDeletions := flag.Bool("sync-deletions", false, "If false, do not delete files in bucket that are not local") - listFileMaxAge := flag.Duration("list-file-max-age", defaultListFileMaxAge, "The maximum age of the list file before falling back to listing the bucket.") flag.Parse() err := os.MkdirAll(*osvOutputPath, 0755) @@ -281,7 +227,7 @@ func main() { go func() { defer listObjWg.Done() logger.Info("Listing CVE5 objects", slog.String("path", *cve5Path)) - cve5Files, cve5ListErr = listObjects(ctx, client, *cve5Path, *listFileMaxAge) + cve5Files, cve5ListErr = listObjects(ctx, client, *cve5Path) if cve5ListErr != nil { logger.Fatal("Failed to list CVE5 objects", slog.Any("err", cve5ListErr)) } @@ -290,7 +236,7 @@ func main() { go func() { defer listObjWg.Done() logger.Info("Listing NVD objects", slog.String("path", *nvdPath)) - nvdFiles, nvdListErr = listObjects(ctx, client, *nvdPath, *listFileMaxAge) + nvdFiles, nvdListErr = listObjects(ctx, client, *nvdPath) if nvdListErr != nil { logger.Fatal("Failed to list NVD objects", slog.Any("err", nvdListErr)) } diff --git a/vulnfeeds/cmd/converters/cve/cve5/bulk-converter/main.go b/vulnfeeds/cmd/converters/cve/cve5/bulk-converter/main.go index ff0f950b2ca..365470e71ce 100644 --- a/vulnfeeds/cmd/converters/cve/cve5/bulk-converter/main.go +++ b/vulnfeeds/cmd/converters/cve/cve5/bulk-converter/main.go @@ -8,7 +8,6 @@ import ( "flag" "log/slog" "os" - "path" "path/filepath" "slices" "strconv" @@ -23,12 +22,10 @@ import ( "github.com/google/osv/vulnfeeds/utility/logger" ) -const defaultStartYear = "2022" - var ( repoDir = flag.String("cve5-repo", "cvelistV5", "CVEListV5 directory path") localOutputDir = flag.String("out-dir", "cve5", "Path to output results.") - startYear = flag.String("start-year", defaultStartYear, "The first in scope year to process.") + startYear = flag.String("start-year", "2022", "The first in scope year to process.") workers = flag.Int("workers", 10, "The number of concurrent workers to use for processing CVEs.") gcsWorkers = flag.Int("gcs-workers", 30, "The number of concurrent workers to use for GCS uploads.") cnaDenyList = flag.String("cna-denylist", "", "A comma-separated list of CNAs to skip. If not provided, defaults to cna_denylist.txt.") @@ -46,20 +43,6 @@ func main() { logger.InitGlobalLogger() defer logger.Close() - startYearInt, err := strconv.Atoi(*startYear) - defaultStartYearInt, _ := strconv.Atoi(defaultStartYear) - isPartial := false - if err != nil { - logger.Error("Invalid start-year, assuming partial", slog.String("start-year", *startYear)) - isPartial = true - } else if startYearInt > defaultStartYearInt { - isPartial = true - } - - if isPartial { - logger.Info("Partial run detected (start-year > " + defaultStartYear + "), will skip files.txt upload") - } - logger.Info("Commencing CVE to OSV conversion run") if err := os.MkdirAll(*localOutputDir, 0755); err != nil { logger.Fatal("Failed to create local output directory", slog.Any("err", err)) @@ -87,29 +70,20 @@ func main() { if err != nil { logger.Fatal("Failed to initialize GCS upload pool", slog.Any("err", err)) } + defer gcsHelper.CloseAndWait() logger.Info("GCS Upload Pool initialized", slog.String("bucket", *outputBucket)) } - outputFilesChan := make(chan string) - var collectorWg sync.WaitGroup - var outputFiles []string - collectorWg.Add(1) - go func() { - defer collectorWg.Done() - for f := range outputFilesChan { - outputFiles = append(outputFiles, f) - } - }() - // Start the worker pool. for range *workers { wg.Add(1) - go worker(&wg, jobs, gcsHelper, *localOutputDir, cnaList, *rejectFailed, outputFilesChan) + go worker(&wg, jobs, gcsHelper, *localOutputDir, cnaList, *rejectFailed) } // Discover files and send them to the workers. logger.Info("Starting conversion of CVEs...") currentYear := time.Now().Year() + startYearInt, _ := strconv.Atoi(*startYear) for year := startYearInt; year <= currentYear; year++ { year := strconv.Itoa(year) @@ -138,37 +112,22 @@ func main() { close(jobs) wg.Wait() - close(outputFilesChan) - collectorWg.Wait() - - if *uploadToGCS && gcsHelper != nil { - gcsHelper.CloseAndWait() - if !isPartial { - logger.Info("Uploading output file list to GCS...") - if err := gcs.UploadFileList(ctx, *outputBucket, *gcsPrefix, outputFiles); err != nil { - logger.Error("Failed to upload output file list", slog.Any("err", err)) - } - } else { - logger.Info("Skipping files.txt upload due to partial run") - } - } - logger.Info("CVE5 Conversion run complete") } // worker is a function that processes CVE files from the jobs channel. -func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDir string, cnas []string, rejectFailed bool, outputFilesChan chan<- string) { +func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDir string, cnas []string, rejectFailed bool) { defer wg.Done() - for cvePath := range jobs { - data, err := os.ReadFile(cvePath) + for path := range jobs { + data, err := os.ReadFile(path) if err != nil { - logger.Info("Failed to read file", slog.String("path", cvePath), slog.Any("err", err)) + logger.Info("Failed to read file", slog.String("path", path), slog.Any("err", err)) continue } var cve models.CVE5 if err := json.Unmarshal(data, &cve); err != nil { - logger.Info("Failed to unmarshal JSON", slog.String("path", cvePath), slog.Any("err", err)) + logger.Info("Failed to unmarshal JSON", slog.String("path", path), slog.Any("err", err)) continue } @@ -180,9 +139,9 @@ func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDi sourceLink := "" baseDirCVEList := "cves/" // The base folder for the CVEListV5 repository. - idx := strings.Index(cvePath, baseDirCVEList) + idx := strings.Index(path, baseDirCVEList) if idx != -1 { - relPath := cvePath[idx:] + relPath := path[idx:] sourceLink = "https://github.com/CVEProject/cvelistV5/tree/main/" + relPath } @@ -194,8 +153,6 @@ func worker(wg *sync.WaitGroup, jobs <-chan string, gcsHelper *gcs.Helper, outDi logger.Info("Queueing OSV record for "+string(cveID), slog.String("cve", string(cveID))) if err := writer.UploadVulnIfChangedAsync(gcsHelper, *gcsPrefix, vuln.Vulnerability); err != nil { logger.Error("Failed to queue vulnerability upload", slog.String("cve", string(cveID)), slog.Any("err", err)) - } else { - outputFilesChan <- path.Join(*gcsPrefix, string(cveID)+".json") } if err := writer.UploadMetricsToGCSAsync(gcsHelper, *gcsPrefix, cveID, metrics); err != nil { diff --git a/vulnfeeds/cmd/converters/cve/nvd-cve-osv/main.go b/vulnfeeds/cmd/converters/cve/nvd-cve-osv/main.go index 56f1172fcab..6301252985e 100644 --- a/vulnfeeds/cmd/converters/cve/nvd-cve-osv/main.go +++ b/vulnfeeds/cmd/converters/cve/nvd-cve-osv/main.go @@ -8,7 +8,6 @@ import ( "log/slog" "net/http" "os" - "path" "path/filepath" "regexp" "runtime/pprof" @@ -27,8 +26,6 @@ import ( "github.com/google/osv/vulnfeeds/vulns" ) -const defaultStartYear = 2016 - var ( jsonPath = flag.String("nvd-json", "", "Path to NVD CVE JSON to examine.") jsonDir = flag.String("nvd-json-dir", "", "Path to directory containing NVD CVE JSON files to examine.") @@ -42,7 +39,7 @@ var ( uploadToGCS = flag.Bool("upload-to-gcs", false, "If true, upload to GCS bucket instead of writing to local disk.") outputBucket = flag.String("output-bucket", "osv-test-cve-osv-conversion", "The GCS bucket to write to.") gcsPrefix = flag.String("gcs-prefix", "nvd-osv", "The prefix within the GCS bucket.") - startYear = flag.Int("start-year", defaultStartYear, "The first in scope year to process. If 0, process all years.") + startYear = flag.Int("start-year", 2016, "The first in scope year to process. If 0, process all years.") ) var ( @@ -74,11 +71,6 @@ func main() { logger.InitGlobalLogger() defer logger.Close() - isPartial := *startYear > defaultStartYear - if isPartial { - logger.Info("Partial run detected (start-year > " + strconv.Itoa(defaultStartYear) + "), will skip files.txt upload") - } - var files []string if *jsonDir != "" { matches, err := filepath.Glob(filepath.Join(*jsonDir, "nvdcve-2.0-*.json")) @@ -138,26 +130,16 @@ func main() { if err != nil { logger.Fatal("Failed to initialize GCS upload pool", slog.Any("err", err)) } + defer gcsHelper.CloseAndWait() logger.Info("GCS Upload Pool initialized", slog.String("bucket", *outputBucket)) } - outputFilesChan := make(chan string) - var collectorWg sync.WaitGroup - var outputFiles []string - collectorWg.Add(1) - go func() { - defer collectorWg.Done() - for f := range outputFilesChan { - outputFiles = append(outputFiles, f) - } - }() - jobs := make(chan models.NVDCVE, *workers) var wg sync.WaitGroup for range *workers { wg.Add(1) - go worker(&wg, jobs, gcsHelper, *outDir, vpRepoCache, repoTagsCache, outputFilesChan) + go worker(&wg, jobs, gcsHelper, *outDir, vpRepoCache, repoTagsCache) } for _, file := range files { @@ -182,21 +164,11 @@ func main() { close(jobs) wg.Wait() - close(outputFilesChan) - collectorWg.Wait() timesBlocked := int64(0) if gcsHelper != nil { timesBlocked = gcsHelper.GetTimesBlocked() gcsHelper.CloseAndWait() - if !isPartial { - logger.Info("Uploading output file list to GCS...") - if err := gcs.UploadFileList(ctx, *outputBucket, *gcsPrefix, outputFiles); err != nil { - logger.Error("Failed to upload output file list", slog.Any("err", err)) - } - } else { - logger.Info("Skipping files.txt upload due to partial run") - } } logger.Info("Conversion Stats", slog.Uint64("total_processed", totalConversionsCount.Load()), @@ -227,7 +199,7 @@ func processCVE(cve models.NVDCVE, vpRepoCache *c.VPRepoCache, repoTagsCache git return nvd.CVEToOSV(cve, repos, vpRepoCache, repoTagsCache, metrics) } -func worker(wg *sync.WaitGroup, jobs <-chan models.NVDCVE, gcsHelper *gcs.Helper, outDir string, vpRepoCache *c.VPRepoCache, repoTagsCache git.RepoTagsCache, outputFilesChan chan<- string) { +func worker(wg *sync.WaitGroup, jobs <-chan models.NVDCVE, gcsHelper *gcs.Helper, outDir string, vpRepoCache *c.VPRepoCache, repoTagsCache git.RepoTagsCache) { defer wg.Done() for cve := range jobs { vuln, metrics, outcome := processCVE(cve, vpRepoCache, repoTagsCache) @@ -273,8 +245,6 @@ func worker(wg *sync.WaitGroup, jobs <-chan models.NVDCVE, gcsHelper *gcs.Helper if vuln != nil { if err := writer.UploadVulnIfChangedAsync(gcsHelper, *gcsPrefix, vuln.Vulnerability); err != nil { logger.Error("Failed to queue vulnerability upload", slog.String("cve", vuln.Id), slog.Any("err", err)) - } else { - outputFilesChan <- path.Join(*gcsPrefix, vuln.Id+".json") } } if *outputMetrics && metrics != nil { diff --git a/vulnfeeds/conversion/writer/writer.go b/vulnfeeds/conversion/writer/writer.go index aeae704131c..a1b835a12cb 100644 --- a/vulnfeeds/conversion/writer/writer.go +++ b/vulnfeeds/conversion/writer/writer.go @@ -30,8 +30,9 @@ import ( const ( // hashMetadataKey is the key for the sha256 hash in the GCS object metadata. - hashMetadataKey = "sha256-hash" - overrideFolder = "osv-output-overrides" // location of overrides within bucket + hashMetadataKey = "sha256-hash" + overrideFolder = "osv-output-overrides" // location of overrides within bucket + parallelStartYear = 2018 ) // ErrUploadSkipped indicates that an upload was intentionally skipped because @@ -309,7 +310,12 @@ func UploadVulnsToGCS( func HandleDeletion(ctx context.Context, outBkt *storage.BucketHandle, osvOutputPath string, validVulnIDs []string) { // Check if any need to be deleted - bucketObjects, err := gcs.ListBucketObjects(ctx, outBkt, osvOutputPath) + currentYear := time.Now().Year() + var breakdowns []string + for year := parallelStartYear; year <= currentYear; year++ { + breakdowns = append(breakdowns, fmt.Sprintf("CVE-%d-", year)) + } + bucketObjects, err := gcs.ListObjectsFast(ctx, outBkt, osvOutputPath, breakdowns) if err != nil { logger.Error("Failed to list bucket objects for deletion check, skipping deletion.", slog.Any("err", err)) return diff --git a/vulnfeeds/gcs-tools/gcs.go b/vulnfeeds/gcs-tools/gcs.go index d54791c6628..a25634a195c 100644 --- a/vulnfeeds/gcs-tools/gcs.go +++ b/vulnfeeds/gcs-tools/gcs.go @@ -9,6 +9,7 @@ import ( "log/slog" "os" "path/filepath" + "slices" "strings" "sync" "sync/atomic" @@ -17,12 +18,10 @@ import ( "github.com/google/osv/vulnfeeds/utility/logger" "golang.org/x/sync/errgroup" "google.golang.org/api/iterator" - "google.golang.org/api/option" ) const ( - hashMetadataKey = "sha256-hash" // hashMetadataKey is the key for the sha256 hash in the GCS object metadata. - FileListFilename = "files.txt" + hashMetadataKey = "sha256-hash" // hashMetadataKey is the key for the sha256 hash in the GCS object metadata. ) type Helper struct { @@ -223,10 +222,16 @@ func DownloadBucket(ctx context.Context, bkt *storage.BucketHandle, prefix strin return nil } -// ListBucketObjects lists the names of all objects in a Google Cloud Storage bucket. -// It does not download the file contents. -func ListBucketObjects(ctx context.Context, bucket *storage.BucketHandle, prefix string) ([]string, error) { - it := bucket.Objects(ctx, &storage.Query{Prefix: prefix}) +// ListBucketObjectsQuery lists the names of all objects in a Google Cloud Storage bucket matching the query. +// It optimizes the request by only retrieving the Name attribute. +func ListBucketObjectsQuery(ctx context.Context, bucket *storage.BucketHandle, query *storage.Query) ([]string, error) { + if query == nil { + query = &storage.Query{} + } + if err := query.SetAttrSelection([]string{"Name"}); err != nil { + return nil, fmt.Errorf("failed to set attribute selection: %w", err) + } + it := bucket.Objects(ctx, query) var filenames []string for { attrs, err := it.Next() @@ -242,28 +247,81 @@ func ListBucketObjects(ctx context.Context, bucket *storage.BucketHandle, prefix return filenames, nil } -// UploadFileList uploads a list of files as a single text file to GCS. -func UploadFileList(ctx context.Context, bucketName string, prefix string, files []string, opts ...option.ClientOption) error { - client, err := storage.NewClient(ctx, opts...) - if err != nil { - return fmt.Errorf("storage.NewClient: %w", err) +// ListObjectsFast lists objects in parallel by splitting the prefix space. +// It takes a globalPrefix and a list of breakdownPrefixes. +// It partitions the space using StartOffset and EndOffset based on the breakdowns. +// breakdownPrefixes should be sorted lexicographically. +func ListObjectsFast(ctx context.Context, bucket *storage.BucketHandle, globalPrefix string, breakdownPrefixes []string) ([]string, error) { + globalPrefix = strings.TrimSuffix(globalPrefix, "/") + + if len(breakdownPrefixes) == 0 { + return ListBucketObjectsQuery(ctx, bucket, &storage.Query{Prefix: globalPrefix}) + } + + // Ensure breakdownPrefixes are sorted + sortedBreakdowns := make([]string, len(breakdownPrefixes)) + copy(sortedBreakdowns, breakdownPrefixes) + slices.Sort(sortedBreakdowns) + + queries := make([]*storage.Query, 0, len(sortedBreakdowns)+1) + globalPrefixWithSlash := globalPrefix + if globalPrefix != "" { + globalPrefixWithSlash = globalPrefix + "/" + } + + // 1. First query: everything before the first breakdown + queries = append(queries, &storage.Query{ + Prefix: globalPrefixWithSlash, + EndOffset: globalPrefixWithSlash + sortedBreakdowns[0], + }) + + // 2. Intermediate queries: ranges between breakdowns + for i := range len(sortedBreakdowns) - 1 { + queries = append(queries, &storage.Query{ + Prefix: globalPrefixWithSlash, + StartOffset: globalPrefixWithSlash + sortedBreakdowns[i], + EndOffset: globalPrefixWithSlash + sortedBreakdowns[i+1], + }) + } + + // 3. Last query: everything from the last breakdown onwards + queries = append(queries, &storage.Query{ + Prefix: globalPrefixWithSlash, + StartOffset: globalPrefixWithSlash + sortedBreakdowns[len(sortedBreakdowns)-1], + }) + + objsChan := make(chan []string, len(queries)) + g, ctx := errgroup.WithContext(ctx) + + for _, q := range queries { + g.Go(func() error { + objs, err := ListBucketObjectsQuery(ctx, bucket, q) + if err != nil { + return err + } + objsChan <- objs + + return nil + }) } - defer client.Close() - bkt := client.Bucket(bucketName) + if err := g.Wait(); err != nil { + return nil, err + } + close(objsChan) - var objectName string - if prefix == "" { - objectName = FileListFilename - } else { - objectName = strings.TrimSuffix(prefix, "/") + "/" + FileListFilename + var allObjs []string + for objs := range objsChan { + allObjs = append(allObjs, objs...) } - content := strings.Join(files, "\n") - if len(files) > 0 { - content += "\n" // Add trailing newline + filtered := make([]string, 0, len(allObjs)) + for _, obj := range allObjs { + if strings.HasSuffix(obj, "/") { + continue + } + filtered = append(filtered, obj) } - reader := strings.NewReader(content) - return UploadToGCS(ctx, bkt, objectName, reader, "text/plain", nil) + return filtered, nil } diff --git a/vulnfeeds/gcs-tools/gcs_test.go b/vulnfeeds/gcs-tools/gcs_test.go index f4e888f9514..65d0514ad77 100644 --- a/vulnfeeds/gcs-tools/gcs_test.go +++ b/vulnfeeds/gcs-tools/gcs_test.go @@ -5,10 +5,10 @@ import ( "context" "os" "path/filepath" + "slices" "testing" "github.com/fsouza/fake-gcs-server/fakestorage" - "google.golang.org/api/option" ) func TestUploadToGCS(t *testing.T) { @@ -210,39 +210,92 @@ func TestDownloadBucket(t *testing.T) { }) } -func TestUploadFileList(t *testing.T) { - server := fakestorage.NewServer([]fakestorage.Object{}) - t.Cleanup(server.Stop) - - t.Log("Server URL:", server.URL()) - - bucketName := "test-bucket" - prefix := "test-prefix" - files := []string{ - "test-prefix/file1.json", - "test-prefix/file2.json", +func TestListObjectsFast(t *testing.T) { + objects := []fakestorage.Object{ + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/CVE-1999-0001.json", + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/CVE-2017-9999.json", + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/CVE-2018-0001.json", + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/CVE-2018-9999.json", + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/CVE-2018/nested.json", // in the gap + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/CVE-2019-0001.json", + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/CVE-2020-0001.json", + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "cve5/non-cve.json", + }, + }, + { + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-bucket", + Name: "other/file.json", // should not be listed + }, + }, } - ctx := context.Background() - // Create bucket first (fake server needs it) + server := fakestorage.NewServer(objects) + t.Cleanup(server.Stop) + client := server.Client() - bkt := client.Bucket(bucketName) - if err := bkt.Create(ctx, "project", nil); err != nil { - t.Fatalf("failed to create bucket: %v", err) - } + bkt := client.Bucket("test-bucket") - err := UploadFileList(ctx, bucketName, prefix, files, option.WithHTTPClient(server.HTTPClient())) + breakdowns := []string{"CVE-2018-", "CVE-2019-"} + got, err := ListObjectsFast(context.Background(), bkt, "cve5", breakdowns) if err != nil { - t.Fatalf("UploadFileList failed: %v", err) + t.Fatalf("ListObjectsFast failed: %v", err) } - obj, err := server.GetObject(bucketName, "test-prefix/files.txt") - if err != nil { - t.Fatalf("failed to get object: %v", err) + expected := []string{ + "cve5/CVE-1999-0001.json", + "cve5/CVE-2017-9999.json", + "cve5/CVE-2018-0001.json", + "cve5/CVE-2018-9999.json", + "cve5/CVE-2018/nested.json", + "cve5/CVE-2019-0001.json", + "cve5/CVE-2020-0001.json", + "cve5/non-cve.json", } - expectedContent := "test-prefix/file1.json\ntest-prefix/file2.json\n" - if string(obj.Content) != expectedContent { - t.Errorf("expected content %q, got %q", expectedContent, string(obj.Content)) + // Sort both to compare + slices.Sort(got) + slices.Sort(expected) + + if !slices.Equal(got, expected) { + t.Errorf("ListObjectsFast returned unexpected list.\ngot: %v\nwant: %v", got, expected) } }