database: only fetch relevant files for conflict check

This should let us avoid dealing with most of the files most of the time.

Signed-off-by: Stephan Renatus <[email protected]>

Author: srenatus

internal/database/database.go

@@ -34,7 +34,6 @@ import (
 	"github.com/open-policy-agent/opa-control-plane/internal/aws"
 	"github.com/open-policy-agent/opa-control-plane/internal/config"
 	"github.com/open-policy-agent/opa-control-plane/internal/database/sourcedatafs"
-	ocp_fs "github.com/open-policy-agent/opa-control-plane/internal/fs"
 	"github.com/open-policy-agent/opa-control-plane/internal/fs/mountfs"
 	"github.com/open-policy-agent/opa-control-plane/internal/jsonpatch"
 	"github.com/open-policy-agent/opa-control-plane/internal/logging"
@@ -404,6 +403,125 @@ WHERE source_name = %s AND path = %s AND (`+conditions+")", d.arg(0), d.arg(1)),
 	}
 }
 
+func (d *Database) checkSourceDataConflict(ctx context.Context, tx *sql.Tx, sourceName, path string, data any, bs []byte, principal string) error {
+	expr, err := authz.Partial(ctx, authz.Access{
+		Principal:  principal,
+		Permission: "sources.data.read",
+		Resource:   "sources",
+		Name:       sourceName,
+	}, nil)
+	if err != nil {
+		return err
+	}
+
+	conditions, args := expr.SQL(d.arg, []any{sourceName, path})
+	offset := len(args)
+	prefixPath := filepath.Dir(path)
+
+	// We're only fetching relevant data "blobs", based on their filesystem location:
+	// 1. anything clashing with the included object/non-object
+	// 2. anything "upwards" the prefixPath.
+	//
+	// For example, to update corp/users/internal/alice/data.json, we'll fetch
+	// 1a. path LIKE corp/users/internal/alice/KEY/% for each KEY in the object (len(keys) > 0)
+	// 1b. path LIKE corp/users/internal/alice/% if data.json is not an object
+	//     --> either non-empty result is a conflict
+	// 2. path IN (corp/users/internal/data.json corp/users/data.json corp/data.json)
+	//    --> these need to be fed into the loader to check
+
+	// deal with (1.), downwards conflicts
+	// NB(sr): We don't need to consult the data itself to determine errors here,
+	// so we only check existing paths.
+	{
+		var keys []string
+		switch d := data.(type) {
+		case map[string]any:
+			keys = slices.Collect(maps.Keys(d))
+			if keys == nil {
+				keys = []string{} // we use keys == nil to signal non-object data
+			}
+		}
+
+		prefixes := make([]any, max(len(keys), 1))
+		if keys == nil { // (1b)
+			prefixes = []any{prefixPath + "/%"} // any prefix path is a conflict
+		} else { // (1a)
+			for i := range keys {
+				prefixes[i] = prefixPath + "/" + keys[i] + "/%"
+			}
+		}
+		prefixArgs := make([]string, len(prefixes))
+		for i := range prefixes {
+			prefixArgs[i] = "path LIKE " + d.arg(offset+i)
+		}
+
+		values := make([]any, 0, len(prefixes)+len(args)+2)
+		values = append(values, sourceName, "true")
+		values = append(values, args[2:]...) // conditions
+		values = append(values, prefixes...)
+
+		query := fmt.Sprintf(`SELECT path FROM sources_data
+WHERE source_name = %[1]s
+  AND (%[2]s = %[2]s)
+  AND (%[3]s)
+  AND (%[4]s)
+ORDER BY path LIMIT 4`,
+			d.arg(0), d.arg(1), strings.Join(prefixArgs, " OR "), conditions)
+
+		files, err := queryPaths(ctx, tx, query, values...)
+		if err != nil {
+			if !errors.Is(err, sql.ErrNoRows) { // no rows, no conflict
+				return err
+			}
+		}
+
+		if len(files) > 0 {
+			if len(files) == 4 {
+				files[3] = "..."
+			}
+			return fmt.Errorf("%w: conflict with %v", ErrDataConflict, files)
+		}
+	}
+
+	// deal with (2.)
+	{
+		upwardsPaths := upwardsPaths(prefixPath)
+		if len(upwardsPaths) == 0 {
+			return nil
+		}
+		inParams := make([]string, len(upwardsPaths))
+		for i := range upwardsPaths {
+			inParams[i] = d.arg(i + offset)
+		}
+
+		query := fmt.Sprintf(`SELECT path FROM sources_data
+ WHERE source_name = %s AND ((path LIKE %s) OR (path in (%s))) AND (%s)`,
+			d.arg(0), d.arg(1), strings.Join(inParams, ","), conditions)
+		values := make([]any, 0, 2+len(upwardsPaths)+len(args))
+		values = append(values, sourceName, prefixPath+"/%")
+		values = append(values, args[2:]...) // conditions
+		values = append(values, upwardsPaths...)
+		files, err := queryPaths(ctx, tx, query, values...)
+		if err != nil {
+			return err
+		}
+		if len(files) == 0 {
+			return nil
+		}
+
+		// Attempt to load, i.e. merge with existing data. If it fails, don't upsert.
+		fs0 := mountfs.New(map[string]fs.FS{filepath.Dir(path): sourcedatafs.NewSingleFS(ctx, func(context.Context) ([]byte, error) { return bs, nil })})
+		fs1 := sourcedatafs.New(ctx, files, func(file string) func(context.Context) ([]byte, error) {
+			return d.sourceData(tx, sourceName, file)
+		})
+		fs2 := merged_fs.NewMergedFS(fs0, fs1)
+		if _, err := loader.NewFileLoader().WithFS(fs2).All([]string{"."}); err != nil {
+			return fmt.Errorf("%w: %w", ErrDataConflict, err)
+		}
+	}
+	return nil
+}
+
 func (d *Database) SourcesDataPut(ctx context.Context, sourceName, path string, data any, principal string) error {
 	path = filepath.ToSlash(path)
 	return tx1(ctx, d, d.sourcesDataPut(ctx, sourceName, path, data, principal))
@@ -430,28 +548,13 @@ func (d *Database) sourcesDataPut(ctx context.Context, sourceName, path string,
 			return err
 		}
 
-		// Attempt to load, i.e. merge with existing data. If it fails, don't upsert.
-		var files []string
-		type Path struct {
-			P string `sql:"path"`
-		}
-		for row, err := range sqlrange.QueryContext[Path](ctx, tx,
-			`SELECT path FROM sources_data WHERE source_name = `+d.arg(0),
-			sourceName) {
-			if err != nil {
-				return err
-			}
-			files = append(files, row.P)
-		}
-
-		fs0 := mountfs.New(map[string]fs.FS{filepath.Dir(path): sourcedatafs.NewSingleFS(ctx, func(context.Context) ([]byte, error) { return bs, nil })})
-		fs1 := sourcedatafs.New(ctx, files, func(file string) func(context.Context) ([]byte, error) {
-			return d.sourceData(tx, sourceName, file)
-		})
-		fs2 := merged_fs.NewMergedFS(fs0, fs1)
-		ocp_fs.Walk(fs2)
-		if _, err := loader.NewFileLoader().WithFS(fs2).All([]string{"."}); err != nil {
-			return fmt.Errorf("%w: %w", ErrDataConflict, err)
+		// NB: We only check for conflicts if the principal has the right to read source data.
+		// (Otherwise, write access to could be abused to guess the data or its layout? Let's
+		// err on the side of caution.)
+		// This is done implicitly: If the conditions are not satisfiable, none of the file
+		// lookups will yield anything.
+		if err := d.checkSourceDataConflict(ctx, tx, sourceName, path, data, bs, principal); err != nil {
+			return err
 		}
 
 		return d.upsert(ctx, tx, "sources_data", []string{"source_name", "path", "data"}, []string{"source_name", "path"}, sourceName, path, bs)
@@ -1308,18 +1411,11 @@ func (d *Database) iterSourceFiles(ctx context.Context, dbish sqlrange.Queryable
 		sourceName)
 }
 
-func (d *Database) iterSourceFilenames(ctx context.Context, dbish sqlrange.Queryable, sourceName string) iter.Seq2[Data, error] {
-	return sqlrange.QueryContext[Data](ctx,
-		dbish,
-		`SELECT path FROM sources_data WHERE source_name = `+d.arg(0),
-		sourceName)
-}
-
 func (d *Database) sourceData(tx *sql.Tx, sourceName, path string) func(context.Context) ([]byte, error) {
 	return func(ctx context.Context) ([]byte, error) {
 		var data []byte
 		err := tx.QueryRowContext(ctx,
-			`SELECT data FROM sources_data WHERE source_name = `+d.arg(0)+`AND path = `+d.arg(1),
+			`SELECT data FROM sources_data WHERE source_name = `+d.arg(0)+` AND path = `+d.arg(1),
 			sourceName, path,
 		).Scan(&data)
 		return data, err
@@ -1766,3 +1862,41 @@ func tx3[T any, U bool | string](ctx context.Context, db *Database, f func(*sql.
 
 	return result, result2, nil
 }
+
+func upwardsPaths(basePath string) []any {
+	prefixes := []any{}
+	parts := strings.Split(basePath, "/")
+	currentPath := ""
+
+	for i := 1; i < len(parts); i++ {
+		if i > 0 {
+			currentPath = strings.Join(parts[:i], "/")
+		}
+		prefixes = append(prefixes, currentPath+"/data.json")
+	}
+	return prefixes
+}
+
+func queryPaths(ctx context.Context, tx *sql.Tx, query string, values ...any) ([]string, error) {
+	rows, err := tx.QueryContext(ctx, query, values...)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var files []string
+	for rows.Next() {
+		var file string
+		if err := rows.Scan(&file); err != nil {
+			return nil, err
+		}
+		files = append(files, file)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return files, nil
+}