send vuln proto in pubsub data

Author: michaelkedar

go/internal/importer/importer.go

@@ -24,6 +24,7 @@ import (
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 	"k8s.io/apimachinery/pkg/util/yaml"
 )
 
@@ -311,7 +312,7 @@ func processUpdate(ctx context.Context, config Config, item WorkItem) {
 			return
 		}
 	}
-	if err := sendToWorker(ctx, config, item, hash, modified); err != nil {
+	if err := sendToWorker(ctx, config, item, hash, modified, &vulnProto); err != nil {
 		logger.ErrorContext(ctx, "Failed to send to worker", slog.Any("error", err), slog.String("source", sourceRepoName), slog.String("path", sourcePath))
 	}
 }
@@ -321,23 +322,32 @@ func computeHash(data []byte) string {
 	return hex.EncodeToString(hash[:])
 }
 
-func sendToWorker(ctx context.Context, config Config, item WorkItem, hash string, modifiedTime time.Time) error {
+func sendToWorker(ctx context.Context, config Config, item WorkItem, hash string, modifiedTime time.Time, vuln *osvschema.Vulnerability) error {
 	var srcTimestamp *time.Time
 	if !item.IsReimport {
 		// Only track the update latency if we're not doing a reimport of the data
 		srcTimestamp = &modifiedTime
 	}
 
-	return publishUpdate(ctx, config.Publisher, item.SourceRepository, item.SourcePath, hash, false, srcTimestamp)
+	return publishUpdate(ctx, config.Publisher, item.SourceRepository, item.SourcePath, hash, false, srcTimestamp, vuln)
 }
 
 func sendDeletionToWorker(ctx context.Context, config Config, item WorkItem) error {
-	return publishUpdate(ctx, config.Publisher, item.SourceRepository, item.SourcePath, "", true, nil)
+	return publishUpdate(ctx, config.Publisher, item.SourceRepository, item.SourcePath, "", true, nil, nil)
 }
 
-func publishUpdate(ctx context.Context, publisher clients.Publisher, source, path, hash string, deleted bool, srcTimestamp *time.Time) error {
+func publishUpdate(ctx context.Context, publisher clients.Publisher, source, path, hash string, deleted bool, srcTimestamp *time.Time, vuln *osvschema.Vulnerability) error {
+	// Send the vulnerability proto in the message data
+	var data []byte
+	if vuln != nil {
+		var err error
+		data, err = proto.Marshal(vuln)
+		if err != nil {
+			return err
+		}
+	}
 	msg := &pubsub.Message{
-		Data: []byte(""),
+		Data: data,
 		Attributes: map[string]string{
 			"type":            "update",
 			"source":          source,

go/internal/importer/importer_test.go

@@ -6,6 +6,9 @@ import (
 	"time"
 
 	"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 TestSendToWorker(t *testing.T) {
@@ -21,8 +24,12 @@ func TestSendToWorker(t *testing.T) {
 	}
 	hash := "some-hash"
 	modifiedTime := time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)
+	vuln := &osvschema.Vulnerability{
+		Id:       "CVE-2023-1234",
+		Modified: timestamppb.New(modifiedTime),
+	}
 
-	err := sendToWorker(ctx, config, item, hash, modifiedTime)
+	err := sendToWorker(ctx, config, item, hash, modifiedTime, vuln)
 	if err != nil {
 		t.Errorf("Expected no error, got %v", err)
 	}
@@ -53,6 +60,19 @@ func TestSendToWorker(t *testing.T) {
 	if msg.Attributes["src_timestamp"] != "1672531200" {
 		t.Errorf("Expected src_timestamp=1672531200, got %s", msg.Attributes["src_timestamp"])
 	}
+	if len(msg.Data) == 0 {
+		t.Errorf("Expected vulnerability data to be present")
+	}
+	var parsedVuln osvschema.Vulnerability
+	if err := proto.Unmarshal(msg.Data, &parsedVuln); err != nil {
+		t.Errorf("Failed to unmarshal vulnerability: %v", err)
+	}
+	if parsedVuln.Id != "CVE-2023-1234" {
+		t.Errorf("Expected vulnerability ID CVE-2023-1234, got %s", parsedVuln.Id)
+	}
+	if parsedVuln.Modified.AsTime() != modifiedTime {
+		t.Errorf("Expected vulnerability modified time %v, got %v", modifiedTime, parsedVuln.Modified.AsTime())
+	}
 }
 
 func TestImporterWorker(t *testing.T) {