udeServer/internal/ws/service.go

// Package ws 提供运维诊断通道(移植自 sdf-serve internal/ws):
// ping/traceroute 起子进程逐行推流,telnet 做 TCP 拨测。
package ws

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"os/exec"
	"regexp"
	"runtime"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"

	"udeServer/internal/types"
)

type Service struct{}

func NewService() *Service {
	return &Service{}
}

// ValidateTarget 校验 ping/traceroute/telnet 的目标地址:仅允许 IP 或
// 字母/数字/点/连字符,阻断空格与特殊字符注入。
func ValidateTarget(target string) bool {
	if net.ParseIP(target) != nil {
		return true
	}
	re := regexp.MustCompile(`^[a-zA-Z0-9.-]+$`)
	return re.MatchString(target)
}

// ValidatePort 校验 telnet 端口范围。
func ValidatePort(port int) bool {
	return port > 0 && port <= 65535
}

// Ping 实时输出(跨平台)
func (s *Service) Ping(conn *websocket.Conn, req types.WSRequest) {
	count := req.Count
	if count == 0 {
		count = 4
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	var args []string
	if runtime.GOOS == "windows" {
		args = []string{req.Target, "/n", fmt.Sprintf("%d", count)}
	} else {
		args = []string{"-c", fmt.Sprintf("%d", count), req.Target}
	}
	cmd := exec.CommandContext(ctx, "ping", args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		_ = conn.WriteJSON(types.WSResponse{Action: "ping", Status: "error", Output: err.Error()})
		return
	}
	if err := cmd.Start(); err != nil {
		_ = conn.WriteJSON(types.WSResponse{Action: "ping", Status: "error", Output: err.Error()})
		return
	}
	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		line := scanner.Text()
		if runtime.GOOS == "windows" {
			line = gbkToUtf8([]byte(line))
		}
		_ = conn.WriteJSON(types.WSResponse{Action: "ping", Status: "running", Output: line})
	}

	if err := cmd.Wait(); err != nil {
		_ = conn.WriteJSON(types.WSResponse{Action: "ping", Status: "error", Output: err.Error()})
		return
	}

	_ = conn.WriteJSON(types.WSResponse{Action: "ping", Status: "done", Output: "Ping finished"})
}

// Traceroute 实时输出(跨平台)
func (s *Service) Traceroute(conn *websocket.Conn, req types.WSRequest) {
	maxHops := req.MaxHops
	if maxHops == 0 {
		maxHops = 30
	}

	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	var cmd *exec.Cmd
	if runtime.GOOS == "windows" {
		// Windows 使用 tracert,/h 指定最大跳数
		cmd = exec.CommandContext(ctx, "tracert", "/h", fmt.Sprintf("%d", maxHops), req.Target)
	} else {
		// Linux/macOS 使用 traceroute
		cmd = exec.CommandContext(ctx, "traceroute", "-m", fmt.Sprintf("%d", maxHops), req.Target)
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		_ = conn.WriteJSON(types.WSResponse{Action: "traceroute", Status: "error", Output: err.Error()})
		return
	}

	if err := cmd.Start(); err != nil {
		_ = conn.WriteJSON(types.WSResponse{Action: "traceroute", Status: "error", Output: err.Error()})
		return
	}

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		line := scanner.Text()
		if runtime.GOOS == "windows" {
			line = gbkToUtf8([]byte(line)) // Windows 输出是 GBK
		}
		_ = conn.WriteJSON(types.WSResponse{Action: "traceroute", Status: "running", Output: line})
	}

	if err := cmd.Wait(); err != nil {
		_ = conn.WriteJSON(types.WSResponse{Action: "traceroute", Status: "error", Output: err.Error()})
		return
	}

	_ = conn.WriteJSON(types.WSResponse{Action: "traceroute", Status: "done", Output: "Traceroute finished"})
}

// Telnet 实时输出
func (s *Service) Telnet(conn *websocket.Conn, req types.WSRequest) {
	address := net.JoinHostPort(req.Host, strconv.Itoa(req.Port))
	_ = conn.WriteJSON(types.WSResponse{Action: "telnet", Status: "running", Output: fmt.Sprintf("尝试连接 %s ...", address)})

	timeout := 5 * time.Second
	start := time.Now()
	type dialResult struct {
		msg string
		ok  bool
	}
	connCh := make(chan dialResult, 1)

	go func() {
		_, err := net.DialTimeout("tcp", address, timeout)
		if err != nil {
			connCh <- dialResult{fmt.Sprintf("连接 %s 失败: %v", address, err), false}
		} else {
			connCh <- dialResult{fmt.Sprintf("连接 %s 成功", address), true}
		}
	}()

	select {
	case res := <-connCh:
		status := "done"
		if !res.ok {
			status = "error"
		}
		_ = conn.WriteJSON(types.WSResponse{Action: "telnet", Status: status, Output: res.msg})
	case <-time.After(timeout + 1*time.Second):
		_ = conn.WriteJSON(types.WSResponse{Action: "telnet", Status: "error", Output: "连接超时"})
	}

	duration := time.Since(start).Milliseconds()
	_ = conn.WriteJSON(types.WSResponse{Action: "telnet", Status: "done", Output: fmt.Sprintf("耗时 %d ms", duration)})
}

// gbkToUtf8 转换 Windows 子进程 GBK 输出。
func gbkToUtf8(data []byte) string {
	reader := transform.NewReader(bytes.NewReader(data), simplifiedchinese.GBK.NewDecoder())
	result, _ := io.ReadAll(reader)
	return string(result)
}