1// Copyright 2016 The OPA Authors.  All rights reserved.
2// Use of this source code is governed by an Apache2
3// license that can be found in the LICENSE file.
4
5package topdown
6
7import (
8	"fmt"
9
10	"github.com/open-policy-agent/opa/ast"
11	"github.com/open-policy-agent/opa/topdown/builtins"
12)
13
14type (
15	// FunctionalBuiltin1 defines an interface for simple functional built-ins.
16	//
17	// Implement this interface if your built-in function takes one input and
18	// produces one output.
19	//
20	// If an error occurs, the functional built-in should return a descriptive
21	// message. The message should not be prefixed with the built-in name as the
22	// framework takes care of this.
23	FunctionalBuiltin1 func(op1 ast.Value) (output ast.Value, err error)
24
25	// FunctionalBuiltin2 defines an interface for simple functional built-ins.
26	//
27	// Implement this interface if your built-in function takes two inputs and
28	// produces one output.
29	//
30	// If an error occurs, the functional built-in should return a descriptive
31	// message. The message should not be prefixed with the built-in name as the
32	// framework takes care of this.
33	FunctionalBuiltin2 func(op1, op2 ast.Value) (output ast.Value, err error)
34
35	// FunctionalBuiltin3 defines an interface for simple functional built-ins.
36	//
37	// Implement this interface if your built-in function takes three inputs and
38	// produces one output.
39	//
40	// If an error occurs, the functional built-in should return a descriptive
41	// message. The message should not be prefixed with the built-in name as the
42	// framework takes care of this.
43	FunctionalBuiltin3 func(op1, op2, op3 ast.Value) (output ast.Value, err error)
44
45	// BuiltinContext contains context from the evaluator that may be used by
46	// built-in functions.
47	BuiltinContext struct {
48		Cache    builtins.Cache
49		Location *ast.Location
50		Tracer   Tracer
51		QueryID  uint64
52		ParentID uint64
53	}
54
55	// BuiltinFunc defines an interface for implementing built-in functions.
56	// The built-in function is called with the plugged operands from the call
57	// (including the output operands.) The implementation should evaluate the
58	// operands and invoke the iteraror for each successful/defined output
59	// value.
60	BuiltinFunc func(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error
61)
62
63// RegisterBuiltinFunc adds a new built-in function to the evaluation engine.
64func RegisterBuiltinFunc(name string, f BuiltinFunc) {
65	builtinFunctions[name] = f
66}
67
68// RegisterFunctionalBuiltin1 adds a new built-in function to the evaluation
69// engine.
70func RegisterFunctionalBuiltin1(name string, fun FunctionalBuiltin1) {
71	builtinFunctions[name] = functionalWrapper1(name, fun)
72}
73
74// RegisterFunctionalBuiltin2 adds a new built-in function to the evaluation
75// engine.
76func RegisterFunctionalBuiltin2(name string, fun FunctionalBuiltin2) {
77	builtinFunctions[name] = functionalWrapper2(name, fun)
78}
79
80// RegisterFunctionalBuiltin3 adds a new built-in function to the evaluation
81// engine.
82func RegisterFunctionalBuiltin3(name string, fun FunctionalBuiltin3) {
83	builtinFunctions[name] = functionalWrapper3(name, fun)
84}
85
86// BuiltinEmpty is used to signal that the built-in function evaluated, but the
87// result is undefined so evaluation should not continue.
88type BuiltinEmpty struct{}
89
90func (BuiltinEmpty) Error() string {
91	return "<empty>"
92}
93
94var builtinFunctions = map[string]BuiltinFunc{}
95
96func functionalWrapper1(name string, fn FunctionalBuiltin1) BuiltinFunc {
97	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
98		result, err := fn(args[0].Value)
99		if err == nil {
100			return iter(ast.NewTerm(result))
101		}
102		if _, empty := err.(BuiltinEmpty); empty {
103			return nil
104		}
105		return handleBuiltinErr(name, bctx.Location, err)
106	}
107}
108
109func functionalWrapper2(name string, fn FunctionalBuiltin2) BuiltinFunc {
110	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
111		result, err := fn(args[0].Value, args[1].Value)
112		if err == nil {
113			return iter(ast.NewTerm(result))
114		}
115		if _, empty := err.(BuiltinEmpty); empty {
116			return nil
117		}
118		return handleBuiltinErr(name, bctx.Location, err)
119	}
120}
121
122func functionalWrapper3(name string, fn FunctionalBuiltin3) BuiltinFunc {
123	return func(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
124		result, err := fn(args[0].Value, args[1].Value, args[2].Value)
125		if err == nil {
126			return iter(ast.NewTerm(result))
127		}
128		if _, empty := err.(BuiltinEmpty); empty {
129			return nil
130		}
131		return handleBuiltinErr(name, bctx.Location, err)
132	}
133}
134
135func handleBuiltinErr(name string, loc *ast.Location, err error) error {
136	switch err := err.(type) {
137	case BuiltinEmpty:
138		return nil
139	case builtins.ErrOperand:
140		return &Error{
141			Code:     TypeErr,
142			Message:  fmt.Sprintf("%v: %v", string(name), err.Error()),
143			Location: loc,
144		}
145	default:
146		return &Error{
147			Code:     InternalErr,
148			Message:  fmt.Sprintf("%v: %v", string(name), err.Error()),
149			Location: loc,
150		}
151	}
152}
153