admin/TaskPool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: update ws security

Browse Source
This commit is contained in:
duorameng
2026-04-27 10:54:00 +08:00
parent f186697116
commit cf9acf3ea0
6 changed files with 162 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/constant/constant.go
+10
View File
@@ -1,5 +1,7 @@
package constant
import "time"
const (
// ConfigPath 配置文件路径
@@ -159,6 +161,14 @@ const (
// Env Type
EnvTypeNormal = "normal"
EnvTypeSecret = "secret"
// WebSocket 安全常量
// PongWait 收到 pong 的超时时间
PongWait = 60 * time.Second
// PingPeriod 发送 ping 的周期
PingPeriod = (PongWait * 9) / 10
// MaxMessageSize 允许的最大消息大小
MaxMessageSize = 1024 * 1024 // 1MB
)
// TablePrefix 表前缀，从配置文件读取
internal/controllers/agent_controller.go
+1 -3
View File
@@ -20,9 +20,7 @@ import (
)
var agentUpgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
CheckOrigin: utils.CheckWSOrigin,
}
// AgentController Agent 控制器
internal/controllers/log_ws_controller.go
+27
View File
@@ -3,6 +3,7 @@ package controllers
import (
"fmt"
"github.com/engigu/baihu-panel/internal/constant"
"github.com/engigu/baihu-panel/internal/database"
"github.com/engigu/baihu-panel/internal/models"
"github.com/engigu/baihu-panel/internal/services/tasks"
@@ -10,6 +11,7 @@ import (
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"time"
)
type LogWSController struct{}
@@ -32,6 +34,23 @@ func (lc *LogWSController) StreamLog(c *gin.Context) {
}
defer conn.Close()
// DoS 保护与心跳设置
conn.SetReadLimit(constant.MaxMessageSize) // 使用与终端一致的限制
conn.SetReadDeadline(time.Now().Add(constant.PongWait))
conn.SetPongHandler(func(string) error {
conn.SetReadDeadline(time.Now().Add(constant.PongWait))
return nil
})
// 启动一个读取循环，用于处理 pong 和检测断开
go func() {
for {
if _, _, err := conn.ReadMessage(); err != nil {
break
}
}
}()
// 1. 检查数据库中是否已结束
var taskLog models.TaskLog
res := database.DB.Where("id = ?", logID).Limit(1).Find(&taskLog)
@@ -68,9 +87,17 @@ func (lc *LogWSController) StreamLog(c *gin.Context) {
sub := tl.Subscribe()
defer tl.Unsubscribe(sub)
ticker := time.NewTicker(constant.PingPeriod)
defer ticker.Stop()
// 推送更新
for {
select {
case <-ticker.C:
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
case data, ok := <-sub:
if !ok {
// 任务结束，尝试刷新最后一次库内完整内容
internal/controllers/terminal_controller.go
+63 -4
View File
@@ -4,11 +4,11 @@ import (
"bufio"
"encoding/json"
"io"
"net/http"
"os"
"path/filepath"
"runtime"
"sync"
"time"
"unicode/utf8"
"github.com/engigu/baihu-panel/internal/constant"
@@ -22,6 +22,7 @@ import (
"golang.org/x/text/transform"
)
type TerminalController struct {
envService *services.EnvService
}
@@ -33,9 +34,7 @@ func NewTerminalController(envService *services.EnvService) *TerminalController
}
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
CheckOrigin: utils.CheckWSOrigin,
}
// toUTF8 将可能是 GBK 编码的字节转换为 UTF-8
@@ -98,6 +97,13 @@ func (tc *TerminalController) HandleWebSocket(c *gin.Context) {
// handlePtyMode 使用 PTY 处理终端（Unix/macOS）
func (tc *TerminalController) handlePtyMode(conn *websocket.Conn, userID string) {
conn.SetReadLimit(constant.MaxMessageSize)
conn.SetReadDeadline(time.Now().Add(constant.PongWait))
conn.SetPongHandler(func(string) error {
conn.SetReadDeadline(time.Now().Add(constant.PongWait))
return nil
})
// 发送 PTY 模式标识
conn.WriteMessage(websocket.TextMessage, []byte("__PTY_MODE__"))
@@ -124,6 +130,7 @@ func (tc *TerminalController) handlePtyMode(conn *websocket.Conn, userID string)
writeMessage := func(data []byte) {
connMu.Lock()
defer connMu.Unlock()
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
conn.WriteMessage(websocket.TextMessage, data)
}
@@ -143,6 +150,27 @@ func (tc *TerminalController) handlePtyMode(conn *websocket.Conn, userID string)
}
}()
// 启动 ping 协程
pingDone := make(chan struct{})
go func() {
ticker := time.NewTicker(constant.PingPeriod)
defer ticker.Stop()
for {
select {
case <-ticker.C:
connMu.Lock()
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
connMu.Unlock()
return
}
connMu.Unlock()
case <-pingDone:
return
}
}
}()
for {
_, message, err := conn.ReadMessage()
if err != nil {
@@ -167,6 +195,7 @@ func (tc *TerminalController) handlePtyMode(conn *websocket.Conn, userID string)
}
}
close(pingDone)
cmd.Process.Kill()
cmd.Wait()
wg.Wait()
@@ -174,6 +203,13 @@ func (tc *TerminalController) handlePtyMode(conn *websocket.Conn, userID string)
// handlePipeMode 使用 pipe 处理终端（Windows）
func (tc *TerminalController) handlePipeMode(conn *websocket.Conn, userID string) {
conn.SetReadLimit(constant.MaxMessageSize)
conn.SetReadDeadline(time.Now().Add(constant.PongWait))
conn.SetPongHandler(func(string) error {
conn.SetReadDeadline(time.Now().Add(constant.PongWait))
return nil
})
// 发送 pipe 模式标识
conn.WriteMessage(websocket.TextMessage, []byte("__PIPE_MODE__"))
@@ -215,6 +251,7 @@ func (tc *TerminalController) handlePipeMode(conn *websocket.Conn, userID string
writeMessage := func(data []byte) {
connMu.Lock()
defer connMu.Unlock()
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
conn.WriteMessage(websocket.TextMessage, data)
}
@@ -238,6 +275,27 @@ func (tc *TerminalController) handlePipeMode(conn *websocket.Conn, userID string
go readOutput(stdout)
go readOutput(stderr)
// 启动 ping 协程
pingDone := make(chan struct{})
go func() {
ticker := time.NewTicker(constant.PingPeriod)
defer ticker.Stop()
for {
select {
case <-ticker.C:
connMu.Lock()
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
connMu.Unlock()
return
}
connMu.Unlock()
case <-pingDone:
return
}
}
}()
for {
_, message, err := conn.ReadMessage()
if err != nil {
@@ -258,6 +316,7 @@ func (tc *TerminalController) handlePipeMode(conn *websocket.Conn, userID string
}
}
close(pingDone)
stdin.Close()
cmd.Process.Kill()
cmd.Wait()
internal/router/router.go
+4 -1
View File
@@ -2,6 +2,7 @@ package router
import (
"strings"
"os"
"github.com/engigu/baihu-panel/internal/controllers"
"github.com/engigu/baihu-panel/internal/middleware"
@@ -31,7 +32,9 @@ type Controllers struct {
}
func Setup(c *Controllers) *gin.Engine {
gin.SetMode(gin.ReleaseMode)
if os.Getenv("GIN_MODE") == "" {
gin.SetMode(gin.ReleaseMode)
}
router := gin.New()
router.Use(middleware.GinLogger(), middleware.GinRecovery())
internal/utils/network.go
+57
View File
@@ -0,0 +1,57 @@
package utils
import (
"net/http"
"net/url"
"os"
"strings"
"github.com/gin-gonic/gin"
)
// CheckWSOrigin 校验 WebSocket 的 Origin 来源是否安全。
// 默认仅允许同源请求，可通过环境变量 BH_ALLOWED_ORIGINS 配置额外的允许列表（逗号分隔）。
func CheckWSOrigin(r *http.Request) bool {
origin := r.Header.Get("Origin")
if origin == "" {
// 非浏览器发起的请求（如直接用脚本连接）通常不带 Origin，默认放行。
return true
}
u, err := url.Parse(origin)
if err != nil {
return false
}
// 0. 开发环境校验：如果是非 Release 模式，默认放行
if gin.Mode() != gin.ReleaseMode {
return true
}
// 1. 同源校验：Origin 的 Host 与请求头中的 Host 一致
if strings.EqualFold(u.Host, r.Host) {
return true
}
// 2. 环境变量配置的允许列表校验
allowedOrigins := os.Getenv("BH_ALLOWED_ORIGINS")
if allowedOrigins != "" {
origins := strings.Split(allowedOrigins, ",")
for _, o := range origins {
o = strings.TrimSpace(o)
if o == "*" {
return true
}
// 匹配完整 Origin (如 http://localhost:5173) 或仅 Host 部分
if strings.EqualFold(o, origin) || strings.EqualFold(o, u.Host) {
return true
}
}
}
// 3. 兜底策略：如果是开发环境常见的 localhost/127.0.0.1，且端口不一致的情况，
// 如果用户没有配置允许列表，我们在非 Release 模式下可以考虑放行，
// 但为了安全，默认应严格限制。建议开发时通过 BH_ALLOWED_ORIGINS=localhost:5173 显式开启。
return false
}