skycfg.go

// Copyright 2018 The Skycfg 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.
//
// SPDX-License-Identifier: Apache-2.0

// Package skycfg is an extension library for the Starlark language that adds support
// for constructing Protocol Buffer messages.
package skycfg

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"go.starlark.net/starlark"
	"go.starlark.net/starlarkjson"
	"go.starlark.net/starlarkstruct"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/reflect/protoregistry"

	"github.com/stripe/skycfg/go/assertmodule"
	"github.com/stripe/skycfg/go/hashmodule"
	"github.com/stripe/skycfg/go/protomodule"
	"github.com/stripe/skycfg/go/urlmodule"
	"github.com/stripe/skycfg/go/yamlmodule"
)

// A FileReader controls how load() calls resolve and read other modules.
type FileReader interface {
	// Resolve parses the "name" part of load("name", "symbol") to a path. This
	// is not required to correspond to a true path on the filesystem, but should
	// be "absolute" within the semantics of this FileReader.
	//
	// fromPath will be empty when loading the root module passed to Load().
	Resolve(ctx context.Context, name, fromPath string) (path string, err error)

	// ReadFile reads the content of the file at the given path, which was
	// returned from Resolve().
	ReadFile(ctx context.Context, path string) ([]byte, error)
}

type localFileReader struct {
	root string
}

// LocalFileReader returns a FileReader that resolves and loads files from
// within a given filesystem directory.
func LocalFileReader(root string) FileReader {
	if root == "" {
		panic("LocalFileReader: empty root path")
	}
	return &localFileReader{root}
}

func (r *localFileReader) Resolve(ctx context.Context, name, fromPath string) (string, error) {
	if fromPath == "" {
		return name, nil
	}
	if filepath.Separator != '/' && strings.ContainsRune(name, filepath.Separator) {
		return "", fmt.Errorf("load(%q): invalid character in module name", name)
	}
	resolved := filepath.Join(r.root, filepath.FromSlash(path.Clean("/"+name)))
	return resolved, nil
}

func (r *localFileReader) ReadFile(ctx context.Context, path string) ([]byte, error) {
	return ioutil.ReadFile(path)
}

// NewProtoMessage returns a Starlark value representing the given Protobuf
// message. It can be returned back to a proto.Message() via AsProtoMessage().
func NewProtoMessage(msg proto.Message) (starlark.Value, error) {
	return protomodule.NewMessage(msg)
}

// AsProtoMessage returns a Protobuf message underlying the given Starlark
// value, which must have been created by NewProtoMessage(). Returns
// (_, false) if the value is not a valid message.
func AsProtoMessage(v starlark.Value) (proto.Message, bool) {
	return protomodule.AsProtoMessage(v)
}

// NewProtoPackage returns a Starlark value representing the given Protobuf
// package. It can be added to global symbols when loading a Skycfg config file.
func NewProtoPackage(r unstableProtoRegistryV2, name string) starlark.Value {
	return protomodule.NewProtoPackage(r.UnstableProtobufTypes(), protoreflect.FullName(name))
}

// A Config is a Skycfg config file that has been fully loaded and is ready
// for execution.
type Config struct {
	filename string
	globals  starlark.StringDict
	locals   starlark.StringDict
	tests    []*Test
}

// A LoadOption adjusts details of how Skycfg configs are loaded.
type LoadOption interface {
	applyLoad(*loadOptions)
}

type loadOptions struct {
	globals       starlark.StringDict
	fileReader    FileReader
	protoRegistry unstableProtoRegistryV2
}

type fnLoadOption func(*loadOptions)

func (fn fnLoadOption) applyLoad(opts *loadOptions) { fn(opts) }

// WithGlobals adds additional global symbols to the Starlark environment
// when loading a Skycfg config.
func WithGlobals(globals starlark.StringDict) LoadOption {
	return fnLoadOption(func(opts *loadOptions) {
		for key, value := range globals {
			opts.globals[key] = value
		}
	})
}

// WithFileReader changes the implementation of load() when loading a
// Skycfg config.
func WithFileReader(r FileReader) LoadOption {
	if r == nil {
		panic("WithFileReader: nil reader")
	}
	return fnLoadOption(func(opts *loadOptions) {
		opts.fileReader = r
	})
}

type unstableProtoRegistryV2 interface {
	// UNSTABLE go-protobuf v2 type registry
	UnstableProtobufTypes() *protoregistry.Types
}

type protoRegistryWrapper struct {
	protoRegistry *protoregistry.Types
}

var _ (unstableProtoRegistryV2) = (*protoRegistryWrapper)(nil)

func (pr *protoRegistryWrapper) UnstableProtobufTypes() *protoregistry.Types {
	return pr.protoRegistry
}

func NewUnstableProtobufRegistryV2(r *protoregistry.Types) unstableProtoRegistryV2 {
	return &protoRegistryWrapper{r}
}

// WithProtoRegistry is an EXPERIMENTAL and UNSTABLE option to override
// how Protobuf message type names are mapped to Go types.
func WithProtoRegistry(r unstableProtoRegistryV2) LoadOption {
	if r == nil {
		panic("WithProtoRegistry: nil registry")
	}
	return fnLoadOption(func(opts *loadOptions) {
		opts.protoRegistry = r
	})
}

// UnstablePredeclaredModules returns a Starlark string dictionary with
// predeclared Skycfg modules which can be used in starlark.ExecFile.
//
// Takes in unstableProtoRegistry as param (if nil will use standard proto
// registry).
//
// Currently provides these modules (see REAMDE for more detailed description):
//  * fail   - interrupts execution and prints a stacktrace.
//  * hash   - supports md5, sha1 and sha245 functions.
//  * json   - marshals plain values (dicts, lists, etc) to JSON.
//  * proto  - package for constructing Protobuf messages.
//  * struct - experimental Starlark struct support.
//  * yaml   - same as "json" package but for YAML.
//  * url    - utility package for encoding URL query string.
func UnstablePredeclaredModules(r unstableProtoRegistryV2) starlark.StringDict {
	return starlark.StringDict{
		"fail":   assertmodule.Fail,
		"hash":   hashmodule.NewModule(),
		"json":   newJsonModule(),
		"proto":  UnstableProtoModule(r),
		"struct": starlark.NewBuiltin("struct", starlarkstruct.Make),
		"yaml":   newYamlModule(),
		"url":    urlmodule.NewModule(),
	}
}

func newJsonModule() starlark.Value {
	// Copy starlarkjson and add compatibility fields
	starlarjsonModule := starlarkjson.Module
	module := &starlarkstruct.Module{
		Name:    starlarjsonModule.Name,
		Members: make(starlark.StringDict),
	}
	for k, v := range starlarjsonModule.Members {
		module.Members[k] = v
	}

	// Aliases for compatibility with pre-v1.0 Skycfg API.
	module.Members["marshal"] = module.Members["encode"]
	module.Members["unmarshal"] = module.Members["decode"]

	return module
}

func newYamlModule() starlark.Value {
	module := yamlmodule.NewModule()

	// Aliases for compatibility with pre-v1.0 Skycfg API.
	module.Members["marshal"] = module.Members["encode"]
	module.Members["unmarshal"] = module.Members["decode"]

	return module
}

func UnstableProtoModule(r unstableProtoRegistryV2) starlark.Value {
	protoTypes := protoregistry.GlobalTypes
	if r != nil {
		protoTypes = r.UnstableProtobufTypes()
	}

	protoModule := protomodule.NewModule(protoTypes)

	// Compatibility aliases
	protoModule.Members["from_json"] = protoModule.Members["decode_json"]
	protoModule.Members["from_text"] = protoModule.Members["decode_text"]
	protoModule.Members["to_any"] = protoModule.Members["encode_any"]
	protoModule.Members["to_json"] = protoModule.Members["encode_json"]
	protoModule.Members["to_text"] = protoModule.Members["encode_text"]

	return protoModule
}

// Load reads a Skycfg config file from the filesystem.
func Load(ctx context.Context, filename string, opts ...LoadOption) (*Config, error) {
	parsedOpts := &loadOptions{
		globals:    starlark.StringDict{},
		fileReader: LocalFileReader(filepath.Dir(filename)),
	}
	for _, opt := range opts {
		opt.applyLoad(parsedOpts)
	}

	overriddenGlobals := parsedOpts.globals
	parsedOpts.globals = UnstablePredeclaredModules(parsedOpts.protoRegistry)
	for key, value := range overriddenGlobals {
		parsedOpts.globals[key] = value
	}
	configLocals, tests, err := loadImpl(ctx, parsedOpts, filename)
	if err != nil {
		return nil, err
	}
	return &Config{
		filename: filename,
		globals:  parsedOpts.globals,
		locals:   configLocals,
		tests:    tests,
	}, nil
}

func loadImpl(ctx context.Context, opts *loadOptions, filename string) (starlark.StringDict, []*Test, error) {
	reader := opts.fileReader

	type cacheEntry struct {
		globals starlark.StringDict
		err     error
	}
	cache := make(map[string]*cacheEntry)
	tests := []*Test{}

	var load func(thread *starlark.Thread, moduleName string) (starlark.StringDict, error)
	load = func(thread *starlark.Thread, moduleName string) (starlark.StringDict, error) {
		var fromPath string
		if thread.CallStackDepth() > 0 {
			fromPath = thread.CallFrame(0).Pos.Filename()
		}
		modulePath, err := reader.Resolve(ctx, moduleName, fromPath)
		if err != nil {
			return nil, err
		}

		e, ok := cache[modulePath]
		if e != nil {
			return e.globals, e.err
		}
		if ok {
			return nil, fmt.Errorf("cycle in load graph")
		}
		moduleSource, err := reader.ReadFile(ctx, modulePath)
		if err != nil {
			cache[modulePath] = &cacheEntry{nil, err}
			return nil, err
		}

		cache[modulePath] = nil
		globals, err := starlark.ExecFile(thread, modulePath, moduleSource, opts.globals)
		cache[modulePath] = &cacheEntry{globals, err}

		for name, val := range globals {
			if !strings.HasPrefix(name, "test_") {
				continue
			}
			if fn, ok := val.(starlark.Callable); ok {
				tests = append(tests, &Test{
					callable: fn,
				})
			}
		}
		return globals, err
	}
	locals, err := load(&starlark.Thread{
		Print: skyPrint,
		Load:  load,
	}, filename)
	return locals, tests, err
}

// Filename returns the original filename passed to Load().
func (c *Config) Filename() string {
	return c.filename
}

// Globals returns the set of variables in the Starlark global namespace,
// including any added to the config loader by WithGlobals().
func (c *Config) Globals() starlark.StringDict {
	return c.globals
}

// Locals returns the set of variables in the Starlark local namespace for
// the top-level module.
func (c *Config) Locals() starlark.StringDict {
	return c.locals
}

// An ExecOption adjusts details of how a Skycfg config's main function is
// executed.
type ExecOption interface {
	applyExec(*execOptions)
}

type execOptions struct {
	vars         *starlark.Dict
	funcName     string
	flattenLists bool
}

type fnExecOption func(*execOptions)

func (fn fnExecOption) applyExec(opts *execOptions) { fn(opts) }

// WithVars adds key:value pairs to the ctx.vars dict passed to main().
func WithVars(vars starlark.StringDict) ExecOption {
	return fnExecOption(func(opts *execOptions) {
		for key, value := range vars {
			opts.vars.SetKey(starlark.String(key), value)
		}
	})
}

// WithEntryPoint changes the name of the Skycfg function to execute.
func WithEntryPoint(name string) ExecOption {
	return fnExecOption(func(opts *execOptions) {
		opts.funcName = name
	})
}

// WithFlattenLists flatten lists one layer deep ([1, [2,3]] -> [1, 2, 3])
func WithFlattenLists() ExecOption {
	return fnExecOption(func(opts *execOptions) {
		opts.flattenLists = true
	})
}

// Main executes main() or a custom entry point function from the top-level Skycfg config
// module, which is expected to return either None or a list of Protobuf messages.
func (c *Config) Main(ctx context.Context, opts ...ExecOption) ([]proto.Message, error) {
	parsedOpts := &execOptions{
		vars:     &starlark.Dict{},
		funcName: "main",
	}
	for _, opt := range opts {
		opt.applyExec(parsedOpts)
	}
	mainVal, ok := c.locals[parsedOpts.funcName]
	if !ok {
		return nil, fmt.Errorf("no %q function found in %q", parsedOpts.funcName, c.filename)
	}
	main, ok := mainVal.(starlark.Callable)
	if !ok {
		return nil, fmt.Errorf("%q must be a function (got a %s)", parsedOpts.funcName, mainVal.Type())
	}

	thread := &starlark.Thread{
		Print: skyPrint,
	}
	thread.SetLocal("context", ctx)
	mainCtx := &starlarkstruct.Module{
		Name: "skycfg_ctx",
		Members: starlark.StringDict(map[string]starlark.Value{
			"vars": parsedOpts.vars,
		}),
	}
	args := starlark.Tuple([]starlark.Value{mainCtx})
	mainVal, err := starlark.Call(thread, main, args, nil)
	if err != nil {
		return nil, err
	}
	mainList, ok := mainVal.(*starlark.List)
	if !ok {
		if _, isNone := mainVal.(starlark.NoneType); isNone {
			return nil, nil
		}
		return nil, fmt.Errorf("%q didn't return a list (got a %s)", parsedOpts.funcName, mainVal.Type())
	}
	var msgs []proto.Message
	for ii := 0; ii < mainList.Len(); ii++ {
		maybeMsg := mainList.Index(ii)
		// Only flatten but not flatten deep. This will flatten out, in order, lists within main list and append the
		// message into msgs
		if maybeMsgList, ok := maybeMsg.(*starlark.List); parsedOpts.flattenLists && ok {
			for iii := 0; iii < maybeMsgList.Len(); iii++ {
				maybeNestedMsg := maybeMsgList.Index(iii)
				msg, ok := AsProtoMessage(maybeNestedMsg)
				if !ok {
					return nil, fmt.Errorf("%q returned something that's not a protobuf (a %s) within a nested list", parsedOpts.funcName, maybeNestedMsg.Type())
				}
				msgs = append(msgs, msg)
			}
		} else {
			msg, ok := AsProtoMessage(maybeMsg)
			if !ok {
				return nil, fmt.Errorf("%q returned something that's not a protobuf (a %s)", parsedOpts.funcName, maybeMsg.Type())
			}
			msgs = append(msgs, msg)
		}
	}
	return msgs, nil
}

// A TestResult is the result of a test run
type TestResult struct {
	TestName string
	Failure  error
	Duration time.Duration
}

// A Test is a test case, which is a skycfg function whose name starts with `test_`.
type Test struct {
	callable starlark.Callable
}

// Name returns the name of the test (the name of the function)
func (t *Test) Name() string {
	return t.callable.Name()
}

// An TestOption adjusts details of how a Skycfg config's test functions are
// executed.
type TestOption interface {
	applyTest(*testOptions)
}

type testOptions struct {
	vars *starlark.Dict
}

type fnTestOption func(*testOptions)

func (fn fnTestOption) applyTest(opts *testOptions) { fn(opts) }

// WithTestVars adds key:value pairs to the ctx.vars dict passed to tests
func WithTestVars(vars starlark.StringDict) TestOption {
	return fnTestOption(func(opts *testOptions) {
		for key, value := range vars {
			opts.vars.SetKey(starlark.String(key), value)
		}
	})
}

// Run actually executes a test. It returns a TestResult if the test completes (even if it fails)
// The error return value will only be non-nil if the test execution itself errors.
func (t *Test) Run(ctx context.Context, opts ...TestOption) (*TestResult, error) {
	parsedOpts := &testOptions{
		vars: &starlark.Dict{},
	}
	for _, opt := range opts {
		opt.applyTest(parsedOpts)
	}

	thread := &starlark.Thread{
		Print: skyPrint,
	}
	thread.SetLocal("context", ctx)

	assertModule := assertmodule.AssertModule()
	testCtx := &starlarkstruct.Module{
		Name: "skycfg_test_ctx",
		Members: starlark.StringDict(map[string]starlark.Value{
			"vars":   parsedOpts.vars,
			"assert": assertModule,
		}),
	}
	args := starlark.Tuple([]starlark.Value{testCtx})

	result := TestResult{
		TestName: t.Name(),
	}

	startTime := time.Now()
	_, err := starlark.Call(thread, t.callable, args, nil)
	result.Duration = time.Since(startTime)
	if err != nil {
		// if there is no assertion error, there was something wrong with the execution itself
		if len(assertModule.Failures) == 0 {
			return nil, err
		}

		// there should only be one failure, because each test run gets its own *TestContext
		// and each assertion failure halts execution.
		if len(assertModule.Failures) > 1 {
			panic("A test run should only have one assertion failure. Something went wrong with the test infrastructure.")
		}
		result.Failure = assertModule.Failures[0]
	}

	return &result, nil
}

// Tests returns all tests defined in the config
func (c *Config) Tests() []*Test {
	return c.tests
}

func skyPrint(t *starlark.Thread, msg string) {
	fmt.Fprintf(os.Stderr, "[%v] %s\n", t.CallFrame(1).Pos, msg)
}