backend/watchdog/tcpClient.go

package watchdog

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

var (
	tcpConn      net.Conn
	tcpConnMutex sync.Mutex
	isConnected  bool
	recvChan     chan string
	stopRecvChan chan struct{}
	recvWg       sync.WaitGroup
)

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

	tcpConnMutex.Lock()
	defer tcpConnMutex.Unlock()

	if recvChan != nil || stopRecvChan != nil {
		return nil
	}

	if tcpConn != nil {
		return fmt.Errorf("TCP client already initialized")
	}

	recvChan = make(chan string, 100)
	stopRecvChan = make(chan struct{})
	isConnected = false

	return nil
}

func ensureWatchdogProcess() error {
	watchdogName, err := getWatchdogBinaryName()
	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(".", "watchdog", 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(".", "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", "comm=")
		output, err := cmd.CombinedOutput()
		if err != nil {
			return false
		}

		for _, line := range strings.Split(string(output), "\n") {
			if strings.TrimSpace(line) == watchdogName {
				return true
			}
		}
		return false
	default:
		return false
	}
}

func tcpConnect(ipaddr string, port int) error {
	tcpConnMutex.Lock()
	defer tcpConnMutex.Unlock()

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

	address := net.JoinHostPort(ipaddr, fmt.Sprintf("%d", port))
	conn, err := net.DialTimeout("tcp", address, 3*time.Second)
	if err != nil {
		return fmt.Errorf("failed to connect to %s: %v", address, err)
	}

	tcpConn = conn
	isConnected = true

	recvWg.Add(1)
	go recvMsg()

	return nil
}

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

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

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

	tcpConnMutex.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()

	tcpConnMutex.Lock()
	if tcpConn == nil {
		tcpConnMutex.Unlock()
		return
	}
	conn := tcpConn
	tcpConnMutex.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
					}
					// Here add logic to handle the message
					if !isResponseMessage(line) {
						// postLog.Debug(fmt.Sprintf("[Watchdog] TCP Socket received message: %s", line))
						parseCommand(line)
					}
				}
			}

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

				tcpConnMutex.Lock()
				if tcpConn != nil {
					tcpConn.Close()
					tcpConn = nil
					isConnected = false
				}
				tcpConnMutex.Unlock()
				return
			}
		}
	}
}

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

func Destroy() error {
	tcpConnMutex.Lock()

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

	tcpConnMutex.Unlock()

	recvWg.Wait()

	tcpConnMutex.Lock()
	defer tcpConnMutex.Unlock()

	if tcpConn != nil {
		err := tcpConn.Close()
		tcpConn = nil
		isConnected = 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 {
	tcpConnMutex.Lock()
	defer tcpConnMutex.Unlock()
	return isConnected && tcpConn != nil
}