Updated based on code review

Signed-off-by: nathannaveen <[email protected]>
guacsec · Aug 12, 2024 · 2980fe7 · 2980fe7
1 parent a9d15ef
commit 2980fe7
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 15 deletions.
diff --git a/pkg/assembler/backends/keyvalue/certifyVuln.go b/pkg/assembler/backends/keyvalue/certifyVuln.go
@@ -18,6 +18,7 @@ package keyvalue
 import (
 	"context"
 	"errors"
+	"fmt"
 	"reflect"
 	"sort"
 	"strings"
@@ -30,6 +31,12 @@ import (
 	"github.com/guacsec/guac/pkg/assembler/kv"
 )
 
+const (
+	certifyVulnLinkType = "certifyVuln"
+	hasSLSALinkType     = "hasSLSA"
+	hasSBOMLinkType     = "hasSBOM"
+)
+
 // Internal data: link between packages and vulnerabilities (certifyVulnerability)
 type certifyVulnerabilityLink struct {
 	ThisID          string
@@ -62,16 +69,62 @@ func (n *certifyVulnerabilityLink) Key() string {
 }
 
 // Helper function to remove vulnerability links. This works by setting all the links expect the specified linkID.
-func (c *demoClient) removeLinks(ctx context.Context, linkID string, links []string, col string, id string) error {
+func (c *demoClient) removeLinks(ctx context.Context, linkID string, linkType string, links []string, col string, id string) error {
 	var newLinks []string
-	for _, id := range links {
-		if id != linkID {
-			newLinks = append(newLinks, id)
+	for _, link := range links {
+		if link != linkID {
+			newLinks = append(newLinks, link)
 		}
 	}
 
-	// Update the entity in the KeyValue store
-	return c.kv.Set(ctx, col, id, newLinks)
+	switch col {
+	case "packages":
+		var pkg pkgVersion
+		if err := c.kv.Get(ctx, col, id, &pkg); err != nil {
+			return fmt.Errorf("error getting package version from keyvalue: %w", err)
+		}
+		switch linkType {
+		case certifyVulnLinkType:
+			pkg.CertifyVulnLinks = newLinks
+		case hasSBOMLinkType:
+			pkg.HasSBOMs = newLinks
+		}
+		return setkv(ctx, col, &pkg, c)
+	case "vulnerabilities":
+		var vuln vulnTypeStruct
+		if err := c.kv.Get(ctx, col, id, &vuln); err != nil {
+			return fmt.Errorf("error getting vulnerability from keyvalue: %w", err)
+		}
+		switch linkType {
+		case certifyVulnLinkType:
+			vuln.VulnIDs = newLinks
+		}
+		return setkv(ctx, col, &vuln, c)
+	case "builders":
+		var builder builderStruct
+		if err := c.kv.Get(ctx, col, id, &builder); err != nil {
+			return fmt.Errorf("error getting builder from keyvalue: %w", err)
+		}
+		switch linkType {
+		case hasSLSALinkType:
+			builder.HasSLSAs = newLinks
+		}
+		return setkv(ctx, col, &builder, c)
+	case "artifacts":
+		var artifact artStruct
+		if err := c.kv.Get(ctx, col, id, &artifact); err != nil {
+			return fmt.Errorf("error getting artifact from keyvalue: %w", err)
+		}
+		switch linkType {
+		case hasSBOMLinkType:
+			artifact.HasSBOMs = newLinks
+		case hasSLSALinkType:
+			artifact.HasSLSAs = newLinks
+		}
+		return setkv(ctx, col, &artifact, c)
+	default:
+		return errors.New("unsupported entity type")
+	}
 }
 
 // DeleteCertifyVuln deletes a specified certifyVuln node along with all associated relationships.
@@ -92,15 +145,15 @@ func (c *demoClient) DeleteCertifyVuln(ctx context.Context, id string) (bool, er
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, foundPackage.CertifyVulnLinks, "packages", foundPackage.ID()); err != nil {
+	if err := c.removeLinks(ctx, link.ThisID, certifyVulnLinkType, foundPackage.CertifyVulnLinks, "packages", foundPackage.ID()); err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 
 	foundVulnNode, err := c.returnFoundVulnerability(ctx, &model.IDorVulnerabilityInput{VulnerabilityNodeID: &link.VulnerabilityID})
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, foundVulnNode.CertifyVulnLinks, "vulnerabilities", foundVulnNode.ID()); err != nil {
+	if err := c.removeLinks(ctx, link.ThisID, certifyVulnLinkType, foundVulnNode.CertifyVulnLinks, "vulnerabilities", foundVulnNode.ID()); err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 

diff --git a/pkg/assembler/backends/keyvalue/hasSBOM.go b/pkg/assembler/backends/keyvalue/hasSBOM.go
@@ -100,15 +100,15 @@ func (c *demoClient) DeleteHasSBOM(ctx context.Context, id string) (bool, error)
 		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
-		if err := c.removeLinks(ctx, link.ThisID, foundPkg.HasSBOMs, "packages", foundPkg.ID()); err != nil {
+		if err := c.removeLinks(ctx, link.ThisID, hasSBOMLinkType, foundPkg.HasSBOMs, "packages", foundPkg.ID()); err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
 	} else if link.Artifact != "" {
 		foundArtifact, err := c.returnFoundArtifact(ctx, &model.IDorArtifactInput{ArtifactID: &link.Artifact})
 		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
-		if err := c.removeLinks(ctx, link.ThisID, foundArtifact.HasSBOMs, "artifacts", foundArtifact.ID()); err != nil {
+		if err := c.removeLinks(ctx, link.ThisID, hasSBOMLinkType, foundArtifact.HasSBOMs, "artifacts", foundArtifact.ID()); err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/hasSLSA.go b/pkg/assembler/backends/keyvalue/hasSLSA.go
@@ -90,7 +90,7 @@ func (c *demoClient) DeleteHasSLSA(ctx context.Context, id string) (bool, error)
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, foundSubject.HasSLSAs, "artifacts", foundSubject.ID()); err != nil {
+	if err := c.removeLinks(ctx, link.ThisID, hasSLSALinkType, foundSubject.HasSLSAs, "artifacts", foundSubject.ID()); err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 
@@ -99,7 +99,7 @@ func (c *demoClient) DeleteHasSLSA(ctx context.Context, id string) (bool, error)
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, foundBuiltBy.HasSLSAs, "builders", foundBuiltBy.ID()); err != nil {
+	if err := c.removeLinks(ctx, link.ThisID, hasSLSALinkType, foundBuiltBy.HasSLSAs, "builders", foundBuiltBy.ID()); err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 
@@ -109,7 +109,7 @@ func (c *demoClient) DeleteHasSLSA(ctx context.Context, id string) (bool, error)
 		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
-		if err := c.removeLinks(ctx, link.ThisID, foundBuiltFrom.HasSLSAs, "artifacts", foundBuiltFrom.ID()); err != nil {
+		if err := c.removeLinks(ctx, link.ThisID, hasSLSALinkType, foundBuiltFrom.HasSLSAs, "artifacts", foundBuiltFrom.ID()); err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/path.go b/pkg/assembler/backends/keyvalue/path.go
@@ -18,7 +18,7 @@ package keyvalue
 import (
 	"context"
 	"fmt"
-	"log"
+	"github.com/guacsec/guac/pkg/logging"
 	"strings"
 
 	"github.com/vektah/gqlparser/v2/gqlerror"
@@ -194,6 +194,8 @@ func (c *demoClient) Nodes(ctx context.Context, ids []string) ([]model.Node, err
 // Delete node and all associated relationships. This functionality is only implemented for
 // certifyVuln, HasSBOM and HasSLSA.
 func (c *demoClient) Delete(ctx context.Context, nodeID string) (bool, error) {
+	logger := logging.FromContext(ctx)
+
 	// Retrieve the node type based on the node ID
 	var k string
 	if err := c.kv.Get(ctx, indexCol, nodeID, &k); err != nil {
@@ -227,7 +229,7 @@ func (c *demoClient) Delete(ctx context.Context, nodeID string) (bool, error) {
 		}
 		return deleted, nil
 	default:
-		log.Printf("Unknown node type: %s", nodeType)
+		logger.Debugf("Delete attempted for node id %s which was of type %s, and that type not supported for delete", nodeID, nodeType)
 	}
 	return false, nil
 }