backend/watchdog/tcpClient.go

package watchdog

import (
	"bufio"
	"fmt"
	"net"
	"strings"
	"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 {
	tcpConnMutex.Lock()
	defer tcpConnMutex.Unlock()

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

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

	return nil
}

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

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

	address := fmt.Sprintf("%s:%d", ipaddr, 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:
					}

					if !isResponseMessage(line) {
						postLog.Debug(fmt.Sprintf("[Watchdog] TCP Socket received message: %s", 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
}