refactor(socket): update socket handling to use a local socket and improve connection management

2026-05-10 14:39:17 +08:00
parent 65f31f28f6
commit 751e6afa83
4 changed files with 124 additions and 121 deletions
@@ -1,87 +1,120 @@
-package socket
-
-import (
-	"Watchdog_Linux-systemd/postLog"
-	"bufio"
-	"fmt"
-	"net"
-	"strings"
-)
-
-var (
-	Conn           net.Conn
-	CommandHandler func(string) error
-)
-
-func BootSocket(networkType, listenAddr string, listenPort int) error {
-	listen, err := net.Listen(networkType, fmt.Sprintf("%s:%d", listenAddr, listenPort))
-	if err != nil {
-		postLog.Fatal(fmt.Sprintf("[Socket] Failed to listen: %v, err: %v, %v", listenAddr, listenPort, err))
-		return fmt.Errorf("failed to listen: %v, err: %v, %v", listenAddr, listenPort, err)
-	}
-	defer listen.Close()
-
-	postLog.Info(fmt.Sprintf("Server is running on %s:%d", listenAddr, listenPort))
-
-	for {
-		Conn, err = listen.Accept()
-		if err != nil {
-			postLog.Error(fmt.Sprintf("Failed to accept: %v, err: %v", Conn, err))
-		}
-		go handleRequest()
-	}
-}
-
-func handleRequest() {
-	defer Conn.Close()
-
-	reader := bufio.NewReader(Conn)
-
-	for {
-		data, err := reader.ReadBytes('\n')
-		if err != nil {
-			return
-		}
-
-		recvMsg := strings.TrimSpace(string(data))
-
-		responseMsg := ""
-		if len(recvMsg) != 0 {
-			postLog.Debug(fmt.Sprintf("Received message: %s", recvMsg))
-			if recvMsg == "watchdogAgentConnectionTest" {
-				responseMsg = "success"
-			} else {
-				if CommandHandler != nil {
-					err := CommandHandler(recvMsg)
-					if err != nil {
-						responseMsg = fmt.Sprintf("error: %v", err)
-					} else {
-						responseMsg = "success"
-					}
-				} else {
-					responseMsg = "error: command handler not initialized"
-				}
-			}
-		}
-
-		Conn.Write([]byte(responseMsg + "\n"))
-	}
-}
-
-func SendMsg(msg string) error {
-	if Conn == nil {
-		return fmt.Errorf("connection is nil")
-	}
-
-	data := []byte(msg + "\n")
-	n, err := Conn.Write(data)
-	if err != nil {
-		return fmt.Errorf("failed to write message: %v", err)
-	}
-
-	if n != len(data) {
-		return fmt.Errorf("incomplete write: wrote %d bytes out of %d", n, len(data))
-	}
-
-	return nil
-}
+package socket
+
+import (
+	"Watchdog_Linux-systemd/postLog"
+	"bufio"
+	"fmt"
+	"net"
+	"os"
+	"strings"
+	"sync"
+)
+
+const SocketPath = "/tmp/super-frpc-watchdog.sock"
+
+var (
+	Conn           net.Conn
+	connMutex      sync.Mutex
+	CommandHandler func(string) error
+)
+
+func BootSocket() error {
+	if err := os.RemoveAll(SocketPath); err != nil {
+		return fmt.Errorf("failed to remove stale socket %s: %w", SocketPath, err)
+	}
+
+	listen, err := net.Listen("unix", SocketPath)
+	if err != nil {
+		postLog.Fatal(fmt.Sprintf("[Socket] Failed to listen on %s: %v", SocketPath, err))
+		return fmt.Errorf("failed to listen on %s: %w", SocketPath, err)
+	}
+	defer func() {
+		listen.Close()
+		os.Remove(SocketPath)
+	}()
+
+	postLog.Info(fmt.Sprintf("Server is running on local socket %s", SocketPath))
+
+	for {
+		conn, err := listen.Accept()
+		if err != nil {
+			postLog.Error(fmt.Sprintf("Failed to accept local socket connection: %v", err))
+			continue
+		}
+		connMutex.Lock()
+		if Conn != nil {
+			Conn.Close()
+		}
+		Conn = conn
+		connMutex.Unlock()
+
+		go handleRequest(conn)
+	}
+}
+
+func handleRequest(conn net.Conn) {
+	defer func() {
+		conn.Close()
+		connMutex.Lock()
+		if Conn == conn {
+			Conn = nil
+		}
+		connMutex.Unlock()
+	}()
+
+	reader := bufio.NewReader(conn)
+
+	for {
+		data, err := reader.ReadBytes('\n')
+		if err != nil {
+			return
+		}
+
+		recvMsg := strings.TrimSpace(string(data))
+
+		responseMsg := ""
+		if len(recvMsg) != 0 {
+			postLog.Debug(fmt.Sprintf("Received message: %s", recvMsg))
+			if recvMsg == "watchdogAgentConnectionTest" {
+				responseMsg = "success"
+			} else {
+				if CommandHandler != nil {
+					err := CommandHandler(recvMsg)
+					if err != nil {
+						responseMsg = fmt.Sprintf("error: %v", err)
+					} else {
+						responseMsg = "success"
+					}
+				} else {
+					responseMsg = "error: command handler not initialized"
+				}
+			}
+		}
+
+		if _, err := conn.Write([]byte(responseMsg + "\n")); err != nil {
+			return
+		}
+	}
+}
+
+func SendMsg(msg string) error {
+	connMutex.Lock()
+	conn := Conn
+	connMutex.Unlock()
+
+	if conn == nil {
+		return fmt.Errorf("connection is nil")
+	}
+
+	data := []byte(msg + "\n")
+	n, err := conn.Write(data)
+	if err != nil {
+		return fmt.Errorf("failed to write message: %v", err)
+	}
+
+	if n != len(data) {
+		return fmt.Errorf("incomplete write: wrote %d bytes out of %d", n, len(data))
+	}
+
+	return nil
+}