distributor/internal/config/secrets.go

package config

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
)

var secretNamePattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`)

type EnvLookup func(string) (string, bool)

type Environment struct {
	values map[string]string
	lookup EnvLookup
}

type SecretConflict struct {
	Name string
}

type SecretLoadResult struct {
	Environment Environment
	Conflicts   []SecretConflict
}

type ResolvedCredentials struct {
	AccessKeyID     string
	SecretAccessKey string
}

func ProcessEnvironment() Environment {
	return NewEnvironment(nil, os.LookupEnv)
}

func NewEnvironment(values map[string]string, lookup EnvLookup) Environment {
	copied := make(map[string]string, len(values))
	for key, value := range values {
		copied[key] = value
	}
	if lookup == nil {
		lookup = os.LookupEnv
	}
	return Environment{values: copied, lookup: lookup}
}

func (e Environment) Lookup(name string) (string, bool) {
	if e.lookup != nil {
		if value, ok := e.lookup(name); ok {
			return value, true
		}
	}
	value, ok := e.values[name]
	return value, ok
}

func (e Environment) ResolveCredentials(creds Credentials) (ResolvedCredentials, error) {
	var resolved ResolvedCredentials
	var err error
	if creds.AccessKeyIDEnv != "" {
		resolved.AccessKeyID, err = e.required(creds.AccessKeyIDEnv)
		if err != nil {
			return ResolvedCredentials{}, err
		}
	}
	if creds.SecretAccessKeyEnv != "" {
		resolved.SecretAccessKey, err = e.required(creds.SecretAccessKeyEnv)
		if err != nil {
			return ResolvedCredentials{}, err
		}
	}
	return resolved, nil
}

func (e Environment) required(name string) (string, error) {
	value, ok := e.Lookup(name)
	if !ok {
		return "", fmt.Errorf("credential environment variable %s is not set", name)
	}
	if value == "" {
		return "", fmt.Errorf("credential environment variable %s is empty", name)
	}
	return value, nil
}

func LoadSecretEnvironment(directory string, lookup EnvLookup) (SecretLoadResult, error) {
	if directory == "" {
		return SecretLoadResult{Environment: NewEnvironment(nil, lookup)}, nil
	}
	values, err := loadSecretValues(directory)
	if err != nil {
		return SecretLoadResult{}, err
	}
	var conflicts []SecretConflict
	for name, value := range values {
		if processValue, ok := lookupValue(lookup, name); ok && processValue != value {
			conflicts = append(conflicts, SecretConflict{Name: name})
		}
	}
	sort.Slice(conflicts, func(i, j int) bool {
		return conflicts[i].Name < conflicts[j].Name
	})
	return SecretLoadResult{
		Environment: NewEnvironment(values, lookup),
		Conflicts:   conflicts,
	}, nil
}

func loadSecretValues(directory string) (map[string]string, error) {
	entries, err := os.ReadDir(directory)
	if err != nil {
		return nil, fmt.Errorf("load secrets directory %q: %w", directory, err)
	}
	values := make(map[string]string)
	for _, entry := range entries {
		path := filepath.Join(directory, entry.Name())
		info, err := os.Stat(path)
		if err != nil {
			return nil, fmt.Errorf("inspect secret file %q: %w", entry.Name(), err)
		}
		if info.IsDir() || !info.Mode().IsRegular() {
			continue
		}
		if !secretNamePattern.MatchString(entry.Name()) {
			return nil, fmt.Errorf("secret filename %q is invalid", entry.Name())
		}
		data, err := os.ReadFile(path)
		if err != nil {
			return nil, fmt.Errorf("read secret file %q: %w", entry.Name(), err)
		}
		values[entry.Name()] = trimOneTrailingLineEnding(string(data))
	}
	return values, nil
}

func trimOneTrailingLineEnding(value string) string {
	if strings.HasSuffix(value, "\r\n") {
		return strings.TrimSuffix(value, "\r\n")
	}
	if strings.HasSuffix(value, "\n") {
		return strings.TrimSuffix(value, "\n")
	}
	return value
}

func lookupValue(lookup EnvLookup, name string) (string, bool) {
	if lookup == nil {
		lookup = os.LookupEnv
	}
	return lookup(name)
}