go-omni/config/references.go

package config

import (
	"errors"
	"fmt"
	"strings"
	"unicode/utf8"
)

// WithLookupEnv supplies an explicit-reference lookup for tests or an embedding
// host. Ordinary projects need no loader options. The function must support
// concurrent loads; each load caches each referenced name independently.
func WithLookupEnv(lookup func(string) (string, bool)) LoaderOption {
	return func(l *Loader) error {
		if l == nil || l.sealed || lookup == nil {
			return errors.New("config: expected a non-nil environment lookup on a new loader")
		}
		l.lookup = lookup
		return nil
	}
}

// resolvedText is textual input bound according to the target scalar type. It
// never becomes an object, collection or null by interpreting environment data.
type resolvedText string

const maxExpandedBytes = 4 << 20

type environmentValue struct {
	value   string
	present bool
}
type referenceResolver struct {
	lookup                   func(string) (string, bool)
	cache                    map[string]environmentValue
	doc                      Document
	nodes, bytes, cacheBytes int
}

func resolveDocument(doc Document, lookup func(string) (string, bool)) (Document, error) {
	r := referenceResolver{lookup: lookup, cache: map[string]environmentValue{}, doc: Document{Positions: doc.Positions, environments: map[string][]string{}}}
	v, err := r.walk(doc.Value, "", 0)
	if err != nil {
		return Document{}, err
	}
	r.doc.Value = v
	return r.doc, nil
}
func (r *referenceResolver) walk(v any, path string, depth int) (any, error) {
	r.nodes++
	if depth > maxConfigDepth || r.nodes > maxConfigNodes {
		return nil, r.doc.diagnostic(path, "configuration nesting or node limit exceeded", nil)
	}
	switch x := v.(type) {
	case string:
		return r.text(x, path, false)
	case Text:
		return r.text(string(x), path, true)
	case map[string]any:
		out := make(map[string]any, len(x))
		for _, key := range sortedKeys(x) {
			value, err := r.walk(x[key], joinField(path, key), depth+1)
			if err != nil {
				return nil, err
			}
			out[key] = value
		}
		return out, nil
	case []any:
		out := make([]any, len(x))
		for i, value := range x {
			var err error
			out[i], err = r.walk(value, fmt.Sprintf("%s[%d]", path, i), depth+1)
			if err != nil {
				return nil, err
			}
		}
		return out, nil
	default:
		return v, nil
	}
}
func (r *referenceResolver) fail(path, name, reason string) error {
	d := r.doc.diagnostic(path, reason, nil)
	d.Env = name
	return d
}
func validEnvironmentName(name string) bool {
	if len(name) == 0 {
		return false
	}
	for i := 0; i < len(name); i++ {
		c := name[i]
		if c == '_' || c >= 'A' && c <= 'Z' || c >= 'a' && c <= 'z' || i > 0 && c >= '0' && c <= '9' {
			continue
		}
		return false
	}
	return true
}
func (r *referenceResolver) lookupValue(name, path string) (environmentValue, error) {
	if v, ok := r.cache[name]; ok {
		return v, nil
	}
	value, present := r.lookup(name)
	if len(value) > maxConfigBytes || r.cacheBytes+len(value) > maxExpandedBytes {
		return environmentValue{}, r.fail(path, name, "environment value exceeds configuration size limit")
	}
	if !utf8.ValidString(value) {
		return environmentValue{}, r.fail(path, name, "environment value is not valid UTF-8")
	}

	r.cacheBytes += len(value)
	v := environmentValue{value, present}
	r.cache[name] = v
	return v, nil
}
func (r *referenceResolver) text(raw, path string, xmlText bool) (any, error) {
	var out strings.Builder
	var names []string
	seen := map[string]bool{}
	if !strings.Contains(raw, "$") {
		if len(raw) > maxConfigBytes || r.bytes+len(raw) > maxExpandedBytes {
			return nil, r.fail(path, "", "expanded configuration exceeds size limit")
		}
		r.bytes += len(raw)
		if xmlText {
			return Text(raw), nil
		}
		return raw, nil
	}
	appendText := func(s string) error {
		if out.Len()+len(s) > maxConfigBytes || r.bytes+len(s) > maxExpandedBytes {
			return r.fail(path, strings.Join(names, ", "), "expanded configuration exceeds size limit")
		}
		out.WriteString(s)
		r.bytes += len(s)
		return nil
	}
	// Scan only the file's decoded scalar. Inserted environment values are never
	// scanned again, including dollars, quotes, separators or newline characters.
	for i := 0; i < len(raw); {
		j := strings.IndexByte(raw[i:], '$')
		if j < 0 {
			if err := appendText(raw[i:]); err != nil {
				return nil, err
			}
			break
		}
		j += i
		if err := appendText(raw[i:j]); err != nil {
			return nil, err
		}
		switch {
		case j+1 < len(raw) && raw[j+1] == '$':
			if err := appendText("$"); err != nil {
				return nil, err
			}
			i = j + 2
		case j+1 < len(raw) && raw[j+1] == '{':
			end := strings.IndexByte(raw[j+2:], '}')
			if end < 0 {
				return nil, r.fail(path, "", "unclosed environment reference; use ${NAME}, ${NAME:-fallback} or $${...} for a literal")
			}
			end += j + 2
			expression := raw[j+2 : end]
			name, fallback, hasDefault := strings.Cut(expression, ":-")
			if !validEnvironmentName(name) || strings.Contains(fallback, "${") {
				return nil, r.fail(path, "", "invalid environment reference; use ${NAME} or ${NAME:-fallback}; nested references are not supported")
			}
			r.nodes++
			if r.nodes > maxConfigNodes {
				return nil, r.fail(path, name, "environment reference count exceeds configuration limit")
			}
			if !seen[name] {
				names = append(names, name)
				seen[name] = true
			}

			value, err := r.lookupValue(name, path)
			if err != nil {
				return nil, err
			}
			selected := value.value
			if !value.present || selected == "" {
				if !hasDefault {
					return nil, r.fail(path, name, "required environment variable is unset or empty")
				}
				selected = fallback
			}
			if err := appendText(selected); err != nil {
				return nil, err
			}
			i = end + 1
		default:
			if err := appendText("$"); err != nil {
				return nil, err
			}
			i = j + 1
		}
	}
	if len(names) > 0 {
		r.doc.environments[path] = names
		return resolvedText(out.String()), nil
	}
	if xmlText {
		return Text(out.String()), nil
	}
	return out.String(), nil
}