daemon.go

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"time"

	"wanctl/internal/client"
	"wanctl/internal/config"
	"wanctl/internal/transport"
)

// cmdSupervise is the restart loop used by the Windows Scheduled Task, whose
// native ONLOGON trigger does not restart a process that exits later. Each
// iteration executes the stable binary path again, so a replacement is picked
// up on the next child start.
func cmdSupervise(ctx context.Context, args []string) error {
	self, err := selfPath()
	if err != nil {
		return err
	}
	// The task runs this under a headless console, so anything printed to it
	// is lost. Write where `wanctl start` and the launchd agent do instead.
	logPath, err := config.LogPath()
	if err != nil {
		return err
	}
	logf, err := os.OpenFile(logPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("open log: %w", err)
	}
	defer logf.Close()
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	exitWithParent(cancel)
	return superviseLoop(ctx, self, append([]string{"agent", "--managed"}, args...), logf)
}

func superviseLoop(ctx context.Context, self string, agentArgs []string, out io.Writer) error {
	for {
		cmd := exec.CommandContext(ctx, self, agentArgs...)
		cmd.Stdout = out
		cmd.Stderr = out
		err := cmd.Run()
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if err != nil {
			fmt.Fprintf(out, "wanctl supervisor: agent exited: %v; restarting in 3s\n", err)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(3 * time.Second):
		}
	}
}

// agentRunning answers "is an agent serving this config dir" from the lock the
// agent holds, not from whether the number in agent.pid happens to name a live
// process. It also clears a pid file that provably belongs to nobody.
//
// The cleanup is deliberately conditional. A pid file whose process is gone is
// unambiguously stale. A pid file whose process is alive while the lock is free
// is either pid reuse or an agent that was just launched and has not reached
// AcquireAgentLock yet -- cmdStart and the Windows update handover both record
// the pid on the new agent's behalf before it locks -- and deleting that file
// would hide a starting agent from `wanctl stop` for good. Reporting it as not
// running is safe in both readings; deleting it is not.
func agentRunning() (int, bool) {
	pid, running := config.AgentRunning()
	if !running && pid > 0 && !processAlive(pid) {
		_ = config.RemovePID()
	}
	return pid, running
}

// cmdStart makes this machine a controlled device: it logs in if there is no
// token yet, then runs the agent detached in the background and records its pid.
//
// The login step lives here rather than in a separate entrypoint because
// enrolling a device and running its agent are one intention. Bare `wanctl`
// used to do both, which meant that someone on a controller-only machine who
// typed `wanctl` to see what it does had their machine registered as a
// controlled device; it prints help now and does nothing.
func cmdStart(ctx context.Context) error {
	if pid, running := agentRunning(); running {
		fmt.Printf("wanctl 服务已在运行 (pid %d)。停止用: wanctl stop\n", pid)
		return nil
	}
	if config.EnvOr("WANCTL_TOKEN", config.StoredToken()) == "" {
		if err := ensureEndpointsConfigured(); err != nil {
			return err
		}
		t, err := enrollForStart(ctx)
		if err != nil {
			return err
		}
		if err := config.SaveToken(t); err != nil {
			return fmt.Errorf("save token: %w", err)
		}
	}
	self, err := selfPath()
	if err != nil {
		return err
	}
	logPath, err := config.LogPath()
	if err != nil {
		return err
	}
	logf, err := os.OpenFile(logPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("open log: %w", err)
	}
	defer logf.Close()

	// The child reads token from <cfg>/token and relay/transport from compile-time
	// defaults, so no flags are needed. detachSysProcAttr() detaches it from this
	// terminal so it survives the parent exiting.
	cmd := selfCommand(self, "agent")
	cmd.Stdout = logf
	cmd.Stderr = logf
	cmd.SysProcAttr = detachSysProcAttr()
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("start agent: %w", err)
	}
	pid := cmd.Process.Pid // capture before Release() zeroes it
	if err := config.WritePID(pid); err != nil {
		return fmt.Errorf("write pid: %w", err)
	}
	_ = cmd.Process.Release() // detach: don't reap; it runs until `wanctl stop`
	fmt.Printf("✓ 服务已转后台 (pid %d)，日志: %s\n  停止用: wanctl stop\n", pid, logPath)
	return nil
}

// enrollForStart is the browser login `wanctl start` performs on a device that
// has no token. It is a variable so a test can drive the rest of start without
// a portal.
var enrollForStart = enroll

// localStatusLine is the one line bare `wanctl` adds under the help text: what
// this machine is right now, so that someone reading the help knows which half
// of it applies to them. It only reads -- printing help must change nothing.
func localStatusLine() string {
	credential := "未登录"
	if config.EnvOr("WANCTL_TOKEN", config.StoredToken()) != "" {
		credential = "已登录"
	}
	agent := "agent 未运行"
	if pid, running := config.AgentRunning(); running {
		agent = "agent 运行中"
		if pid > 0 {
			agent = fmt.Sprintf("agent 运行中 (pid %d)", pid)
		}
	}
	return fmt.Sprintf("本机: %s · %s", credential, agent)
}

// terminateAgent is terminatePID behind a variable so a test can drive the
// stop-then-start sequence without signalling a real process.
var terminateAgent = terminatePID

// agentStopTimeout bounds the wait for a signalled agent to release the
// config-dir lock. An agent shutting down closes relay connections and finishes
// whatever request it was serving first; ten seconds is far longer than that
// takes and still short enough to report rather than hang.
// Variables, not constants, so a test can exercise the give-up path without
// waiting out the real deadline.
var (
	agentStopTimeout = 10 * time.Second
	agentStopPoll    = 50 * time.Millisecond
)

// cmdStop terminates the background agent and waits for it to let go.
//
// It refuses to signal a pid whose lock nobody holds. Trusting the pid file
// alone meant that a stale one, plus the pid reuse that follows sooner or
// later, made `wanctl stop` (and the stop that `wanctl update` performs on its
// way to a restart) kill whatever process had inherited the number.
//
// Returning as soon as the signal was delivered was the other half of the same
// bug. Terminating is asynchronous: `wanctl update` stopped the old agent,
// swapped the binary and started a new one while the old process was still
// shutting down and still holding the lock, so the new agent could not acquire
// it and exited. The parent had already printed "✓ 服务已转后台" and nothing
// restarted it -- a device dropped off the relay for fifty minutes that way.
func cmdStop() error {
	recorded := config.ReadPID()
	pid, running := agentRunning()
	if !running {
		if recorded > 0 {
			fmt.Printf("wanctl 服务未在运行（agent.pid 里记的 %d 已失效%s）\n", recorded, staleNote(recorded))
			return nil
		}
		fmt.Println("wanctl 服务未在运行")
		return nil
	}
	if pid <= 0 {
		fmt.Println("wanctl 服务在运行,但没有记录 pid,无法从这里停止它。请停止启动它的那个服务(计划任务/systemd/launchd)")
		return nil
	}
	if err := terminateAgent(pid); err != nil {
		return fmt.Errorf("stop pid %d: %w", pid, err)
	}
	if !awaitAgentLockRelease(config.AgentRunning, agentStopTimeout, agentStopPoll, time.Sleep) {
		return fmt.Errorf("已向 agent (pid %d) 发出停止信号,但它在 %s 内没有释放 %s 的锁 —— 没有重启它,以免两个 agent 抢同一个配置目录。等它退出后再跑 wanctl start",
			pid, agentStopTimeout, configDirForMessage())
	}
	_ = config.RemovePID()
	fmt.Printf("✓ 已停止 wanctl 服务 (pid %d)\n", pid)
	return nil
}

// awaitAgentLockRelease polls until nobody holds the config-dir lock, and
// reports whether that happened before the deadline. The clock and the probe
// are injected so the decision is testable without a real agent.
func awaitAgentLockRelease(running func() (int, bool), timeout, poll time.Duration, sleep func(time.Duration)) bool {
	for waited := time.Duration(0); ; waited += poll {
		if _, held := running(); !held {
			return true
		}
		if waited >= timeout {
			return false
		}
		sleep(poll)
	}
}

// configDirForMessage names the directory in an error, or says nothing useful
// rather than failing when it cannot be resolved.
func configDirForMessage() string {
	if dir, err := transport.ConfigDir(); err == nil {
		return dir
	}
	return "配置目录"
}

// staleNote distinguishes the two ways a pid file outlives its agent, because
// the second one is the one that used to make wanctl kill a stranger.
func staleNote(pid int) string {
	if processAlive(pid) {
		return "，该 pid 现在属于别的进程，不会去动它"
	}
	return "，已清理"
}

// cmdStatus reports either the local background agent or a specified remote
// agent. Keeping argument parsing here prevents unknown status arguments from
// being silently discarded by the top-level dispatcher.
func cmdStatus(ctx context.Context, args []string) error {
	target, err := parseStatusArgs(args)
	if err != nil {
		return err
	}
	if target != "" {
		c, err := client.New()
		if err != nil {
			return err
		}
		status, err := c.Status(ctx, target)
		if err != nil {
			return fmt.Errorf("remote device %q: %w", target, err)
		}
		out, flush := deviceOutput(os.Stdout)
		defer flush()
		fmt.Fprintf(out, "● 远端设备 %s 在线，agent 运行中\n", target)
		if !status.Detailed {
			fmt.Fprintln(out, "  详情: 该设备版本较旧，无法报告 policy mode 或 agent 版本")
			return nil
		}
		fmt.Fprintf(out, "  Policy mode: %s\n", status.Mode)
		fmt.Fprintf(out, "  Agent 版本: %s\n", status.Version)
		return nil
	}
	return printLocalStatus()
}

func parseStatusArgs(args []string) (string, error) {
	fs := withHelp(flag.NewFlagSet("status", flag.ContinueOnError))
	fs.SetOutput(io.Discard)
	target := fs.String("target", "", "device (NS/DEV or DEV)")
	if err := fs.Parse(args); err != nil {
		return "", err
	}
	if fs.NArg() != 0 {
		return "", fmt.Errorf("usage: wanctl status [-target NS/DEV]")
	}
	if len(args) != 0 && strings.TrimSpace(*target) == "" {
		return "", fmt.Errorf("status target must not be empty")
	}
	return *target, nil
}

func printLocalStatus() error {
	fmt.Println("本机 agent:")
	if pid, running := agentRunning(); running {
		fmt.Printf("● 运行中 (pid %d)\n", pid)
		// Worth saying before an upgrade rather than after: `wanctl update` can
		// swap the binary here but cannot restart this agent, so the old build
		// keeps serving until whoever owns it restarts the service.
		if !canTerminatePID(pid) {
			fmt.Println("  托管: 由另一个账户运行（supervisor），本用户无法停止或重启它")
			fmt.Println("        升级后需以管理员/root 重启该服务才能生效")
		}
	} else {
		fmt.Println("○ 未运行（运行 `wanctl start` 启动）")
	}
	// The endpoints, so a front-end that cannot run `wanctl config` (the
	// Android app) can still show which instance this device belongs to.
	for _, k := range []string{"relay", "portal"} {
		if v, _ := config.Setting(k); v != "" {
			fmt.Printf("  %s: %s\n", k, v)
		} else {
			fmt.Printf("  %s: 未配置\n", k)
		}
	}
	// Whether this device keeps itself current, said where someone chasing a
	// version mismatch will look first.
	if config.AutoUpdateEnabled() {
		fmt.Println("  自动更新: 开启")
	} else {
		fmt.Println("  自动更新: 关闭")
	}
	if config.StoredToken() != "" {
		fmt.Println("  凭证: 已登录")
	} else if os.Getenv("WANCTL_TOKEN") != "" {
		fmt.Println("  凭证: 来自 WANCTL_TOKEN 环境变量")
	} else {
		fmt.Println("  凭证: 未登录")
	}
	return nil
}

// cmdLogout stops the agent and clears the stored token.
func cmdLogout() error {
	_ = cmdStop()
	if err := config.ClearToken(); err != nil {
		return err
	}
	fmt.Println("✓ 已登出（已清除本地凭证）")
	return nil
}