kra-new/internal/config/runtime.go

package config

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/go-viper/mapstructure/v2"
	"github.com/spf13/viper"
)

var ErrConfigPathRequired = errors.New("configuration path is required")

// Store owns an immutable configuration snapshot and its Viper-backed file
// watcher.
// Every accessor returns a deep copy, so readers can safely use it without
// coordinating with a reload in another goroutine.
type Store struct {
	snapshot  atomic.Pointer[Config]
	path      string
	mu        sync.RWMutex
	listeners map[uint64]func(*Config)
	nextID    uint64
	watcher   *fsnotify.Watcher
	stop      chan struct{}
	done      chan struct{}
}

// Load reads a YAML/JSON/TOML configuration file or a directory containing
// config.yaml. Environment variables with KRA_ prefix override file values;
// nested keys use underscores (for example KRA_ADMIN_JWT_SIGNING_KEY).
func Load(path string) (*Config, error) {
	resolved, err := resolvePath(path)
	if err != nil {
		return nil, err
	}
	v := newViper(resolved)
	if err = v.ReadInConfig(); err != nil {
		return nil, fmt.Errorf("read configuration %q: %w", resolved, err)
	}
	config, err := decode(v)
	if err != nil {
		return nil, fmt.Errorf("decode configuration %q: %w", resolved, err)
	}
	setConfigPath(config, resolved)
	return config, nil
}

// NewStore creates a store from an already decoded snapshot. The input is
// copied immediately and may be reused or mutated by the caller afterwards.
func NewStore(config *Config) *Store {
	store := &Store{listeners: make(map[uint64]func(*Config))}
	if config != nil && config.Admin != nil {
		store.path = config.Admin.ConfigPath
	}
	store.Replace(config)
	return store
}

// LoadStore loads the initial snapshot and starts the single process-wide
// Viper-backed watcher.
func LoadStore(path string) (*Store, error) {
	config, err := Load(path)
	if err != nil {
		return nil, err
	}
	store := NewStore(config)
	if err = store.Watch(); err != nil {
		store.Close()
		return nil, err
	}
	return store, nil
}

func (r *Store) Snapshot() *Config {
	if r == nil {
		return nil
	}
	return cloneConfig(r.snapshot.Load())
}

func (r *Store) Server() *Server {
	config := r.Snapshot()
	if config == nil {
		return nil
	}
	return config.Server
}

func (r *Store) Data() *Data {
	config := r.Snapshot()
	if config == nil {
		return nil
	}
	return config.Data
}

func (r *Store) Admin() *Admin {
	config := r.Snapshot()
	if config == nil {
		return nil
	}
	return config.Admin
}

func (r *Store) Values() (*Data, *Admin) {
	config := r.Snapshot()
	if config == nil {
		return nil, nil
	}
	return config.Data, config.Admin
}

func (r *Store) UpdateDatabase(database *Database) {
	if r == nil {
		return
	}
	config := r.Snapshot()
	if config == nil {
		config = &Config{}
	}
	if config.Data == nil {
		config.Data = &Data{}
	}
	config.Data.Database = cloneDatabase(database)
	r.replace(config, false)
}

func (r *Store) ConfigPath() string {
	if r == nil {
		return ""
	}
	r.mu.RLock()
	path := r.path
	r.mu.RUnlock()
	return path
}

func (r *Store) Replace(config *Config) {
	r.replace(config, true)
}

func (r *Store) replace(config *Config, notify bool) {
	if r == nil {
		return
	}
	next := cloneConfig(config)
	if next == nil {
		next = &Config{}
	}
	if path := r.ConfigPath(); path != "" {
		setConfigPath(next, path)
	}
	r.snapshot.Store(next)
	if !notify {
		return
	}
	r.mu.RLock()
	listeners := make([]func(*Config), 0, len(r.listeners))
	for _, listener := range r.listeners {
		listeners = append(listeners, listener)
	}
	r.mu.RUnlock()
	for _, listener := range listeners {
		listener(cloneConfig(next))
	}
}

// Subscribe registers a callback after each successful replacement. The
// callback is invoked outside the runtime lock, and the returned function is
// idempotent so cleanup can safely be deferred by multiple owners.
func (r *Store) Subscribe(listener func(*Config)) func() {
	if r == nil || listener == nil {
		return func() {}
	}
	r.mu.Lock()
	r.nextID++
	id := r.nextID
	r.listeners[id] = listener
	r.mu.Unlock()
	var once sync.Once
	return func() {
		once.Do(func() {
			r.mu.Lock()
			delete(r.listeners, id)
			r.mu.Unlock()
		})
	}
}

// Watch starts watching the source directory. Repeated calls are harmless.
func (r *Store) Watch() error {
	if r == nil {
		return ErrConfigPathRequired
	}
	path := r.ConfigPath()
	if path == "" {
		return ErrConfigPathRequired
	}
	r.mu.Lock()
	if r.watcher != nil {
		r.mu.Unlock()
		return nil
	}
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		r.mu.Unlock()
		return fmt.Errorf("create configuration watcher: %w", err)
	}
	if err = watcher.Add(filepath.Dir(path)); err != nil {
		_ = watcher.Close()
		r.mu.Unlock()
		return fmt.Errorf("watch configuration directory: %w", err)
	}
	r.watcher = watcher
	r.stop = make(chan struct{})
	r.done = make(chan struct{})
	stop, done := r.stop, r.done
	r.mu.Unlock()
	go r.watchLoop(watcher, stop, done, path)
	return nil
}

func (r *Store) watchLoop(watcher *fsnotify.Watcher, stop <-chan struct{}, done chan<- struct{}, path string) {
	defer close(done)
	const debounce = 100 * time.Millisecond
	var timer *time.Timer
	var timerC <-chan time.Time
	events := watcher.Events
	errors := watcher.Errors
	for {
		select {
		case event, ok := <-events:
			if !ok {
				events = nil
				if errors == nil {
					return
				}
				continue
			}
			if filepath.Clean(event.Name) != path || event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Rename) == 0 {
				continue
			}
			if timer == nil {
				timer = time.NewTimer(debounce)
			} else {
				if !timer.Stop() {
					select {
					case <-timer.C:
					default:
					}
				}
				timer.Reset(debounce)
			}
			timerC = timer.C
		case <-timerC:
			timerC = nil
			next, err := Load(path)
			if err != nil {
				// Keep the last valid snapshot. A partially-written file must not
				// take down a running process or publish invalid state.
				continue
			}
			current := r.Snapshot()
			if next.Admin == nil {
				next.Admin = &Admin{}
			}
			// Storage and Email move to database-backed integration settings after
			// first initialization. External config file edits must not erase the
			// active values when those sections are absent from config.yaml.
			if current != nil && current.Admin != nil {
				if next.Admin.Storage == nil {
					next.Admin.Storage = cloneStorage(current.Admin.Storage)
				}
				if next.Admin.Email == nil && current.Admin.Email != nil {
					email := *current.Admin.Email
					next.Admin.Email = &email
				}
			}
			r.Replace(next)
		case _, ok := <-errors:
			// fsnotify errors are intentionally non-fatal; the watcher remains
			// useful for subsequent events and Close always terminates it.
			if !ok {
				errors = nil
				if events == nil {
					return
				}
			}
		case <-stop:
			if timer != nil {
				timer.Stop()
			}
			return
		}
	}
}

func (r *Store) Close() {
	if r == nil {
		return
	}
	r.mu.Lock()
	watcher, stop, done := r.watcher, r.stop, r.done
	r.watcher = nil
	r.stop = nil
	r.done = nil
	r.mu.Unlock()
	if watcher == nil {
		return
	}
	close(stop)
	_ = watcher.Close()
	<-done
}

func resolvePath(path string) (string, error) {
	path = strings.TrimSpace(path)
	if path == "" {
		return "", ErrConfigPathRequired
	}
	absolute, err := filepath.Abs(path)
	if err != nil {
		return "", fmt.Errorf("resolve configuration path: %w", err)
	}
	if info, statErr := os.Stat(absolute); statErr == nil && info.IsDir() {
		absolute = filepath.Join(absolute, "config.yaml")
	}
	return filepath.Clean(absolute), nil
}

func newViper(path string) *viper.Viper {
	v := viper.New()
	v.SetConfigFile(path)
	v.SetConfigType(strings.TrimPrefix(filepath.Ext(path), "."))
	v.SetEnvPrefix("KRA")
	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_", "-", "_"))
	v.AutomaticEnv()
	return v
}

func decode(v *viper.Viper) (*Config, error) {
	var config Config
	hook := mapstructure.ComposeDecodeHookFunc(
		mapstructure.StringToTimeDurationHookFunc(),
		mapstructure.StringToSliceHookFunc(","),
	)
	if err := v.Unmarshal(&config, viper.DecodeHook(hook)); err != nil {
		return nil, err
	}
	return &config, nil
}

func setConfigPath(config *Config, path string) {
	if config != nil && config.Admin != nil {
		config.Admin.ConfigPath = path
	}
}

func cloneConfig(value *Config) *Config {
	if value == nil {
		return nil
	}
	// Viper's mapstructure output contains only value fields and pointers;
	// round-tripping through a temporary map is unnecessarily lossy for
	// durations. Explicit copy via YAML gives us a robust deep copy while
	// preserving time.Duration values through the text hook.
	copy := *value
	if value.Server != nil {
		server := *value.Server
		copy.Server = &server
		if value.Server.HTTP != nil {
			http := *value.Server.HTTP
			copy.Server.HTTP = &http
		}
	}
	if value.Data != nil {
		data := *value.Data
		copy.Data = &data
		data.Database = cloneDatabase(value.Data.Database)
		data.Redis = cloneRedis(value.Data.Redis)
		data.Mongo = cloneMongo(value.Data.Mongo)
		data.DatabaseList = make([]*Database, len(value.Data.DatabaseList))
		for i, item := range value.Data.DatabaseList {
			data.DatabaseList[i] = cloneDatabase(item)
		}
		data.RedisList = make([]*Redis, len(value.Data.RedisList))
		for i, item := range value.Data.RedisList {
			data.RedisList[i] = cloneRedis(item)
		}
	}
	if value.Admin != nil {
		copy.Admin = cloneAdmin(value.Admin)
	}
	return &copy
}

func cloneDatabase(value *Database) *Database {
	if value == nil {
		return nil
	}
	copy := *value
	return &copy
}
func cloneRedis(value *Redis) *Redis {
	if value == nil {
		return nil
	}
	copy := *value
	copy.ClusterAddrs = append([]string(nil), value.ClusterAddrs...)
	return &copy
}
func cloneMongo(value *Mongo) *Mongo {
	if value == nil {
		return nil
	}
	copy := *value
	copy.Hosts = make([]*MongoHost, len(value.Hosts))
	for i, item := range value.Hosts {
		if item != nil {
			host := *item
			copy.Hosts[i] = &host
		}
	}
	return &copy
}

func cloneAdmin(value *Admin) *Admin {
	if value == nil {
		return nil
	}
	copy := *value
	if value.JWT != nil {
		item := *value.JWT
		copy.JWT = &item
	}
	if value.Captcha != nil {
		item := *value.Captcha
		copy.Captcha = &item
	}
	if value.Local != nil {
		item := *value.Local
		copy.Local = &item
	}
	if value.Email != nil {
		item := *value.Email
		copy.Email = &item
	}
	if value.Media != nil {
		item := *value.Media
		copy.Media = &item
	}
	if value.System != nil {
		item := *value.System
		copy.System = &item
	}
	if value.Zap != nil {
		item := *value.Zap
		item.FileOnlyModules = append([]string(nil), value.Zap.FileOnlyModules...)
		copy.Zap = &item
	}
	if value.App != nil {
		item := *value.App
		copy.App = &item
	}
	if value.CORS != nil {
		item := *value.CORS
		item.Whitelist = make([]*CORSRule, len(value.CORS.Whitelist))
		for i, rule := range value.CORS.Whitelist {
			if rule != nil {
				next := *rule
				item.Whitelist[i] = &next
			}
		}
		copy.CORS = &item
	}
	if value.Storage != nil {
		copy.Storage = cloneStorage(value.Storage)
	}
	copy.DiskList = make([]*Disk, len(value.DiskList))
	for i, item := range value.DiskList {
		if item != nil {
			next := *item
			copy.DiskList[i] = &next
		}
	}
	return &copy
}

func cloneStorage(value *Storage) *Storage {
	if value == nil {
		return nil
	}
	copy := *value
	if value.Qiniu != nil {
		item := *value.Qiniu
		copy.Qiniu = &item
	}
	copy.AliyunOSS = cloneObjectStore(value.AliyunOSS)
	copy.HuaweiOBS = cloneObjectStore(value.HuaweiOBS)
	copy.TencentCOS = cloneObjectStore(value.TencentCOS)
	copy.AWSS3 = cloneObjectStore(value.AWSS3)
	copy.CloudflareR2 = cloneObjectStore(value.CloudflareR2)
	copy.Minio = cloneObjectStore(value.Minio)
	return &copy
}
func cloneObjectStore(value *ObjectStore) *ObjectStore {
	if value == nil {
		return nil
	}
	copy := *value
	return &copy
}