[优化] 增强 WebSocket 处理逻辑，添加上下文取消支持和关闭机制

[优化] 重构 WebSocket 处理逻辑，添加消息处理接口和心跳机制

openflared/internal/flared/runner.go

@@ -49,31 +49,31 @@ func (r *Runner) Run(ctx context.Context) error {
 	}
 }
 
-func (r *Runner) handleConnection(ctx context.Context, conn *wsclient.Connection) {
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		default:
-		}
+type flaredWSHandler struct {
+	runner *Runner
+}
 
-		msg, err := conn.Receive()
-		if err != nil {
-			slog.Error("flared ws receive failed", "error", err)
-			return
-		}
+func (h *flaredWSHandler) OnConnect(ctx context.Context) error {
+	return nil
+}
 
-		switch msg.Type {
-		case "ping":
-			_ = conn.SendPong()
-		case "active_config":
-			// Server notifies there is a new config available
-			slog.Info("received config update notification from server")
-			r.SyncService.Trigger()
-		default:
-			slog.Debug("ignored unknown ws message type", "type", msg.Type)
-		}
+func (h *flaredWSHandler) HandleMessage(ctx context.Context, msg wsclient.WSMessage) error {
+	switch msg.Type {
+	case "active_config":
+		slog.Info("received config update notification from server")
+		h.runner.SyncService.Trigger()
+	default:
+		slog.Debug("ignored unknown ws message type", "type", msg.Type)
 	}
+	return nil
+}
+
+func (h *flaredWSHandler) OnClose(err error) {
+	slog.Error("flared ws receive failed", "error", err)
+}
+
+func (r *Runner) handleConnection(ctx context.Context, conn *wsclient.Connection) {
+	_ = conn.RunReceiveLoop(ctx, &flaredWSHandler{runner: r})
 }
 
 func (r *Runner) sleepContext(ctx context.Context, d time.Duration) {

openflared/internal/wsclient/client.go

@@ -9,6 +9,9 @@ import (
 	shared "openflare/utils/wsclient"
 )
 
+type WSMessage = shared.WSMessage
+type MessageHandler = shared.MessageHandler
+
 type Client struct {
 	sharedClient *shared.Client
 }
@@ -41,6 +44,10 @@ func (c *Client) Connect(ctx context.Context) (*Connection, error) {
 	return &Connection{sharedConn: conn}, nil
 }
 
+func (conn *Connection) SendPing() error {
+	return conn.sharedConn.SendMessage("ping", nil)
+}
+
 func (conn *Connection) SendPong() error {
 	return conn.sharedConn.SendMessage("pong", nil)
 }
@@ -59,6 +66,10 @@ func (conn *Connection) Receive() (service.WSMessage, error) {
 	}, nil
 }
 
+func (conn *Connection) RunReceiveLoop(ctx context.Context, handler shared.MessageHandler) error {
+	return conn.sharedConn.RunReceiveLoop(ctx, handler)
+}
+
 func (conn *Connection) Close() error {
 	return conn.sharedConn.Close()
 }