LokiStack/backend/services/image_registry.go

package services

// 镜像仓库源 — 全局设置(参照 kuboard 的国内/海外镜像选择):
// overseas = docker.io 官方源;cn = daocloud 国内镜像;custom = 用户自填 registry 前缀。
// 影响平台拉取的全部镜像:Loki / Grafana / Promtail / Fluent Bit。
import (
	"encoding/json"
	"fmt"
	"strings"

	"opsxloki/models"
	"opsxloki/repository"
)

const (
	imageRegistrySettingKey = "image_registry"
	RegistryOverseas        = "overseas" // docker.io
	RegistryCN              = "cn"       // docker.m.daocloud.io(国内加速)
	RegistryCustom          = "custom"   // 用户私有仓库前缀
	daocloudPrefix          = "docker.m.daocloud.io"
)

// ImageRegistryConfig 镜像源配置
type ImageRegistryConfig struct {
	Mode       string `json:"mode"`       // overseas | cn | custom
	CustomHost string `json:"customHost"` // custom 模式下的仓库前缀,如 swr.cn-east-2.myhuaweicloud.com/xxx
}

type ImageRegistryPublic struct {
	Mode       string `json:"mode"`
	CustomHost string `json:"customHost"`
	// 示例预览(前端展示生效镜像)
	SampleLoki      string `json:"sampleLoki"`
	SampleGrafana   string `json:"sampleGrafana"`
	SamplePromtail  string `json:"samplePromtail"`
	SampleFluentBit string `json:"sampleFluentBit"`
}

func GetImageRegistry() *ImageRegistryConfig {
	db := repository.GetDB()
	if db == nil {
		return &ImageRegistryConfig{Mode: RegistryOverseas}
	}
	db.RLock()
	var raw string
	for i := range db.Settings {
		if db.Settings[i].Key == imageRegistrySettingKey {
			raw = db.Settings[i].Value
			break
		}
	}
	db.RUnlock()
	cfg := &ImageRegistryConfig{Mode: RegistryOverseas}
	if raw != "" {
		_ = json.Unmarshal([]byte(raw), cfg)
	}
	if cfg.Mode == "" {
		cfg.Mode = RegistryOverseas
	}
	return cfg
}

func SaveImageRegistry(mode, customHost string) error {
	if mode != RegistryOverseas && mode != RegistryCN && mode != RegistryCustom {
		return fmt.Errorf("无效的镜像源模式")
	}
	if mode == RegistryCustom {
		customHost = strings.TrimSuffix(strings.TrimSpace(customHost), "/")
		if customHost == "" || !strings.Contains(customHost, ".") {
			return fmt.Errorf("自定义镜像仓库地址无效(需含域名,如 swr.cn-east-2.myhuaweicloud.com/your-ns)")
		}
	}
	data, _ := json.Marshal(&ImageRegistryConfig{Mode: mode, CustomHost: customHost})
	db := repository.GetDB()
	newID := db.NextID("settings")
	db.Lock()
	defer db.Unlock()
	for i := range db.Settings {
		if db.Settings[i].Key == imageRegistrySettingKey {
			db.Settings[i].Value = string(data)
			db.RequestSave()
			return nil
		}
	}
	db.Settings = append(db.Settings, models.SystemSetting{ID: newID, Key: imageRegistrySettingKey, Value: string(data)})
	db.RequestSave()
	return nil
}

// ResolveImage 按当前镜像源解析镜像地址:docker.io/grafana/loki:3.4.0 形式的官方镜像
// 会按 mode 重写前缀;已带私有仓库前缀(含镜像源域名)的地址原样返回。
func ResolveImage(official string) string {
	cfg := GetImageRegistry()
	switch cfg.Mode {
	case RegistryCN:
		if strings.HasPrefix(official, "docker.io/") {
			return daocloudPrefix + "/" + strings.TrimPrefix(official, "docker.io/")
		}
		if !strings.Contains(strings.SplitN(official, "/", 2)[0], ".") {
			// 无仓库前缀的官方镜像(grafana/loki:tag)
			return daocloudPrefix + "/" + official
		}
		return official
	case RegistryCustom:
		if strings.HasPrefix(official, "docker.io/") {
			return cfg.CustomHost + "/" + strings.TrimPrefix(official, "docker.io/")
		}
		if !strings.Contains(strings.SplitN(official, "/", 2)[0], ".") {
			return cfg.CustomHost + "/" + official
		}
		return official
	default:
		if strings.HasPrefix(official, "docker.io/") {
			return official // 官方模式保留全限定名
		}
		return official
	}
}

// ImageRegistryPublicConfig 对外输出(含示例预览)
func ImageRegistryPublicConfig() *ImageRegistryPublic {
	cfg := GetImageRegistry()
	return &ImageRegistryPublic{
		Mode:            cfg.Mode,
		CustomHost:      cfg.CustomHost,
		SampleLoki:      ResolveImage("docker.io/grafana/loki:3.4.0"),
		SampleGrafana:   ResolveImage("docker.io/grafana/grafana:11.3.0"),
		SamplePromtail:  ResolveImage("docker.io/grafana/promtail:3.4.0"),
		SampleFluentBit: ResolveImage("docker.io/fluent/fluent-bit:5.1.5"),
	}
}