feature opsflow添加认证

cmd/hytop/lib/lib.go

@@ -2,7 +2,7 @@ package lib
 
 /*
 #cgo CFLAGS: -I.
-#cgo LDFLAGS: -L/opt/hyhal/lib -lamd_smi
+#cgo LDFLAGS: -L/opt/hyhal/lib -lrocm_smi64
 
 #include <stdlib.h>
 #include <string.h>

cmd/opsflow/auth/auth_test.go

+package auth
+
+import (
+	"encoding/base64"
+	"testing"
+	"time"
+
+	"github.com/pquerna/otp"
+	"github.com/pquerna/otp/totp"
+)
+
+func TestOTP(t *testing.T) {
+	key, err := totp.Generate(totp.GenerateOpts{
+
+		Issuer:      "System",
+		AccountName: "Opsflow",
+	})
+	if err != nil {
+		t.Error(err)
+	}
+	t.Log(key.Secret())
+}
+
+// OA5BDGLJ2DYGAWNCLJYSNZFAESPQ7BRL
+
+func verify(sec, input string) bool {
+	return totp.Validate(input, sec)
+}
+
+func TestAuth(t *testing.T) {
+	sec := "OA5BDGLJ2DYGAWNCLJYSNZFAESPQ7BRL"
+	out, err := totp.GenerateCode(sec, time.Now())
+	if err != nil {
+		t.Error(err)
+	}
+	t.Log(out)
+
+	if verify(sec, out) {
+		t.Log("auth ok")
+	} else {
+		t.Log("auth fail")
+	}
+
+	ok, err := totp.ValidateCustom(out, sec, time.Now(), totp.ValidateOpts{
+		Period:    30, // 每 30 秒更新一次
+		Skew:      1,  // 允许前后偏移 1 个周期（即允许 30 秒的时间误差）
+		Digits:    otp.DigitsSix,
+		Algorithm: otp.AlgorithmSHA1,
+	})
+	if err != nil {
+		t.Error(err)
+	}
+	t.Logf("auth result: %v\n", ok)
+}
+
+func TestGetCode(t *testing.T) {
+	sec := "OA5BDGLJ2DYGAWNCLJYSNZFAESPQ7BRL"
+	out, err := totp.GenerateCode(sec, time.Now())
+	if err != nil {
+		t.Error(err)
+	}
+	t.Log(base64.StdEncoding.EncodeToString([]byte(out)))
+}

cmd/opsflow/main.go

@@ -53,6 +53,7 @@ OPSFLOW_RCCL_ALL_REDUCE_PERF_ARGS: set rccl all reduce perf args`)
 	cfg.SetDefault("rccl_test_path", "/opt/rccl-tests/build")
 	cfg.SetEnvPrefix("OPSFLOW")
 	cfg.AutomaticEnv()
+	cfg.SetDefault("auth_key", "OA5BDGLJ2DYGAWNCLJYSNZFAESPQ7BRL")
 	cfg.BindPFlag("debug_mode", pflag.Lookup("debug"))
 	cfg.BindPFlag("rccl_test_path", pflag.Lookup("rccl-test-path"))
 	cfg.BindPFlag("rccl_all_reduce_perf_args", pflag.Lookup("rccl-all-reduce-perf-args"))

cmd/opsflow/opsflow.yaml

 rccl_test_path: /home/panyq/wangx/rccl-tests/build-dan
 rccl_all_reduce_perf_args: "-b 8 -e 128M -f 2 -g 8 -d half"
 debug_mode: false
+# Google的基于时间的验证码认证
+auth_key: OA5BDGLJ2DYGAWNCLJYSNZFAESPQ7BRL
\ No newline at end of file

cmd/opsflow/web/main.go

 package web
 
 import (
+	"encoding/base64"
 	"get-container/cmd/opsflow/backend"
+	"net/http"
 	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/pquerna/otp"
+	"github.com/pquerna/otp/totp"
 	"github.com/spf13/viper"
 
 	swaggerFiles "github.com/swaggo/files"
@@ -68,12 +73,57 @@ func ReturnGinList[T any](ctx *gin.Context, data []T, err error) {
 	})
 }
 
+func webAuth(ctx *gin.Context) {
+	authHeader := ctx.GetHeader("Authorization")
+	if authHeader == "" {
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "Not find header Authorization"})
+		ctx.Abort()
+		return
+	}
+	fields := strings.SplitN(authHeader, " ", 2)
+	if len(fields) != 2 || fields[0] != "Bearer" {
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "Authorization error"})
+		ctx.Abort()
+		return
+	}
+	code, err := base64.StdEncoding.DecodeString(fields[1])
+	if err != nil {
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
+		ctx.Abort()
+		return
+	}
+	sec := globalCfg.GetString("auth_key")
+	ok, err := totp.ValidateCustom(string(code), sec, time.Now(), totp.ValidateOpts{
+		Period:    30, // 每 30 秒更新一次
+		Skew:      1,  // 允许前后偏移 1 个周期（即允许 30 秒的时间误差）
+		Digits:    otp.DigitsSix,
+		Algorithm: otp.AlgorithmSHA1,
+	})
+	if err != nil {
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
+		ctx.Abort()
+		return
+	}
+	if ok {
+		ctx.Next()
+	} else {
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "Authorization failed"})
+		ctx.Abort()
+		return
+	}
+}
+
 func WebServer(addr string) error {
 	engine := gin.Default()
+	cmdGroup := engine.Group("/api/cmd")
 	if globalCfg.GetBool("debug_mode") {
+		// 调试模式
 		engine.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))
+	} else {
+		// 非调试模式，添加认证
+		cmdGroup.Use(webAuth)
 	}
-	cmdGroup := engine.Group("/api/cmd")
+
 	cmdGroup.GET("/all", _controller.GetAllInfo)
 	cmdGroup.GET("/loginUser", _controller.GetOnlineUser)
 	cmdGroup.GET("/sysload", _controller.GetSysLoad)

go.mod

@@ -17,6 +17,7 @@ require (
 
 require (
 	github.com/KyleBanks/depth v1.2.1 // indirect
+	github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect
 	github.com/bytedance/sonic v1.14.0 // indirect
 	github.com/bytedance/sonic/loader v0.3.0 // indirect
 	github.com/cloudwego/base64x v0.1.6 // indirect
@@ -47,6 +48,7 @@ require (
 	github.com/muesli/clusters v0.0.0-20200529215643-2700303c1762 // indirect
 	github.com/muesli/kmeans v0.3.1 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.4 // indirect
+	github.com/pquerna/otp v1.5.0 // indirect
 	github.com/quic-go/qpack v0.5.1 // indirect
 	github.com/quic-go/quic-go v0.54.0 // indirect
 	github.com/sagikazarmark/locafero v0.11.0 // indirect