Repository: cerbos/cerbos
Pull request #2169 (2024-05-29 16:10:32)

feat(repl): Autocomplete for REPL directives, and a small fix for filenames
#### Description

There are two commits here, but one is a simple one liner.
- fix `:load` with a plain filename, these were rejected by the parser as they were treated as identifiers
- adds some autocomplete support for `:` directives:
  - `:...` get completed with directives
  - `:load` completes filenames
  - `:let` completes variables
  - `:exec` completes rule numbers


tcolgate: Added tests

Patch:
diff --git a/cmd/cerbos/repl/internal/directive.go b/cmd/cerbos/repl/internal/directive.go
index 12e96d73..d25b801c 100644
--- a/cmd/cerbos/repl/internal/directive.go
+++ b/cmd/cerbos/repl/internal/directive.go
@@ -57,7 +57,7 @@ type LetDirective struct {
 }
 
 type LoadDirective struct {
-	Path string `parser:"'load' @Path"`
+	Path string `parser:"'load' @(Path|Ident)"`
 }
 
 type ExecDirective struct {
diff --git a/cmd/cerbos/repl/internal/repl.go b/cmd/cerbos/repl/internal/repl.go
index 63cd4de7..7f2b93ea 100644
--- a/cmd/cerbos/repl/internal/repl.go
+++ b/cmd/cerbos/repl/internal/repl.go
@@ -9,11 +9,16 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/fs"
 	"os"
+	"path/filepath"
 	"reflect"
+	"regexp"
 	"sort"
+	"strconv"
 	"strings"
 	"time"
+	"unicode"
 
 	"github.com/alecthomas/participle/v2"
 	"github.com/google/cel-go/cel"
@@ -688,6 +693,153 @@ func (r *REPL) mkEnv() (*cel.Env, error) {
 	return conditions.StdEnv.Extend(cel.Declarations(decls...))
 }
 
+func (r *REPL) Complete(line string) []string {
+	parts := strings.Fields(line)
+	if len(parts) == 0 {
+		parts = []string{""}
+	}
+
+	chars := []rune(line)
+
+	if len(chars) != 0 && unicode.IsSpace(chars[len(chars)-1]) {
+		parts = append(parts, "")
+	}
+
+	lastTokIndx := len(parts) - 1
+	lastTok := parts[lastTokIndx]
+
+	if lastTokIndx == 0 {
+		return r.completeCmd(lastTok)
+	}
+
+	var compls []string
+	// This isn't ideal, but the directives are baked into tags
+	// making this trickier to implement neatly
+	switch parts[0] {
+	case ":load":
+		compls = r.completeFile(lastTok)
+	case ":let":
+		compls = r.completeVar(lastTok)
+	case ":exec":
+		compls = r.completeRule(lastTok)
+	}
+
+	res := make([]string, len(compls))
+	for i, c := range compls {
+		parts[lastTokIndx] = c
+		res[i] = strings.Join(parts, " ")
+	}
+
+	return res
+}
+
+var directives = []string{
+	":help",
+	":let",
+	":load",
+	":vars",
+	":rules",
+	":exec",
+	":reset",
+	":quit",
+}
+
+func (r *REPL) completeCmd(prefix string) (c []string) {
+	var matches []string
+
+	for _, d := range directives {
+		if strings.HasPrefix(d, prefix) {
+			matches = append(matches, d)
+		}
+	}
+
+	return matches
+}
+
+var (
+	skippableDirs  = regexp.MustCompile(`^(_schemas|testdata|derived_roles)$`)
+	skippableFiles = regexp.MustCompile(`(^\.)|(_test\.(yaml|yml|json)$)`)
+	matchingFiles  = regexp.MustCompile(`\.(yaml|yml|json)$`)
+)
+
+func (r *REPL) completeFile(prefix string) (c []string) {
+	path := filepath.Dir(prefix)
+	if path == "" {
+		path = "."
+	}
+
+	// This is subtle, but the first path may be a complete directory the user explicitly
+	// passed, so should not be skipped. This is also needed for the tests to pass (since
+	// we skip testdata
+	top := true
+
+	var matches []string
+
+	// TODO(tcm): this should probably limit the depth of the search
+	_ = filepath.WalkDir(path, func(path string, d fs.DirEntry, err error) error {
+		if top {
+			defer func() { top = false }()
+		}
+
+		if err != nil {
+			return err
+		}
+		dn := d.Name()
+		if !top && d.IsDir() && skippableDirs.MatchString(dn) {
+			return fs.SkipDir
+		}
+		if d.IsDir() {
+			return nil
+		}
+
+		if strings.HasPrefix(path, prefix) && !skippableFiles.MatchString(dn) && matchingFiles.MatchString(dn) {
+			matches = append(matches, path)
+		}
+		return nil
+	})
+	return matches
+}
+
+func (r *REPL) completeVar(prefix string) (c []string) {
+	var matches []string
+
+	for k := range r.vars {
+		if strings.HasPrefix(k, prefix) {
+			matches = append(matches, k)
+		}
+	}
+	sort.Strings(matches)
+
+	return matches
+}
+
+func (r *REPL) completeRule(prefix string) (c []string) {
+	var matches []string
+
+	if len(prefix) == 0 {
+		prefix = "#"
+	}
+
+	if r.policy == nil || !strings.HasPrefix(prefix, "#") {
+		return matches
+	}
+
+	numStr := prefix[1:]
+
+	if _, err := strconv.Atoi(numStr); numStr != "" && err != nil {
+		return matches
+	}
+
+	for idx := range r.policy.rules {
+		idxStr := strconv.Itoa(idx)
+		if strings.HasPrefix(idxStr, numStr) {
+			matches = append(matches, "#"+idxStr)
+		}
+	}
+
+	return matches
+}
+
 func isTerminated(line string, stack *runeStack) (string, bool) {
 	if line[len(line)-1] == '\\' {
 		return line[:len(line)-1], false
diff --git a/cmd/cerbos/repl/internal/testdata/complete/.hidden.yaml b/cmd/cerbos/repl/internal/testdata/complete/.hidden.yaml
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/internal/testdata/complete/_schemas/schema.json b/cmd/cerbos/repl/internal/testdata/complete/_schemas/schema.json
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/internal/testdata/complete/file.json b/cmd/cerbos/repl/internal/testdata/complete/file.json
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/internal/testdata/complete/file.yaml b/cmd/cerbos/repl/internal/testdata/complete/file.yaml
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/internal/testdata/complete/file.yml b/cmd/cerbos/repl/internal/testdata/complete/file.yml
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/internal/testdata/complete/sub/file.json b/cmd/cerbos/repl/internal/testdata/complete/sub/file.json
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/internal/testdata/complete/sub/file.yaml b/cmd/cerbos/repl/internal/testdata/complete/sub/file.yaml
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/internal/testdata/complete/sub/file.yml b/cmd/cerbos/repl/internal/testdata/complete/sub/file.yml
new file mode 100644
index 00000000..e69de29b
diff --git a/cmd/cerbos/repl/repl.go b/cmd/cerbos/repl/repl.go
index cb5a1992..f46e6a2e 100644
--- a/cmd/cerbos/repl/repl.go
+++ b/cmd/cerbos/repl/repl.go
@@ -51,6 +51,8 @@ func (c *Cmd) Run(k *kong.Kong) error {
 		return fmt.Errorf("failed to initialize the REPL: %w", err)
 	}
 
+	reader.SetCompleter(r.Complete)
+
 	return r.Loop()
 }
 


Test patch:
diff --git a/cmd/cerbos/repl/internal/repl_test.go b/cmd/cerbos/repl/internal/repl_test.go
index bf0eb9b7..21be7462 100644
--- a/cmd/cerbos/repl/internal/repl_test.go
+++ b/cmd/cerbos/repl/internal/repl_test.go
@@ -5,6 +5,7 @@ package internal
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 	"testing"
 
@@ -544,3 +545,121 @@ func loadConditionsFromPolicy(t *testing.T, path string) []*runtimev1.Condition
 
 	return conds
 }
+
+func TestComplete(t *testing.T) {
+	testRulesProtos := make([]proto.Message, 3)
+
+	wd, _ := os.Getwd()
+
+	expFilesStrs := []string{
+		"testdata/complete/file.json",
+		"testdata/complete/file.yaml",
+		"testdata/complete/file.yml",
+		"testdata/complete/sub/file.json",
+		"testdata/complete/sub/file.yaml",
+		"testdata/complete/sub/file.yml",
+	}
+
+	expFileLoads := make([]string, len(expFilesStrs))
+	for i, f := range expFilesStrs {
+		expFileLoads[i] = fmt.Sprintf(":load %s", filepath.FromSlash(f))
+	}
+
+	expFileAbsLoads := make([]string, len(expFilesStrs))
+	for i, f := range expFilesStrs {
+		expFileAbsLoads[i] = fmt.Sprintf(":load %s", filepath.Join(wd, filepath.FromSlash(f)))
+	}
+
+	testCases := []struct {
+		input    string
+		expected []string
+		rules    []proto.Message
+	}{
+		{"1 ", []string{}, nil},
+		{":l", []string{":let", ":load"}, nil},
+		{":lo", []string{":load"}, nil},
+
+		{":load", []string{":load"}, nil},
+		{
+			fmt.Sprintf(":load %s", filepath.Join("testdata", "complete")),
+			expFileLoads,
+			nil,
+		},
+		{
+			fmt.Sprintf(":load %s", filepath.Join(wd, "testdata", "complete")),
+			expFileAbsLoads,
+			nil,
+		},
+
+		{
+			":let ",
+			[]string{
+				":let G",
+				":let P",
+				":let R",
+				":let V",
+				":let _",
+				":let globals",
+				":let request",
+				":let runtime",
+				":let variables",
+			},
+			nil,
+		},
+
+		{
+			":exec ",
+			[]string{
+				":exec #0",
+				":exec #1",
+				":exec #2",
+			},
+			testRulesProtos,
+		},
+		{
+			":exec #",
+			[]string{
+				":exec #0",
+				":exec #1",
+				":exec #2",
+			},
+			testRulesProtos,
+		},
+		{
+			":exec #2",
+			[]string{
+				":exec #2",
+			},
+			testRulesProtos,
+		},
+		{
+			":exec #3",
+			[]string{},
+			testRulesProtos,
+		},
+		{
+			":exec #abc",
+			[]string{},
+			testRulesProtos,
+		},
+		{
+			":exec abc",
+			[]string{},
+			testRulesProtos,
+		},
+	}
+
+	for idx, tc := range testCases {
+		t.Run(fmt.Sprintf("TestCase_%d", idx), func(t *testing.T) {
+			mockOut := &mockOutput{}
+			repl, err := NewREPL(nil, mockOut)
+			require.NoError(t, err)
+
+			if tc.rules != nil {
+				repl.policy = &policyHolder{rules: tc.rules}
+			}
+
+			require.Equal(t, tc.expected, repl.Complete(tc.input))
+		})
+	}
+}