perf: 支持多个 guacd 地址

jumpserver · Mar 19, 2024 · b29902d · b29902d
1 parent 6717132
commit b29902d
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -3,6 +3,8 @@ package config
 import (
 	"fmt"
 	"log"
+	"math/rand"
+	"net"
 	"os"
 	"path/filepath"
 	"strings"
@@ -29,6 +31,8 @@ type Config struct {
 	HTTPPort       string `mapstructure:"HTTPD_PORT"`
 	LogLevel       string `mapstructure:"LOG_LEVEL"`
 
+	GuacdAddrs string `mapstructure:"GUACD_ADDRS"`
+
 	GuaHost                   string `mapstructure:"GUA_HOST"`
 	GuaPort                   string `mapstructure:"GUA_PORT"`
 	DisableAllCopyPaste       bool   `mapstructure:"JUMPSERVER_DISABLE_ALL_COPY_PASTE"`
@@ -54,6 +58,14 @@ type Config struct {
 	EnablePanda       bool   `mapstructure:"ENABLE_PANDA"`
 }
 
+func (c *Config) SelectGuacdAddr() string {
+	if len(c.GuacdAddrs) == 0 {
+		return net.JoinHostPort(c.GuaHost, c.GuaPort)
+	}
+	addresses := strings.Split(c.GuacdAddrs, ",")
+	return addresses[rand.Intn(len(addresses))]
+}
+
 func Setup(configPath string) {
 	var conf = getDefaultConfig()
 	loadConfigFromEnv(&conf)

diff --git a/pkg/tunnel/conn.go b/pkg/tunnel/conn.go
@@ -3,7 +3,6 @@ package tunnel
 import (
 	"encoding/json"
 	"fmt"
-	"net"
 	"sort"
 	"strings"
 	"sync"
@@ -14,7 +13,6 @@ import (
 	"github.com/gorilla/websocket"
 
 	"lion/pkg/common"
-	"lion/pkg/config"
 	"lion/pkg/guacd"
 	"lion/pkg/jms-sdk-go/model"
 	"lion/pkg/logger"
@@ -48,6 +46,8 @@ type Connection struct {
 	guacdTunnel *guacd.Tunnel
 	Service     *session.Server
 
+	guacdAddr string
+
 	ws *websocket.Conn
 
 	wsLock    sync.Mutex
@@ -370,8 +370,7 @@ func (t *Connection) CloneMonitorTunnel() (*guacd.Tunnel, error) {
 	info := guacd.NewClientInformation()
 	conf := guacd.NewConfiguration()
 	conf.ConnectionID = t.guacdTunnel.UUID
-	guacdAddr := net.JoinHostPort(config.GlobalConfig.GuaHost,
-		config.GlobalConfig.GuaPort)
+	guacdAddr := t.guacdAddr
 	monitorTunnel, err := guacd.NewTunnel(guacdAddr, conf, info)
 	if err != nil {
 		return nil, err

diff --git a/pkg/tunnel/server.go b/pkg/tunnel/server.go
@@ -180,7 +180,8 @@ func (g *GuacamoleTunnelServer) Connect(ctx *gin.Context) {
 	}
 
 	var tunnel *guacd.Tunnel
-	guacdAddr := net.JoinHostPort(config.GlobalConfig.GuaHost, config.GlobalConfig.GuaPort)
+	guacdAddr := config.GlobalConfig.SelectGuacdAddr()
+
 	tunnel, err = guacd.NewTunnel(guacdAddr, conf, info)
 	if err != nil {
 		logger.Errorf("Connect tunnel err: %+v", err)
@@ -204,6 +205,7 @@ func (g *GuacamoleTunnelServer) Connect(ctx *gin.Context) {
 		sessionId, info.OptimalScreenWidth, info.OptimalScreenHeight)
 
 	conn := Connection{
+		guacdAddr:   guacdAddr,
 		Sess:        tunnelSession,
 		guacdTunnel: tunnel,
 		Service:     g.SessionService,