feat(update): keep the monitoring agent correct after fly update

.github/workflows/ci.yml

@@ -33,5 +33,11 @@ jobs:
     - name: Run the checks
       run: make check
 
+    # make check runs as the runner user, so the tests that need root skip.
+    # They cover "sudo fly update": the owner of the binary, and a link that
+    # a hostile owner of the binary folder puts in place of the new file.
+    - name: Run the tests that need root
+      run: sudo env "PATH=$PATH" "GOCACHE=$(go env GOCACHE)" "GOMODCACHE=$(go env GOMODCACHE)" go test -count=1 -run 'TestReplaceBinary' ./internal/release/
+
     - name: Build the release archives
       run: make release

cmd/version.go

@@ -1,12 +1,16 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 
 	"github.com/flywp/server-cli/internal/release"
+	"github.com/flywp/server-cli/internal/service"
 	"github.com/flywp/server-cli/internal/version"
+	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
 )
 
@@ -25,6 +29,9 @@ var versionCmd = &cobra.Command{
 var updateCmd = &cobra.Command{
 	Use:   "update",
 	Short: "Update fly-cli to the latest version",
+	Long: `Update fly-cli to the latest release. When fly already runs the latest
+release, the command does nothing. On a server with the monitoring agent, the
+command also restarts the agent, so that the agent runs the new binary.`,
 	RunE: func(cmd *cobra.Command, args []string) error {
 		if os.Geteuid() != 0 {
 			return errors.New("the update command must be run as root, please run 'sudo fly update'")
@@ -41,17 +48,17 @@ var updateCmd = &cobra.Command{
 			fmt.Printf("This is not a release build (version %s). Latest release: %s\n", version.Version, latest)
 		case !update.Available:
 			fmt.Println("You are already running the latest version.")
-			return nil
+			return restartStaleAgent(cmd.Context())
 		default:
 			fmt.Printf("New version available: %s\n", latest)
 		}
 
 		if !yesFlag {
-			fmt.Printf("Do you want to install %s? (y/n): ", latest)
-			var response string
-			// An empty or unreadable answer cancels the update.
-			_, _ = fmt.Scanln(&response)
-			if response != "y" && response != "Y" {
+			ok, err := confirm(os.Stdin, os.Stdout, latest)
+			if err != nil {
+				return err
+			}
+			if !ok {
 				fmt.Println("Update cancelled.")
 				return nil
 			}
@@ -61,12 +68,61 @@ var updateCmd = &cobra.Command{
 		if err := release.SelfUpdate(cmd.Context(), update.Release); err != nil {
 			return fmt.Errorf("updating: %w", err)
 		}
-
 		fmt.Printf("Updated to %s.\n", latest)
-		return nil
+
+		return restartAgent(cmd.Context())
 	},
 }
 
+// confirm asks whether to install latest. Without a terminal nobody can
+// answer, so it returns an error: a script must not read "cancelled" as done.
+func confirm(in *os.File, out io.Writer, latest string) (bool, error) {
+	if !isatty.IsTerminal(in.Fd()) && !isatty.IsCygwinTerminal(in.Fd()) {
+		return false, errors.New("there is no terminal to confirm the update: run 'sudo fly update --yes'")
+	}
+
+	_, _ = fmt.Fprintf(out, "Do you want to install %s? (y/n): ", latest)
+	var response string
+	// An empty or unreadable answer cancels the update.
+	_, _ = fmt.Fscanln(in, &response)
+
+	return response == "y" || response == "Y", nil
+}
+
+// restartAgent restarts the monitoring agent, if the server has it, so that it
+// runs the new binary.
+func restartAgent(ctx context.Context) error {
+	if !service.Installed() {
+		return nil
+	}
+
+	fmt.Println("Restarting the monitoring agent...")
+	if err := service.Restart(ctx); err != nil {
+		return fmt.Errorf("the update is installed, but the monitoring agent did not restart: %w", err)
+	}
+
+	return nil
+}
+
+// restartStaleAgent restarts the monitoring agent when it still runs a binary
+// that an earlier update replaced.
+func restartStaleAgent(ctx context.Context) error {
+	if !service.Installed() {
+		return nil
+	}
+
+	stale, err := service.Stale(ctx)
+	if err != nil {
+		return fmt.Errorf("checking the monitoring agent: %w", err)
+	}
+	if !stale {
+		return nil
+	}
+
+	fmt.Println("The monitoring agent runs an older binary.")
+	return restartAgent(ctx)
+}
+
 func init() {
 	updateCmd.Flags().BoolVarP(&yesFlag, "yes", "y", false, "Automatically answer yes to update confirmation")
 	rootCmd.AddCommand(versionCmd)

cmd/version_test.go

@@ -0,0 +1,117 @@
+package cmd
+
+import (
+	"context"
+	"io"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/flywp/server-cli/internal/service"
+)
+
+func TestConfirmNeedsATerminal(t *testing.T) {
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = r.Close() }()
+	_, _ = w.WriteString("y\n")
+	_ = w.Close()
+
+	// A script that pipes "y" still needs --yes: the pipe is not a terminal.
+	ok, err := confirm(r, io.Discard, "v0.2.0")
+	if ok || err == nil || !strings.Contains(err.Error(), "--yes") {
+		t.Errorf("confirm() = %v, %v; want an error that tells to use --yes", ok, err)
+	}
+}
+
+// fakeAgentService makes the agent unit exist and puts a fake systemctl in
+// PATH. It returns the log of the systemctl calls.
+func fakeAgentService(t *testing.T, installed bool) string {
+	return fakeAgentServiceExit(t, installed, 0)
+}
+
+// fakeAgentServiceExit is fakeAgentService with a systemctl that exits with
+// exit.
+func fakeAgentServiceExit(t *testing.T, installed bool, exit int) string {
+	t.Helper()
+
+	unit := filepath.Join(t.TempDir(), "fly-agent.service")
+	if installed {
+		if err := os.WriteFile(unit, []byte("[Unit]\n"), 0o644); err != nil {
+			t.Fatal(err)
+		}
+	}
+	old := service.UnitPath
+	service.UnitPath = unit
+	t.Cleanup(func() { service.UnitPath = old })
+
+	dir := t.TempDir()
+	log := filepath.Join(t.TempDir(), "calls")
+	script := "#!/bin/sh\nprintf '%s\\n' \"$*\" >> " + log + "\n[ \"$1\" = show ] && echo 0\n[ " + strconv.Itoa(exit) + " -ne 0 ] && echo 'Failed to connect to bus' >&2\nexit " + strconv.Itoa(exit) + "\n"
+	if err := os.WriteFile(filepath.Join(dir, "systemctl"), []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH"))
+
+	return log
+}
+
+func systemctlCalls(t *testing.T, log string) string {
+	t.Helper()
+	data, err := os.ReadFile(log)
+	if err != nil && !os.IsNotExist(err) {
+		t.Fatal(err)
+	}
+	return strings.TrimSpace(string(data))
+}
+
+func TestRestartAgentAfterAnUpdate(t *testing.T) {
+	log := fakeAgentService(t, true)
+	if err := restartAgent(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if got := systemctlCalls(t, log); got != "try-restart fly-agent" {
+		t.Errorf("systemctl calls = %q, want try-restart fly-agent", got)
+	}
+}
+
+func TestNoRestartWithoutTheAgent(t *testing.T) {
+	log := fakeAgentService(t, false)
+	if err := restartAgent(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if err := restartStaleAgent(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if got := systemctlCalls(t, log); got != "" {
+		t.Errorf("systemctl calls = %q, want none on a server without the agent", got)
+	}
+}
+
+func TestNoRestartWhenTheAgentDoesNotRun(t *testing.T) {
+	// The fake systemctl shows MainPID 0: the agent does not run, so systemd
+	// starts the binary on the disk.
+	log := fakeAgentService(t, true)
+	if err := restartStaleAgent(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if got := systemctlCalls(t, log); got != "show --property=MainPID --value fly-agent" {
+		t.Errorf("systemctl calls = %q, want only the check", got)
+	}
+}
+
+func TestASystemctlFailureShowsItsMessage(t *testing.T) {
+	fakeAgentServiceExit(t, true, 1)
+
+	for name, f := range map[string]func(context.Context) error{"restartAgent": restartAgent, "restartStaleAgent": restartStaleAgent} {
+		var stderr strings.Builder
+		code := exitCode(f(context.Background()), &stderr)
+		if code != 1 || !strings.Contains(stderr.String(), "Failed to connect to bus") {
+			t.Errorf("%s: exit %d, stderr %q; want exit 1 and the systemctl message", name, code, stderr.String())
+		}
+	}
+}

go.mod

@@ -6,6 +6,7 @@ toolchain go1.27.1
 
 require (
 	github.com/fatih/color v1.19.0
+	github.com/mattn/go-isatty v0.0.24
 	github.com/oklog/ulid/v2 v2.1.2
 	github.com/spf13/cobra v1.10.2
 	golang.org/x/mod v0.41.0
@@ -16,6 +17,5 @@ require (
 require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/mattn/go-colorable v0.1.15 // indirect
-	github.com/mattn/go-isatty v0.0.24 // indirect
 	github.com/spf13/pflag v1.0.10 // indirect
 )

internal/release/release.go

@@ -16,6 +16,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"runtime"
+	"syscall"
 	"time"
 
 	"github.com/flywp/server-cli/internal/version"
@@ -226,6 +227,10 @@ func writeBinary(exe string, r io.Reader) (err error) {
 		_ = tmp.Close()
 		return fmt.Errorf("making binary executable: %w", err)
 	}
+	if err = keepOwner(tmp, exe); err != nil {
+		_ = tmp.Close()
+		return fmt.Errorf("keeping the owner of the binary: %w", err)
+	}
 	// Put the binary on the disk before the rename: after a power loss, a
 	// renamed but empty binary would not start.
 	if err = tmp.Sync(); err != nil {
@@ -247,3 +252,21 @@ func writeBinary(exe string, r io.Reader) (err error) {
 
 	return nil
 }
+
+// keepOwner gives the open file f the owner of exe. "sudo fly update" then
+// keeps the binary of the agent with the server user, not with root. Only
+// root can give a file to a different user; for other users the owner is
+// already correct. It changes the open file, not a path: a path in a
+// directory of an other user can be replaced by a link to a root file.
+func keepOwner(f *os.File, exe string) error {
+	info, err := os.Stat(exe)
+	if err != nil {
+		return nil
+	}
+	st, ok := info.Sys().(*syscall.Stat_t)
+	if !ok || os.Geteuid() != 0 {
+		return nil
+	}
+
+	return f.Chown(int(st.Uid), int(st.Gid))
+}