AetherForge/agent/deploy/tunnel_manager_windows.go

//go:build windows

package deploy

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"sync"
	"syscall"
)

// TunnelKind identifies a tracked tunnel process.
type TunnelKind string

const (
	TunnelCloudflared TunnelKind = "cloudflared"
	TunnelSSHForward  TunnelKind = "ssh_forward"
)

// SSHForwardMeta describes one local port forward on the agent.
type SSHForwardMeta struct {
	LocalPort  int    `json:"local_port"`
	RemoteHost string `json:"remote_host"`
	RemotePort int    `json:"remote_port"`
	SSHUser    string `json:"ssh_user,omitempty"`
	JumpHost   string `json:"jump_host,omitempty"`
}

type tunnelEntry struct {
	kind TunnelKind
	pid  int
	meta string // JSON metadata for ssh forwards; URL for cloudflared
}

var (
	tunnelMu        sync.Mutex
	trackedTunnels  []tunnelEntry
)

// RegisterTunnelPID records a background tunnel process for later stop/status.
func RegisterTunnelPID(kind TunnelKind, pid int, meta string) {
	if pid <= 0 {
		return
	}
	tunnelMu.Lock()
	defer tunnelMu.Unlock()
	if kind == TunnelCloudflared {
		next := make([]tunnelEntry, 0, len(trackedTunnels)+1)
		for _, t := range trackedTunnels {
			if t.kind != TunnelCloudflared {
				next = append(next, t)
			}
		}
		trackedTunnels = append(next, tunnelEntry{kind: kind, pid: pid, meta: meta})
		return
	}
	trackedTunnels = append(trackedTunnels, tunnelEntry{kind: kind, pid: pid, meta: meta})
}

const processQueryLimitedInformation = 0x1000

var processAliveFn = processAliveImpl

func processAlive(pid int) bool {
	return processAliveFn(pid)
}

func processAliveImpl(pid int) bool {
	if pid <= 0 {
		return false
	}
	h, err := syscall.OpenProcess(processQueryLimitedInformation, false, uint32(pid))
	if err != nil {
		return false
	}
	defer syscall.CloseHandle(h)
	var code uint32
	err = syscall.GetExitCodeProcess(h, &code)
	if err != nil {
		return false
	}
	return code == 259 // STILL_ACTIVE
}

func killPID(pid int) error {
	if pid <= 0 {
		return fmt.Errorf("invalid pid")
	}
	return HiddenRun("taskkill", "/F", "/PID", fmt.Sprintf("%d", pid))
}

func pruneDead() {
	alive := trackedTunnels[:0]
	for _, t := range trackedTunnels {
		if processAlive(t.pid) {
			alive = append(alive, t)
		}
	}
	trackedTunnels = alive
}

// TunnelStatus returns JSON describing tracked tunnel processes.
func TunnelStatus() string {
	tunnelMu.Lock()
	pruneDead()
	snap := append([]tunnelEntry(nil), trackedTunnels...)
	tunnelMu.Unlock()

	st := TunnelStatusJSON{SSHForwards: []SSHForwardLive{}}
	for _, t := range snap {
		switch t.kind {
		case TunnelCloudflared:
			st.CloudflaredRunning = true
			st.CloudflaredURL = t.meta
			st.CloudflaredPID = t.pid
		case TunnelSSHForward:
			var meta SSHForwardMeta
			_ = json.Unmarshal([]byte(t.meta), &meta)
			st.SSHForwards = append(st.SSHForwards, SSHForwardLive{
				LocalPort:  meta.LocalPort,
				RemoteHost: meta.RemoteHost,
				RemotePort: meta.RemotePort,
				SSHUser:    meta.SSHUser,
				JumpHost:   meta.JumpHost,
				PID:        t.pid,
				Running:    true,
			})
		}
	}
	b, _ := json.Marshal(st)
	return string(b)
}

// StopTunnels stops tracked tunnel kinds. Empty kinds stops all.
func StopTunnels(kinds ...TunnelKind) (stopped int, msgs []string) {
	tunnelMu.Lock()
	defer tunnelMu.Unlock()
	pruneDead()

	wantAll := len(kinds) == 0
	want := map[TunnelKind]bool{}
	for _, k := range kinds {
		want[k] = true
	}

	remaining := trackedTunnels[:0]
	for _, t := range trackedTunnels {
		if wantAll || want[t.kind] {
			if err := killPID(t.pid); err != nil {
				msgs = append(msgs, fmt.Sprintf("%s pid %d: %v", t.kind, t.pid, err))
			} else {
				stopped++
				msgs = append(msgs, fmt.Sprintf("stopped %s pid %d", t.kind, t.pid))
			}
			continue
		}
		remaining = append(remaining, t)
	}
	trackedTunnels = remaining
	return stopped, msgs
}

// ResetTrackedTunnels clears registry without killing (tests only).
func ResetTrackedTunnels() {
	tunnelMu.Lock()
	trackedTunnels = nil
	tunnelMu.Unlock()
}

// CloudflaredTargetFromEnv returns trimmed URL or empty.
func CloudflaredTargetFromEnv() string {
	return strings.TrimSpace(os.Getenv("AETHERFORGE_TUNNEL_URL"))
}