mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-09-29 22:06:38 +08:00
ryan
119 lines
3.2 KiB
Go
119 lines
3.2 KiB
Go
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
|
|
}
|
|
payload.IP = service.ResolveReportedNodeIP(payload.IP, c.Request.RemoteAddr)
|
|
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
|
|
}
|
|
}
|
|
}
|
|
}
|