Create indices in guacone (guacsec#175)

* Create indices in guacone Solves guacsec#133. I'll try to find a solution where the indices are created automatically so we don't need to keep track of this functionality when adding new node types. Signed-off-by: Mihai Maruseac <[email protected]> * Fix format Signed-off-by: Mihai Maruseac <[email protected]> Signed-off-by: Mihai Maruseac <[email protected]>
Qiwen-Zhou · Oct 21, 2022 · 4b98664 · 4b98664
1 parent b515708
commit 4b98664
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 0 deletions.
diff --git a/cmd/guacone/cmd/files.go b/cmd/guacone/cmd/files.go
@@ -175,6 +175,12 @@ func getAssembler(opts options) (func([]assembler.Graph) error, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	err = createIndices(client)
+	if err != nil {
+		return nil, err
+	}
+
 	return func(gs []assembler.Graph) error {
 		combined := assembler.Graph{
 			Nodes: []assembler.GuacNode{},
@@ -190,3 +196,23 @@ func getAssembler(opts options) (func([]assembler.Graph) error, error) {
 		return nil
 	}, nil
 }
+
+func createIndices(client graphdb.Client) error {
+	indices := map[string][]string{
+		"Artifact":    {"digest", "name"},
+		"Package":     {"purl", "name"},
+		"Metadata":    {"id"},
+		"Attestation": {"digest"},
+	}
+
+	for label, attributes := range indices {
+		for _, attribute := range attributes {
+			err := assembler.CreateIndexOn(client, label, attribute)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
diff --git a/pkg/assembler/graphdb.go b/pkg/assembler/graphdb.go
@@ -86,6 +86,26 @@ func StoreGraph(g Graph, client graphdb.Client) error {
 	return err
 }
 
+// CreateIndexOn creates database indixes in the graph database given by Client
+// to optimize performance.
+func CreateIndexOn(client graphdb.Client, nodeLabel string, nodeAttribute string) error {
+	session := client.NewSession(neo4j.SessionConfig{})
+	defer session.Close()
+
+	var sb strings.Builder
+	sb.WriteString("CREATE INDEX IF NOT EXISTS FOR (n:")
+	sb.WriteString(nodeLabel) // not user controlled
+	sb.WriteString(") ON n.")
+	sb.WriteString(nodeAttribute) // not user controlled
+
+	_, err := session.WriteTransaction(
+		func(tx graphdb.Transaction) (interface{}, error) {
+			return tx.Run(sb.String(), nil)
+		})
+
+	return err
+}
+
 // Creates the "MERGE (n:${NODE_TYPE} {${ATTR}:${VALUE}, ...})" part of the query
 func queryPartForMergeNode(sb *strings.Builder, n GuacNode, label string) error {
 	node_data := n.Properties()