Fix #2738 - undefined function for custom-builtins

The build and check commands were retrieving registered builtins too
early during initial execution, which did not give enough time for
custom binaries to register custom builtin functions. This change
sets the initial capabilities flag to nil instead, and then retrieves built-ins
immediately prior to instantiating a new compiler. This prevents an
undefined function rego_type_error during the build and check commands.
Fixes #2738.

This is my commit message

Signed-off-by: Grant Shively <gshively@godaddy.com>

cmd/build.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/bundle"
 
 	"github.com/spf13/cobra"
@@ -248,9 +249,17 @@ func dobuild(params buildParams, args []string) error {
 			return fmt.Errorf("enable bundle mode (ie. --bundle) to verify or sign bundle files or directories")
 		}
 	}
-
+	var capabilities *ast.Capabilities
+	// if capabilities are not provided as a cmd flag,
+	// then ast.CapabilitiesForThisVersion must be called
+	// within dobuild to ensure custom builtins are properly captured
+	if checkParams.capabilities.C != nil {
+		capabilities = checkParams.capabilities.C
+	} else {
+		capabilities = ast.CapabilitiesForThisVersion()
+	}
 	compiler := compile.New().
-		WithCapabilities(params.capabilities.C).
+		WithCapabilities(capabilities).
 		WithTarget(params.target.String()).
 		WithAsBundle(params.bundleMode).
 		WithOptimizationLevel(params.optimizationLevel).

cmd/check.go

@@ -85,10 +85,18 @@ func checkModules(args []string) int {
 			modules[m.Name] = m.Parsed
 		}
 	}
-
+	var capabilities *ast.Capabilities
+	// if capabilities are not provided as a cmd flag,
+	// then ast.CapabilitiesForThisVersion must be called
+	// within checkModules to ensure custom builtins are properly captured
+	if checkParams.capabilities.C != nil {
+		capabilities = checkParams.capabilities.C
+	} else {
+		capabilities = ast.CapabilitiesForThisVersion()
+	}
 	compiler := ast.NewCompiler().
 		SetErrorLimit(checkParams.errLimit).
-		WithCapabilities(checkParams.capabilities.C)
+		WithCapabilities(capabilities)
 
 	compiler.Compile(modules)
 

cmd/flags.go

@@ -144,7 +144,9 @@ type capabilitiesFlag struct {
 
 func newcapabilitiesFlag() *capabilitiesFlag {
 	return &capabilitiesFlag{
-		C: ast.CapabilitiesForThisVersion(),
+		// cannot call ast.CapabilitiesForThisVersion here because
+		// custom builtins cannot be registered by this point in execution
+		C: nil,
 	}
 }