feat: 添加安全的WebSocket连接管理，支持ping任务功能

server/task.go

@@ -3,14 +3,18 @@ package server
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"log"
+	"net"
 	"net/http"
 	"os/exec"
 	"runtime"
 	"strings"
 	"time"
 
+	ping "github.com/go-ping/ping"
 	"github.com/komari-monitor/komari-agent/cmd/flags"
+	"github.com/komari-monitor/komari-agent/ws"
 )
 
 func NewTask(task_id, command string) {
@@ -79,3 +83,79 @@ func uploadTaskResult(taskID, result string, exitCode int, finishedAt time.Time)
 		}
 	}
 }
+func icmpPing(target string, timeout time.Duration) error {
+	pinger, err := getPinger(target)
+	if err != nil {
+		return err
+	}
+	pinger.Count = 1
+	pinger.Timeout = timeout
+	pinger.SetPrivileged(true)
+	return pinger.Run()
+}
+
+func getPinger(target string) (*ping.Pinger, error) {
+	return ping.NewPinger(target)
+}
+
+func tcpPing(target string, timeout time.Duration) error {
+	if !strings.Contains(target, ":") {
+		target += ":80"
+	}
+	conn, err := net.DialTimeout("tcp", target, timeout)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	return nil
+}
+
+func httpPing(target string, timeout time.Duration) (int64, error) {
+	client := http.Client{
+		Timeout: timeout,
+	}
+	start := time.Now()
+	resp, err := client.Get(target)
+	latency := time.Since(start).Milliseconds()
+	if err != nil {
+		return 0, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode >= 200 && resp.StatusCode < 400 {
+		return latency, nil
+	}
+	return latency, errors.New("http status not ok")
+}
+
+func NewPingTask(conn *ws.SafeConn, taskID uint, pingType, pingTarget string) {
+	if taskID == 0 {
+		return
+	}
+	pingResult := 0
+	timeout := 3 * time.Second
+	switch pingType {
+	case "icmp":
+		start := time.Now()
+		if err := icmpPing(pingTarget, timeout); err == nil {
+			pingResult = int(time.Since(start).Milliseconds())
+		}
+	case "tcp":
+		start := time.Now()
+		if err := tcpPing(pingTarget, timeout); err == nil {
+			pingResult = int(time.Since(start).Milliseconds())
+		}
+	case "http":
+		if latency, err := httpPing(pingTarget, timeout); err == nil {
+			pingResult = int(latency)
+		}
+	default:
+		return
+	}
+	payload := map[string]interface{}{
+		"type":        "ping_result",
+		"task_id":     taskID,
+		"value":       pingResult,
+		"finished_at": time.Now(),
+	}
+	_ = conn.WriteJSON(payload)
+}

server/websocket.go

@@ -11,6 +11,7 @@ import (
 	"github.com/komari-monitor/komari-agent/cmd/flags"
 	"github.com/komari-monitor/komari-agent/monitoring"
 	"github.com/komari-monitor/komari-agent/terminal"
+	"github.com/komari-monitor/komari-agent/ws"
 )
 
 func EstablishWebSocketConnection() {
@@ -18,7 +19,7 @@ func EstablishWebSocketConnection() {
 	websocketEndpoint := strings.TrimSuffix(flags.Endpoint, "/") + "/api/clients/report?token=" + flags.Token
 	websocketEndpoint = "ws" + strings.TrimPrefix(websocketEndpoint, "http")
 
-	var conn *websocket.Conn
+	var conn *ws.SafeConn
 	defer func() {
 		if conn != nil {
 			conn.Close()
@@ -73,7 +74,7 @@ func EstablishWebSocketConnection() {
 	}
 }
 
-func connectWebSocket(websocketEndpoint string) (*websocket.Conn, error) {
+func connectWebSocket(websocketEndpoint string) (*ws.SafeConn, error) {
 	dialer := &websocket.Dialer{
 		HandshakeTimeout: 5 * time.Second,
 	}
@@ -85,11 +86,10 @@ func connectWebSocket(websocketEndpoint string) (*websocket.Conn, error) {
 		return nil, err
 	}
 
-	return conn, nil
+	return ws.NewSafeConn(conn), nil
 }
 
-func handleWebSocketMessages(conn *websocket.Conn, done chan<- struct{}) {
-
+func handleWebSocketMessages(conn *ws.SafeConn, done chan<- struct{}) {
 	defer close(done)
 	for {
 		_, message_raw, err := conn.ReadMessage()
@@ -104,6 +104,10 @@ func handleWebSocketMessages(conn *websocket.Conn, done chan<- struct{}) {
 			// Remote Exec
 			ExecCommand string `json:"command,omitempty"`
 			ExecTaskID  string `json:"task_id,omitempty"`
+			// Ping
+			PingTaskID uint   `json:"ping_task_id,omitempty"`
+			PingType   string `json:"ping_type,omitempty"`
+			PingTarget string `json:"ping_target,omitempty"`
 		}
 		err = json.Unmarshal(message_raw, &message)
 		if err != nil {
@@ -119,7 +123,10 @@ func handleWebSocketMessages(conn *websocket.Conn, done chan<- struct{}) {
 			go NewTask(message.ExecTaskID, message.ExecCommand)
 			continue
 		}
-
+		if message.Message == "ping" || message.PingTaskID != 0 || message.PingType != "" || message.PingTarget != "" {
+			go NewPingTask(conn, message.PingTaskID, message.PingType, message.PingTarget)
+			continue
+		}
 	}
 }