- Implement TCP client for watchdog service communication - Add watchdog configuration in config.json - Update main.go to initialize and connect to watchdog - Add watchdog related functions (connect, command handling) - Update README.md with new watchdog feature - Improve route setup logging in router.go
165 lines
2.8 KiB
Go
165 lines
2.8 KiB
Go
package watchdog
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"net"
|
|
"super-frpc/postLog"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
tcpConn net.Conn
|
|
tcpConnMutex sync.Mutex
|
|
isConnected bool
|
|
recvChan chan string
|
|
stopRecvChan chan struct{}
|
|
)
|
|
|
|
func notifyMessage(message string) {
|
|
postLog.Info(fmt.Sprintf("%s", message))
|
|
}
|
|
|
|
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
|
|
|
|
go recvMsg()
|
|
|
|
return nil
|
|
}
|
|
|
|
func sendMsg(message string, target 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:
|
|
return response, nil
|
|
case <-time.After(time.Duration(target) * time.Second):
|
|
return "", fmt.Errorf("timeout waiting for response")
|
|
}
|
|
}
|
|
|
|
func recvMsg() {
|
|
reader := bufio.NewReader(tcpConn)
|
|
for {
|
|
select {
|
|
case <-stopRecvChan:
|
|
return
|
|
default:
|
|
tcpConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
|
|
line, err := reader.ReadString('\n')
|
|
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
|
|
}
|
|
|
|
line = line[:len(line)-1]
|
|
if len(line) > 0 {
|
|
select {
|
|
case recvChan <- line:
|
|
default:
|
|
select {
|
|
case <-recvChan:
|
|
recvChan <- line
|
|
default:
|
|
}
|
|
}
|
|
|
|
if len(line) > 0 && line != "success" && !isResponseMessage(line) {
|
|
notifyMessage(line)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func isResponseMessage(msg string) bool {
|
|
return msg == "success" || msg == "failed"
|
|
}
|
|
|
|
|
|
|
|
func Destroy() error {
|
|
tcpConnMutex.Lock()
|
|
defer tcpConnMutex.Unlock()
|
|
|
|
if stopRecvChan != nil {
|
|
close(stopRecvChan)
|
|
stopRecvChan = nil
|
|
}
|
|
|
|
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
|
|
}
|