remote.go

package viper

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"reflect"
	"slices"
	"sync"

	"dario.cat/mergo"
)

// SupportedRemoteProviders are universally supported remote providers.
var SupportedRemoteProviders = []string{"etcd", "etcd3", "consul", "firestore", "nats"}

func resetRemote() {
	SupportedRemoteProviders = []string{"etcd", "etcd3", "consul", "firestore", "nats"}
}

type remoteConfigFactory interface {
	Get(rp RemoteProvider) (io.Reader, error)
	Watch(rp RemoteProvider) (io.Reader, error)
	WatchChannel(rp RemoteProvider) (<-chan *RemoteResponse, chan bool)
}

type RemoteResponse struct {
	Value []byte
	Error error
}

// RemoteConfig is optional, see the remote package.
var RemoteConfig remoteConfigFactory

// UnsupportedRemoteProviderError denotes encountering an unsupported remote
// provider. Currently only etcd and Consul are supported.
type UnsupportedRemoteProviderError string

// Error returns the formatted remote provider error.
func (str UnsupportedRemoteProviderError) Error() string {
	return fmt.Sprintf("Unsupported Remote Provider Type %q", string(str))
}

// RemoteConfigError denotes encountering an error while trying to
// pull the configuration from the remote provider.
type RemoteConfigError string

// Error returns the formatted remote provider error.
func (rce RemoteConfigError) Error() string {
	return fmt.Sprintf("Remote Configurations Error: %s", string(rce))
}

type defaultRemoteProvider struct {
	provider      string
	endpoint      string
	endpoints     []string
	path          string
	secretKeyring string
}

func (rp defaultRemoteProvider) Provider() string {
	return rp.provider
}

func (rp defaultRemoteProvider) Endpoint() string {
	return rp.endpoint
}

func (rp defaultRemoteProvider) Endpoints() []string {
	return rp.endpoints
}

func (rp defaultRemoteProvider) Path() string {
	return rp.path
}

func (rp defaultRemoteProvider) SecretKeyring() string {
	return rp.secretKeyring
}

// RemoteProvider stores the configuration necessary
// to connect to a remote key/value store.
// Optional secretKeyring to unencrypt encrypted values
// can be provided.
type RemoteProvider interface {
	Provider() string
	Endpoint() string
	Endpoints() []string
	Path() string
	SecretKeyring() string
}

// AddRemoteProvider adds a remote configuration source.
// Remote Providers are searched in the order they are added.
// provider is a string value: "etcd", "etcd3", "consul", "firestore" or "nats" are currently supported.
// endpoint is the url.  etcd requires http://ip:port, consul requires ip:port, nats requires nats://ip:port
// path is the path in the k/v store to retrieve configuration
// To retrieve a config file called myapp.json from /configs/myapp.json
// you should set path to /configs and set config name (SetConfigName()) to
// "myapp".
func AddRemoteProvider(provider, endpoint, path string) error {
	return v.AddRemoteProvider(provider, endpoint, path)
}

func (v *Viper) AddRemoteProvider(provider, endpoint, path string) error {
	if !slices.Contains(SupportedRemoteProviders, provider) {
		return UnsupportedRemoteProviderError(provider)
	}
	if provider != "" && endpoint != "" {
		v.logger.Info("adding remote provider", "provider", provider, "endpoint", endpoint)

		rp := &defaultRemoteProvider{
			endpoint: endpoint,
			provider: provider,
			path:     path,
		}
		if !v.providerPathExists(rp) {
			v.remoteProviders = append(v.remoteProviders, rp)
		}
	}
	return nil
}

func (v *Viper) AddRemoteProviderCluster(provider string, endpoints []string, path string) error {
	if !slices.Contains(SupportedRemoteProviders, provider) {
		return UnsupportedRemoteProviderError(provider)
	}
	if provider != "" && len(endpoints) != 0 {
		v.logger.Info("adding remote provider", "provider", provider, "endpoints", endpoints)

		rp := &defaultRemoteProvider{
			endpoints: endpoints,
			provider:  provider,
			path:      path,
		}
		if !v.providerPathExists(rp) {
			v.remoteProviders = append(v.remoteProviders, rp)
		}
	}
	return nil
}

// AddSecureRemoteProvider adds a remote configuration source.
// Secure Remote Providers are searched in the order they are added.
// provider is a string value: "etcd", "etcd3", "consul", "firestore" or "nats" are currently supported.
// endpoint is the url.  etcd requires http://ip:port  consul requires ip:port
// secretkeyring is the filepath to your openpgp secret keyring.  e.g. /etc/secrets/myring.gpg
// path is the path in the k/v store to retrieve configuration
// To retrieve a config file called myapp.json from /configs/myapp.json
// you should set path to /configs and set config name (SetConfigName()) to
// "myapp".
// Secure Remote Providers are implemented with github.com/sagikazarmark/crypt.
func AddSecureRemoteProvider(provider, endpoint, path, secretkeyring string) error {
	return v.AddSecureRemoteProvider(provider, endpoint, path, secretkeyring)
}

func (v *Viper) AddSecureRemoteProvider(provider, endpoint, path, secretkeyring string) error {
	if !slices.Contains(SupportedRemoteProviders, provider) {
		return UnsupportedRemoteProviderError(provider)
	}
	if provider != "" && endpoint != "" {
		v.logger.Info("adding remote provider", "provider", provider, "endpoint", endpoint)

		rp := &defaultRemoteProvider{
			endpoint:      endpoint,
			provider:      provider,
			path:          path,
			secretKeyring: secretkeyring,
		}
		if !v.providerPathExists(rp) {
			v.remoteProviders = append(v.remoteProviders, rp)
		}
	}
	return nil
}

func (v *Viper) providerPathExists(p *defaultRemoteProvider) bool {
	for _, y := range v.remoteProviders {
		if reflect.DeepEqual(y, p) {
			return true
		}
	}
	return false
}

// ReadRemoteConfig attempts to get configuration from a remote source
// and read it in the remote configuration registry.
func ReadRemoteConfig() error {
	return v.getKeyValueConfig(false)
}
func ReadRemoteConfigWithMerged(merged bool) error {
	return v.getKeyValueConfig(merged)
}

func (v *Viper) ReadRemoteConfig() error {
	return v.getKeyValueConfig(false)
}
func (v *Viper) ReadRemoteConfigWithMerged(merged bool) error {
	return v.getKeyValueConfig(merged)
}

func WatchRemoteConfig() error { return v.WatchRemoteConfig() }
func (v *Viper) WatchRemoteConfig() error {
	return v.watchKeyValueConfig()
}

func (v *Viper) WatchRemoteConfigOnChannel() error {
	return v.watchKeyValueConfigOnChannel()
}

// WatchRemoteConfigWithChannel WatchRemoteConfigOnChannel 的增强实现,多了channel回调
//
//	@receiver v
//	@param ctx 要取消请传入 context.WithCancel
//	@param receiver
//	@param deepMerge
//	@return error
func (v *Viper) WatchRemoteConfigWithChannel(ctx context.Context, receiver chan *RemoteResponse, deepMerge bool) error {
	return v.watchKeyValueConfigWithChannel(ctx, receiver, deepMerge)
}

// Retrieve the first found remote configuration.
func (v *Viper) getKeyValueConfig(deepMerge bool) error {
	if RemoteConfig == nil {
		return RemoteConfigError("Enable the remote features by doing a blank import of the viper/remote package: '_ github.com/spf13/viper/remote'")
	}

	if len(v.remoteProviders) == 0 {
		return RemoteConfigError("No Remote Providers")
	}

	var found = false
	for _, rp := range v.remoteProviders {
		val, err := v.getRemoteConfig(rp)
		if err != nil {
			v.logger.Error(fmt.Errorf("get remote config: %w", err).Error())

			continue
		}
		found = true
		v.mergeRemoteConfigSnapshot(v.kvstore, val, deepMerge)
	}
	if found {
		return nil
	}
	return RemoteConfigError("No Files Found")
}

func (v *Viper) getRemoteConfig(provider RemoteProvider) (map[string]any, error) {
	reader, err := RemoteConfig.Get(provider)
	if err != nil {
		return nil, err
	}
	val := map[string]any{}
	err = v.unmarshalReader(reader, val)
	return val, err
}

// mergeRemoteConfigSnapshots rebuilds the effective remote configuration from
// all known provider snapshots.
//
// remoteProviders are registered from low priority to high priority, so this
// function intentionally merges snapshots in slice order. Later snapshots
// overwrite earlier snapshots for the same key, while lower-priority snapshots
// still contribute fields that no higher-priority snapshot defines.
func (v *Viper) mergeRemoteConfigSnapshots(snapshots []map[string]any, deepMerge bool) map[string]any {
	merged := make(map[string]any)

	for _, snapshot := range snapshots {
		if snapshot == nil {
			continue
		}
		v.mergeRemoteConfigSnapshot(merged, snapshot, deepMerge)
	}

	return merged
}

func (v *Viper) mergeRemoteConfigSnapshot(dst, snapshot map[string]any, deepMerge bool) {
	if deepMerge {
		// mergo merges nested maps by mutating the destination map in place.
		// Copy the source snapshot first so the effective kvstore never shares
		// nested map or slice values with a provider snapshot. The watch path
		// later replaces individual snapshots and then rebuilds kvstore from
		// scratch, so keeping those two data sets independent avoids stale
		// references between generations of the merged config.
		if err := mergo.Merge(&dst, copyRemoteConfigMap(snapshot), mergo.WithOverride); err != nil {
			v.logger.Error(fmt.Errorf("remote config merge error: %w", err).Error())
			for k, v_ := range snapshot {
				dst[k] = copyRemoteConfigValue(v_)
			}
		}
		return
	}

	for k, v_ := range snapshot {
		dst[k] = copyRemoteConfigValue(v_)
	}
}

func copyRemoteConfigMap(src map[string]any) map[string]any {
	dst := make(map[string]any, len(src))
	for k, v := range src {
		dst[k] = copyRemoteConfigValue(v)
	}
	return dst
}

func copyRemoteConfigValue(src any) any {
	switch v := src.(type) {
	case map[string]any:
		return copyRemoteConfigMap(v)
	case []any:
		dst := make([]any, len(v))
		for i, item := range v {
			dst[i] = copyRemoteConfigValue(item)
		}
		return dst
	default:
		return v
	}
}

// Retrieve the first found remote configuration.
func (v *Viper) watchKeyValueConfigOnChannel() error {
	if len(v.remoteProviders) == 0 {
		return RemoteConfigError("No Remote Providers")
	}

	for _, rp := range v.remoteProviders {
		respc, _ := RemoteConfig.WatchChannel(rp)
		// Todo: Add quit channel
		go func(rc <-chan *RemoteResponse) {
			for {
				b, ok := <-rc
				// 校验异常
				if !ok {
					break
				}
				if b.Error != nil {
					v.logger.Error(fmt.Errorf("viper watchKeyValueConfigWithChannel watch remote config: %w", b.Error).Error())
					break
				}
				reader := bytes.NewReader(b.Value)
				v.unmarshalReader(reader, v.kvstore)
			}
		}(respc)
		return nil
	}
	return RemoteConfigError("No Files Found")
}

// watchKeyValueConfigWithChannel Retrieve the first found remote configuration.
//
//	比 watchKeyValueConfigOnChannel 多了channel回调
func (v *Viper) watchKeyValueConfigWithChannel(ctx context.Context, receiver chan *RemoteResponse, deepMerge bool) error {
	if RemoteConfig == nil {
		return RemoteConfigError("Enable the remote features by doing a blank import of the viper/remote package: '_ github.com/spf13/viper/remote'")
	}

	if len(v.remoteProviders) == 0 {
		return RemoteConfigError("No Remote Providers")
	}
	providers := append([]*defaultRemoteProvider(nil), v.remoteProviders...)
	snapshots := make([]map[string]any, len(providers))
	snapshotMu := sync.Mutex{}

	for i, rp := range providers {
		val, err := v.getRemoteConfig(rp)
		if err != nil {
			return fmt.Errorf("viper watchKeyValueConfigWithChannel get remote config from %s %s: %w", rp.Provider(), rp.Path(), err)
		}
		snapshots[i] = val
	}
	v.kvstore = v.mergeRemoteConfigSnapshots(snapshots, deepMerge)

	responseChans := make([]<-chan *RemoteResponse, len(providers))
	quitChans := make([]chan bool, len(providers))
	for i, rp := range providers {
		respc, quit := RemoteConfig.WatchChannel(rp)
		if respc == nil || quit == nil {
			for _, startedQuit := range quitChans[:i] {
				if startedQuit != nil {
					close(startedQuit)
				}
			}
			return fmt.Errorf("viper watchKeyValueConfigWithChannel watch channel unavailable for %s %s", rp.Provider(), rp.Path())
		}
		responseChans[i] = respc
		quitChans[i] = quit
	}

	for i := range providers {
		respc := responseChans[i]
		quit := quitChans[i]
		// 去掉todo 已经加上quit channel
		go func(providerIndex int, rc <-chan *RemoteResponse, quit chan<- bool) {
			// 关闭quit避免crypt库里goroutine泄漏,TODO 更好的方式是使用传入CancelContext来控制退出,这里受限于接口定义,不好改动
			defer close(quit)
			for {
				select {
				case <-ctx.Done():
					quit <- true
					return
				case b, ok := <-rc:
					if !ok {
						return
					}
					if b.Error != nil {
						v.logger.Error(fmt.Errorf("viper watchKeyValueConfigWithChannel watch remote config: %w", b.Error).Error())
						receiver <- b
						return
					}
					reader := bytes.NewReader(b.Value)
					val := map[string]any{}
					err := v.unmarshalReader(reader, val)
					if err != nil {
						v.logger.Error(fmt.Errorf("viper watchKeyValueConfigWithChannel watch remote config: %w", err).Error())
						continue
					}
					snapshotMu.Lock()
					snapshots[providerIndex] = val
					// 注册顺序就是远程配置优先级，从低到高。
					//
					// 这里不能把本次 watch 到的 val 直接 Merge 到 v.kvstore：
					// 如果低优先级配置后变化，而同一个字段在高优先级配置里也存在，
					// “最后一次事件”就会覆盖“更高优先级”，最终结果和优先级语义相反。
					//
					// 所以每个 provider 只维护自己的最新快照；任意快照变化后，都从空 map
					// 按注册顺序重新合并全部快照。这样低优先级源只能提供高优先级源没有覆盖
					// 的字段，高优先级源始终在最后合并并覆盖同名字段。
					v.kvstore = v.mergeRemoteConfigSnapshots(snapshots, deepMerge)
					snapshotMu.Unlock()
					receiver <- b
				}
			}
		}(i, respc, quit)
		// 官方bug 造成只监听一个 provider
		// return nil
	}
	return nil
}

// Retrieve the first found remote configuration.
func (v *Viper) watchKeyValueConfig() error {
	if len(v.remoteProviders) == 0 {
		return RemoteConfigError("No Remote Providers")
	}

	for _, rp := range v.remoteProviders {
		val, err := v.watchRemoteConfig(rp)
		if err != nil {
			v.logger.Error(fmt.Errorf("watch remote config: %w", err).Error())

			continue
		}
		v.kvstore = val
		return nil
	}
	return RemoteConfigError("No Files Found")
}

func (v *Viper) watchRemoteConfig(provider RemoteProvider) (map[string]any, error) {
	reader, err := RemoteConfig.Watch(provider)
	if err != nil {
		return nil, err
	}
	err = v.unmarshalReader(reader, v.kvstore)
	return v.kvstore, err
}