LokiStack/backend/services/loki_docker.go

package services

import (
	"context"
	"encoding/base64"
	"fmt"
	"strings"
	"time"

	"lokistack/models"
)

func dockerContainerName(id uint) string { return fmt.Sprintf("lokistack-loki-%d", id) }
func dockerConfigDir(id uint) string     { return fmt.Sprintf("/etc/lokistack/loki-%d", id) }
func dockerConfigPath(id uint) string    { return dockerConfigDir(id) + "/loki.yaml" }
func dockerDataDir(id uint) string       { return fmt.Sprintf("/var/lib/lokistack/loki-%d", id) }

// 历史命名的容器与目录 — 仅用于迁移与清理,
// 避免已部署实例在升级后留下孤儿容器/目录。
func legacyDockerContainerName(id uint) string { return fmt.Sprintf("lokideploy-loki-%d", id) }
func legacyDockerConfigDir(id uint) string     { return fmt.Sprintf("/etc/lokideploy/loki-%d", id) }
func legacyDockerDataDir(id uint) string       { return fmt.Sprintf("/var/lib/lokideploy/loki-%d", id) }

func dockerLokiImage(version string) (string, error) {
	if version == "" {
		version = "3.0.0"
	}
	for _, supported := range LokiVersionList {
		if supported == version && strings.HasPrefix(version, "3.") {
			return ResolveImageFor("loki", "grafana/loki", version), nil
		}
	}
	return "", fmt.Errorf("主机 Docker 部署请选择支持的 Loki 3.x 版本")
}

func stopDockerLoki(ctx context.Context, deployment *models.LokiDeployment) error {
	if deployment.Host == nil {
		return fmt.Errorf("部署未关联主机")
	}
	client, err := NewSSHClientWithAuth(HostSSHAuth(deployment.Host))
	if err != nil {
		return fmt.Errorf("连接部署主机失败: %w", err)
	}
	defer client.Close()
	if _, stderr, err := client.Execute(ctx, "docker stop "+dockerContainerName(deployment.ID)); err != nil {
		return fmt.Errorf("停止 Loki 容器失败: %v: %s", err, strings.TrimSpace(stderr))
	}
	// 历史命名的同名容器一并停掉(可能不存在,尽力而为)
	client.Execute(ctx, "docker stop "+legacyDockerContainerName(deployment.ID)+" >/dev/null 2>&1 || true")
	return nil
}

func deleteDockerLoki(ctx context.Context, deployment *models.LokiDeployment) error {
	if deployment.Host == nil {
		return fmt.Errorf("部署未关联主机")
	}
	client, err := NewSSHClientWithAuth(HostSSHAuth(deployment.Host))
	if err != nil {
		return fmt.Errorf("连接部署主机失败: %w", err)
	}
	defer client.Close()
	name := dockerContainerName(deployment.ID)
	legacyName := legacyDockerContainerName(deployment.ID)
	if _, _, err := client.Execute(ctx, "docker info >/dev/null 2>&1"); err != nil {
		return fmt.Errorf("目标主机 Docker 不可用，无法安全清理部署")
	}
	// 仅删除该部署的容器和配置(含历史命名)；日志数据目录保留，以便人工恢复。
	cmd := "for n in " + name + " " + legacyName + "; do if docker container inspect $n >/dev/null 2>&1; then docker rm -f $n >/dev/null || exit 1; fi; done; rm -rf -- " +
		dockerConfigDir(deployment.ID) + " " + legacyDockerConfigDir(deployment.ID)
	if _, stderr, err := client.Execute(ctx, cmd); err != nil {
		return fmt.Errorf("清理 Loki 容器与配置失败: %v: %s", err, strings.TrimSpace(stderr))
	}
	return nil
}

// uploadDockerFile 通过 SSH 写入固定路径。内容仅以 base64 传给远程 shell，避免 YAML 或凭证参与 shell 解析。
func uploadDockerFile(ctx context.Context, client *SSHClient, path, content, mode string) error {
	encoded := base64.StdEncoding.EncodeToString([]byte(content))
	cmd := fmt.Sprintf("umask 077 && printf '%%s' '%s' | base64 -d > %s && chmod %s %s", encoded, path, mode, path)
	_, _, err := client.Execute(ctx, cmd)
	if err != nil {
		return fmt.Errorf("写入主机配置失败: %w", err)
	}
	return nil
}

func (d *LokiDeployer) deployDocker(ctx context.Context) error {
	deployment := d.deployment
	image, err := dockerLokiImage(deployment.Version)
	if err != nil {
		return err
	}
	if deployment.Host == nil {
		return fmt.Errorf("请选择目标主机")
	}
	client, err := NewSSHClientWithAuth(HostSSHAuth(deployment.Host))
	if err != nil {
		return err
	}
	defer client.Close()
	if _, _, err := client.Execute(ctx, "docker info >/dev/null 2>&1"); err != nil {
		return fmt.Errorf("目标主机 Docker 不可用，请检查安装状态及 SSH 用户权限")
	}
	name := dockerContainerName(deployment.ID)
	configDir := dockerConfigDir(deployment.ID)
	dataDir := dockerDataDir(deployment.ID)
	// 历史命名目录原子迁移到新命名;存在才迁,不覆盖已有数据
	migrate := fmt.Sprintf(
		"if [ -d %s ] && [ ! -d %s ]; then mkdir -p /etc/lokistack && mv %s %s; fi; "+
			"if [ -d %s ] && [ ! -d %s ]; then mkdir -p /var/lib/lokistack && mv %s %s; fi",
		legacyDockerConfigDir(deployment.ID), configDir, legacyDockerConfigDir(deployment.ID), configDir,
		legacyDockerDataDir(deployment.ID), dataDir, legacyDockerDataDir(deployment.ID), dataDir,
	)
	if _, stderr, err := client.Execute(ctx, migrate); err != nil {
		return fmt.Errorf("迁移历史部署目录失败(可手动执行: %s): %v: %s", migrate, err, strings.TrimSpace(stderr))
	}
	if _, _, err := client.Execute(ctx, "mkdir -p "+configDir+" "+dataDir+" && chmod 755 "+configDir+" && chown 10001:10001 "+dataDir); err != nil {
		return fmt.Errorf("创建主机持久化目录失败: %w", err)
	}
	config := d.generateLokiConfig()
	if err := uploadDockerFile(ctx, client, dockerConfigPath(deployment.ID), config, "644"); err != nil {
		return err
	}
	// env 文件由 Docker daemon 读取；Loki 配置只保留变量引用。
	env := ""
	if sc := deployment.StorageConfig; sc != nil {
		for _, value := range []string{sc.AccessKey, sc.SecretKey} {
			if strings.ContainsAny(value, "\r\n") {
				return fmt.Errorf("存储凭证包含不支持的换行字符")
			}
		}
		env = "LOKI_STORAGE_ACCESS_KEY=" + sc.AccessKey + "\nLOKI_STORAGE_SECRET_KEY=" + sc.SecretKey + "\n"
	}
	if err := uploadDockerFile(ctx, client, configDir+"/storage.env", env, "600"); err != nil {
		return err
	}
	d.log("info", "拉取 Docker 镜像 "+image, "")
	if _, stderr, err := client.Execute(ctx, "docker pull "+image); err != nil {
		return fmt.Errorf("拉取 Loki 镜像失败: %v: %s", err, strings.TrimSpace(stderr))
	}
	// 同一部署 ID 的重部署只替换容器(含历史命名)，配置和日志数据目录保留。
	client.Execute(ctx, "docker rm -f "+name+" "+legacyDockerContainerName(deployment.ID)+" >/dev/null 2>&1 || true")
	port := deployment.ListenPort
	if port == 0 {
		port = 3100
	}
	listenIP := deployment.ListenIP
	if listenIP == "" {
		listenIP = deployment.Host.HostIP
	}
	cmd := fmt.Sprintf("docker run -d --name %s --restart unless-stopped --env-file %s/storage.env -p %s:%d:3100 -v %s:/etc/loki:ro -v %s:/var/lib/loki %s -config.file=/etc/loki/loki.yaml -config.expand-env=true", name, configDir, listenIP, port, configDir, dataDir, image)
	if _, stderr, err := client.Execute(ctx, cmd); err != nil {
		return fmt.Errorf("启动 Loki 容器失败: %v: %s", err, strings.TrimSpace(stderr))
	}
	d.log("info", "等待 Loki 容器就绪", "")
	readyCmd := fmt.Sprintf("curl -fsS --max-time 2 http://%s:%d/ready >/dev/null", listenIP, port)
	for i := 0; i < 30; i++ {
		if _, _, err := client.Execute(ctx, readyCmd); err == nil {
			d.log("success", fmt.Sprintf("Loki 已启动：http://%s:%d", listenIP, port), "")
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
	return fmt.Errorf("Loki 容器启动后未就绪，请在目标主机执行 docker logs %s", name)
}