From 418100c2f525773e95193a6e1f8b32b4d12bfe7b Mon Sep 17 00:00:00 2001
From: knqyf263 <knqyf263@gmail.com>
Date: Fri, 13 Sep 2024 12:00:07 +0400
Subject: [PATCH] refactor: reuse the parsed SBOM

Signed-off-by: knqyf263 <knqyf263@gmail.com>
---
 pkg/sbom/core/bom.go  |  8 ++++++++
 pkg/sbom/io/encode.go | 39 +++++++++++++++++++++++----------------
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/pkg/sbom/core/bom.go b/pkg/sbom/core/bom.go
index c0a082d013b5..d6ee429a0576 100644
--- a/pkg/sbom/core/bom.go
+++ b/pkg/sbom/core/bom.go
@@ -269,6 +269,10 @@ func (b *BOM) AddRelationship(parent, child *Component, relationshipType Relatio
 	}
 }
 
+func (b *BOM) AddVulnerability(c *Component, vuln Vulnerability) {
+	b.vulnerabilities[c.id] = append(b.vulnerabilities[c.id], vuln)
+}
+
 func (b *BOM) AddVulnerabilities(c *Component, vulns []Vulnerability) {
 	if c.id == uuid.Nil {
 		b.AddComponent(c)
@@ -279,6 +283,10 @@ func (b *BOM) AddVulnerabilities(c *Component, vulns []Vulnerability) {
 	b.vulnerabilities[c.id] = vulns
 }
 
+func (b *BOM) ClearVulnerabilities() {
+	b.vulnerabilities = make(map[uuid.UUID][]Vulnerability)
+}
+
 func (b *BOM) Root() *Component {
 	root, ok := b.components[b.rootID]
 	if !ok {
diff --git a/pkg/sbom/io/encode.go b/pkg/sbom/io/encode.go
index 975b98d40b10..a91ae2965036 100644
--- a/pkg/sbom/io/encode.go
+++ b/pkg/sbom/io/encode.go
@@ -20,9 +20,8 @@ import (
 )
 
 type Encoder struct {
-	bom        *core.BOM
-	opts       core.Options
-	components map[uuid.UUID]*core.Component
+	bom  *core.BOM
+	opts core.Options
 }
 
 func NewEncoder(opts core.Options) *Encoder {
@@ -31,7 +30,7 @@ func NewEncoder(opts core.Options) *Encoder {
 
 func (e *Encoder) Encode(report types.Report) (*core.BOM, error) {
 	if report.BOM != nil {
-		e.components = report.BOM.Components()
+		return e.reuseBOM(report)
 	}
 	// Metadata component
 	root, err := e.rootComponent(report)
@@ -262,16 +261,6 @@ func (e *Encoder) encodePackages(parent *core.Component, result types.Result) {
 	}
 }
 
-// existedPkgIdentifier tries to look for package identifier (BOM-ref, PURL) by component name and component type
-func (e *Encoder) existedPkgIdentifier(name string, componentType core.ComponentType) ftypes.PkgIdentifier {
-	for _, c := range e.components {
-		if c.Name == name && c.Type == componentType {
-			return c.PkgIdentifier
-		}
-	}
-	return ftypes.PkgIdentifier{}
-}
-
 func (e *Encoder) resultComponent(root *core.Component, r types.Result, osFound *ftypes.OS) *core.Component {
 	component := &core.Component{
 		Name: r.Target,
@@ -294,10 +283,8 @@ func (e *Encoder) resultComponent(root *core.Component, r types.Result, osFound
 			component.Version = osFound.Name
 		}
 		component.Type = core.TypeOS
-		component.PkgIdentifier = e.existedPkgIdentifier(component.Name, component.Type)
 	case types.ClassLangPkg:
 		component.Type = core.TypeApplication
-		component.PkgIdentifier = e.existedPkgIdentifier(component.Name, component.Type)
 	}
 
 	e.bom.AddRelationship(root, component, core.RelationshipContains)
@@ -446,6 +433,26 @@ func (*Encoder) belongToParent(pkg ftypes.Package, parents map[string]ftypes.Pac
 	}
 }
 
+func (e *Encoder) reuseBOM(report types.Report) (*core.BOM, error) {
+	report.BOM.ClearVulnerabilities()
+
+	// Group components by BOM-Ref
+	components := lo.MapKeys(report.BOM.Components(), func(value *core.Component, _ uuid.UUID) string {
+		return value.PkgIdentifier.BOMRef
+	})
+
+	for _, result := range report.Results {
+		for _, vuln := range result.Vulnerabilities {
+			c, ok := components[vuln.PkgIdentifier.BOMRef]
+			if !ok || c == nil {
+				continue
+			}
+			report.BOM.AddVulnerability(c, e.vulnerability(vuln))
+		}
+	}
+	return report.BOM, nil
+}
+
 func filterProperties(props []core.Property) []core.Property {
 	return lo.Filter(props, func(property core.Property, _ int) bool {
 		return !(property.Value == "" || (property.Name == core.PropertySrcEpoch && property.Value == "0"))