Merge pull request #84 from daveshanley/custom-functions

Created new custom functions code.

cmd/lint.go

@@ -35,6 +35,8 @@ func GetLintCommand() *cobra.Command {
 			rulesetFlag, _ := cmd.Flags().GetString("ruleset")
 			silent, _ := cmd.Flags().GetBool("silent")
 
+			functionsFlag, _ := cmd.Flags().GetString("functions")
+
 			if !silent {
 				PrintBanner()
 			}
@@ -65,6 +67,8 @@ func GetLintCommand() *cobra.Command {
 			// default is recommended rules, based on spectral (for now anyway)
 			selectedRS := defaultRuleSets.GenerateOpenAPIRecommendedRuleSet()
 
+			customFunctions, _ := LoadCustomFunctions(functionsFlag)
+
 			// if ruleset has been supplied, lets make sure it exists, then load it in
 			// and see if it's valid. If so - let's go!
 			if rulesetFlag != "" {
@@ -75,6 +79,7 @@ func GetLintCommand() *cobra.Command {
 					pterm.Println()
 					return rsErr
 				}
+
 				selectedRS, rsErr = BuildRuleSetFromUserSuppliedSet(rsBytes, defaultRuleSets)
 				if rsErr != nil {
 					return rsErr
@@ -84,8 +89,9 @@ func GetLintCommand() *cobra.Command {
 			pterm.Info.Printf("Linting against %d rules: %s\n", len(selectedRS.Rules), selectedRS.DocumentationURI)
 			start := time.Now()
 			result := motor.ApplyRulesToRuleSet(&motor.RuleSetExecution{
-				RuleSet: selectedRS,
-				Spec:    specBytes,
+				RuleSet:         selectedRS,
+				Spec:            specBytes,
+				CustomFunctions: customFunctions,
 			})
 
 			results := result.Results

cmd/root.go

@@ -50,6 +50,7 @@ func init() {
 	cobra.OnInitialize(initConfig)
 	rootCmd.PersistentFlags().BoolP("time", "t", false, "Show how long vacuum took to run")
 	rootCmd.PersistentFlags().StringP("ruleset", "r", "", "Path to a spectral ruleset configuration")
+	rootCmd.PersistentFlags().StringP("functions", "f", "", "Path to custom functions")
 
 	rootCmd.AddCommand(GetLintCommand())
 	rootCmd.AddCommand(GetVacuumReportCommand())

cmd/shared_functions.go

@@ -5,6 +5,8 @@ package cmd
 
 import (
 	"fmt"
+	"github.com/daveshanley/vacuum/model"
+	"github.com/daveshanley/vacuum/plugin"
 	"github.com/daveshanley/vacuum/rulesets"
 	"github.com/pterm/pterm"
 	"github.com/pterm/pterm/putils"
@@ -43,3 +45,19 @@ func PrintBanner() {
 	pterm.Printf("version: %s | compiled: %s\n\n", Version, Date)
 	pterm.Println()
 }
+
+// LoadCustomFunctions will scan for (and load) custom functions defined as vacuum plugins.
+func LoadCustomFunctions(functionsFlag string) (map[string]model.RuleFunction, error) {
+	// check custom functions
+	if functionsFlag != "" {
+		pm, err := plugin.LoadFunctions(functionsFlag)
+		if err != nil {
+			pterm.Error.Printf("Unable to open custom functions: %v\n", err)
+			pterm.Println()
+			return nil, err
+		}
+		pterm.Info.Printf("Loaded %d custom function(s) successfully.\n", pm.LoadedFunctionCount())
+		return pm.GetCustomFunctions(), nil
+	}
+	return nil, nil
+}

functions/functions.go

@@ -7,11 +7,18 @@ import (
 	"github.com/daveshanley/vacuum/functions/core"
 	openapi_functions "github.com/daveshanley/vacuum/functions/openapi"
 	"github.com/daveshanley/vacuum/model"
+	"github.com/daveshanley/vacuum/plugin"
 	"sync"
 )
 
+type customFunction struct {
+	functionHook plugin.FunctionHook
+	schemaHook   plugin.FunctionSchema
+}
+
 type functionsModel struct {
-	functions map[string]model.RuleFunction
+	functions       map[string]model.RuleFunction
+	customFunctions map[string]customFunction
 }
 
 // Functions is used to Query available functions loaded into vacuum
@@ -32,13 +39,17 @@ func MapBuiltinFunctions() Functions {
 
 	coreFunctionGrab.Do(func() {
 		var funcs map[string]model.RuleFunction
+		var customFuncs map[string]customFunction
 
 		if functionsSingleton != nil {
 			funcs = functionsSingleton.functions
+			customFuncs = functionsSingleton.customFunctions
 		} else {
 			funcs = make(map[string]model.RuleFunction)
+			customFuncs = make(map[string]customFunction)
 			functionsSingleton = &functionsModel{
-				functions: funcs,
+				functions:       funcs,
+				customFunctions: customFuncs,
 			}
 		}
 
@@ -90,6 +101,10 @@ func MapBuiltinFunctions() Functions {
 	return functionsSingleton
 }
 
+func (fm functionsModel) RegisterCustomFunction(name string, function plugin.FunctionHook, schema plugin.FunctionSchema) {
+	fm.customFunctions[name] = customFunction{functionHook: function, schemaHook: schema}
+}
+
 func (fm functionsModel) GetAllFunctions() map[string]model.RuleFunction {
 	return fm.functions
 }

motor/rule_applicator.go

@@ -23,14 +23,15 @@ type ruleContext struct {
 	errors           *[]error
 	index            *model.SpecIndex
 	specInfo         *model.SpecInfo
+	customFunctions  map[string]model.RuleFunction
 }
 
 // RuleSetExecution is an instruction set for executing a ruleset. It's a convenience structure to allow the signature
 // of ApplyRules to change, without a huge refactor. The ApplyRules function only returns a single error also.
 type RuleSetExecution struct {
-	RuleSet *rulesets.RuleSet // The RuleSet in which to apply
-	Spec    []byte            // The raw bytes of the OpenAPI specification.
-
+	RuleSet         *rulesets.RuleSet             // The RuleSet in which to apply
+	Spec            []byte                        // The raw bytes of the OpenAPI specification.
+	CustomFunctions map[string]model.RuleFunction // custom functions loaded from plugin.
 }
 
 // RuleSetExecutionResult returns the results of running the ruleset against the supplied spec.
@@ -136,6 +137,7 @@ func ApplyRulesToRuleSet(execution *RuleSetExecution) *RuleSetExecutionResult {
 				errors:           &errors,
 				index:            index,
 				specInfo:         specInfo,
+				customFunctions:  execution.CustomFunctions,
 			}
 			go runRule(ctx)
 		}
@@ -316,6 +318,14 @@ var lock sync.Mutex
 func buildResults(ctx ruleContext, ruleAction model.RuleAction, nodes []*yaml.Node) *[]model.RuleFunctionResult {
 
 	ruleFunction := ctx.builtinFunctions.FindFunction(ruleAction.Function)
+	// not found, check if it's been registered as a custom function
+	if ruleFunction == nil {
+		if ctx.customFunctions != nil {
+			if ctx.customFunctions[ruleAction.Function] != nil {
+				ruleFunction = ctx.customFunctions[ruleAction.Function]
+			}
+		}
+	}
 
 	if ruleFunction != nil {
 

plugin/plugin_loader.go

@@ -0,0 +1,51 @@
+package plugin
+
+import (
+	"github.com/pterm/pterm"
+	"io/ioutil"
+	"path/filepath"
+	"plugin"
+	"strings"
+)
+
+// LoadFunctions will load custom functions found in the supplied path
+func LoadFunctions(path string) (*Manager, error) {
+
+	dirEntries, err := ioutil.ReadDir(path)
+	if err != nil {
+		return nil, err
+	}
+
+	pm := createPluginManager()
+
+	for _, entry := range dirEntries {
+		if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".so") {
+			fPath := filepath.Join(path, entry.Name())
+
+			// found something
+			pterm.Info.Printf("Located custom function plugin: %s\n", fPath)
+
+			// let's try and open it.
+			p, e := plugin.Open(fPath)
+			if e != nil {
+				return nil, e
+			}
+
+			// look up the Boot function and store as a Symbol
+			var bootFunc plugin.Symbol
+			bootFunc, e = p.Lookup("Boot")
+			if err != nil {
+				return nil, err
+			}
+
+			// lets go pedro!
+			if bootFunc != nil {
+				bootFunc.(func(*Manager))(pm)
+			} else {
+				pterm.Error.Printf("Unable to boot plugin")
+			}
+		}
+	}
+
+	return pm, nil
+}

plugin/plugin_loader_test.go

@@ -0,0 +1,29 @@
+package plugin
+
+import (
+	"github.com/stretchr/testify/assert"
+	"runtime"
+	"testing"
+)
+
+func TestLoadFunctions_Nowhere(t *testing.T) {
+	pm, err := LoadFunctions("nowhere")
+	assert.Nil(t, pm)
+	assert.Error(t, err)
+}
+
+func TestLoadFunctions(t *testing.T) {
+	pm, err := LoadFunctions("../model/test_files")
+	assert.NotNil(t, pm)
+	assert.NoError(t, err)
+	assert.Equal(t, 0, pm.LoadedFunctionCount())
+}
+
+func TestLoadFunctions_Sample(t *testing.T) {
+	pm, err := LoadFunctions("sample")
+	if runtime.GOOS != "windows" { // windows does not support this feature, at all.
+		assert.NotNil(t, pm)
+		assert.NoError(t, err)
+		assert.Equal(t, 0, pm.LoadedFunctionCount())
+	}
+}

plugin/plugin_manager.go

@@ -0,0 +1,33 @@
+package plugin
+
+import (
+	"github.com/daveshanley/vacuum/model"
+	"gopkg.in/yaml.v3"
+)
+
+type FunctionSchema func() model.RuleFunctionSchema
+type FunctionHook func(nodes []*yaml.Node, context model.RuleFunctionContext) []model.RuleFunctionResult
+
+type Manager struct {
+	customFunctions map[string]model.RuleFunction
+}
+
+func createPluginManager() *Manager {
+	return &Manager{
+		customFunctions: make(map[string]model.RuleFunction),
+	}
+}
+
+// RegisterFunction allows a custom function to be hooked in
+func (pm *Manager) RegisterFunction(name string, ruleFunction model.RuleFunction) {
+	pm.customFunctions[name] = ruleFunction
+}
+
+// LoadedFunctionCount returns the number of available and ready to use functions.
+func (pm *Manager) LoadedFunctionCount() int {
+	return len(pm.customFunctions)
+}
+
+func (pm *Manager) GetCustomFunctions() map[string]model.RuleFunction {
+	return pm.customFunctions
+}

plugin/plugin_manager_test.go

@@ -0,0 +1,16 @@
+package plugin
+
+import (
+	"github.com/daveshanley/vacuum/functions/core"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestPluginManager_RegisterFunction(t *testing.T) {
+
+	pm := createPluginManager()
+
+	pm.RegisterFunction("defined", core.Defined{})
+	assert.Len(t, pm.GetCustomFunctions(), 1)
+
+}

plugin/sample/boot.go

@@ -0,0 +1,15 @@
+package main
+
+import "github.com/daveshanley/vacuum/plugin"
+
+// Boot is called by the Manager when the module is located.
+// all custom functions should be registered here.
+func Boot(pm *plugin.Manager) {
+
+	sampleA := SampleRuleFunction_A{}
+	sampleB := SampleRuleFunction_B{}
+
+	// register custom functions with vacuum plugin manager.
+	pm.RegisterFunction(sampleA.GetSchema().Name, sampleA)
+	pm.RegisterFunction(sampleB.GetSchema().Name, sampleB)
+}

plugin/sample/sample_a.go

@@ -0,0 +1,32 @@
+package main
+
+import (
+	"github.com/daveshanley/vacuum/model"
+	"gopkg.in/yaml.v3"
+)
+
+// SampleRuleFunction_A is an example custom rule that does nothing.
+type SampleRuleFunction_A struct {
+}
+
+// GetSchema returns a model.RuleFunctionSchema defining the schema of the Defined rule.
+func (s SampleRuleFunction_A) GetSchema() model.RuleFunctionSchema {
+	return model.RuleFunctionSchema{
+		Name: "uselessFunc",
+	}
+}
+
+// RunRule will execute the Sample rule, based on supplied context and a supplied []*yaml.Node slice.
+func (s SampleRuleFunction_A) RunRule(nodes []*yaml.Node, context model.RuleFunctionContext) []model.RuleFunctionResult {
+
+	// return a single result, for a made up linting failure.
+	return []model.RuleFunctionResult{
+		{
+			Message:   "this is a useless function that will always error out.",
+			StartNode: &yaml.Node{Line: 1, Column: 0},
+			EndNode:   &yaml.Node{Line: 2, Column: 0},
+			Path:      "$.i.do.not.exist",
+			Rule:      context.Rule,
+		},
+	}
+}

plugin/sample/sample_b.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	"fmt"
+	"github.com/daveshanley/vacuum/model"
+	"gopkg.in/yaml.v3"
+)
+
+// SampleRuleFunction_B is an example custom rule that checks only a single path exists.
+type SampleRuleFunction_B struct {
+}
+
+// GetSchema returns a model.RuleFunctionSchema defining the schema of the Defined rule.
+func (s SampleRuleFunction_B) GetSchema() model.RuleFunctionSchema {
+	return model.RuleFunctionSchema{
+		Name: "checkSinglePathExists",
+	}
+}
+
+// RunRule will execute the Sample rule, based on supplied context and a supplied []*yaml.Node slice.
+func (s SampleRuleFunction_B) RunRule(nodes []*yaml.Node, context model.RuleFunctionContext) []model.RuleFunctionResult {
+
+	// get the index https://quobix.com/vacuum/api/spec-index/
+	index := context.Index
+
+	// get the paths node from the index.
+	paths := index.GetPathsNode()
+
+	// checks if there are more than two nodes present in the paths node, if so, more than one path is present.
+	if len(paths.Content) > 2 {
+		return []model.RuleFunctionResult{
+			{
+				Message:   fmt.Sprintf("more than a single path exists, there are %v", len(paths.Content)/2),
+				StartNode: paths,
+				EndNode:   paths,
+				Path:      "$.paths",
+				Rule:      context.Rule,
+			},
+		}
+	}
+	return nil
+}