Merge pull request #1484 from bensmrs/main

Add function checks to scriptlet validation

internal/server/scriptlet/load/load.go

@@ -3,6 +3,7 @@ package load
 import (
 	"fmt"
 	"slices"
+	"sort"
 	"sync"
 
 	"go.starlark.net/starlark"
@@ -15,8 +16,8 @@ const nameInstancePlacement = "instance_placement"
 // prefixQEMU is the prefix used in Starlark for the QEMU scriptlet.
 const prefixQEMU = "qemu"
 
-// prefixAuthorization is the prefix used in Starlark for the Authorization scriptlet.
-const prefixAuthorization = "authorization"
+// nameAuthorization is the name used in Starlark for the Authorization scriptlet.
+const nameAuthorization = "authorization"
 
 // compile compiles a scriptlet.
 func compile(programName string, src string, preDeclared []string) (*starlark.Program, error) {
@@ -33,6 +34,72 @@ func compile(programName string, src string, preDeclared []string) (*starlark.Pr
 	return mod, nil
 }
 
+// validate validates a scriptlet by compiling it and checking the presence of required functions.
+func validate(compiler func(string, string) (*starlark.Program, error), programName string, src string, requiredFunctions map[string][]string) error {
+	prog, err := compiler(programName, src)
+	if err != nil {
+		return err
+	}
+
+	thread := &starlark.Thread{Name: programName}
+	globals, err := prog.Init(thread, nil)
+	if err != nil {
+		return err
+	}
+
+	globals.Freeze()
+
+	var notFound []string
+	for funName, requiredArgs := range requiredFunctions {
+		// The function is missing if its name is not found in the globals.
+		funv := globals[funName]
+		if funv == nil {
+			notFound = append(notFound, funName)
+			continue
+		}
+
+		// The function is missing if its name is not bound to a function.
+		fun, ok := funv.(*starlark.Function)
+		if !ok {
+			notFound = append(notFound, funName)
+		}
+
+		// Get the function arguments.
+		argc := fun.NumParams()
+		var args []string
+		for i := range argc {
+			arg, _ := fun.Param(i)
+			args = append(args, arg)
+		}
+
+		// Return an error early if the function does not have the right arguments.
+		match := len(args) == len(requiredArgs)
+		if match {
+			sort.Strings(args)
+			sort.Strings(requiredArgs)
+			for i := range args {
+				if args[i] != requiredArgs[i] {
+					match = false
+					break
+				}
+			}
+		}
+
+		if !match {
+			return fmt.Errorf("The function %q defines arguments %q (expected: %q)", funName, args, requiredArgs)
+		}
+	}
+
+	switch len(notFound) {
+	case 0:
+		return nil
+	case 1:
+		return fmt.Errorf("The function %q is required but has not been found in the scriptlet", notFound[0])
+	default:
+		return fmt.Errorf("The functions %q are required but have not been found in the scriptlet", notFound)
+	}
+}
+
 var programsMu sync.Mutex
 var programs = make(map[string]*starlark.Program)
 
@@ -89,8 +156,9 @@ func InstancePlacementCompile(name string, src string) (*starlark.Program, error
 
 // InstancePlacementValidate validates the instance placement scriptlet.
 func InstancePlacementValidate(src string) error {
-	_, err := InstancePlacementCompile(nameInstancePlacement, src)
-	return err
+	return validate(InstancePlacementCompile, nameInstancePlacement, src, map[string][]string{
+		"instance_placement": {"request", "candidate_members"},
+	})
 }
 
 // InstancePlacementSet compiles the instance placement scriptlet into memory for use with InstancePlacementRun.
@@ -131,8 +199,9 @@ func QEMUCompile(name string, src string) (*starlark.Program, error) {
 
 // QEMUValidate validates the QEMU scriptlet.
 func QEMUValidate(src string) error {
-	_, err := QEMUCompile(prefixQEMU, src)
-	return err
+	return validate(QEMUCompile, prefixQEMU, src, map[string][]string{
+		"qemu_hook": {"hook_name"},
+	})
 }
 
 // QEMUSet compiles the QEMU scriptlet into memory for use with QEMURun.
@@ -157,17 +226,18 @@ func AuthorizationCompile(name string, src string) (*starlark.Program, error) {
 
 // AuthorizationValidate validates the authorization scriptlet.
 func AuthorizationValidate(src string) error {
-	_, err := AuthorizationCompile(prefixAuthorization, src)
-	return err
+	return validate(AuthorizationCompile, nameAuthorization, src, map[string][]string{
+		"authorize": {"details", "object", "entitlement"},
+	})
 }
 
 // AuthorizationSet compiles the authorization scriptlet into memory for use with AuthorizationRun.
 // If empty src is provided the current program is deleted.
 func AuthorizationSet(src string) error {
-	return set(AuthorizationCompile, prefixAuthorization, src)
+	return set(AuthorizationCompile, nameAuthorization, src)
 }
 
 // AuthorizationProgram returns the precompiled authorization scriptlet program.
 func AuthorizationProgram() (*starlark.Program, *starlark.Thread, error) {
-	return program("Authorization", prefixAuthorization)
+	return program("Authorization", nameAuthorization)
 }