worker/internal/installer/install.go

package installer

import (
	"errors"
	"fmt"
	"io"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

var dockerImagePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._/:@-]*$`)

const (
	DefaultURL   = "https://rsmon.ru"
	DefaultImage = "reg.rsxx.ru/rsmon/rsmon-worker:latest"
	binaryPath   = "/usr/local/bin/rsmon-worker"
	envPath      = "/etc/rsmon-worker/worker.env"
	unitPath     = "/etc/systemd/system/rsmon-worker.service"
)

const systemdUnit = `[Unit]
Description=RSMon distributed monitoring worker
After=network-online.target

[Service]
Type=simple
User=root
EnvironmentFile=/etc/rsmon-worker/worker.env
ExecStart=/usr/local/bin/rsmon-worker
Restart=on-failure

[Install]
WantedBy=multi-user.target
`

const dockerSystemdUnit = `[Unit]
Description=RSMon distributed monitoring worker (Docker)
After=network-online.target docker.service
Requires=docker.service

[Service]
Type=simple
User=root
ExecStartPre=-docker rm -f rsmon-worker
ExecStart=docker run --rm --name rsmon-worker --network host --cap-add NET_RAW --env-file /etc/rsmon-worker/worker.env -v rsmon-worker-data:/var/lib/rsmon-worker %s
ExecStop=docker stop rsmon-worker
Restart=on-failure

[Install]
WantedBy=multi-user.target
`

type InstallOptions struct {
	Binary  string
	EnvFile string
	Token   string
	URL     string
	Docker  bool
	Image   string
	NoStart bool
}

func Install(opts InstallOptions) error {
	if os.Geteuid() != 0 {
		return errors.New("install must be run as root")
	}
	if opts.URL == "" {
		opts.URL = DefaultURL
	}
	if err := ValidateURL(opts.URL); err != nil {
		return err
	}
	if opts.EnvFile == "" {
		if err := ValidateToken(opts.Token); err != nil {
			return err
		}
	}
	if opts.Image == "" {
		opts.Image = DefaultImage
	}
	if opts.Docker {
		if err := ValidateImage(opts.Image); err != nil {
			return err
		}
	} else if opts.Binary == "" {
		var err error
		opts.Binary, err = os.Executable()
		if err != nil {
			return fmt.Errorf("locate worker executable: %w", err)
		}
	}
	unit := systemdUnit
	if opts.Docker {
		if _, err := exec.LookPath("docker"); err != nil {
			return errors.New("docker is required for --docker installation")
		}
		if err := command("docker", "pull", opts.Image); err != nil {
			return err
		}
		unit = DockerUnit(opts.Image)
	} else if err := copyAtomic(opts.Binary, binaryPath, 0755); err != nil {
		return fmt.Errorf("install binary: %w", err)
	}
	if err := writeAtomic(unitPath, []byte(unit), 0644); err != nil {
		return fmt.Errorf("install systemd unit: %w", err)
	}
	if opts.EnvFile != "" {
		if err := copyAtomic(opts.EnvFile, envPath, 0600); err != nil {
			return fmt.Errorf("install environment: %w", err)
		}
	} else {
		if err := writeAtomic(envPath, Environment(opts.URL, opts.Token), 0600); err != nil {
			return fmt.Errorf("install environment: %w", err)
		}
	}
	if err := command("systemctl", "daemon-reload"); err != nil {
		return err
	}
	if err := command("systemctl", "enable", "rsmon-worker.service"); err != nil {
		return err
	}
	if !opts.NoStart {
		if err := command("systemctl", "restart", "rsmon-worker.service"); err != nil {
			return err
		}
		if err := command("systemctl", "is-active", "--quiet", "rsmon-worker.service"); err != nil {
			return err
		}
	}
	return nil
}

func ValidateURL(raw string) error {
	if strings.ContainsAny(raw, "\r\n") {
		return fmt.Errorf("server URL must be an absolute HTTP(S) URL")
	}
	u, err := url.Parse(raw)
	if err != nil || u.Host == "" || (u.Scheme != "https" && u.Scheme != "http") {
		return fmt.Errorf("server URL must be an absolute HTTP(S) URL")
	}
	return nil
}

func Environment(serverURL, token string) []byte {
	return []byte("RSMON_URL=" + serverURL + "\nRSMON_TOKEN=" + token + "\nWORKER_HOST=127.0.0.1\n")
}

func ValidateToken(token string) error {
	if strings.TrimSpace(token) == "" || strings.ContainsAny(token, "\r\n") {
		return errors.New("worker token must be non-empty and contain no newlines")
	}
	return nil
}

func ValidateImage(image string) error {
	if !dockerImagePattern.MatchString(image) {
		return errors.New("Docker image must be one non-option argument")
	}
	return nil
}

func DockerUnit(image string) string {
	return fmt.Sprintf(dockerSystemdUnit, image)
}

func copyAtomic(src, dst string, mode os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	return atomicFile(dst, mode, func(out *os.File) error {
		_, err := io.Copy(out, in)
		return err
	})
}

func writeAtomic(dst string, data []byte, mode os.FileMode) error {
	return atomicFile(dst, mode, func(out *os.File) error {
		_, err := out.Write(data)
		return err
	})
}

func atomicFile(dst string, mode os.FileMode, write func(*os.File) error) error {
	if err := os.MkdirAll(filepath.Dir(dst), 0755); err != nil {
		return err
	}
	tmp, err := os.CreateTemp(filepath.Dir(dst), ".rsmon-worker-*")
	if err != nil {
		return err
	}
	name := tmp.Name()
	defer os.Remove(name)
	if err := tmp.Chmod(mode); err != nil {
		tmp.Close()
		return err
	}
	if err := write(tmp); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	return os.Rename(name, dst)
}

func command(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("%s failed: %w", name, err)
	}
	return nil
}