LokiStack/backend/services/image_registry.go

package services

// 镜像仓库源 — 全局设置(参照 kuboard 的国内/海外镜像选择):
// overseas = docker.io 官方源;cn = daocloud 国内镜像;custom = 用户自填 registry 前缀。
// 影响平台拉取的全部镜像:Loki / Grafana / Promtail / Fluent Bit。
// 每个组件可单独配置仓库与 tag(Loki/Grafana 的 tag 留空则跟随部署版本)。
import (
	"encoding/json"
	"fmt"
	"strings"

	"lokistack/models"
	"lokistack/repository"
)

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

// ComponentImage 单个组件的镜像配置
type ComponentImage struct {
	Repository string `json:"repository"` // 如 grafana/loki
	Tag        string `json:"tag"`        // 留空表示跟随部署版本 / 使用默认
}

// ImageRegistryConfig 镜像源配置
type ImageRegistryConfig struct {
	Mode       string                    `json:"mode"`       // overseas | cn | custom
	CustomHost string                    `json:"customHost"` // custom 模式下的仓库前缀
	Images     map[string]ComponentImage `json:"images"`     // loki/grafana/promtail/fluentbit
}

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

var defaultComponentImages = map[string]ComponentImage{
	"loki":      {Repository: "grafana/loki", Tag: ""},    // 跟随部署版本
	"grafana":   {Repository: "grafana/grafana", Tag: ""}, // 跟随部署版本
	"promtail":  {Repository: "grafana/promtail", Tag: "3.4.0"},
	"fluentbit": {Repository: "fluent/fluent-bit", Tag: "5.1.5"},
}

func (c *ImageRegistryConfig) fillDefaults() {
	if c.Images == nil {
		c.Images = map[string]ComponentImage{}
	}
	for name, def := range defaultComponentImages {
		img, ok := c.Images[name]
		if !ok || (img.Repository == "" && img.Tag == "") {
			c.Images[name] = def
			continue
		}
		if img.Repository == "" {
			img.Repository = def.Repository
		}
		c.Images[name] = img
	}
}

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

func SaveImageRegistry(mode, customHost string, images map[string]ComponentImage) 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)")
		}
	}
	cfg := &ImageRegistryConfig{Mode: mode, CustomHost: customHost, Images: images}
	cfg.fillDefaults()
	// 清理空白输入
	clean := map[string]ComponentImage{}
	for name, img := range cfg.Images {
		clean[name] = ComponentImage{Repository: strings.TrimSpace(img.Repository), Tag: strings.TrimSpace(img.Tag)}
	}
	cfg.Images = clean
	data, _ := json.Marshal(cfg)
	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], ".") {
			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:
		return official
	}
}

// ResolveImageFor 按组件配置解析镜像；组件 tag 为空时使用 fallbackTag（通常是部署版本）。
func ResolveImageFor(component, fallbackRepo, fallbackTag string) string {
	cfg := GetImageRegistry()
	img, ok := cfg.Images[component]
	repo := fallbackRepo
	tag := fallbackTag
	if ok {
		if strings.TrimSpace(img.Repository) != "" {
			repo = strings.TrimSpace(img.Repository)
		}
		if strings.TrimSpace(img.Tag) != "" {
			tag = strings.TrimSpace(img.Tag)
		}
	}
	official := repo
	if !strings.Contains(repo, ":") && tag != "" {
		official = repo + ":" + tag
	}
	return ResolveImage(official)
}

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