srv/cmd/gendb,srv/internal/database: refactor

Logic for srv/cmd/gendb is is moved into internal/database.

Change-Id: I1ef7bd9e15969570d00d97555657417edd0fd1ca
Reviewed-on: https://go-review.googlesource.com/c/vuln/+/372996
Trust: Julie Qiu <julie@golang.org>
Run-TryBot: Julie Qiu <julie@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
diff --git a/srv/cmd/gendb/main.go b/srv/cmd/gendb/main.go
index 4d58c8a..73d4e50 100644
--- a/srv/cmd/gendb/main.go
+++ b/srv/cmd/gendb/main.go
@@ -7,128 +7,24 @@
 package main
 
 import (
-	"encoding/json"
 	"flag"
-	"fmt"
-	"io/ioutil"
-	"os"
-	"path/filepath"
-	"strings"
+	"log"
 
-	"golang.org/x/vuln/client"
-	"golang.org/x/vuln/osv"
-	"golang.org/x/vuln/srv/internal"
 	"golang.org/x/vuln/srv/internal/database"
-	"golang.org/x/vuln/srv/internal/report"
-	"gopkg.in/yaml.v2"
 )
 
-func failf(format string, args ...interface{}) {
-	why := fmt.Sprintf(format, args...)
-	fmt.Fprintln(os.Stderr, why)
-	os.Exit(1)
-}
-
 // TODO(rolandshoemaker): once we have the HTML representation ready this should
 // be the prefix for that.
 const dbURL = "https://go.googlesource.com/vuln/+/refs/heads/master/reports/"
 
+var (
+	yamlDir = flag.String("reports", "reports", "Directory containing yaml reports")
+	jsonDir = flag.String("out", "out", "Directory to write JSON database to")
+)
+
 func main() {
-	yamlDir := flag.String("reports", "reports", "Directory containing yaml reports")
-	jsonDir := flag.String("out", "out", "Directory to write JSON database to")
 	flag.Parse()
-
-	yamlFiles, err := ioutil.ReadDir(*yamlDir)
-	if err != nil {
-		failf("can't read %q: %s", *yamlDir, err)
-	}
-
-	jsonVulns := map[string][]osv.Entry{}
-	var entries []osv.Entry
-	for _, f := range yamlFiles {
-		if !strings.HasSuffix(f.Name(), ".yaml") {
-			continue
-		}
-		content, err := ioutil.ReadFile(filepath.Join(*yamlDir, f.Name()))
-		if err != nil {
-			failf("can't read %q: %s", f.Name(), err)
-		}
-		var vuln report.Report
-		if err := yaml.UnmarshalStrict(content, &vuln); err != nil {
-			failf("unable to unmarshal %q: %s", f.Name(), err)
-		}
-		if lints := vuln.Lint(); len(lints) > 0 {
-			fmt.Fprintf(os.Stderr, "invalid vulnerability file %q:\n", f.Name())
-			for _, lint := range lints {
-				fmt.Fprintf(os.Stderr, "\t%s\n", lint)
-			}
-			os.Exit(1)
-		}
-
-		name := strings.TrimSuffix(filepath.Base(f.Name()), filepath.Ext(f.Name()))
-
-		// TODO(rolandshoemaker): once the HTML representation is ready this should be
-		// the link to the HTML page.
-		linkName := fmt.Sprintf("%s%s.yaml", dbURL, name)
-		entry, paths := database.Generate(name, linkName, vuln)
-		for _, path := range paths {
-			jsonVulns[path] = append(jsonVulns[path], entry)
-		}
-		entries = append(entries, entry)
-	}
-
-	index := make(client.DBIndex, len(jsonVulns))
-	for path, vulns := range jsonVulns {
-		outPath := filepath.Join(*jsonDir, path)
-		content, err := json.Marshal(vulns)
-		if err != nil {
-			failf("failed to marshal json: %s", err)
-		}
-		if err := os.MkdirAll(filepath.Dir(outPath), 0700); err != nil {
-			failf("failed to create directory %q: %s", filepath.Dir(outPath), err)
-		}
-		if err := ioutil.WriteFile(outPath+".json", content, 0644); err != nil {
-			failf("failed to write %q: %s", outPath+".json", err)
-		}
-		for _, v := range vulns {
-			if v.Modified.After(index[path]) || v.Published.After(index[path]) {
-				index[path] = v.Modified
-			}
-		}
-	}
-
-	indexJSON, err := json.Marshal(index)
-	if err != nil {
-		failf("failed to marshal index json: %s", err)
-	}
-	if err := ioutil.WriteFile(filepath.Join(*jsonDir, "index.json"), indexJSON, 0644); err != nil {
-		failf("failed to write index: %s", err)
-	}
-
-	// Write a directory containing entries by ID.
-	idDir := filepath.Join(*jsonDir, internal.IDDirectory)
-	if err := os.MkdirAll(idDir, 0700); err != nil {
-		failf("failed to create directory %q: %v", idDir, err)
-	}
-	var idIndex []string
-	for _, e := range entries {
-		outPath := filepath.Join(idDir, e.ID+".json")
-		content, err := json.Marshal(e)
-		if err != nil {
-			failf("failed to marshal json: %v", err)
-		}
-		if err := ioutil.WriteFile(outPath, content, 0644); err != nil {
-			failf("failed to write %q: %v", outPath, err)
-		}
-		idIndex = append(idIndex, e.ID)
-	}
-
-	// Write an index.json in the ID directory with a list of all the IDs.
-	idIndexJSON, err := json.Marshal(idIndex)
-	if err != nil {
-		failf("failed to marshal index json: %s", err)
-	}
-	if err := ioutil.WriteFile(filepath.Join(idDir, "index.json"), idIndexJSON, 0644); err != nil {
-		failf("failed to write index: %s", err)
+	if err := database.Generate(*yamlDir, *jsonDir, dbURL); err != nil {
+		log.Fatal(err)
 	}
 }
diff --git a/srv/internal/database/generate.go b/srv/internal/database/generate.go
index 7ca8542..e231774 100644
--- a/srv/internal/database/generate.go
+++ b/srv/internal/database/generate.go
@@ -6,14 +6,117 @@
 package database
 
 import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
 	"strings"
 
 	"golang.org/x/mod/semver"
+	"golang.org/x/vuln/client"
 	"golang.org/x/vuln/osv"
+	"golang.org/x/vuln/srv/internal"
+	"golang.org/x/vuln/srv/internal/derrors"
 	"golang.org/x/vuln/srv/internal/report"
+	"gopkg.in/yaml.v2"
 )
 
-func Generate(id string, url string, r report.Report) (osv.Entry, []string) {
+func Generate(yamlDir, jsonDir, dbURL string) (err error) {
+	defer derrors.Wrap(&err, "Generate(%q)", yamlDir)
+	yamlFiles, err := ioutil.ReadDir(yamlDir)
+	if err != nil {
+		return fmt.Errorf("can't read %q: %s", yamlDir, err)
+	}
+
+	jsonVulns := map[string][]osv.Entry{}
+	var entries []osv.Entry
+	for _, f := range yamlFiles {
+		if !strings.HasSuffix(f.Name(), ".yaml") {
+			continue
+		}
+		content, err := ioutil.ReadFile(filepath.Join(yamlDir, f.Name()))
+		if err != nil {
+			return fmt.Errorf("can't read %q: %s", f.Name(), err)
+		}
+		var vuln report.Report
+		if err := yaml.UnmarshalStrict(content, &vuln); err != nil {
+			return fmt.Errorf("unable to unmarshal %q: %s", f.Name(), err)
+		}
+		if lints := vuln.Lint(); len(lints) > 0 {
+			return fmt.Errorf("vuln.Lint: %v", lints)
+		}
+
+		name := strings.TrimSuffix(filepath.Base(f.Name()), filepath.Ext(f.Name()))
+
+		// TODO(rolandshoemaker): once the HTML representation is ready this should be
+		// the link to the HTML page.
+		linkName := fmt.Sprintf("%s%s.yaml", dbURL, name)
+		entry, paths := generateOSVEntry(name, linkName, vuln)
+		for _, path := range paths {
+			jsonVulns[path] = append(jsonVulns[path], entry)
+		}
+		entries = append(entries, entry)
+	}
+
+	index := make(client.DBIndex, len(jsonVulns))
+	for path, vulns := range jsonVulns {
+		outPath := filepath.Join(jsonDir, path)
+		content, err := json.Marshal(vulns)
+		if err != nil {
+			return fmt.Errorf("failed to marshal json: %s", err)
+		}
+		if err := os.MkdirAll(filepath.Dir(outPath), 0700); err != nil {
+			return fmt.Errorf("failed to create directory %q: %s", filepath.Dir(outPath), err)
+		}
+		if err := ioutil.WriteFile(outPath+".json", content, 0644); err != nil {
+			return fmt.Errorf("failed to write %q: %s", outPath+".json", err)
+		}
+		for _, v := range vulns {
+			if v.Modified.After(index[path]) || v.Published.After(index[path]) {
+				index[path] = v.Modified
+			}
+		}
+	}
+
+	indexJSON, err := json.Marshal(index)
+	if err != nil {
+		return fmt.Errorf("failed to marshal index json: %s", err)
+	}
+	if err := ioutil.WriteFile(filepath.Join(jsonDir, "index.json"), indexJSON, 0644); err != nil {
+		return fmt.Errorf("failed to write index: %s", err)
+	}
+
+	// Write a directory containing entries by ID.
+	idDir := filepath.Join(jsonDir, internal.IDDirectory)
+	if err := os.MkdirAll(idDir, 0700); err != nil {
+		return fmt.Errorf("failed to create directory %q: %v", idDir, err)
+	}
+	var idIndex []string
+	for _, e := range entries {
+		outPath := filepath.Join(idDir, e.ID+".json")
+		content, err := json.Marshal(e)
+		if err != nil {
+			return fmt.Errorf("failed to marshal json: %v", err)
+		}
+		if err := ioutil.WriteFile(outPath, content, 0644); err != nil {
+			return fmt.Errorf("failed to write %q: %v", outPath, err)
+		}
+		idIndex = append(idIndex, e.ID)
+	}
+
+	// Write an index.json in the ID directory with a list of all the IDs.
+	idIndexJSON, err := json.Marshal(idIndex)
+	if err != nil {
+		return fmt.Errorf("failed to marshal index json: %s", err)
+	}
+	if err := ioutil.WriteFile(filepath.Join(idDir, "index.json"), idIndexJSON, 0644); err != nil {
+		return fmt.Errorf("failed to write index: %s", err)
+	}
+	return nil
+}
+
+func generateOSVEntry(id string, url string, r report.Report) (osv.Entry, []string) {
 	importPath := r.Module
 	if r.Package != "" {
 		importPath = r.Package
@@ -68,7 +171,6 @@
 	for module := range moduleMap {
 		modules = append(modules, module)
 	}
-
 	return entry, modules
 }
 
diff --git a/srv/internal/database/generate_test.go b/srv/internal/database/generate_test.go
index 2108e66..2150a85 100644
--- a/srv/internal/database/generate_test.go
+++ b/srv/internal/database/generate_test.go
@@ -166,7 +166,7 @@
 	wantModules := []string{"example.com/vulnerable/v2", "vanity.host/vulnerable", "example.com/also-vulnerable"}
 	sort.Strings(wantModules)
 
-	gotEntry, gotModules := Generate("GO-1991-0001", url, r)
+	gotEntry, gotModules := generateOSVEntry("GO-1991-0001", url, r)
 	if diff := cmp.Diff(wantEntry, gotEntry, cmp.Comparer(func(a, b time.Time) bool { return a.Equal(b) })); diff != "" {
 		t.Errorf("Generate returned unexpected entry (-want +got):\n%s", diff)
 	}