vulnerability_test.go

package datastore

import (
	"bytes"
	"context"
	"errors"
	"strings"
	"sync"
	"testing"
	"time"

	"cloud.google.com/go/datastore"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"github.com/google/osv.dev/go/internal/models"
	"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/timestamppb"
)

func TestVulnerabilityStore_ListBySource(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    testutils.NewMockStorage(),
	})

	vulns := []Vulnerability{
		{
			SourceID:    "source-a:path/1.json",
			IsWithdrawn: false,
		},
		{
			SourceID:    "source-a:path/2.json",
			IsWithdrawn: true,
		},
		{
			SourceID:    "source-b:path/3.json",
			IsWithdrawn: false,
		},
	}

	keys := []*datastore.Key{
		datastore.NameKey("Vulnerability", "ID-1", nil),
		datastore.NameKey("Vulnerability", "ID-2", nil),
		datastore.NameKey("Vulnerability", "ID-3", nil),
	}

	if _, err := dsClient.PutMulti(ctx, keys, vulns); err != nil {
		t.Fatalf("Failed to setup test data: %v", err)
	}

	tests := []struct {
		name          string
		source        string
		skipWithdrawn bool
		want          []*models.VulnSourceRef
	}{
		{
			name:          "Source A - Skip Withdrawn",
			source:        "source-a",
			skipWithdrawn: true,
			want: []*models.VulnSourceRef{
				{ID: "ID-1", Source: "source-a", Path: "path/1.json"},
			},
		},
		{
			name:          "Source A - Include Withdrawn",
			source:        "source-a",
			skipWithdrawn: false,
			want: []*models.VulnSourceRef{
				{ID: "ID-1", Source: "source-a", Path: "path/1.json"},
				{ID: "ID-2", Source: "source-a", Path: "path/2.json"},
			},
		},
		{
			name:          "Source B",
			source:        "source-b",
			skipWithdrawn: true,
			want: []*models.VulnSourceRef{
				{ID: "ID-3", Source: "source-b", Path: "path/3.json"},
			},
		},
		{
			name:          "Non-existent Source",
			source:        "source-c",
			skipWithdrawn: true,
			want:          []*models.VulnSourceRef{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var got []*models.VulnSourceRef
			for entry, err := range store.ListBySource(ctx, tt.source, tt.skipWithdrawn) {
				if err != nil {
					t.Fatalf("ListBySource() error: %v", err)
				}
				got = append(got, entry)
			}

			sortOpt := cmpopts.SortSlices(func(a, b *models.VulnSourceRef) bool {
				return a.ID < b.ID
			})

			if diff := cmp.Diff(tt.want, got, cmpopts.EquateEmpty(), sortOpt); diff != "" {
				t.Errorf("ListBySource() mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

func TestVulnerabilityStore_Get(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	gcsStore := testutils.NewMockStorage()
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    gcsStore,
	})

	v := &osvschema.Vulnerability{
		Id: "GHSA-xxxx-yyyy-zzzz",
	}
	data, err := proto.Marshal(v)
	if err != nil {
		t.Fatalf("failed to marshal vulnerability: %v", err)
	}

	if err := gcsStore.WriteObject(ctx, "all/pb/GHSA-xxxx-yyyy-zzzz.pb", data, nil); err != nil {
		t.Fatalf("failed to write mock GCS object: %v", err)
	}

	got, err := store.GetFull(ctx, "GHSA-xxxx-yyyy-zzzz")
	if err != nil {
		t.Fatalf("Get() error = %v", err)
	}

	if got.GetId() != v.GetId() {
		t.Errorf("Get() id mismatch, got %s, want %s", got.GetId(), v.GetId())
	}
}

func TestVulnerabilityStore_Get_NotFound(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	gcsStore := testutils.NewMockStorage()
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    gcsStore,
	})

	_, err := store.GetFull(ctx, "GHSA-not-exists")
	if !errors.Is(err, models.ErrNotFound) {
		t.Errorf("want ErrNotFound, got %v", err)
	}
}

func TestVulnerabilityStore_GetWithMetadata(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	gcsStore := testutils.NewMockStorage()
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    gcsStore,
	})

	v := &osvschema.Vulnerability{
		Id: "GHSA-xxxx-yyyy-zzzz",
	}
	data, err := proto.Marshal(v)
	if err != nil {
		t.Fatalf("failed to marshal vulnerability: %v", err)
	}

	if err := gcsStore.WriteObject(ctx, "all/pb/GHSA-xxxx-yyyy-zzzz.pb", data, nil); err != nil {
		t.Fatalf("failed to write mock GCS object: %v", err)
	}

	key := datastore.NameKey("Vulnerability", "GHSA-xxxx-yyyy-zzzz", nil)
	dv := Vulnerability{
		SourceID: "source-x:path/1.json",
	}
	if _, err := dsClient.Put(ctx, key, &dv); err != nil {
		t.Fatalf("failed to setup datastore test data: %v", err)
	}

	gotV, gotRef, err := store.GetWithMetadata(ctx, "GHSA-xxxx-yyyy-zzzz")
	if err != nil {
		t.Fatalf("GetWithMetadata() error = %v", err)
	}

	if gotV.GetId() != v.GetId() {
		t.Errorf("GetWithMetadata() id mismatch, got %s, want %s", gotV.GetId(), v.GetId())
	}

	if gotRef.Source != "source-x" {
		t.Errorf("gotRef.Source mismatch, got %s, want source-x", gotRef.Source)
	}

	if gotRef.Path != "path/1.json" {
		t.Errorf("gotRef.Path mismatch, got %s, want path/1.json", gotRef.Path)
	}
}

func TestVulnerabilityStore_GetWithMetadata_NotFound(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	gcsStore := testutils.NewMockStorage()
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    gcsStore,
	})

	_, _, err := store.GetWithMetadata(ctx, "GHSA-not-exists")
	if !errors.Is(err, models.ErrNotFound) {
		t.Errorf("want ErrNotFound, got %v", err)
	}
}

func TestVulnerabilityStore_Write(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	gcsStore := testutils.NewMockStorage()
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    gcsStore,
	})

	v := &osvschema.Vulnerability{
		Id:        "TEST-WRITE-123",
		Published: timestamppb.Now(),
		Affected: []*osvschema.Affected{
			{
				Package: &osvschema.Package{
					Name:      "test-pkg",
					Ecosystem: "npm",
				},
				Ranges: []*osvschema.Range{
					{
						Type: osvschema.Range_ECOSYSTEM,
						Events: []*osvschema.Event{
							{Introduced: "1.0.0"},
							{Fixed: "2.0.0"},
						},
					},
				},
			},
		},
	}

	req := models.WriteRequest{
		ID:       "TEST-WRITE-123",
		Source:   "test-source",
		Path:     "TEST-WRITE-123.json",
		Raw:      v,
		Enriched: v,
		AffectedCommits: models.AffectedCommitsResult{
			Skip:    false,
			Commits: [][]byte{[]byte("test-commit")},
		},
	}

	if err := store.Write(ctx, req); err != nil {
		t.Fatalf("Write() failed: %v", err)
	}

	// 1. Verify Vulnerability was written
	var dv Vulnerability
	vKey := datastore.NameKey("Vulnerability", req.ID, nil)
	if err := dsClient.Get(ctx, vKey, &dv); err != nil {
		t.Fatalf("failed to fetch Vulnerability after Write: %v", err)
	}
	if dv.SourceID != "test-source:TEST-WRITE-123.json" {
		t.Errorf("Vulnerability.SourceID mismatch, got %s", dv.SourceID)
	}

	// 2. Verify ListedVulnerability was written
	var lv ListedVulnerability
	lvKey := datastore.NameKey("ListedVulnerability", req.ID, nil)
	if err := dsClient.Get(ctx, lvKey, &lv); err != nil {
		t.Fatalf("failed to fetch ListedVulnerability after Write: %v", err)
	}
	if lv.Summary != v.GetSummary() {
		t.Errorf("ListedVulnerability Summary mismatch")
	}

	// 3. Verify AffectedCommits was written
	var ac AffectedCommits
	acKey := datastore.NameKey("AffectedCommits", "TEST-WRITE-123-0", nil)
	if err := dsClient.Get(ctx, acKey, &ac); err != nil {
		t.Fatalf("failed to fetch AffectedCommits after Write: %v", err)
	}
	if len(ac.Commits) != 1 || string(ac.Commits[0]) != "test-commit" {
		t.Errorf("AffectedCommits mismatch, got %v", ac.Commits)
	}

	// 4. Verify AffectedVersions was written
	var avs []AffectedVersions
	q := datastore.NewQuery("AffectedVersions").FilterField("vuln_id", "=", req.ID)
	if _, err := dsClient.GetAll(ctx, q, &avs); err != nil {
		t.Fatalf("failed to query AffectedVersions after Write: %v", err)
	}
	if len(avs) == 0 {
		t.Errorf("Expected AffectedVersions to be written, got 0")
	}

	// 5. Verify GCS object was written
	path := "all/pb/TEST-WRITE-123.pb"
	gcsData, err := gcsStore.ReadObject(ctx, path)
	if err != nil {
		t.Fatalf("failed to read GCS object after Write: %v", err)
	}
	gotV := new(osvschema.Vulnerability)
	if err := proto.Unmarshal(gcsData, gotV); err != nil {
		t.Fatalf("failed to unmarshal vulnerability from GCS: %v", err)
	}
	if gotV.GetId() != v.GetId() {
		t.Errorf("GCS Vulnerability ID mismatch, got %s, want %s", gotV.GetId(), v.GetId())
	}
}

func TestVulnerabilityStore_Write_GCSFailure_Retry(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	gcsStore := testutils.NewMockStorage()
	mockPublisher := &testutils.MockPublisher{}

	store := NewVulnerabilityStore(VulnStoreConfig{
		Client:               dsClient,
		GCS:                  gcsStore,
		FailedWritePublisher: mockPublisher,
	})

	// Make GCS write fail
	gcsStore.WriteError = errors.New("GCS write failed")

	v := &osvschema.Vulnerability{
		Id: "TEST-WRITE-RETRY",
	}

	req := models.WriteRequest{
		ID:       "TEST-WRITE-RETRY",
		Source:   "test-source",
		Path:     "TEST-WRITE-RETRY.json",
		Enriched: v,
	}

	if err := store.Write(ctx, req); err != nil {
		t.Fatalf("Write() failed: %v", err)
	}

	// Verify that it attempted to publish to the retry topic
	if len(mockPublisher.Messages) != 1 {
		t.Fatalf("Expected 1 message published, got %d", len(mockPublisher.Messages))
	}

	msg := mockPublisher.Messages[0]
	if msg.Attributes["type"] != "gcs_retry" {
		t.Errorf("Expected message type 'gcs_retry', got %q", msg.Attributes["type"])
	}

	expectedData, _ := proto.Marshal(v)
	if !bytes.Equal(msg.Data, expectedData) {
		t.Errorf("Message data mismatch")
	}
}

func TestVulnerabilityStore_GetModified(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
	})

	now := time.Now().UTC().Truncate(time.Millisecond) // Datastore might lose precision, truncate to ms
	vuln := Vulnerability{
		Modified: now,
	}
	key := datastore.NameKey("Vulnerability", "TEST-MOD-1", nil)
	if _, err := dsClient.Put(ctx, key, &vuln); err != nil {
		t.Fatalf("Failed to setup test data: %v", err)
	}

	got, err := store.GetModified(ctx, "TEST-MOD-1")
	if err != nil {
		t.Fatalf("GetModified() error = %v", err)
	}

	if !got.Equal(now) {
		t.Errorf("GetModified() got = %v, want %v", got, now)
	}
}

func TestVulnerabilityStore_GetModified_NotFound(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
	})

	_, err := store.GetModified(ctx, "TEST-MOD-NOT-EXISTS")
	if !errors.Is(err, models.ErrNotFound) {
		t.Errorf("want ErrNotFound, got %v", err)
	}
}

func TestVulnerabilityStore_GetModified_Batch(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	// Set a slightly longer timeout to ensure they all get batched
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client:       dsClient,
		BatchTimeout: 10 * time.Millisecond,
	})

	now := time.Now().UTC().Truncate(time.Millisecond)
	vulns := []Vulnerability{
		{Modified: now},
		{Modified: now.Add(time.Hour)},
	}
	keys := []*datastore.Key{
		datastore.NameKey("Vulnerability", "TEST-BATCH-1", nil),
		datastore.NameKey("Vulnerability", "TEST-BATCH-2", nil),
	}
	if _, err := dsClient.PutMulti(ctx, keys, vulns); err != nil {
		t.Fatalf("Failed to setup test data: %v", err)
	}

	var wg sync.WaitGroup
	wg.Add(3) // 2 exists, 1 not exists

	var got1, got2 time.Time
	var err1, err2, err3 error

	go func() {
		defer wg.Done()
		got1, err1 = store.GetModified(ctx, "TEST-BATCH-1")
	}()

	go func() {
		defer wg.Done()
		got2, err2 = store.GetModified(ctx, "TEST-BATCH-2")
	}()

	go func() {
		defer wg.Done()
		_, err3 = store.GetModified(ctx, "TEST-BATCH-3") // Should be NotFound
	}()

	wg.Wait()

	if err1 != nil {
		t.Errorf("Request 1 failed: %v", err1)
	} else if !got1.Equal(now) {
		t.Errorf("Request 1 got %v, want %v", got1, now)
	}

	want2 := now.Add(time.Hour)
	if err2 != nil {
		t.Errorf("Request 2 failed: %v", err2)
	} else if !got2.Equal(want2) {
		t.Errorf("Request 2 got %v, want %v", got2, want2)
	}

	if !errors.Is(err3, models.ErrNotFound) {
		t.Errorf("Request 3 expected ErrNotFound, got %v", err3)
	}
}

func TestVulnerabilityStore_GetModified_BatchTriggerEarly(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	// Set a very long timeout but small max elements
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client:           dsClient,
		BatchTimeout:     10 * time.Second, // Long timeout
		BatchMaxElements: 2,                // Trigger immediately on 2 elements
	})

	now := time.Now().UTC().Truncate(time.Millisecond)
	vulns := []Vulnerability{
		{Modified: now},
		{Modified: now.Add(time.Hour)},
	}
	keys := []*datastore.Key{
		datastore.NameKey("Vulnerability", "TEST-TRIGGER-1", nil),
		datastore.NameKey("Vulnerability", "TEST-TRIGGER-2", nil),
	}
	if _, err := dsClient.PutMulti(ctx, keys, vulns); err != nil {
		t.Fatalf("Failed to setup test data: %v", err)
	}

	var wg sync.WaitGroup
	wg.Add(2)

	var got1, got2 time.Time
	var err1, err2 error

	start := time.Now()

	go func() {
		defer wg.Done()
		got1, err1 = store.GetModified(ctx, "TEST-TRIGGER-1")
	}()

	// Give it a tiny bit of time to make sure the first request is registered as leader
	time.Sleep(10 * time.Millisecond)

	go func() {
		defer wg.Done()
		got2, err2 = store.GetModified(ctx, "TEST-TRIGGER-2") // This should trigger the batch early
	}()

	wg.Wait()

	duration := time.Since(start)
	if duration > 2*time.Second {
		t.Errorf("Batch took too long to complete (%v), expected early trigger", duration)
	}

	if err1 != nil {
		t.Errorf("Request 1 failed: %v", err1)
	} else if !got1.Equal(now) {
		t.Errorf("Request 1 got %v, want %v", got1, now)
	}

	want2 := now.Add(time.Hour)
	if err2 != nil {
		t.Errorf("Request 2 failed: %v", err2)
	} else if !got2.Equal(want2) {
		t.Errorf("Request 2 got %v, want %v", got2, want2)
	}
}

func TestVulnerabilityStore_Ping(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	mockGCS := testutils.NewMockStorage()

	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    mockGCS,
	})

	// 1. Test Happy Path: Both GCS and Datastore are healthy (but empty).
	// Ping should return nil because it correctly filters out ErrNotFound.
	if err := store.Ping(ctx); err != nil {
		t.Errorf("Expected Ping to succeed on healthy stores, got: %v", err)
	}

	// 2. Test GCS Outage: Mock GCS returning a real connection error.
	mockGCS.ReadError = errors.New("GCS connection timeout")
	if err := store.Ping(ctx); err == nil {
		t.Errorf("Expected Ping to fail when GCS is unreachable, got nil")
	} else if !strings.Contains(err.Error(), "gcs unreachable") {
		t.Errorf("Expected GCS unreachable error, got: %v", err)
	}

	// 3. Test Datastore Outage: If we pass a nil client, Ping should fail
	// cleanly without panicking (due to our batcher panic-safety).
	brokenStore := NewVulnerabilityStore(VulnStoreConfig{
		Client: nil,
		GCS:    mockGCS,
	})
	// Reset the GCS error so only Datastore is failing
	mockGCS.ReadError = nil

	if err := brokenStore.Ping(ctx); err == nil {
		t.Errorf("Expected Ping to fail when Datastore client is nil, got nil")
	} else if !strings.Contains(err.Error(), "datastore/batcher unreachable") {
		t.Errorf("Expected Datastore/Batcher unreachable error, got: %v", err)
	}
}

func TestVulnerabilityStore_ListModifiedSinceAndMutate(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	mockGCS := testutils.NewMockStorage()
	store := NewVulnerabilityStore(VulnStoreConfig{
		Client: dsClient,
		GCS:    mockGCS,
	})

	baseTime := time.Date(2026, 9, 23, 1, 0, 0, 0, time.UTC)
	v1 := &osvschema.Vulnerability{
		Id:       "VULN-OLD",
		Modified: timestamppb.New(baseTime),
		Summary:  "Old vuln",
	}
	v2 := &osvschema.Vulnerability{
		Id:       "VULN-NEW",
		Modified: timestamppb.New(baseTime.Add(2 * time.Hour)),
		Summary:  "New vuln",
	}

	for _, v := range []*osvschema.Vulnerability{v1, v2} {
		err := store.Write(ctx, models.WriteRequest{
			ID:       v.GetId(),
			Source:   "test-source",
			Path:     v.GetId() + ".json",
			Raw:      v,
			Enriched: v,
			AffectedCommits: models.AffectedCommitsResult{
				Skip: true,
			},
		})
		if err != nil {
			t.Fatalf("store.Write(%s) failed: %v", v.GetId(), err)
		}
	}

	since := baseTime.Add(1 * time.Hour)
	var gotIDs []string
	for v, err := range store.ListModifiedSince(ctx, &since) {
		if err != nil {
			t.Fatalf("ListModifiedSince failed: %v", err)
		}
		gotIDs = append(gotIDs, v.GetId())
	}
	if diff := cmp.Diff([]string{"VULN-NEW"}, gotIDs); diff != "" {
		t.Errorf("ListModifiedSince mismatch (-want +got):\n%s", diff)
	}

	mutatedTime := baseTime.Add(3 * time.Hour)
	updated, err := store.Mutate(ctx, "VULN-NEW", func(v *osvschema.Vulnerability) (bool, error) {
		v.Aliases = []string{"CVE-2026-9999"}
		v.Modified = timestamppb.New(mutatedTime)

		return true, nil
	})
	if err != nil {
		t.Fatalf("store.Mutate failed: %v", err)
	}
	if !updated {
		t.Fatalf("expected updated=true from store.Mutate")
	}

	full, err := store.GetFull(ctx, "VULN-NEW")
	if err != nil {
		t.Fatalf("store.GetFull failed: %v", err)
	}
	if diff := cmp.Diff([]string{"CVE-2026-9999"}, full.GetAliases()); diff != "" {
		t.Errorf("Aliases mismatch after Mutate (-want +got):\n%s", diff)
	}
	if !full.GetModified().AsTime().Equal(mutatedTime) {
		t.Errorf("Modified after Mutate = %v, want %v", full.GetModified().AsTime(), mutatedTime)
	}
}

func TestJobDataStore(t *testing.T) {
	ctx := context.Background()
	dsClient := testutils.MustNewDatastoreClientForTesting(t)
	jobStore := NewJobDataStore(dsClient)

	_, err := models.GetJobData[time.Time](ctx, jobStore, "test_last_run")
	if !errors.Is(err, models.ErrNotFound) {
		t.Fatalf("GetJobData on missing key err = %v, want models.ErrNotFound", err)
	}

	now := time.Date(2026, 9, 23, 4, 0, 0, 0, time.UTC)
	if err := jobStore.Set(ctx, "test_last_run", now); err != nil {
		t.Fatalf("Set failed: %v", err)
	}

	lr, err := models.GetJobData[time.Time](ctx, jobStore, "test_last_run")
	if err != nil || !lr.Equal(now) {
		t.Fatalf("GetJobData[time.Time] = (%v, %v), want %v", lr, err, now)
	}

	if err := jobStore.Set(ctx, "test_retry_list", []string{"A", "B"}); err != nil {
		t.Fatalf("Set failed: %v", err)
	}
	list, err := models.GetJobData[[]string](ctx, jobStore, "test_retry_list")
	if err != nil {
		t.Fatalf("GetJobData[[]string] failed: %v", err)
	}
	if diff := cmp.Diff([]string{"A", "B"}, list); diff != "" {
		t.Errorf("GetJobData[[]string] mismatch (-want +got):\n%s", diff)
	}
}