bouncer/internal/config/bootstrap.go

package config

import (
	"os"
	"path/filepath"
	"strings"
	"time"

	"forge.cadoles.com/cadoles/bouncer/internal/store"
	"github.com/pkg/errors"
	"gopkg.in/yaml.v3"
)

type BootstrapConfig struct {
	Proxies              map[store.ProxyName]BootstrapProxyConfig `yaml:"proxies"`
	Dir                  InterpolatedString                       `yaml:"dir"`
	LockTimeout          InterpolatedDuration                     `yaml:"lockTimeout"`
	MaxConnectionRetries InterpolatedInt                          `yaml:"maxRetries"`
}

func (c *BootstrapConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
	src := struct {
		Proxies map[store.ProxyName]BootstrapProxyConfig `yaml:"proxies"`
		Dir     InterpolatedString                       `yaml:"dir"`
	}{
		Proxies: make(map[store.ProxyName]BootstrapProxyConfig),
		Dir:     "",
	}

	if err := unmarshal(&src); err != nil {
		return errors.WithStack(err)
	}

	c.Proxies = src.Proxies
	c.Dir = src.Dir

	if src.Dir != "" {
		proxies, err := loadBootstrapDir(string(src.Dir))
		if err != nil {
			return errors.Wrapf(err, "could not load bootstrap dir '%s'", src.Dir)
		}

		c.Proxies = overrideProxies(c.Proxies, proxies)
	}

	return nil
}

type BootstrapProxyConfig struct {
	Enabled  InterpolatedBool                         `yaml:"enabled"`
	Weight   InterpolatedInt                          `yaml:"weight"`
	To       InterpolatedString                       `yaml:"to"`
	From     InterpolatedStringSlice                  `yaml:"from"`
	Layers   map[store.LayerName]BootstrapLayerConfig `yaml:"layers"`
	Recreate InterpolatedBool                         `yaml:"recreate"`
}

type BootstrapLayerConfig struct {
	Enabled InterpolatedBool   `yaml:"enabled"`
	Type    InterpolatedString `yaml:"type"`
	Weight  InterpolatedInt    `yaml:"weight"`
	Options InterpolatedMap    `yaml:"options"`
}

func NewDefaultBootstrapConfig() BootstrapConfig {
	return BootstrapConfig{
		Dir:                  "",
		LockTimeout:          *NewInterpolatedDuration(30 * time.Second),
		MaxConnectionRetries: 10,
	}
}

func loadBootstrapDir(dir string) (map[store.ProxyName]BootstrapProxyConfig, error) {
	pattern := filepath.Join(dir, "*.yml")

	files, err := filepath.Glob(pattern)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	proxies := make(map[store.ProxyName]BootstrapProxyConfig)
	for _, f := range files {
		proxy, err := loadBootstrappedProxyConfig(f)
		if err != nil {
			return nil, errors.Wrapf(err, "could not load proxy bootstrap file '%s'", f)
		}

		name := store.ProxyName(strings.TrimSuffix(filepath.Base(f), filepath.Ext(f)))
		proxies[name] = *proxy
	}

	return proxies, nil
}

func loadBootstrappedProxyConfig(filename string) (*BootstrapProxyConfig, error) {
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, errors.Wrapf(err, "could not read file '%s'", filename)
	}

	proxy := BootstrapProxyConfig{}

	if err := yaml.Unmarshal(data, &proxy); err != nil {
		return nil, errors.Wrapf(err, "could not unmarshal proxy")
	}

	return &proxy, nil
}

func overrideProxies(base map[store.ProxyName]BootstrapProxyConfig, proxies map[store.ProxyName]BootstrapProxyConfig) map[store.ProxyName]BootstrapProxyConfig {
	for name, proxy := range proxies {
		base[name] = proxy
	}

	return base
}