backend/watchdog/client.go

package watchdog

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"super-frpc/database"
	"super-frpc/global"
	"super-frpc/postLog"
	"super-frpc/sys"
	"super-frpc/utils"
	"sync"
	"time"
)

const socketPath = "/tmp/super-frpc-watchdog.sock"

var (
	localConn      net.Conn
	localConnMutex sync.Mutex
	recvChan       chan string
	stopRecvChan   chan struct{}
	recvWg         sync.WaitGroup
)

func Init() error {
	localConnMutex.Lock()
	if recvChan != nil || stopRecvChan != nil {
		localConnMutex.Unlock()
		return nil
	}
	if localConn != nil {
		localConnMutex.Unlock()
		return fmt.Errorf("local socket client already initialized")
	}
	localConnMutex.Unlock()

	if global.CurrentConfig.Watchdog.Enabled {
		if err := ensureWatchdogProcess(); err != nil {
			postLog.Error(fmt.Sprintf("Failed to boot watchdog program: %v", err))
			return err
		}
	}

	localConnMutex.Lock()
	recvChan = make(chan string, 100)
	stopRecvChan = make(chan struct{})
	global.Is.WatchdogConnected = false
	localConnMutex.Unlock()

	if global.CurrentConfig.Watchdog.Enabled {
		if !Connect() {
			postLog.Warning("[watchdog] initial connect failed; caller may retry with Connect")
		}
	}

	return nil
}

func syncRunningInstancesToWatchdog() {
	instances, err := database.DBListFrpcInstances()
	if err != nil {
		postLog.Error(fmt.Sprintf("[watchdog] failed to list frpc instances: %v", err))
		return
	}

	added := 0
	for _, instance := range instances {
		if err := sys.IsInstanceRunning(instance.ID); err != nil {
			postLog.Debug(fmt.Sprintf("[watchdog] instance %d is not running: %v", instance.ID, err))
			continue
		}

		serviceName, err := database.GetServiceNameByInstanceID(instance.ID)
		if err != nil {
			postLog.Warning(fmt.Sprintf("[watchdog] failed to get service name for instance %d: %v", instance.ID, err))
			continue
		}

		if !AddInstance(serviceName) {
			postLog.Warning(fmt.Sprintf("[watchdog] failed to add running instance %s to monitor list", serviceName))
			continue
		}

		added++
	}

	postLog.Info(fmt.Sprintf("[watchdog] synced %d running instance(s) to monitor list", added))
}

func ensureWatchdogProcess() error {
	watchdogName, err := getWatchdogBinaryName()
	exec.Command("chmod", "+x", filepath.Join(global.CurrentPath, watchdogName)).Run()
	if err != nil {
		return err
	}

	if isWatchdogProcessRunning(watchdogName) {
		postLog.Info(fmt.Sprintf("Watchdog process %s is already running, skipping initialize...", watchdogName))
		return nil
	}
	postLog.Info(fmt.Sprintf("Watchdog process %s is not running, starting it...", watchdogName))

	watchdogPath := filepath.Join(global.CurrentPath, watchdogName)
	if !utils.IsFileExist(watchdogPath) {
		return fmt.Errorf("watchdog binary not found: %s", watchdogPath)
	}

	cmd := exec.Command(watchdogPath)
	devNull, err := os.OpenFile(os.DevNull, os.O_RDWR, 0)
	if err != nil {
		return fmt.Errorf("failed to open null device: %w", err)
	}
	defer devNull.Close()

	cmd.Stdin = devNull
	cmd.Stdout = devNull
	cmd.Stderr = devNull
	configureWatchdogCommand(cmd)

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("failed to start watchdog process %s: %w", watchdogPath, err)
	}

	return nil
}

func getWatchdogBinaryName() (string, error) {
	switch sys.GetInitSystem() {
	case "systemd":
		return "watchdog_systemd", nil
	case "init.d":
		return "watchdog_initd", nil
	case "windows":
		if utils.IsFileExist(filepath.Join(global.CurrentPath, "watchdog", "watchdog_windows.exe")) {
			return "watchdog_windows.exe", nil
		}
		return "watchdog_windows", nil
	default:
		return "", fmt.Errorf("unsupported init system: %s", sys.GetInitSystem())
	}
}

func isWatchdogProcessRunning(watchdogName string) bool {
	switch sys.GetInitSystem() {
	case "windows":
		cmd := exec.Command("tasklist", "/FI", fmt.Sprintf("IMAGENAME eq %s", watchdogName))
		output, err := cmd.CombinedOutput()
		if err != nil {
			return false
		}
		outputStr := strings.ToLower(string(output))
		return strings.Contains(outputStr, strings.ToLower(watchdogName)) && !strings.Contains(outputStr, "no tasks are running")
	case "systemd", "init.d":
		cmd := exec.Command("ps", "-A", "-o", "args=")
		output, err := cmd.CombinedOutput()
		if err != nil {
			return false
		}

		for _, line := range strings.Split(string(output), "\n") {
			fields := strings.Fields(strings.TrimSpace(line))
			if len(fields) > 0 && filepath.Base(fields[0]) == watchdogName {
				return true
			}
		}
		return false
	default:
		return false
	}
}

func localSocketConnect() error {
	localConnMutex.Lock()
	defer localConnMutex.Unlock()

	if localConn != nil {
		return fmt.Errorf("already connected")
	}

	conn, err := net.DialTimeout("unix", socketPath, 3*time.Second)
	if err != nil {
		return fmt.Errorf("failed to connect to %s: %v", socketPath, err)
	}

	localConn = conn
	global.Is.WatchdogConnected = true

	recvWg.Add(1)
	go recvMsg()

	return nil
}

func sendMsg(message string, timeout int) (string, error) {
	localConnMutex.Lock()

	if localConn == nil {
		localConnMutex.Unlock()
		return "", fmt.Errorf("not connected")
	}

	_, err := localConn.Write([]byte(message + "\n"))
	if err != nil {
		localConnMutex.Unlock()
		return "", fmt.Errorf("failed to send message: %v", err)
	}

	localConnMutex.Unlock()

	select {
	case response := <-recvChan:
		for len(recvChan) > 0 {
			<-recvChan
		}
		return strings.TrimSpace(response), nil
	case <-time.After(time.Duration(timeout) * time.Second):
		return "", fmt.Errorf("timeout waiting for response")
	}
}

func recvMsg() {
	defer recvWg.Done()

	localConnMutex.Lock()
	if localConn == nil {
		localConnMutex.Unlock()
		return
	}
	conn := localConn
	localConnMutex.Unlock()

	reader := bufio.NewReader(conn)

	for {
		select {
		case <-stopRecvChan:
			return
		default:
			conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))

			data, err := reader.ReadBytes('\n')

			if len(data) > 0 {
				line := strings.TrimSpace(string(data))
				if len(line) > 0 {
					select {
					case recvChan <- line:
					default:
						// drop the message
					}
					if !isResponseMessage(line) {
						parseCommand(line)
					}
				}
			}

			if err != nil {
				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
					continue
				}

				localConnMutex.Lock()
				if localConn != nil {
					localConn.Close()
					localConn = nil
					global.Is.WatchdogConnected = false
				}
				localConnMutex.Unlock()
				return
			}
		}
	}
}

func isResponseMessage(msg string) bool {
	return msg == "success" || msg == "failed"
}

func Destroy() error {
	localConnMutex.Lock()

	if stopRecvChan != nil {
		close(stopRecvChan)
		stopRecvChan = nil
	}

	localConnMutex.Unlock()

	recvWg.Wait()

	localConnMutex.Lock()
	defer localConnMutex.Unlock()

	if localConn != nil {
		err := localConn.Close()
		localConn = nil
		global.Is.WatchdogConnected = false
		if err != nil {
			return fmt.Errorf("failed to close connection: %v", err)
		}
	}

	if recvChan != nil {
		close(recvChan)
		recvChan = nil
	}

	return nil
}

func IsConnected() bool {
	localConnMutex.Lock()
	defer localConnMutex.Unlock()
	return global.Is.WatchdogConnected && localConn != nil
}