proxy-pool/internal/adapters/providerapi/template_validation.go

package providerapi

import (
	"fmt"
	"reflect"
	"text/template"
	"text/template/parse"
)

const (
	maxTemplateASTDepth   = 32
	maxTemplateASTNodes   = 4096
	maxTemplateRangeDepth = 1
)

func validateTemplateComplexity(parsed *template.Template) error {
	nodes := 0
	for _, definition := range parsed.Templates() {
		if definition.Tree == nil || definition.Tree.Root == nil {
			continue
		}
		if err := walkTemplateNode(definition.Tree.Root, 0, 0, &nodes); err != nil {
			return err
		}
	}
	return nil
}

func walkTemplateNode(node parse.Node, depth, rangeDepth int, nodes *int) error {
	if node == nil {
		return nil
	}
	value := reflect.ValueOf(node)
	if value.Kind() == reflect.Pointer && value.IsNil() {
		return nil
	}
	*nodes++
	if *nodes > maxTemplateASTNodes {
		return fmt.Errorf("%w: node count exceeds %d", ErrTemplateTooComplex, maxTemplateASTNodes)
	}
	if depth > maxTemplateASTDepth {
		return fmt.Errorf("%w: nesting depth exceeds %d", ErrTemplateTooComplex, maxTemplateASTDepth)
	}

	walk := func(child parse.Node, childDepth, childRangeDepth int) error {
		return walkTemplateNode(child, childDepth, childRangeDepth, nodes)
	}
	switch current := node.(type) {
	case *parse.ListNode:
		for _, child := range current.Nodes {
			if err := walk(child, depth, rangeDepth); err != nil {
				return err
			}
		}
	case *parse.ActionNode:
		return walk(current.Pipe, depth, rangeDepth)
	case *parse.CommandNode:
		for _, argument := range current.Args {
			if err := walk(argument, depth, rangeDepth); err != nil {
				return err
			}
		}
	case *parse.PipeNode:
		for _, declaration := range current.Decl {
			if err := walk(declaration, depth, rangeDepth); err != nil {
				return err
			}
		}
		for _, command := range current.Cmds {
			if err := walk(command, depth, rangeDepth); err != nil {
				return err
			}
		}
	case *parse.ChainNode:
		return walk(current.Node, depth, rangeDepth)
	case *parse.TemplateNode:
		return walk(current.Pipe, depth, rangeDepth)
	case *parse.IfNode:
		if err := walk(current.Pipe, depth, rangeDepth); err != nil {
			return err
		}
		if err := walk(current.List, depth+1, rangeDepth); err != nil {
			return err
		}
		return walk(current.ElseList, depth+1, rangeDepth)
	case *parse.RangeNode:
		if rangeDepth >= maxTemplateRangeDepth {
			return fmt.Errorf("%w: nested range exceeds depth %d", ErrTemplateTooComplex, maxTemplateRangeDepth)
		}
		if err := walk(current.Pipe, depth, rangeDepth); err != nil {
			return err
		}
		if err := walk(current.List, depth+1, rangeDepth+1); err != nil {
			return err
		}
		return walk(current.ElseList, depth+1, rangeDepth+1)
	case *parse.WithNode:
		if err := walk(current.Pipe, depth, rangeDepth); err != nil {
			return err
		}
		if err := walk(current.List, depth+1, rangeDepth); err != nil {
			return err
		}
		return walk(current.ElseList, depth+1, rangeDepth)
	}
	return nil
}

func rejectRecursiveTemplates(parsed *template.Template) error {
	graph := make(map[string][]string)
	for _, definition := range parsed.Templates() {
		if definition.Tree == nil || definition.Tree.Root == nil {
			continue
		}
		var calls []string
		collectTemplateCalls(definition.Tree.Root, &calls)
		graph[definition.Name()] = calls
	}

	const (
		visiting = 1
		visited  = 2
	)
	state := make(map[string]int, len(graph))
	var visit func(string) bool
	visit = func(name string) bool {
		switch state[name] {
		case visiting:
			return true
		case visited:
			return false
		}
		state[name] = visiting
		for _, called := range graph[name] {
			if visit(called) {
				return true
			}
		}
		state[name] = visited
		return false
	}
	for name := range graph {
		if visit(name) {
			return ErrRecursiveTemplate
		}
	}
	return nil
}

func collectTemplateCalls(node parse.Node, calls *[]string) {
	if node == nil {
		return
	}
	switch current := node.(type) {
	case *parse.ListNode:
		if current == nil {
			return
		}
		for _, child := range current.Nodes {
			collectTemplateCalls(child, calls)
		}
	case *parse.TemplateNode:
		if current == nil {
			return
		}
		*calls = append(*calls, current.Name)
	case *parse.IfNode:
		if current == nil {
			return
		}
		collectTemplateCalls(current.List, calls)
		collectTemplateCalls(current.ElseList, calls)
	case *parse.RangeNode:
		if current == nil {
			return
		}
		collectTemplateCalls(current.List, calls)
		collectTemplateCalls(current.ElseList, calls)
	case *parse.WithNode:
		if current == nil {
			return
		}
		collectTemplateCalls(current.List, calls)
		collectTemplateCalls(current.ElseList, calls)
	}
}