121 lines
2.3 KiB
Go
121 lines
2.3 KiB
Go
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
|
|
}
|