udeServer/internal/handler/ops/opshandler.go

package ops

import (
	"encoding/json"
	"net/http"

	"github.com/gorilla/websocket"

	"udeServer/internal/svc"
	"udeServer/internal/types"
	"udeServer/internal/ws"
)

// upGrader 允许跨源升级;ws/ops 在审计中间件 skip 名单中,
// 升级后按 WSRequest/WSResponse 消息驱动,不走 goctl 常规 logic 链。
var upGrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// OpsHandler 是 WebSocket 运维诊断入口(移植自 sdf-serve ops 模块):
// 1. HTTP 升级成 WebSocket;2. 循环读 WSRequest;3. 校验目标后分发到 ws.Service。
// 动作定义: ping(逐行回推)/traceroute(逐行回推)/telnet(连接成败+耗时)。
func OpsHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		conn, err := upGrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				break
			}
			var req types.WSRequest
			if err := json.Unmarshal(message, &req); err != nil {
				sendErr(conn, req.Action, "invalid request")
				continue
			}

			// 安全校验:防目标地址注入
			if req.Action == "ping" || req.Action == "traceroute" {
				if !ws.ValidateTarget(req.Target) {
					sendErr(conn, req.Action, "非法目标地址")
					continue
				}
			}
			if req.Action == "telnet" {
				if !ws.ValidateTarget(req.Host) || !ws.ValidatePort(req.Port) {
					sendErr(conn, "telnet", "非法 host 或 port")
					continue
				}
			}

			// 执行操作
			switch req.Action {
			case "ping":
				go svcCtx.WsServer.Ping(conn, req)
			case "traceroute":
				go svcCtx.WsServer.Traceroute(conn, req)
			case "telnet":
				go svcCtx.WsServer.Telnet(conn, req)
			default:
				sendErr(conn, req.Action, "unsupported action")
			}
		}
	}
}

func sendErr(conn *websocket.Conn, action, msg string) {
	_ = conn.WriteJSON(types.WSResponse{
		Action: action,
		Status: "error",
		Output: msg,
	})
}