internal/{vuln,frontend}: add another in-memory test client for vulndb-v1

This change adds an additional test client for reading from the vulndb
v1 databases. It is easier to use than the existing test client because
it takes in osv entries instead of txtar files. The logic to generate
the test client is mostly copied from x/vulndb.

This change also enables the vulndb-v1 experiment in all tests.

Change-Id: I386c48a1b90614ac1ca86899e9f1af0fe4f7ece2
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/481780
Run-TryBot: Tatiana Bradley <tatianabradley@google.com>
Reviewed-by: Tatiana Bradley <tatianabradley@google.com>
Reviewed-by: Jamal Carvalho <jamal@golang.org>
TryBot-Result: kokoro <noreply+kokoro@google.com>
diff --git a/internal/frontend/search_test.go b/internal/frontend/search_test.go
index c0af85d..4c99bee 100644
--- a/internal/frontend/search_test.go
+++ b/internal/frontend/search_test.go
@@ -18,6 +18,7 @@
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"golang.org/x/pkgsite/internal"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/fetchdatasource"
 	"golang.org/x/pkgsite/internal/licenses"
 	"golang.org/x/pkgsite/internal/postgres"
@@ -39,7 +40,10 @@
 	for _, v := range modules {
 		postgres.MustInsertModule(ctx, t, testDB, v)
 	}
-	vc := vuln.NewTestClient(testEntries)
+	vc, err := vuln.NewTestClient(testEntries)
+	if err != nil {
+		t.Fatal(err)
+	}
 	for _, test := range []struct {
 		name         string
 		method       string
@@ -225,7 +229,8 @@
 }
 
 func TestFetchSearchPage(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	ctx := experiment.NewContext(context.Background(), internal.ExperimentVulndbV1)
+	ctx, cancel := context.WithTimeout(ctx, testTimeout)
 	defer cancel()
 	defer postgres.ResetTestDB(testDB, t)
 
@@ -311,10 +316,13 @@
 				}},
 			}},
 		}}
-
-		vc = vuln.NewTestClient(vulnEntries)
 	)
 
+	vc, err := vuln.NewTestClient(vulnEntries)
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	for _, m := range []*internal.Module{moduleFoo, moduleBar} {
 		postgres.MustInsertModule(ctx, t, testDB, m)
 	}
@@ -547,7 +555,10 @@
 }
 
 func TestSearchVulnAlias(t *testing.T) {
-	vc := vuln.NewTestClient(testEntries)
+	vc, err := vuln.NewTestClient(testEntries)
+	if err != nil {
+		t.Fatal(err)
+	}
 	for _, test := range []struct {
 		name     string
 		mode     string
@@ -619,7 +630,10 @@
 }
 
 func TestSearchVulnModulePath(t *testing.T) {
-	vc := vuln.NewTestClient(testEntries)
+	vc, err := vuln.NewTestClient(testEntries)
+	if err != nil {
+		t.Fatal(err)
+	}
 	for _, test := range []struct {
 		name     string
 		mode     string
diff --git a/internal/frontend/versions_test.go b/internal/frontend/versions_test.go
index 70c6212..3c0a9aa 100644
--- a/internal/frontend/versions_test.go
+++ b/internal/frontend/versions_test.go
@@ -10,6 +10,7 @@
 
 	"github.com/google/go-cmp/cmp"
 	"golang.org/x/pkgsite/internal"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/postgres"
 	"golang.org/x/pkgsite/internal/stdlib"
 	"golang.org/x/pkgsite/internal/testing/sample"
@@ -91,6 +92,7 @@
 
 	vulnFixedVersion := "1.2.3"
 	vulnEntry := &osv.Entry{
+		ID:      "GO-1999-0001",
 		Details: "vuln",
 		Affected: []osv.Affected{{
 			Package: osv.Package{
@@ -107,7 +109,10 @@
 			},
 		}},
 	}
-	vc := vuln.NewTestClient([]*osv.Entry{vulnEntry})
+	vc, err := vuln.NewTestClient([]*osv.Entry{vulnEntry})
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	for _, tc := range []struct {
 		name        string
@@ -146,6 +151,7 @@
 					func() *VersionList {
 						vl := makeList(v1Path, modulePath1, "v1", []string{"v1.3.0", "v1.2.3", "v1.2.1"}, false)
 						vl.Versions[2].Vulns = []vuln.Vuln{{
+							ID:      vulnEntry.ID,
 							Details: vulnEntry.Details,
 						}}
 						return vl
@@ -188,7 +194,8 @@
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			ctx, cancel := context.WithTimeout(context.Background(), testTimeout*2)
+			ctx := experiment.NewContext(context.Background(), internal.ExperimentVulndbV1)
+			ctx, cancel := context.WithTimeout(ctx, testTimeout*2)
 			defer cancel()
 			defer postgres.ResetTestDB(testDB, t)
 
diff --git a/internal/frontend/vulns_test.go b/internal/frontend/vulns_test.go
index f211227..f19116a 100644
--- a/internal/frontend/vulns_test.go
+++ b/internal/frontend/vulns_test.go
@@ -10,6 +10,8 @@
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+	"golang.org/x/pkgsite/internal"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/vuln"
 	"golang.org/x/vuln/osv"
 )
@@ -38,8 +40,11 @@
 }
 
 func TestNewVulnListPage(t *testing.T) {
-	ctx := context.Background()
-	c := vuln.NewTestClient(testEntries)
+	ctx := experiment.NewContext(context.Background(), internal.ExperimentVulndbV1)
+	c, err := vuln.NewTestClient(testEntries)
+	if err != nil {
+		t.Fatal(err)
+	}
 	got, err := newVulnListPage(ctx, c)
 	if err != nil {
 		t.Fatal(err)
@@ -56,8 +61,11 @@
 }
 
 func TestNewVulnPage(t *testing.T) {
-	ctx := context.Background()
-	c := vuln.NewTestClient(testEntries)
+	ctx := experiment.NewContext(context.Background(), internal.ExperimentVulndbV1)
+	c, err := vuln.NewTestClient(testEntries)
+	if err != nil {
+		t.Fatal(err)
+	}
 	got, err := newVulnPage(ctx, c, "GO-1990-02")
 	if err != nil {
 		t.Fatal(err)
diff --git a/internal/vuln/client_test.go b/internal/vuln/client_test.go
index f58bba8..b4188b3 100644
--- a/internal/vuln/client_test.go
+++ b/internal/vuln/client_test.go
@@ -341,7 +341,7 @@
 // Test that Client can pick the right underlying client, based
 // on whether the v1 experiment is active.
 func TestCli(t *testing.T) {
-	v1, err := newTestV1Client(dbTxtar)
+	v1, err := newTestClientFromTxtar(dbTxtar)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -396,7 +396,7 @@
 
 // Run the test for both the v1 and legacy clients.
 func runClientTest(t *testing.T, test func(*testing.T, cli)) {
-	v1, err := newTestV1Client(dbTxtar)
+	v1, err := newTestClientFromTxtar(dbTxtar)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/vuln/client_v1.go b/internal/vuln/client_v1.go
index b4060e3..6f0ef9c 100644
--- a/internal/vuln/client_v1.go
+++ b/internal/vuln/client_v1.go
@@ -234,12 +234,6 @@
 	return dec, nil
 }
 
-var (
-	idDir           = "ID"
-	modulesEndpoint = "index/modules"
-	vulnsEndpoint   = "index/vulns"
-)
-
 func (c *client) modules(ctx context.Context) ([]byte, error) {
 	return c.src.get(ctx, modulesEndpoint)
 }
diff --git a/internal/vuln/schema.go b/internal/vuln/schema.go
index 3113f5b..6e38bf4 100644
--- a/internal/vuln/schema.go
+++ b/internal/vuln/schema.go
@@ -6,6 +6,20 @@
 
 import "time"
 
+var (
+	idDir           = "ID"
+	dbEndpoint      = "index/db"
+	modulesEndpoint = "index/modules"
+	vulnsEndpoint   = "index/vulns"
+)
+
+// DBMeta contains metadata about the database itself.
+type DBMeta struct {
+	// Modified is the time the database was last modified, calculated
+	// as the most recent time any single OSV entry was modified.
+	Modified time.Time `json:"modified"`
+}
+
 // ModuleMeta contains metadata about a Go module that has one
 // or more vulnerabilities in the database.
 //
diff --git a/internal/vuln/source.go b/internal/vuln/source.go
index ed699f6..3e1638e 100644
--- a/internal/vuln/source.go
+++ b/internal/vuln/source.go
@@ -7,12 +7,16 @@
 import (
 	"compress/gzip"
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
 	"os"
 	"path/filepath"
+	"sort"
+
+	"golang.org/x/vuln/osv"
 )
 
 type source interface {
@@ -94,6 +98,92 @@
 	return os.ReadFile(filepath.Join(db.dir, endpoint+".json"))
 }
 
+// Create a new in-memory source for testing.
+// Adapted from x/vulndb/internal/database.go.
+func newInMemorySource(entries []*osv.Entry) (*inMemorySource, error) {
+	data := make(map[string][]byte)
+	db := DBMeta{}
+	modulesMap := make(map[string]*ModuleMeta)
+	vulnsMap := make(map[string]*VulnMeta)
+	for _, entry := range entries {
+		if entry.ID == "" {
+			return nil, fmt.Errorf("entry %v has no ID", entry)
+		}
+		if _, ok := vulnsMap[entry.ID]; ok {
+			return nil, fmt.Errorf("id %q appears twice", entry.ID)
+		}
+		if entry.Modified.After(db.Modified) {
+			db.Modified = entry.Modified
+		}
+		for _, affected := range entry.Affected {
+			modulePath := affected.Package.Name
+			if _, ok := modulesMap[modulePath]; !ok {
+				modulesMap[modulePath] = &ModuleMeta{
+					Path:  modulePath,
+					Vulns: []ModuleVuln{},
+				}
+			}
+			module := modulesMap[modulePath]
+			module.Vulns = append(module.Vulns, ModuleVuln{
+				ID:       entry.ID,
+				Modified: entry.Modified,
+				Fixed:    latestFixedVersion(affected.Ranges),
+			})
+		}
+		vulnsMap[entry.ID] = &VulnMeta{
+			ID:       entry.ID,
+			Modified: entry.Modified,
+			Aliases:  entry.Aliases,
+		}
+		b, err := json.Marshal(entry)
+		if err != nil {
+			return nil, err
+		}
+		data[idDir+"/"+entry.ID] = b
+	}
+
+	b, err := json.Marshal(db)
+	if err != nil {
+		return nil, err
+	}
+	data[dbEndpoint] = b
+
+	// Add the modules endpoint.
+	modules := make([]*ModuleMeta, 0, len(modulesMap))
+	for _, module := range modulesMap {
+		modules = append(modules, module)
+	}
+	sort.SliceStable(modules, func(i, j int) bool {
+		return modules[i].Path < modules[j].Path
+	})
+	for _, module := range modules {
+		sort.SliceStable(module.Vulns, func(i, j int) bool {
+			return module.Vulns[i].ID < module.Vulns[j].ID
+		})
+	}
+	b, err = json.Marshal(modules)
+	if err != nil {
+		return nil, err
+	}
+	data[modulesEndpoint] = b
+
+	// Add the vulns endpoint.
+	vulns := make([]*VulnMeta, 0, len(vulnsMap))
+	for _, vuln := range vulnsMap {
+		vulns = append(vulns, vuln)
+	}
+	sort.SliceStable(vulns, func(i, j int) bool {
+		return vulns[i].ID < vulns[j].ID
+	})
+	b, err = json.Marshal(vulns)
+	if err != nil {
+		return nil, err
+	}
+	data[vulnsEndpoint] = b
+
+	return &inMemorySource{data: data}, nil
+}
+
 // inMemorySource reads databases from an in-memory map.
 // Intended for use in unit tests.
 type inMemorySource struct {
@@ -107,3 +197,18 @@
 	}
 	return b, nil
 }
+
+func latestFixedVersion(ranges []osv.AffectsRange) string {
+	var latestFixed string
+	for _, r := range ranges {
+		if r.Type == "SEMVER" {
+			for _, e := range r.Events {
+				fixed := e.Fixed
+				if fixed != "" && less(latestFixed, fixed) {
+					latestFixed = fixed
+				}
+			}
+		}
+	}
+	return latestFixed
+}
diff --git a/internal/vuln/source_test.go b/internal/vuln/source_test.go
index 28ee485..8aade71 100644
--- a/internal/vuln/source_test.go
+++ b/internal/vuln/source_test.go
@@ -13,6 +13,8 @@
 	"os"
 	"path/filepath"
 	"testing"
+
+	"golang.org/x/vuln/osv"
 )
 
 func TestNewSource(t *testing.T) {
@@ -110,6 +112,34 @@
 	}
 }
 
+func TestNewInMemorySource(t *testing.T) {
+	fromTxtar, err := newTestClientFromTxtar(dbTxtar)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fromEntries, err := newInMemorySource([]*osv.Entry{&testOSV1, &testOSV2, &testOSV3})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx := context.Background()
+	endpoints := []string{dbEndpoint, modulesEndpoint, vulnsEndpoint, idDir + "/" + testOSV1.ID, idDir + "/" + testOSV2.ID, idDir + "/" + testOSV3.ID}
+	for _, endpoint := range endpoints {
+		got, err := fromEntries.get(ctx, endpoint)
+		if err != nil {
+			t.Fatal(err)
+		}
+		want, err := fromTxtar.src.get(ctx, endpoint)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if string(got) != string(want) {
+			t.Errorf("newInMemorySource().get(%q) = %s, want %s", endpoint, got, want)
+		}
+	}
+}
+
 func gzipped(data []byte) ([]byte, error) {
 	var b bytes.Buffer
 	w := gzip.NewWriter(&b)
diff --git a/internal/vuln/test_client.go b/internal/vuln/test_client.go
index c2b2f5a..5753c5d 100644
--- a/internal/vuln/test_client.go
+++ b/internal/vuln/test_client.go
@@ -5,27 +5,28 @@
 package vuln
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 
 	"golang.org/x/tools/txtar"
 	vulnc "golang.org/x/vuln/client"
 	"golang.org/x/vuln/osv"
 )
 
-// NewTestClient creates an in-memory client for use in tests,
-// It's logic is different from the real client, so it should not be used to
-// test the client itself, but can be used to test code that depends on the
-// client.
-func NewTestClient(entries []*osv.Entry) *Client {
-	return &Client{legacy: newTestLegacyClient(entries)}
+// NewTestClient creates an in-memory client for use in tests.
+func NewTestClient(entries []*osv.Entry) (*Client, error) {
+	inMemory, err := newInMemorySource(entries)
+	if err != nil {
+		return nil, err
+	}
+	return &Client{legacy: newTestLegacyClient(entries), v1: &client{inMemory}}, nil
 }
 
-// newTestV1Client creates an in-memory client for use in tests.
-// It uses all the logic of the real v1 client, except that it reads
-// raw database data from the given txtar file instead of making HTTP
-// requests.
-// It can be used to test core functionality of the v1 client.
-func newTestV1Client(txtarFile string) (*client, error) {
+// newTestClientFromTxtar creates an in-memory client for use in tests.
+// It reads test data from a txtar file which must follow the
+// v1 database schema.
+func newTestClientFromTxtar(txtarFile string) (*client, error) {
 	data := make(map[string][]byte)
 
 	ar, err := txtar.ParseFile(txtarFile)
@@ -34,12 +35,24 @@
 	}
 
 	for _, f := range ar.Files {
-		data[f.Name] = f.Data
+		fdata, err := removeWhitespace(f.Data)
+		if err != nil {
+			return nil, err
+		}
+		data[f.Name] = fdata
 	}
 
 	return &client{&inMemorySource{data: data}}, nil
 }
 
+func removeWhitespace(data []byte) ([]byte, error) {
+	var b bytes.Buffer
+	if err := json.Compact(&b, data); err != nil {
+		return nil, err
+	}
+	return b.Bytes(), nil
+}
+
 func newTestLegacyClient(entries []*osv.Entry) *legacyClient {
 	c := &testVulnClient{
 		entries:          entries,
diff --git a/internal/vuln/vulns_test.go b/internal/vuln/vulns_test.go
index 41a9c1b..c39ccec 100644
--- a/internal/vuln/vulns_test.go
+++ b/internal/vuln/vulns_test.go
@@ -79,7 +79,7 @@
 	}
 
 	legacyClient := newTestLegacyClient([]*osv.Entry{&e, &e2, &stdlib})
-	v1Client, err := newTestV1Client("testdata/db2.txtar")
+	v1Client, err := newTestClientFromTxtar("testdata/db2.txtar")
 	if err != nil {
 		t.Fatal(err)
 	}