[优化] Phase2

openflare_server/controller/agent.go

@@ -191,31 +191,26 @@ func agentWSReadTimeout() time.Duration {
 	return timeout
 }
 
-func streamAgentWSMessages(c *gin.Context, conn *websocket.Conn, client *service.AgentWSClient) {
-	pingTicker := time.NewTicker(30 * time.Second)
-	defer pingTicker.Stop()
+func agentWSWriteTimeout() time.Duration {
+	return 10 * time.Second
+}
+
+func streamAgentWSMessages(c *gin.Context, conn *websocket.Conn, client *service.WSClient) {
 	for {
 		select {
-		case message := <-client.Messages():
-			slog.Debug("agent ws sending message", "node_id", client.NodeID(), "type", message.Type)
-			if err := websocket.JSON.Send(conn, message); err != nil {
-				slog.Debug("agent ws send failed", "node_id", client.NodeID(), "type", message.Type, "error", err)
-				client.Close()
-				return
-			}
-		case <-pingTicker.C:
-			message := service.AgentWSOutboundMessage{Type: service.AgentWSMessageTypePing}
-			slog.Debug("agent ws sending ping", "node_id", client.NodeID())
-			if err := websocket.JSON.Send(conn, message); err != nil {
-				slog.Debug("agent ws ping failed", "node_id", client.NodeID(), "error", err)
-				client.Close()
-				return
-			}
+		case <-c.Request.Context().Done():
+			return
 		case <-client.Done():
 			return
-		case <-c.Request.Context().Done():
-			client.Close()
-			return
+		case message, ok := <-client.Messages():
+			if !ok {
+				return
+			}
+			_ = conn.SetWriteDeadline(time.Now().Add(agentWSWriteTimeout()))
+			if err := websocket.JSON.Send(conn, message); err != nil {
+				slog.Debug("agent ws send failed", "node_id", client.ID(), "error", err)
+				return
+			}
 		}
 	}
 }

openflare_server/controller/flared.go

@@ -0,0 +1,171 @@
+package controller
+
+import (
+	"log/slog"
+	"net"
+	"openflare/model"
+	"openflare/service"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"golang.org/x/net/websocket"
+)
+
+// FlaredHeartbeat godoc
+// @Summary Report OpenFlared client heartbeat
+// @Tags Flared
+// @Accept json
+// @Produce json
+// @Security TunnelTokenAuth
+// @Param payload body service.FlaredHeartbeatPayload true "Flared heartbeat payload"
+// @Success 200 {object} map[string]interface{}
+// @Failure 400 {object} map[string]interface{}
+// @Router /api/flared/heartbeat [post]
+func FlaredHeartbeat(c *gin.Context) {
+	var payload service.FlaredHeartbeatPayload
+	if !bindJSON(c, &payload) {
+		return
+	}
+	authTunnel, ok := c.Get("tunnel")
+	if !ok {
+		respondUnauthorized(c, "无权进行此操作")
+		return
+	}
+	tunnel := authTunnel.(*model.Tunnel)
+	result, err := service.HeartbeatFlared(tunnel, payload)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, result)
+}
+
+// FlaredGetActiveConfig godoc
+// @Summary Get active tunnel config for OpenFlared client
+// @Tags Flared
+// @Produce json
+// @Security TunnelTokenAuth
+// @Success 200 {object} map[string]interface{}
+// @Router /api/flared/config/active [get]
+func FlaredGetActiveConfig(c *gin.Context) {
+	authTunnel, ok := c.Get("tunnel")
+	if !ok {
+		respondUnauthorized(c, "无权进行此操作")
+		return
+	}
+	tunnel := authTunnel.(*model.Tunnel)
+	config, err := service.GetFlaredTunnelConfig(tunnel)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, config)
+}
+
+// FlaredReportApplyLog godoc
+// @Summary Report apply log for OpenFlared client
+// @Tags Flared
+// @Accept json
+// @Produce json
+// @Security TunnelTokenAuth
+// @Param payload body service.ApplyLogPayload true "Apply log payload"
+// @Success 200 {object} map[string]interface{}
+// @Failure 400 {object} map[string]interface{}
+// @Router /api/flared/apply-log [post]
+func FlaredReportApplyLog(c *gin.Context) {
+	var payload service.ApplyLogPayload
+	if !bindJSON(c, &payload) {
+		return
+	}
+	authTunnel, ok := c.Get("tunnel")
+	if !ok {
+		respondUnauthorized(c, "无权进行此操作")
+		return
+	}
+	tunnel := authTunnel.(*model.Tunnel)
+	payload.NodeID = tunnel.TunnelID
+	log, err := service.ReportApplyLog(payload)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, log)
+}
+
+// FlaredWebSocket godoc
+// @Summary Upgrade OpenFlared connection to websocket
+// @Tags Flared
+// @Security TunnelTokenAuth
+// @Router /api/flared/ws [get]
+func FlaredWebSocket(c *gin.Context) {
+	authTunnel, ok := c.Get("tunnel")
+	if !ok {
+		respondUnauthorized(c, "无权进行此操作")
+		return
+	}
+	tunnel := authTunnel.(*model.Tunnel)
+	slog.Debug("flared ws upgrade requested", "tunnel_id", tunnel.TunnelID, "remote", c.Request.RemoteAddr)
+	websocket.Handler(func(conn *websocket.Conn) {
+		client := service.RegisterFlaredWSClient(tunnel.TunnelID)
+		defer service.UnregisterFlaredWSClient(client)
+		defer func() {
+			_ = conn.Close()
+			slog.Debug("flared ws connection closed", "tunnel_id", tunnel.TunnelID)
+		}()
+
+		slog.Debug("flared ws upgrade succeeded", "tunnel_id", tunnel.TunnelID, "remote", c.Request.RemoteAddr)
+
+		go func() {
+			<-client.Done()
+			_ = conn.Close()
+		}()
+
+		go streamFlaredWSMessages(c, conn, client)
+
+		for {
+			var message service.WSMessage
+			_ = conn.SetReadDeadline(time.Now().Add(agentWSReadTimeout()))
+			if err := websocket.JSON.Receive(conn, &message); err != nil {
+				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+					slog.Debug("flared ws receive timeout", "tunnel_id", tunnel.TunnelID)
+					return
+				}
+				slog.Debug("flared ws receive failed", "tunnel_id", tunnel.TunnelID, "error", err)
+				return
+			}
+			slog.Debug("flared ws message received", "tunnel_id", tunnel.TunnelID, "type", message.Type)
+			switch message.Type {
+			case "status":
+				// Handle status if needed for flared
+			case "ping":
+				if !service.SendFlaredWSPong(tunnel.TunnelID) {
+					slog.Debug("flared ws pong enqueue failed", "tunnel_id", tunnel.TunnelID)
+				}
+			case "pong":
+				slog.Debug("flared ws pong received", "tunnel_id", tunnel.TunnelID)
+			default:
+				slog.Debug("flared ws unsupported message type", "tunnel_id", tunnel.TunnelID, "type", message.Type)
+			}
+		}
+	}).ServeHTTP(c.Writer, c.Request)
+}
+
+func streamFlaredWSMessages(c *gin.Context, conn *websocket.Conn, client *service.WSClient) {
+	for {
+		select {
+		case <-c.Request.Context().Done():
+			return
+		case <-client.Done():
+			return
+		case message, ok := <-client.Messages():
+			if !ok {
+				return
+			}
+			_ = conn.SetWriteDeadline(time.Now().Add(agentWSWriteTimeout()))
+			if err := websocket.JSON.Send(conn, message); err != nil {
+				slog.Debug("flared ws send failed", "tunnel_id", client.ID(), "error", err)
+				return
+			}
+		}
+	}
+}

openflare_server/controller/relay.go

@@ -0,0 +1,117 @@
+package controller
+
+import (
+	"log/slog"
+	"net"
+	"openflare/model"
+	"openflare/service"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"golang.org/x/net/websocket"
+)
+
+// RelayHeartbeat godoc
+// @Summary Report relay heartbeat
+// @Tags Relay
+// @Accept json
+// @Produce json
+// @Security AgentTokenAuth
+// @Param payload body service.RelayHeartbeatPayload true "Relay heartbeat payload"
+// @Success 200 {object} map[string]interface{}
+// @Failure 400 {object} map[string]interface{}
+// @Router /api/relay/heartbeat [post]
+func RelayHeartbeat(c *gin.Context) {
+	var payload service.RelayHeartbeatPayload
+	if !bindJSON(c, &payload) {
+		return
+	}
+	authNode, ok := c.Get("relay_node")
+	if !ok {
+		respondUnauthorized(c, "无权进行此操作")
+		return
+	}
+	node := authNode.(*model.Node)
+	result, err := service.HeartbeatRelay(node, payload)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, result)
+}
+
+// RelayWebSocket godoc
+// @Summary Upgrade relay connection to websocket
+// @Tags Relay
+// @Security AgentTokenAuth
+// @Router /api/relay/ws [get]
+func RelayWebSocket(c *gin.Context) {
+	authNode, ok := c.Get("relay_node")
+	if !ok {
+		respondUnauthorized(c, "无权进行此操作")
+		return
+	}
+	node := authNode.(*model.Node)
+	slog.Debug("relay ws upgrade requested", "node_id", node.NodeID, "remote", c.Request.RemoteAddr)
+	websocket.Handler(func(conn *websocket.Conn) {
+		client := service.RegisterRelayWSClient(node.NodeID)
+		defer service.UnregisterRelayWSClient(client)
+		defer func() {
+			_ = conn.Close()
+			slog.Debug("relay ws connection closed", "node_id", node.NodeID)
+		}()
+
+		slog.Debug("relay ws upgrade succeeded", "node_id", node.NodeID, "remote", c.Request.RemoteAddr)
+
+		go func() {
+			<-client.Done()
+			_ = conn.Close()
+		}()
+
+		go streamRelayWSMessages(c, conn, client)
+
+		for {
+			var message service.WSMessage
+			_ = conn.SetReadDeadline(time.Now().Add(agentWSReadTimeout()))
+			if err := websocket.JSON.Receive(conn, &message); err != nil {
+				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+					slog.Debug("relay ws receive timeout", "node_id", node.NodeID)
+					return
+				}
+				slog.Debug("relay ws receive failed", "node_id", node.NodeID, "error", err)
+				return
+			}
+			slog.Debug("relay ws message received", "node_id", node.NodeID, "type", message.Type)
+			switch message.Type {
+			case "ping":
+				if !service.SendRelayWSPong(node.NodeID) {
+					slog.Debug("relay ws pong enqueue failed", "node_id", node.NodeID)
+				}
+			case "pong":
+				slog.Debug("relay ws pong received", "node_id", node.NodeID)
+			default:
+				slog.Debug("relay ws unsupported message type", "node_id", node.NodeID, "type", message.Type)
+			}
+		}
+	}).ServeHTTP(c.Writer, c.Request)
+}
+
+func streamRelayWSMessages(c *gin.Context, conn *websocket.Conn, client *service.WSClient) {
+	for {
+		select {
+		case <-c.Request.Context().Done():
+			return
+		case <-client.Done():
+			return
+		case message, ok := <-client.Messages():
+			if !ok {
+				return
+			}
+			_ = conn.SetWriteDeadline(time.Now().Add(agentWSWriteTimeout()))
+			if err := websocket.JSON.Send(conn, message); err != nil {
+				slog.Debug("relay ws send failed", "node_id", client.ID(), "error", err)
+				return
+			}
+		}
+	}
+}

openflare_server/controller/tunnel.go

@@ -0,0 +1,84 @@
+package controller
+
+import (
+	"openflare/service"
+
+	"github.com/gin-gonic/gin"
+)
+
+func GetTunnels(c *gin.Context) {
+	tunnels, err := service.ListTunnels()
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, tunnels)
+}
+
+func GetTunnel(c *gin.Context) {
+	id, ok := parseIDParam(c)
+	if !ok {
+		return
+	}
+	tunnel, err := service.GetTunnel(id)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, tunnel)
+}
+
+func CreateTunnel(c *gin.Context) {
+	var input service.TunnelInput
+	if !bindJSON(c, &input) {
+		return
+	}
+	tunnel, err := service.CreateTunnel(input)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, tunnel)
+}
+
+func UpdateTunnel(c *gin.Context) {
+	id, ok := parseIDParam(c)
+	if !ok {
+		return
+	}
+	var input service.TunnelInput
+	if !bindJSON(c, &input) {
+		return
+	}
+	tunnel, err := service.UpdateTunnel(id, input)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, tunnel)
+}
+
+func DeleteTunnel(c *gin.Context) {
+	id, ok := parseIDParam(c)
+	if !ok {
+		return
+	}
+	if err := service.DeleteTunnel(id); err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, nil)
+}
+
+func RotateTunnelToken(c *gin.Context) {
+	id, ok := parseIDParam(c)
+	if !ok {
+		return
+	}
+	tunnel, err := service.RotateTunnelToken(id)
+	if err != nil {
+		respondFailure(c, err.Error())
+		return
+	}
+	respondSuccess(c, tunnel)
+}