[engine] Disallow calling register_check during load()

See rationale in runtime_shac.go.

If a valid reason for calling register_check from loaded files arises we
can reconsider, but it's easier if we're strict from the beginning.

Change-Id: Ic344d0b596b9b5f2d1be9c38e91d68cb339daf61
Reviewed-on: https://fuchsia-review.googlesource.com/c/shac-project/shac/+/909933
Commit-Queue: Oliver Newman <olivernewman@google.com>
Fuchsia-Auto-Submit: Oliver Newman <olivernewman@google.com>
Reviewed-by: Marc-Antoine Ruel <maruel@google.com>

doc/stdlib.md

@@ -748,6 +748,11 @@ def fail_often(ctx):
 shac.register_check(cb, fail_often)
 ```
 
+register_check cannot be called while a file is being imported by `load()`.
+It can only be called from the top-level execution path of the main
+shac.star file, either directly or by importing a function that calls
+register_check from another file.
+
 ### Arguments
 
 * **check**: `shac.check()` object or Starlark function that is called back to implement the check.

doc/stdlib.star

@@ -743,6 +743,11 @@ def _shac_register_check(check):
     shac.register_check(cb, fail_often)
     ```
 
+    register_check cannot be called while a file is being imported by `load()`.
+    It can only be called from the top-level execution path of the main
+    shac.star file, either directly or by importing a function that calls
+    register_check from another file.
+
   Args:
     check: `shac.check()` object or Starlark function that is called back to
       implement the check.

internal/engine/env.go

@@ -108,9 +108,12 @@ type loadedSource struct {
 // files in parallel.
 type starlarkEnv struct {
 	// Immutable.
-	// globals is available to all load() statements. They must be frozen via
-	// Freeze().
-	globals starlark.StringDict
+	// execedFileGlobals is available to the top-level Starlark file. They must be frozen
+	// via Freeze().
+	execedFileGlobals starlark.StringDict
+	// loadedFileGlobals is available to any loaded Starlark files. They must be
+	// frozen via Freeze().
+	loadedFileGlobals starlark.StringDict
 	// packages are all the available packages. It must include __main__.
 	packages map[string]fs.FS
 
@@ -145,14 +148,14 @@ func (s *starlarkEnv) load(ctx context.Context, sk sourceKey, pi printImpl) (sta
 		if err != nil {
 			return nil, err
 		}
-		return s.loadInner(ctx, th, skn, pi)
+		return s.loadInner(ctx, th, skn, pi, false)
 	}
 	t.SetLocal("shac.top", sk)
 	t.SetLocal("shac.pkg", sk)
-	return s.loadInner(ctx, t, sk, pi)
+	return s.loadInner(ctx, t, sk, pi, true)
 }
 
-func (s *starlarkEnv) loadInner(ctx context.Context, th *starlark.Thread, sk sourceKey, pi printImpl) (starlark.StringDict, error) {
+func (s *starlarkEnv) loadInner(ctx context.Context, th *starlark.Thread, sk sourceKey, pi printImpl, execed bool) (starlark.StringDict, error) {
 	key := sk.String()
 	s.mu.Lock()
 	ss, ok := s.sources[key]
@@ -194,7 +197,11 @@ func (s *starlarkEnv) loadInner(ctx context.Context, th *starlark.Thread, sk sou
 				oldsk := th.Local("shac.pkg").(sourceKey)
 				th.SetLocal("shac.pkg", sk)
 				fp := syntax.FilePortion{Content: d, FirstLine: 1, FirstCol: 1}
-				ss.globals, ss.err = starlark.ExecFile(th, sk.String(), fp, s.globals)
+				globals := s.execedFileGlobals
+				if !execed {
+					globals = s.loadedFileGlobals
+				}
+				ss.globals, ss.err = starlark.ExecFile(th, sk.String(), fp, globals)
 				th.SetLocal("shac.pkg", oldsk)
 				var errl resolve.ErrorList
 				if errors.As(ss.err, &errl) {

internal/engine/run.go

@@ -234,9 +234,10 @@ func runInner(ctx context.Context, root, tmpdir, main string, r Report, allowNet
 		return err
 	}
 	env := starlarkEnv{
-		globals:  getPredeclared(),
-		sources:  map[string]*loadedSource{},
-		packages: packages,
+		execedFileGlobals: getPredeclared(true),
+		loadedFileGlobals: getPredeclared(false),
+		sources:           map[string]*loadedSource{},
+		packages:          packages,
 	}
 
 	newState := func(scm scmCheckout, subdir string, idx int) *shacState {

internal/engine/run_test.go

@@ -1283,6 +1283,11 @@ func TestTestDataFailOrThrow(t *testing.T) {
 			"cannot load ./load-recurse.star: //load-recurse.star was loaded in a cycle dependency graph",
 			"  //load-recurse.star:15:1: in <toplevel>\n",
 		},
+		{
+			"load-register_check.star",
+			"cannot load ./fail-check.star: shac.register_check: cannot be called from a loaded file",
+			"  //load-register_check.star:16:1: in <toplevel>\n",
+		},
 		{
 			"shac-immutable.star",
 			"can't assign to .key field of struct",

internal/engine/runtime.go

@@ -31,9 +31,9 @@ import (
 // The upstream starlark interpreter includes all the symbols described at
 // https://github.com/google/starlark-go/blob/HEAD/doc/spec.md#built-in-constants-and-functions
 // See https://pkg.go.dev/go.starlark.net/starlark#Universe for the default list.
-func getPredeclared() starlark.StringDict {
+func getPredeclared(execed bool) starlark.StringDict {
 	return starlark.StringDict{
-		"shac": toValue("shac", getShac()),
+		"shac": toValue("shac", getShac(execed)),
 
 		// Add https://bazel.build/rules/lib/json so it feels more natural to bazel
 		// users.

internal/engine/runtime_shac.go

@@ -33,15 +33,30 @@ var (
 // getShac returns the global shac object.
 //
 // Make sure to update //doc/stdlib.star whenever this function is modified.
-func getShac() starlark.StringDict {
-	return starlark.StringDict{
+func getShac(execed bool) starlark.StringDict {
+	ret := starlark.StringDict{
 		"check":          newBuiltin("shac.check", shacCheck),
 		"commit_hash":    starlark.String(getCommitHash()),
 		"register_check": newBuiltinNone("shac.register_check", shacRegisterCheck),
 		"version": starlark.Tuple{
 			starlark.MakeInt(Version[0]), starlark.MakeInt(Version[1]), starlark.MakeInt(Version[2]),
 		},
 	}
+	if !execed {
+		// Disallow loaded files from calling `shac.register_check`. Only the
+		// top-level shac.star file can call `shac.register_check`; outside
+		// shac.star, `shac.register_check` cannot be called at the top level,
+		// only in functions that are loaded and called by shac.star.
+		//
+		// This prevents libraries from forcibly registering checks for all
+		// users, and forces users to be somewhat explicit about which checks
+		// are registered.
+		ret["register_check"] = newBuiltinNone("shac.register_check",
+			func(ctx context.Context, s *shacState, name string, args starlark.Tuple, kwargs []starlark.Tuple) error {
+				return fmt.Errorf("cannot be called from a loaded file")
+			})
+	}
+	return ret
 }
 
 // shacRegisterCheck implements native function shac.register_check().

internal/engine/testdata/fail_or_throw/load-register_check.star

@@ -0,0 +1,16 @@
+# Copyright 2023 The Shac Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Loading a file that registers a check is not allowed.
+load("./fail-check.star", "cb")