path: root/internal/control/bootstrap.go

package control

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

type BootstrapOptions struct {
	StateDir string
	DryRun   bool
}

func BootstrapNode(ctx context.Context, runner SSHExecutor, node Node, opts BootstrapOptions) (*NodeState, error) {
	for idx := range node.Protocols {
		if err := ensureRealityProfile(&node.Protocols[idx]); err != nil {
			return nil, err
		}
		if err := ensureHysteria2Profile(&node.Protocols[idx]); err != nil {
			return nil, err
		}
	}
	if err := ValidateNode(node); err != nil {
		return nil, err
	}

	now := time.Now().UTC()
	state := &NodeState{
		NodeID:          node.ID,
		BootstrapStatus: "pending",
		PublicHost:      publicHost(node),
		Services:        serviceStatuses(node.Protocols, "configured"),
		Metadata: map[string]any{
			"provider": node.Provider,
			"region":   node.Region,
			"dry_run":  opts.DryRun,
		},
	}

	if opts.DryRun {
		state.BootstrapStatus = "planned"
		state.LastBootstrapAt = &now
		state.Metadata["release_id"] = buildReleaseID(now)
		if err := SaveNodeState(opts.StateDir, *state); err != nil {
			return nil, err
		}
		return state, nil
	}

	relID := buildReleaseID(now)
	bundleDir, tarballPath, err := buildRuntimeBundle(node, relID)
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(bundleDir)
	defer os.Remove(tarballPath)

	result, err := runner.Run(ctx, node, RenderBootstrapPrepareScript())
	if err != nil {
		state.BootstrapStatus = "failed"
		state.LastBootstrapAt = &now
		state.Metadata["stderr"] = strings.TrimSpace(result.Stderr)
		state.Metadata["stdout"] = strings.TrimSpace(result.Stdout)
		if saveErr := SaveNodeState(opts.StateDir, *state); saveErr != nil {
			return nil, fmt.Errorf("%w; save state: %v", err, saveErr)
		}
		return nil, err
	}

	remoteTarballPath := "/tmp/vpnem-node-" + node.ID + ".tar.gz"
	if err := runner.CopyFile(ctx, node, tarballPath, remoteTarballPath); err != nil {
		state.BootstrapStatus = "failed"
		state.LastBootstrapAt = &now
		state.Metadata["release_id"] = relID
		state.Metadata["copy_error"] = err.Error()
		if saveErr := SaveNodeState(opts.StateDir, *state); saveErr != nil {
			return nil, fmt.Errorf("%w; save state: %v", err, saveErr)
		}
		return nil, err
	}

	result, err = runner.Run(ctx, node, RenderBootstrapFinalizeScript(node, relID, remoteTarballPath))
	if err != nil {
		state.BootstrapStatus = "failed"
		state.LastBootstrapAt = &now
		state.Metadata["stderr"] = strings.TrimSpace(result.Stderr)
		state.Metadata["stdout"] = strings.TrimSpace(result.Stdout)
		if saveErr := SaveNodeState(opts.StateDir, *state); saveErr != nil {
			return nil, fmt.Errorf("%w; save state: %v", err, saveErr)
		}
		return nil, err
	}

	state.BootstrapStatus = "ready"
	state.LastBootstrapAt = &now
	state.Metadata["release_id"] = relID
	state.Metadata["stdout"] = strings.TrimSpace(result.Stdout)
	if err := SaveNodeState(opts.StateDir, *state); err != nil {
		return nil, err
	}
	return state, nil
}

func CheckNode(ctx context.Context, runner SSHExecutor, node Node, stateDir string) (*NodeState, error) {
	now := time.Now().UTC()
	result, err := runner.Check(ctx, node)
	state := &NodeState{
		NodeID:            node.ID,
		PublicHost:        publicHost(node),
		LastHealthCheckAt: &now,
		Services:          serviceStatuses(node.Protocols, "unknown"),
		Metadata:          map[string]any{},
	}

	if err != nil {
		state.BootstrapStatus = "unreachable"
		state.Metadata["stderr"] = strings.TrimSpace(result.Stderr)
		if saveErr := SaveNodeState(stateDir, *state); saveErr != nil {
			return nil, fmt.Errorf("%w; save state: %v", err, saveErr)
		}
		return nil, err
	}

	state.BootstrapStatus = "reachable"
	state.Metadata["stdout"] = strings.TrimSpace(result.Stdout)

	runtimeResult, runtimeErr := runner.Run(ctx, node, RenderHealthCheckScript(node))
	if runtimeErr != nil {
		state.Metadata["runtime_stderr"] = strings.TrimSpace(runtimeResult.Stderr)
		state.Metadata["runtime_stdout"] = strings.TrimSpace(runtimeResult.Stdout)
	} else {
		services, metadata := parseHealthCheckOutput(runtimeResult.Stdout, node.Protocols)
		if len(services) > 0 {
			state.Services = services
		}
		for k, v := range metadata {
			state.Metadata[k] = v
		}
		if healthy, ok := metadata["healthz_http_code"].(int); ok && healthy == 200 {
			state.BootstrapStatus = "healthy"
		} else if allServicesRunning(state.Services) {
			state.BootstrapStatus = "ready"
		}
	}
	if err := SaveNodeState(stateDir, *state); err != nil {
		return nil, err
	}
	return state, nil
}

func RenderBootstrapPrepareScript() string {
	var b strings.Builder
	b.WriteString("set -eu\n")
	b.WriteString("export DEBIAN_FRONTEND=noninteractive\n")
	b.WriteString("mkdir -p /opt/vpnem-node/releases\n")
	b.WriteString("if command -v apt-get >/dev/null 2>&1; then\n")
	b.WriteString("  apt-get update\n")
	b.WriteString("  apt-get install -y ca-certificates curl tar gzip openssl docker.io docker-compose || true\n")
	b.WriteString("elif command -v dnf >/dev/null 2>&1; then\n")
	b.WriteString("  dnf install -y ca-certificates curl tar gzip openssl docker docker-compose-plugin docker-compose || true\n")
	b.WriteString("elif command -v pacman >/dev/null 2>&1; then\n")
	b.WriteString("  pacman -Sy --noconfirm ca-certificates curl tar gzip openssl docker docker-compose || true\n")
	b.WriteString("elif command -v apk >/dev/null 2>&1; then\n")
	b.WriteString("  apk add --no-cache ca-certificates curl tar gzip openssl docker-cli-compose || true\n")
	b.WriteString("fi\n")
	b.WriteString("if command -v systemctl >/dev/null 2>&1; then systemctl enable --now docker || true; fi\n")
	b.WriteString("if ! command -v docker >/dev/null 2>&1; then\n")
	b.WriteString("  echo 'docker is not installed after bootstrap prepare' >&2\n")
	b.WriteString("  exit 1\n")
	b.WriteString("fi\n")
	b.WriteString("printf 'vpnem-node bootstrap prepared\\n'\n")
	return b.String()
}

func RenderBootstrapFinalizeScript(node Node, releaseID, remoteTarballPath string) string {
	var b strings.Builder
	releaseDir := "/opt/vpnem-node/releases/" + releaseID
	b.WriteString("set -eu\n")
	b.WriteString("mkdir -p " + releaseDir + "\n")
	b.WriteString("tar -xzf " + remoteTarballPath + " -C " + releaseDir + "\n")
	b.WriteString("ln -sfn " + releaseDir + " /opt/vpnem-node/current\n")
	b.WriteString("rm -f " + remoteTarballPath + "\n")
	b.WriteString("if ! command -v docker >/dev/null 2>&1; then\n")
	b.WriteString("  echo 'docker is not installed on target node' >&2\n")
	b.WriteString("  exit 1\n")
	b.WriteString("fi\n")
	b.WriteString("if docker compose version >/dev/null 2>&1; then\n")
	b.WriteString("  docker compose -f /opt/vpnem-node/current/docker-compose.yml up -d --force-recreate\n")
	b.WriteString("elif command -v docker-compose >/dev/null 2>&1; then\n")
	b.WriteString("  docker-compose -f /opt/vpnem-node/current/docker-compose.yml up -d --force-recreate\n")
	b.WriteString("else\n")
	b.WriteString("  echo 'docker compose is not available on target node' >&2\n")
	b.WriteString("  exit 1\n")
	b.WriteString("fi\n")
	b.WriteString("printf 'vpnem-node release ")
	b.WriteString(shellQuoteValue(releaseID))
	b.WriteString(" ready for ")
	b.WriteString(shellQuoteValue(node.ID))
	b.WriteString("\\n'\n")
	return b.String()
}

func RenderHealthCheckScript(node Node) string {
	var b strings.Builder
	b.WriteString("set -eu\n")
	b.WriteString("if [ -f /opt/vpnem-node/current/docker-compose.yml ]; then\n")
	b.WriteString("  if command -v docker >/dev/null 2>&1 && docker compose version >/dev/null 2>&1; then\n")
	b.WriteString("    docker compose -f /opt/vpnem-node/current/docker-compose.yml ps --format json 2>/dev/null || true\n")
	b.WriteString("  elif command -v docker-compose >/dev/null 2>&1; then\n")
	b.WriteString("    docker-compose -f /opt/vpnem-node/current/docker-compose.yml ps --format json 2>/dev/null || true\n")
	b.WriteString("  fi\n")
	b.WriteString("  if command -v docker >/dev/null 2>&1; then\n")
	b.WriteString("    docker ps --format '{{json .}}' 2>/dev/null || true\n")
	b.WriteString("  fi\n")
	b.WriteString("fi\n")
	if needsEdgeProxy(node) {
		b.WriteString("printf 'HEALTHZ_HTTP_CODE='; ")
		b.WriteString("curl -ks --resolve ")
		b.WriteString(shellQuoteValue(node.Domain))
		b.WriteString(":443:127.0.0.1 -o /dev/null -w '%{http_code}' https://")
		b.WriteString(shellQuoteValue(node.Domain))
		b.WriteString("/healthz || true\n")
	}
	if needsHysteria2HealthInbound(node) {
		b.WriteString("printf 'HY2_MIXED_PORT='; ")
		b.WriteString("curl -sS --max-time 5 --proxy socks5h://127.0.0.1:1080 https://ifconfig.me/ip || true\n")
	}
	return b.String()
}

func serviceStatuses(protocols []ProtocolProfile, status string) []ServiceStatus {
	services := make([]ServiceStatus, 0, len(protocols))
	for _, protocol := range protocols {
		if !protocol.Enabled {
			continue
		}
		services = append(services, ServiceStatus{
			Type:   protocol.Type,
			Status: status,
			Port:   protocol.Port,
		})
	}
	return services
}

func parseHealthCheckOutput(stdout string, protocols []ProtocolProfile) ([]ServiceStatus, map[string]any) {
	services := serviceStatuses(protocols, "unknown")
	metadata := map[string]any{}
	lines := strings.Split(stdout, "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if strings.HasPrefix(line, "HEALTHZ_HTTP_CODE=") {
			codeStr := strings.TrimPrefix(line, "HEALTHZ_HTTP_CODE=")
			if code, err := strconv.Atoi(codeStr); err == nil {
				metadata["healthz_http_code"] = code
			}
			continue
		}
		if strings.HasPrefix(line, "HY2_MIXED_PORT=") {
			value := strings.TrimSpace(strings.TrimPrefix(line, "HY2_MIXED_PORT="))
			metadata["hy2_mixed_port"] = value
			if value != "" {
				markServicesByTypes(services, []string{"hysteria2"}, "running")
			}
			continue
		}

		var entry map[string]any
		if err := jsonUnmarshalLine(line, &entry); err != nil {
			continue
		}
		serviceName, _ := entry["Service"].(string)
		state, _ := entry["State"].(string)
		if serviceName == "" {
			if labels, _ := entry["Labels"].(string); strings.Contains(labels, "com.docker.compose.service=sing-box") {
				serviceName = "sing-box"
			} else if names, _ := entry["Names"].(string); strings.Contains(names, "sing-box") {
				serviceName = "sing-box"
			}
		}
		if state == "" {
			if status, _ := entry["Status"].(string); strings.HasPrefix(strings.ToLower(status), "up") {
				state = "running"
			}
		}
		if serviceName == "" || state == "" {
			continue
		}
		metadata["docker_"+serviceName] = state
		switch serviceName {
		case "sing-box":
			markServicesByTypes(services, []string{"vless", "vless-reality", "shadowsocks", "socks", "socks5", "vmess", "hysteria2"}, state)
		case "caddy":
			markServicesByTypes(services, []string{"vless", "vmess"}, state)
		}
	}
	return services, metadata
}

func allServicesRunning(services []ServiceStatus) bool {
	if len(services) == 0 {
		return false
	}
	for _, service := range services {
		if service.Status != "running" {
			return false
		}
	}
	return true
}

func markServicesByTypes(services []ServiceStatus, kinds []string, state string) {
	set := make(map[string]struct{}, len(kinds))
	for _, kind := range kinds {
		set[kind] = struct{}{}
	}
	for idx := range services {
		if _, ok := set[services[idx].Type]; ok {
			services[idx].Status = state
		}
	}
}

func jsonUnmarshalLine(line string, out *map[string]any) error {
	decoder := strings.NewReader(line)
	return json.NewDecoder(decoder).Decode(out)
}

func publicHost(node Node) string {
	if strings.TrimSpace(node.Domain) != "" {
		return node.Domain
	}
	return node.Host
}

func shellQuoteValue(value string) string {
	value = strings.ReplaceAll(value, "\n", "")
	return value
}

func buildRuntimeBundle(node Node, releaseID string) (string, string, error) {
	rootDir, err := os.MkdirTemp("", "vpnem-node-bundle-*")
	if err != nil {
		return "", "", err
	}
	bundleDir := filepath.Join(rootDir, "bundle")
	if err := RenderRuntimeBundle(bundleDir, node, releaseID); err != nil {
		os.RemoveAll(rootDir)
		return "", "", err
	}
	tarballPath := filepath.Join(rootDir, "bundle.tar.gz")
	if err := CreateTarGzFromDir(bundleDir, tarballPath); err != nil {
		os.RemoveAll(rootDir)
		return "", "", err
	}
	return rootDir, tarballPath, nil
}

func buildReleaseID(now time.Time) string {
	return now.UTC().Format("20060102-150405")
}