feat(watchdog): implement keepalive mechanism and enhance connection handling
This commit is contained in:
@@ -0,0 +1,284 @@
|
||||
package watchdog
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"super-frpc/global"
|
||||
"super-frpc/postLog"
|
||||
"super-frpc/sys"
|
||||
"super-frpc/utils"
|
||||
"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()
|
||||
if recvChan != nil || stopRecvChan != nil {
|
||||
tcpConnMutex.Unlock()
|
||||
return nil
|
||||
}
|
||||
if tcpConn != nil {
|
||||
tcpConnMutex.Unlock()
|
||||
return fmt.Errorf("TCP client already initialized")
|
||||
}
|
||||
tcpConnMutex.Unlock()
|
||||
|
||||
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()
|
||||
|
||||
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", "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 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
|
||||
}
|
||||
Reference in New Issue
Block a user