FlowStack/backend/internal/handler/disktest.go

package handler

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"flowstack/internal/crypto"
	"flowstack/internal/database"
	"flowstack/internal/model"
	"flowstack/internal/sshexec"
)

const (
	ddFileSize = ".fs-disktest"
	ddSizeMB   = 256
	ddTimeout  = 60 * time.Second
)

type cmdRunner func(cmd string) (string, error)

func localRunner(cmd string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), ddTimeout)
	defer cancel()
	out, err := exec.CommandContext(ctx, "sh", "-c", cmd).CombinedOutput()
	return string(out), err
}

// parseDDMBps 兼容两种 dd 输出：
// Linux: "268435456 bytes (268 MB, 256 MiB) copied, 1.2 s, 224 MB/s"
// macOS: "268435456 bytes transferred in 1.234 secs (217543210 bytes/sec)"
var (
	linuxSpeedRe = regexp.MustCompile(`([\d.]+)\s*([kMGT]?)B/s`)
	macSpeedRe   = regexp.MustCompile(`\(([\d,]+)\s*bytes/sec\)`)
)

func parseDDMBps(out string) float64 {
	if ms := linuxSpeedRe.FindAllStringSubmatch(out, -1); len(ms) > 0 {
		m := ms[len(ms)-1]
		v, err := strconv.ParseFloat(m[1], 64)
		if err != nil {
			return 0
		}
		switch m[2] {
		case "k", "K":
			v /= 1024
		case "G":
			v *= 1024
		case "T":
			v *= 1024 * 1024
		}
		return v
	}
	if m := macSpeedRe.FindStringSubmatch(out); m != nil {
		v, err := strconv.ParseFloat(strings.ReplaceAll(m[1], ",", ""), 64)
		if err == nil {
			return v / 1e6
		}
	}
	return 0
}

// TestBuildNodeDisk POST /api/build-nodes/:id/disk-test —— dd 实测磁盘读写
func TestBuildNodeDisk(c *gin.Context) {
	var n model.BuildNode
	if err := database.DB.First(&n, c.Param("id")).Error; err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "build node not found"})
		return
	}

	// 本地节点直接 exec，否则 SSH（需要凭证）
	isLocal := n.IP == "127.0.0.1" || n.IP == "localhost" || n.IP == "::1"
	var run cmdRunner
	mode := "ssh"
	if isLocal {
		run = localRunner
		mode = "local"
	} else {
		if n.CredentialID == nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "远程节点需先在编辑中关联 SSH 凭证"})
			return
		}
		var cred model.Credential
		if err := database.DB.First(&cred, *n.CredentialID).Error; err != nil || cred.Content == "" {
			c.JSON(http.StatusBadRequest, gin.H{"error": "SSH 凭证不存在或未设置内容"})
			return
		}
		plain, err := crypto.DecryptField(cred.Content)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "凭证解密失败"})
			return
		}
		var cc struct {
			PrivateKey string `json:"private_key"`
			Password   string `json:"password"`
			SSHUser    string `json:"ssh_user"`
			SSHPort    int    `json:"ssh_port"`
		}
		if err := json.Unmarshal([]byte(plain), &cc); err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "凭证内容格式错误"})
			return
		}
		cfg := sshexec.Config{Host: n.IP, Port: cc.SSHPort, User: cc.SSHUser, Password: cc.Password, PrivateKey: cc.PrivateKey}
		run = func(cmd string) (string, error) {
			return sshexec.Run(cfg, cmd, ddTimeout)
		}
	}

	ws := n.WorkspacePath
	if ws == "" {
		ws = "/data/flowstack-workspace"
	}
	file := ws + "/" + ddFileSize
	start := time.Now()

	if _, err := run(fmt.Sprintf("mkdir -p '%s'", ws)); err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "创建工作空间失败: " + err.Error()})
		return
	}
	defer func() { _, _ = run(fmt.Sprintf("rm -f '%s'", file)) }()

	// 写测试：direct 优先（绕过缓存），失败降级普通模式
	writeCmd := fmt.Sprintf("dd if=/dev/zero of='%s' bs=1m count=%d oflag=direct 2>&1", file, ddSizeMB)
	writeOut, err := run(writeCmd)
	direct := err == nil
	if err != nil || parseDDMBps(writeOut) == 0 {
		writeOut, err = run(fmt.Sprintf("dd if=/dev/zero of='%s' bs=1m count=%d 2>&1", file, ddSizeMB))
		direct = false
	}
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "写入测试失败: " + firstLine(writeOut, err)})
		return
	}
	writeMBps := parseDDMBps(writeOut)

	// 读测试
	readCmd := fmt.Sprintf("dd if='%s' of=/dev/null bs=1m 2>&1", file)
	if direct {
		readCmd = fmt.Sprintf("dd if='%s' of=/dev/null bs=1m iflag=direct 2>&1", file)
	}
	readOut, err := run(readCmd)
	if err != nil && direct {
		readOut, err = run(fmt.Sprintf("dd if='%s' of=/dev/null bs=1m 2>&1", file))
	}
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "读取测试失败: " + firstLine(readOut, err)})
		return
	}
	readMBps := parseDDMBps(readOut)

	if direct {
		mode += "/direct"
	} else {
		mode += "/cache"
	}

	// 回写节点
	now := time.Now()
	database.DB.Model(&n).Updates(map[string]any{
		"disk_write_mbps": writeMBps, "disk_read_mbps": readMBps,
		"disk_test_mode": mode, "disk_tested_at": &now,
	})
	c.JSON(http.StatusOK, gin.H{
		"success": true, "writeMBps": writeMBps, "readMBps": readMBps,
		"mode": mode, "elapsedMs": time.Since(start).Milliseconds(),
		"message": fmt.Sprintf("写入 %.0f MB/s · 读取 %.0f MB/s（%s）", writeMBps, readMBps, map[bool]string{true: "direct 直测", false: "含缓存"}[direct]),
	})
}