From e3e9cae4f6bc21b8af939ce16cb873093c51975e Mon Sep 17 00:00:00 2001 From: Anders Eknert Date: Mon, 27 May 2024 10:27:40 +0200 Subject: [PATCH] lsp: Add code lens support for evaluating rules - This is currently limited to VS Code clients, but may be extended to other clients too in the furure, if they can support it on their side. - Implementation wil recursively look for an `input.json` file and pick the one closest to the file evaluated, down to the level of the workspace root (if any file is found). Signed-off-by: Anders Eknert --- README.md | 2 +- bundle/regal/ast/search.rego | 5 -- bundle/regal/main.rego | 4 +- internal/lsp/eval.go | 113 ++++++++++++++++++++++++++++ internal/lsp/eval_test.go | 108 ++++++++++++++++++++++++++ internal/lsp/server.go | 142 ++++++++++++++++++++++++++++++++++- internal/lsp/types/types.go | 15 ++++ 7 files changed, 378 insertions(+), 11 deletions(-) create mode 100644 internal/lsp/eval.go create mode 100644 internal/lsp/eval_test.go diff --git a/README.md b/README.md index c5400291..f04a22af 100644 --- a/README.md +++ b/README.md @@ -702,7 +702,7 @@ in the near future: - [ ] Make "Check on save" unnecessary by allowing diagnostics to include [compilation errors](https://github.com/StyraInc/regal/issues/745) -- [ ] Add Code Lens to "Evaluate" any rule or package (VS Code only, initially) +- [x] Add Code Lens to "Evaluate" any rule or package (VS Code only, initially) - [ ] Implement [Signature Help](https://github.com/StyraInc/regal/issues/695) feature The roadmap is updated when all the current items have been completed. diff --git a/bundle/regal/ast/search.rego b/bundle/regal/ast/search.rego index 6199a781..4f118350 100644 --- a/bundle/regal/ast/search.rego +++ b/bundle/regal/ast/search.rego @@ -287,11 +287,6 @@ find_some_decl_names_in_scope(rule, location) := {some_var.value | _before_location(rule, some_var, location) } -# _rules_with_bodies[rule_index] := rule if { -# some rule_index, rule in input.rules -# not generated_body(rule) -# } - exprs[rule_index][expr_index] := expr if { some rule_index, rule in input.rules some expr_index, expr in rule.body diff --git a/bundle/regal/main.rego b/bundle/regal/main.rego index 0307a900..8ccf2935 100644 --- a/bundle/regal/main.rego +++ b/bundle/regal/main.rego @@ -12,10 +12,10 @@ lint.aggregates := aggregate lint.ignore_directives[input.regal.file.name] := ast.ignore_directives -lint.violations := report - lint_aggregate.violations := aggregate_report +lint.violations := report + rules_to_run[category][title] if { some category, title config.merged_config.rules[category][title] diff --git a/internal/lsp/eval.go b/internal/lsp/eval.go new file mode 100644 index 00000000..1d2d504f --- /dev/null +++ b/internal/lsp/eval.go @@ -0,0 +1,113 @@ +package lsp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path" + "path/filepath" + "strings" + + "github.com/open-policy-agent/opa/ast" + "github.com/open-policy-agent/opa/bundle" + "github.com/open-policy-agent/opa/rego" + "github.com/open-policy-agent/opa/topdown" + + "github.com/styrainc/regal/internal/lsp/uri" + "github.com/styrainc/regal/pkg/builtins" +) + +func (l *LanguageServer) Eval(ctx context.Context, query string, input string) (rego.ResultSet, error) { + modules := l.cache.GetAllModules() + moduleFiles := make([]bundle.ModuleFile, 0, len(modules)) + + for fileURI, module := range modules { + moduleFiles = append(moduleFiles, bundle.ModuleFile{ + URL: fileURI, + Parsed: module, + Path: uri.ToPath(l.clientIdentifier, fileURI), + }) + } + + bd := bundle.Bundle{ + Data: make(map[string]any), + Manifest: bundle.Manifest{ + Roots: &[]string{""}, + Metadata: map[string]any{"name": "workspace"}, + }, + Modules: moduleFiles, + } + + regoArgs := prepareRegoArgs(ast.MustParseBody(query), bd) + + pq, err := rego.New(regoArgs...).PrepareForEval(ctx) + if err != nil { + return nil, fmt.Errorf("failed preparing query: %w", err) + } + + if input != "" { + inputMap := make(map[string]any) + + err = json.Unmarshal([]byte(input), &inputMap) + if err != nil { + return nil, fmt.Errorf("failed unmarshalling input: %w", err) + } + + return pq.Eval(ctx, rego.EvalInput(inputMap)) //nolint:wrapcheck + } + + return pq.Eval(ctx) //nolint:wrapcheck +} + +type EvalPathResult struct { + Value any `json:"value"` + IsUndefined bool `json:"isUndefined"` +} + +func FindInput(file string, workspacePath string) string { + relative := strings.TrimPrefix(file, workspacePath) + components := strings.Split(path.Dir(relative), string(filepath.Separator)) + + for i := range len(components) { + inputPath := path.Join(workspacePath, path.Join(components[:len(components)-i]...), "input.json") + + if input, err := os.ReadFile(inputPath); err == nil { + return string(input) + } + } + + return "" +} + +func (l *LanguageServer) EvalWorkspacePath(ctx context.Context, query string, input string) (EvalPathResult, error) { + resultQuery := "result := " + query + + result, err := l.Eval(ctx, resultQuery, input) + if err != nil { + return EvalPathResult{}, fmt.Errorf("failed evaluating query: %w", err) + } + + if len(result) == 0 { + return EvalPathResult{IsUndefined: true}, nil + } + + res, ok := result[0].Bindings["result"] + if !ok { + return EvalPathResult{}, errors.New("expected result in bindings, didn't get it") + } + + return EvalPathResult{Value: res}, nil +} + +func prepareRegoArgs(query ast.Body, bd bundle.Bundle) []func(*rego.Rego) { + return []func(*rego.Rego){ + rego.ParsedQuery(query), + rego.ParsedBundle("workspace", &bd), + rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule), + rego.Function1(builtins.RegalLastMeta, builtins.RegalLast), + rego.EnablePrintStatements(true), + rego.PrintHook(topdown.NewPrintHook(os.Stderr)), + } +} diff --git a/internal/lsp/eval_test.go b/internal/lsp/eval_test.go new file mode 100644 index 00000000..db47785d --- /dev/null +++ b/internal/lsp/eval_test.go @@ -0,0 +1,108 @@ +package lsp + +import ( + "context" + "os" + "testing" + + "github.com/styrainc/regal/internal/parse" +) + +func TestEvalWorkspacePath(t *testing.T) { + t.Parallel() + + ls := NewLanguageServer(&LanguageServerOptions{ErrorLog: os.Stderr}) + + policy1 := `package policy1 + + import rego.v1 + + import data.policy2 + + default allow := false + + allow if policy2.allow + ` + + policy2 := `package policy2 + + import rego.v1 + + allow if input.exists + ` + + module1 := parse.MustParseModule(policy1) + module2 := parse.MustParseModule(policy2) + + ls.cache.SetFileContents("file://policy1.rego", policy1) + ls.cache.SetFileContents("file://policy2.rego", policy2) + ls.cache.SetModule("file://policy1.rego", module1) + ls.cache.SetModule("file://policy2.rego", module2) + + res, err := ls.EvalWorkspacePath(context.TODO(), "data.policy1.allow", `{"exists": true}`) + if err != nil { + t.Fatal(err) + } + + empty := EvalPathResult{} + + if res == empty { + t.Fatal("expected result, got nil") + } + + if val, ok := res.Value.(bool); !ok || val != true { + t.Fatalf("expected true, got false") + } +} + +func TestFindInput(t *testing.T) { + t.Parallel() + + tmpDir := t.TempDir() + + workspacePath := tmpDir + "/workspace" + file := tmpDir + "/workspace/foo/bar/baz.rego" + + if err := os.MkdirAll(workspacePath+"/foo/bar", 0o755); err != nil { + t.Fatal(err) + } + + if FindInput(file, workspacePath) != "" { + t.Fatalf("did not expect to find input.json") + } + + content := `{"x": 1}` + + createWithContent(t, tmpDir+"/workspace/foo/bar/input.json", content) + + if res := FindInput(file, workspacePath); res != content { + t.Errorf("expected input at %s, got %s", content, res) + } + + err := os.Remove(tmpDir + "/workspace/foo/bar/input.json") + if err != nil { + t.Fatal(err) + } + + createWithContent(t, tmpDir+"/workspace/input.json", content) + + if res := FindInput(file, workspacePath); res != content { + t.Errorf("expected input at %s, got %s", content, res) + } +} + +func createWithContent(t *testing.T, path string, content string) { + t.Helper() + + f, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + + defer f.Close() + + _, err = f.WriteString(content) + if err != nil { + t.Fatal(err) + } +} diff --git a/internal/lsp/server.go b/internal/lsp/server.go index bef3033f..4ebe1045 100644 --- a/internal/lsp/server.go +++ b/internal/lsp/server.go @@ -109,6 +109,7 @@ type fileUpdateEvent struct { Content string } +//nolint:gocyclo func (l *LanguageServer) Handle( ctx context.Context, conn *jsonrpc2.Conn, @@ -148,6 +149,8 @@ func (l *LanguageServer) Handle( return l.handleTextDocumentHover(ctx, conn, req) case "textDocument/inlayHint": return l.handleTextDocumentInlayHint(ctx, conn, req) + case "textDocument/codeLens": + return l.handleTextDocumentCodeLens(ctx, conn, req) case "textDocument/completion": return l.handleTextDocumentCompletion(ctx, conn, req) case "workspace/didChangeWatchedFiles": @@ -447,6 +450,59 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) { commands.ParseOptions{TargetArgIndex: 0, RowArgIndex: 1, ColArgIndex: 2}, params, ) + case "regal.eval": + fmt.Fprintf(os.Stderr, "regal.eval called with params: %v\n", params) + + if len(params.Arguments) != 3 { + l.logError(fmt.Errorf("expected three arguments, got %d", len(params.Arguments))) + + break + } + + file, ok := params.Arguments[0].(string) + if !ok { + l.logError(fmt.Errorf("expected first argument to be a string, got %T", params.Arguments[0])) + + break + } + + path, ok := params.Arguments[1].(string) + if !ok { + l.logError(fmt.Errorf("expected second argument to be a string, got %T", params.Arguments[1])) + + break + } + + line, ok := params.Arguments[2].(float64) + if !ok { + l.logError(fmt.Errorf("expected third argument to be a number, got %T", params.Arguments[2])) + + break + } + + input := FindInput( + uri.ToPath(l.clientIdentifier, file), + uri.ToPath(l.clientIdentifier, l.workspaceRootURI), + ) + + result, err := l.EvalWorkspacePath(ctx, path, input) + if err != nil { + fmt.Fprintf(os.Stderr, "failed to evaluate workspace path: %v\n", err) + + break + } + + responseParams := map[string]any{ + "result": result, + "line": line, + } + + responseResult := map[string]any{} + + err = l.conn.Call(ctx, "regal/showEvalResult", responseParams, &responseResult) + if err != nil { + l.logError(fmt.Errorf("failed %s notify: %v", "regal/hello", err.Error())) + } } if err != nil { @@ -934,6 +990,78 @@ func (l *LanguageServer) handleTextDocumentInlayHint( return inlayHints, nil } +func (l *LanguageServer) handleTextDocumentCodeLens( + _ context.Context, + _ *jsonrpc2.Conn, + req *jsonrpc2.Request, +) (result any, err error) { + var params types.CodeLensParams + if err := json.Unmarshal(*req.Params, ¶ms); err != nil { + return nil, fmt.Errorf("failed to unmarshal params: %w", err) + } + + if l.clientIdentifier != clients.IdentifierVSCode { + // only VSCode has the client side capability to handle the callback request + // to handle the result of evaluation, so displaying code lenses for any other + // editor is likely just going to result in a bad experience + return nil, nil // return a null response, as per the spec + } + + module, ok := l.cache.GetModule(params.TextDocument.URI) + if !ok { + l.logError(fmt.Errorf("failed to get module for uri %q", params.TextDocument.URI)) + + // return a null response, as per the spec + return nil, nil + } + + codeLenses := make([]types.CodeLens, 0) + + // Package + + pkgLens := types.CodeLens{ + Range: locationToRange(module.Package.Location), + Command: &types.Command{ + Title: "Evaluate", + Command: "regal.eval", + Arguments: &[]any{ + module.Package.Location.File, + module.Package.Path.String(), + module.Package.Location.Row, + }, + }, + } + + codeLenses = append(codeLenses, pkgLens) + + // Rules + + for _, rule := range module.Rules { + if rule.Head.Args != nil { + // Skip functions for now, as it's not clear how to best + // provide inputs for them. + continue + } + + ruleLens := types.CodeLens{ + Range: locationToRange(rule.Location), + Command: &types.Command{ + Title: "Evaluate", + Command: "regal.eval", + Arguments: &[]any{ + module.Package.Location.File, + module.Package.Path.String() + "." + string(rule.Head.Name), + rule.Head.Location.Row, + }, + }, + } + + codeLenses = append(codeLenses, ruleLens) + } + + return codeLenses, nil +} + func (l *LanguageServer) handleTextDocumentCompletion( _ context.Context, _ *jsonrpc2.Conn, @@ -1527,7 +1655,7 @@ func (l *LanguageServer) handleInitialize( ctx context.Context, _ *jsonrpc2.Conn, req *jsonrpc2.Request, -) (result any, err error) { +) (any, error) { var params types.InitializeParams if err := json.Unmarshal(*req.Params, ¶ms); err != nil { return nil, fmt.Errorf("failed to unmarshal params: %w", err) @@ -1554,7 +1682,7 @@ func (l *LanguageServer) handleInitialize( }, } - result = types.InitializeResult{ + initializeResult := types.InitializeResult{ Capabilities: types.ServerCapabilities{ TextDocumentSyncOptions: types.TextDocumentSyncOptions{ OpenClose: true, @@ -1590,6 +1718,7 @@ func (l *LanguageServer) handleInitialize( }, ExecuteCommandProvider: types.ExecuteCommandOptions{ Commands: []string{ + "regal.eval", "regal.fix.opa-fmt", "regal.fix.use-rego-v1", "regal.fix.use-assignment-operator", @@ -1610,6 +1739,13 @@ func (l *LanguageServer) handleInitialize( }, } + // Since evaluation requires some client side handling, this can't be supported + // purely by the LSP. Clients that are capable of handling the code lens callback + // should be added here though. + if l.clientIdentifier == clients.IdentifierVSCode { + initializeResult.Capabilities.CodeLensProvider = &types.CodeLensOptions{} + } + if l.workspaceRootURI != "" { configFile, err := config.FindConfig(uri.ToPath(l.clientIdentifier, l.workspaceRootURI)) if err == nil { @@ -1624,7 +1760,7 @@ func (l *LanguageServer) handleInitialize( l.diagnosticRequestWorkspace <- "server initialize" } - return result, nil + return initializeResult, nil } func (l *LanguageServer) loadWorkspaceContents(ctx context.Context, newOnly bool) ([]string, error) { diff --git a/internal/lsp/types/types.go b/internal/lsp/types/types.go index a8266b76..5d2283ef 100644 --- a/internal/lsp/types/types.go +++ b/internal/lsp/types/types.go @@ -105,6 +105,7 @@ type ServerCapabilities struct { WorkspaceSymbolProvider bool `json:"workspaceSymbolProvider"` DefinitionProvider bool `json:"definitionProvider"` CompletionProvider CompletionOptions `json:"completionProvider"` + CodeLensProvider *CodeLensOptions `json:"codeLensProvider,omitempty"` } type CompletionOptions struct { @@ -189,6 +190,20 @@ type CodeAction struct { Command Command `json:"command"` } +type CodeLensOptions struct { + ResolveProvider *bool `json:"resolveProvider,omitempty"` +} + +type CodeLensParams struct { + TextDocument TextDocumentIdentifier `json:"textDocument"` +} + +type CodeLens struct { + Range Range `json:"range"` + Command *Command `json:"command,omitempty"` + Data *any `json:"data,omitempty"` +} + type Command struct { Title string `json:"title"` Tooltip string `json:"tooltip"`