lpra/internal/agent/update.go

package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

// doUpdate downloads the latest agent binary for this platform, replaces the
// running executable and restarts the service. It runs as root when it can:
// the agent's own user if that is root, otherwise sudo when it needs no
// password. It never uninstalls - only updates.
func (a *Agent) doUpdate() proto.TaskResult {
	exe, err := os.Executable()
	if err != nil {
		return proto.TaskResult{Error: "не удалось определить путь бинарника: " + err.Error()}
	}
	exe, _ = filepath.EvalSymlinks(exe)
	sudo := ""
	if os.Geteuid() != 0 {
		sudo = sudoPath()
	}

	url := fmt.Sprintf("%s/download/agent-%s-%s", a.cfg.Host, runtime.GOOS, runtime.GOARCH)
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil)
	resp, err := a.stream.Do(req)
	if err != nil {
		return proto.TaskResult{Error: "скачивание: " + err.Error()}
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return proto.TaskResult{Error: "скачивание: сервер вернул " + resp.Status}
	}

	// The new binary lands next to the old one, or, when root's hand is
	// needed to put it there, in a temporary folder first.
	dir := filepath.Dir(exe)
	if sudo != "" {
		dir = ""
	}
	f, err := os.CreateTemp(dir, ".lpagent-new-*")
	if err != nil {
		return proto.TaskResult{Error: "создание файла: " + err.Error() + userNote()}
	}
	tmp := f.Name()
	defer os.Remove(tmp)
	n, err := io.Copy(f, io.LimitReader(resp.Body, 256<<20))
	if cerr := f.Close(); err == nil {
		err = cerr
	}
	if err != nil {
		return proto.TaskResult{Error: "запись: " + err.Error()}
	}
	if err := checkBinary(tmp, n); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	if err := os.Chmod(tmp, 0o755); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}

	how := "от " + currentUser()
	if sudo != "" {
		// install copies it in place under a new inode, as a rename would.
		if out, err := exec.Command(sudo, "-n", "install", "-m", "0755", tmp, exe).CombinedOutput(); err != nil {
			return proto.TaskResult{Error: "замена бинарника через sudo: " + strings.TrimSpace(string(out)+" "+err.Error())}
		}
		how = "через sudo"
	} else if err := os.Rename(tmp, exe); err != nil {
		// Renaming over a running executable is allowed; the live process
		// keeps the old inode until it restarts.
		return proto.TaskResult{Error: "замена бинарника: " + err.Error() + userNote()}
	}

	// Restart shortly, after the result has been reported.
	go func() {
		time.Sleep(1500 * time.Millisecond)
		restartSelf(exe, sudo)
	}()

	data, _ := json.Marshal(map[string]string{"status": "обновлено " + how + ", перезапуск"})
	return proto.TaskResult{OK: true, Data: data}
}

// sudoPath is sudo when it lets this user run commands without a password.
func sudoPath() string {
	p, err := exec.LookPath("sudo")
	if err != nil {
		return ""
	}
	if exec.Command(p, "-n", "true").Run() != nil {
		return ""
	}
	return p
}

func userNote() string {
	if os.Geteuid() == 0 {
		return ""
	}
	return " (агент запущен от " + currentUser() + ", без root и sudo без пароля - обновите установщиком)"
}

// checkBinary makes sure what came is a program for this system, not an
// error page: a binary of a sensible size with the right header.
func checkBinary(path string, size int64) error {
	if size < 1<<20 {
		return fmt.Errorf("скачанный файл слишком мал (%d байт) - не похоже на агента", size)
	}
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	var magic [4]byte
	if _, err := io.ReadFull(f, magic[:]); err != nil {
		return err
	}
	ok := false
	switch runtime.GOOS {
	case "linux":
		ok = string(magic[:]) == "\x7fELF"
	case "darwin":
		ok = magic == [4]byte{0xcf, 0xfa, 0xed, 0xfe} || magic == [4]byte{0xca, 0xfe, 0xba, 0xbe}
	default:
		ok = true
	}
	if !ok {
		return errors.New("скачанный файл - не программа для " + runtime.GOOS)
	}
	return nil
}

// restartSelf starts the new binary. Under systemd the service is restarted
// (queued, so it goes on after this process is stopped); otherwise the
// process replaces itself with the new binary.
func restartSelf(exe, sudo string) {
	if os.Getenv("INVOCATION_ID") != "" {
		if path, err := exec.LookPath("systemctl"); err == nil {
			args := []string{"--no-block", "restart", "lpagent"}
			cmd := exec.Command(path, args...)
			if sudo != "" {
				cmd = exec.Command(sudo, append([]string{"-n", path}, args...)...)
			}
			if cmd.Run() == nil {
				return
			}
		}
		// systemd brings a service that stopped back up (Restart=always).
		os.Exit(0)
	}
	CloseTerminals()
	if err := syscall.Exec(exe, os.Args, os.Environ()); err != nil {
		log.Printf("update: restart: %v", err)
		os.Exit(0)
	}
}