mirror of
https://github.com/veops/oneterm.git
synced 2025-10-18 21:24:46 +08:00
feat(backend): Add Redis protocol support
This commit is contained in:
pycook
@@ -4,19 +4,23 @@ go 1.21.3
|
||||
|
||||
require (
|
||||
github.com/BurntSushi/toml v1.4.0
|
||||
github.com/atotto/clipboard v0.1.4
|
||||
github.com/charmbracelet/bubbles v0.19.0
|
||||
github.com/charmbracelet/bubbletea v0.27.1
|
||||
github.com/charmbracelet/lipgloss v0.13.0
|
||||
github.com/creack/pty v1.1.24
|
||||
github.com/fatih/color v1.17.0
|
||||
github.com/getwe/figlet4go v0.0.0-20160909034824-bc879344e874
|
||||
github.com/gin-gonic/gin v1.10.0
|
||||
github.com/go-resty/resty/v2 v2.14.0
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/gorilla/websocket v1.5.3
|
||||
github.com/mattn/go-runewidth v0.0.16
|
||||
github.com/nicksnyder/go-i18n/v2 v2.4.0
|
||||
github.com/oklog/run v1.1.0
|
||||
github.com/pkg/sftp v1.13.6
|
||||
github.com/redis/go-redis/v9 v9.6.1
|
||||
github.com/rivo/uniseg v0.4.7
|
||||
github.com/samber/lo v1.47.0
|
||||
github.com/spf13/cast v1.7.0
|
||||
github.com/spf13/pflag v1.0.5
|
||||
@@ -30,20 +34,19 @@ require (
|
||||
golang.org/x/text v0.17.0
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1
|
||||
gorm.io/driver/mysql v1.5.7
|
||||
gorm.io/driver/postgres v1.5.11
|
||||
gorm.io/gorm v1.25.11
|
||||
gorm.io/plugin/soft_delete v1.2.1
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
|
||||
github.com/atotto/clipboard v0.1.4 // indirect
|
||||
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
|
||||
github.com/charmbracelet/x/ansi v0.1.4 // indirect
|
||||
github.com/charmbracelet/x/input v0.1.0 // indirect
|
||||
github.com/charmbracelet/x/term v0.1.1 // indirect
|
||||
github.com/charmbracelet/x/windows v0.1.0 // indirect
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
|
||||
github.com/golang/snappy v0.0.1 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
|
||||
github.com/jackc/pgx/v5 v5.5.5 // indirect
|
||||
@@ -51,13 +54,11 @@ require (
|
||||
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
|
||||
github.com/mattn/go-colorable v0.1.13 // indirect
|
||||
github.com/mattn/go-localereader v0.0.1 // indirect
|
||||
github.com/mattn/go-runewidth v0.0.16 // indirect
|
||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
|
||||
github.com/muesli/cancelreader v0.2.2 // indirect
|
||||
github.com/muesli/termenv v0.15.2 // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||
gorm.io/driver/postgres v1.5.11 // indirect
|
||||
golang.org/x/term v0.23.0 // indirect
|
||||
)
|
||||
|
||||
require (
|
||||
|
||||
@@ -41,6 +41,8 @@ github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJ
|
||||
github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg=
|
||||
github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
|
||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||
github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
|
||||
github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
|
||||
@@ -91,8 +93,6 @@ github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ
|
||||
github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
|
||||
github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
|
||||
github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
|
||||
github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
|
||||
github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
|
||||
github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
|
||||
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
@@ -169,7 +169,6 @@ github.com/muesli/termenv v0.15.2 h1:GohcuySI0QmI3wN8Ok9PtKGkgkFIk7y6Vpb5PvrY+Wo
|
||||
github.com/muesli/termenv v0.15.2/go.mod h1:Epx+iuz8sNs7mNKhxzH4fWXGNpZwUaJKRS1noLXviQ8=
|
||||
github.com/nicksnyder/go-i18n/v2 v2.4.0 h1:3IcvPOAvnCKwNm0TB0dLDTuawWEj+ax/RERNC+diLMM=
|
||||
github.com/nicksnyder/go-i18n/v2 v2.4.0/go.mod h1:nxYSZE9M0bf3Y70gPQjN9ha7XNHX7gMc814+6wVyEI4=
|
||||
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
|
||||
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
|
||||
github.com/oklog/run v1.1.0 h1:GEenZ1cK0+q0+wsJew9qUg/DyD8k3JzYsZAi5gYi2mA=
|
||||
github.com/oklog/run v1.1.0/go.mod h1:sVPdnTZT1zYwAJeCMu2Th4T21pA3FPOQRfWjQlk7DVU=
|
||||
@@ -311,7 +310,6 @@ golang.org/x/term v0.22.0/go.mod h1:F3qCibpT5AMpCRfhfT53vVJwhLtIVHhB9XDjfFvnMI4=
|
||||
golang.org/x/term v0.23.0 h1:F6D4vR+EHoL9/sWAWgAR1H2DcHr4PareCbAaCo1RpuU=
|
||||
golang.org/x/term v0.23.0/go.mod h1:DgV24QBUrK6jhZXl+20l6UWznPlwAHm1Q1mGHtydmSk=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
@@ -338,8 +336,9 @@ google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFW
|
||||
google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f h1:BLraFXnmrev5lT+xlilqcH8XK9/i0At2xKjWk4p6zsU=
|
||||
gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
gopkg.in/ini.v1 v1.67.0 h1:Dgnx+6+nfE+IfzjUEISNeydPJh9AXNNsWbGP9KzCsOA=
|
||||
gopkg.in/ini.v1 v1.67.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k=
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=
|
||||
|
||||
@@ -181,6 +181,7 @@ func DoConnect(ctx *gin.Context, ws *websocket.Conn) (sess *gsession.Session, er
|
||||
if !sess.IsGuacd() {
|
||||
w, h := cast.ToInt(ctx.Query("w")), cast.ToInt(ctx.Query("h"))
|
||||
sess.SshParser = gsession.NewParser(sess.SessionId, w, h)
|
||||
sess.SshParser.Protocol = sess.Protocol
|
||||
|
||||
sessionService := service.NewSessionService()
|
||||
cmds, err := sessionService.GetSshParserCommands(ctx, []int(asset.AccessAuth.CmdIds))
|
||||
@@ -238,7 +239,5 @@ func DoConnect(ctx *gin.Context, ws *websocket.Conn) (sess *gsession.Session, er
|
||||
|
||||
// This is an external function in authorization.go, but needed here to avoid circular imports
|
||||
func hasAuthorization(ctx *gin.Context, sess *gsession.Session) bool {
|
||||
// Implementation handled externally by the imported authorization package
|
||||
// This is just a placeholder to satisfy the compiler
|
||||
return true // Always return true for now, will be handled by proper authorization checks
|
||||
return service.DefaultAuthService.HasAuthorization(ctx, sess)
|
||||
}
|
||||
|
||||
@@ -2,28 +2,25 @@ package connect
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"time"
|
||||
"sync/atomic"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
"github.com/charmbracelet/lipgloss/table"
|
||||
"github.com/creack/pty"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/samber/lo"
|
||||
"github.com/spf13/cast"
|
||||
mysqlDriver "gorm.io/driver/mysql"
|
||||
"gorm.io/gorm"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/veops/oneterm/internal/model"
|
||||
gsession "github.com/veops/oneterm/internal/session"
|
||||
"github.com/veops/oneterm/internal/tunneling"
|
||||
"github.com/veops/oneterm/pkg/logger"
|
||||
)
|
||||
|
||||
// connectOther connects to other protocols (MySQL, Redis, etc.)
|
||||
// connectOther connects to other protocols (Redis, MySQL, etc.)
|
||||
func connectOther(ctx *gin.Context, sess *gsession.Session, asset *model.Asset, account *model.Account, gateway *model.Gateway) (err error) {
|
||||
chs := sess.Chans
|
||||
defer func() {
|
||||
@@ -32,138 +29,189 @@ func connectOther(ctx *gin.Context, sess *gsession.Session, asset *model.Asset,
|
||||
}
|
||||
}()
|
||||
|
||||
protocol := strings.Split(sess.Protocol, ":")[0]
|
||||
ip, port, err := tunneling.Proxy(false, sess.SessionId, protocol, asset, gateway)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
// Handle Redis connections
|
||||
if sess.IsRedis() {
|
||||
logger.L().Info("Starting Redis connection", zap.String("sessionId", sess.SessionId))
|
||||
|
||||
var (
|
||||
rdb *redis.Client
|
||||
db *gorm.DB
|
||||
)
|
||||
switch protocol {
|
||||
case "redis":
|
||||
rdb = redis.NewClient(&redis.Options{
|
||||
Addr: fmt.Sprintf("%s:%d", ip, port),
|
||||
Password: account.Password,
|
||||
DialTimeout: time.Second,
|
||||
// Setup proxy and connection parameters
|
||||
protocol := strings.Split(sess.Protocol, ":")[0]
|
||||
ip, port, err := tunneling.Proxy(false, sess.SessionId, protocol, asset, gateway)
|
||||
if err != nil {
|
||||
logger.L().Error("Failed to setup tunnel", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
// Build redis-cli command
|
||||
args := []string{"-h", ip, "-p", fmt.Sprintf("%d", port)}
|
||||
if account.Password != "" {
|
||||
args = append(args, "-a", account.Password)
|
||||
}
|
||||
logger.L().Info("Starting redis-cli", zap.String("host", ip), zap.Int("port", port))
|
||||
|
||||
// Create command and pseudo-terminal
|
||||
cmd := exec.CommandContext(sess.Gctx, "redis-cli", args...)
|
||||
cmd.Env = append(os.Environ(), "TERM=xterm-256color")
|
||||
ptmx, err := pty.Start(cmd)
|
||||
if err != nil {
|
||||
logger.L().Error("Failed to start redis-cli with pty", zap.Error(err))
|
||||
return fmt.Errorf("failed to start redis-cli: %w", err)
|
||||
}
|
||||
|
||||
// Set standard terminal size
|
||||
_ = pty.Setsize(ptmx, &pty.Winsize{
|
||||
Cols: 80,
|
||||
Rows: 24,
|
||||
})
|
||||
_, err = rdb.Ping(ctx).Result()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
case "mysql":
|
||||
dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?charset=utf8mb4&parseTime=True&loc=Local", account.Account, account.Password, ip, port)
|
||||
db, err = gorm.Open(mysqlDriver.Open(dsn))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
chs.ErrChan <- err
|
||||
// Simplified IO channel setup - direct connection
|
||||
chs.Rin, chs.Win = io.Pipe()
|
||||
|
||||
sess.G.Go(func() error {
|
||||
reader := bufio.NewReader(chs.Rin)
|
||||
buf := &bytes.Buffer{}
|
||||
pt := ""
|
||||
ss := strings.Split(sess.Protocol, ":")
|
||||
if len(ss) == 2 {
|
||||
pt = ss[1]
|
||||
}
|
||||
sess.Prompt = fmt.Sprintf("%s@%s:%s> ", account.Account, asset.Name, pt)
|
||||
chs.OutChan <- append(byteRN, []byte(sess.Prompt)...)
|
||||
for {
|
||||
select {
|
||||
case <-sess.Gctx.Done():
|
||||
return nil
|
||||
default:
|
||||
rn, size, err := reader.ReadRune()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if size <= 0 || rn == utf8.RuneError {
|
||||
continue
|
||||
}
|
||||
p := make([]byte, utf8.RuneLen(rn))
|
||||
utf8.EncodeRune(p, rn)
|
||||
p = bytes.ReplaceAll(p, byteT, byteS)
|
||||
for bytes.HasSuffix(p, byteDel) {
|
||||
p = p[:len(p)-1]
|
||||
if buf.Len() > 0 {
|
||||
var dels []byte
|
||||
last, ok := lo.Last([]rune(buf.String()))
|
||||
for i := 0; ok && i < lipgloss.Width(string(last)); i++ {
|
||||
dels = append(dels, byteClearCur...)
|
||||
}
|
||||
chs.OutChan <- dels
|
||||
buf.Truncate(buf.Len() - len([]byte(string(last))))
|
||||
}
|
||||
}
|
||||
if len(p) <= 0 {
|
||||
continue
|
||||
}
|
||||
chs.OutChan <- p
|
||||
buf.Write(p)
|
||||
bs := buf.Bytes()
|
||||
if idx := bytes.LastIndex(bs, byteClearAll); idx >= 0 {
|
||||
buf.Reset()
|
||||
continue
|
||||
}
|
||||
if idx := bytes.LastIndex(bs, byteR); idx < 0 {
|
||||
continue
|
||||
}
|
||||
bs = bs[:len(bs)-1]
|
||||
if bytes.Equal(bs, []byte("exit")) {
|
||||
sess.Once.Do(func() { close(chs.AwayChan) })
|
||||
// Create a reader to read PTY output
|
||||
ptmxReader := bufio.NewReader(ptmx)
|
||||
|
||||
// Add an atomic variable to track if exit message has been sent
|
||||
var exitMessageSent int32
|
||||
|
||||
// Monitor process exit
|
||||
sess.G.Go(func() error {
|
||||
err := cmd.Wait()
|
||||
logger.L().Info("Redis cli process exited", zap.Error(err))
|
||||
|
||||
// Only send termination message if not already sent
|
||||
if atomic.CompareAndSwapInt32(&exitMessageSent, 0, 1) {
|
||||
// Send termination message
|
||||
terminationMsg := "\r\n\033[31mThe connection is closed!\033[0m\r\n"
|
||||
chs.OutBuf.WriteString(terminationMsg)
|
||||
}
|
||||
|
||||
sess.Once.Do(func() {
|
||||
logger.L().Info("Closing AwayChan from Redis process monitor")
|
||||
close(chs.AwayChan)
|
||||
})
|
||||
return fmt.Errorf("redis-cli process terminated: %w", err)
|
||||
})
|
||||
|
||||
// Goroutine 1: Process input, detect exit command
|
||||
sess.G.Go(func() error {
|
||||
defer ptmx.Close()
|
||||
buf := make([]byte, 1024)
|
||||
var inputBuffer string
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-sess.Gctx.Done():
|
||||
return nil
|
||||
}
|
||||
buf.Reset()
|
||||
var (
|
||||
res any
|
||||
rows *sql.Rows
|
||||
)
|
||||
if len(bs) > 0 {
|
||||
switch protocol {
|
||||
case "redis":
|
||||
parts := lo.Map(reRedis.FindAllString(string(bs), -1), func(p string, _ int) any { return p })
|
||||
res, err = rdb.Do(ctx, parts...).Result()
|
||||
case "mysql":
|
||||
if rows, err = db.WithContext(ctx).Raw(string(bs)).Rows(); err == nil {
|
||||
heads, _ := rows.Columns()
|
||||
n := len(heads)
|
||||
rs := make([][]string, 0)
|
||||
for rows.Next() {
|
||||
r := make([]any, n)
|
||||
r = lo.Map(r, func(v any, _ int) any { return new(any) })
|
||||
if err = rows.Scan(r...); err != nil {
|
||||
default:
|
||||
n, err := chs.Rin.Read(buf)
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if n > 0 {
|
||||
input := string(buf[:n])
|
||||
|
||||
// Accumulate user input to detect complete commands
|
||||
if input != "\r" {
|
||||
// Check if all characters are printable
|
||||
allPrintable := true
|
||||
for _, ch := range input {
|
||||
if ch < 32 || ch > 126 {
|
||||
allPrintable = false
|
||||
break
|
||||
}
|
||||
rs = append(rs, lo.Map(r, func(v any, i int) string { return cast.ToString(v) }))
|
||||
}
|
||||
res = strings.ReplaceAll(table.New().Border(border).Headers(heads...).Rows(rs...).String(), "\n", "\r\n")
|
||||
if allPrintable {
|
||||
inputBuffer += input
|
||||
}
|
||||
}
|
||||
|
||||
// Detect command end (enter key)
|
||||
if input == "\r" {
|
||||
processCmd := strings.TrimSpace(inputBuffer)
|
||||
|
||||
// Check for exit command
|
||||
if strings.EqualFold(processCmd, "exit") || strings.EqualFold(processCmd, "quit") {
|
||||
// Send command to Redis CLI for normal exit
|
||||
if _, err := ptmx.Write(buf[:n]); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Mark exit message as sent, but don't send it here. Let the process exit handler send it.
|
||||
atomic.StoreInt32(&exitMessageSent, 1)
|
||||
|
||||
inputBuffer = ""
|
||||
continue
|
||||
}
|
||||
|
||||
// Reset command buffer
|
||||
inputBuffer = ""
|
||||
}
|
||||
|
||||
// Forward input to Redis CLI
|
||||
if _, err := ptmx.Write(buf[:n]); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
chs.OutChan <- []byte(fmt.Sprintf("\n%s\r\n%s", lo.Ternary[any](err == nil, lo.Ternary(res == nil, "", res), err), sess.Prompt))
|
||||
err = nil
|
||||
}
|
||||
}
|
||||
})
|
||||
sess.G.Go(func() (err error) {
|
||||
for {
|
||||
select {
|
||||
case <-sess.Gctx.Done():
|
||||
return
|
||||
case <-chs.AwayChan:
|
||||
return
|
||||
case <-chs.WindowChan:
|
||||
continue
|
||||
})
|
||||
|
||||
// Goroutine 2: Read Redis CLI output and send to OutChan
|
||||
sess.G.Go(func() error {
|
||||
for {
|
||||
select {
|
||||
case <-sess.Gctx.Done():
|
||||
return nil
|
||||
default:
|
||||
rn, size, err := ptmxReader.ReadRune()
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if size <= 0 || rn == utf8.RuneError {
|
||||
continue
|
||||
}
|
||||
|
||||
p := make([]byte, utf8.RuneLen(rn))
|
||||
utf8.EncodeRune(p, rn)
|
||||
|
||||
// Send to OutChan for HandleTerm processing
|
||||
chs.OutChan <- p
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
sess.G.Wait()
|
||||
// Goroutine 3: Handle window size changes
|
||||
sess.G.Go(func() error {
|
||||
for {
|
||||
select {
|
||||
case <-sess.Gctx.Done():
|
||||
return nil
|
||||
case <-chs.AwayChan:
|
||||
return fmt.Errorf("away")
|
||||
case window := <-chs.WindowChan:
|
||||
// Adjust Redis terminal size
|
||||
_ = pty.Setsize(ptmx, &pty.Winsize{
|
||||
Cols: uint16(window.Width),
|
||||
Rows: uint16(window.Height),
|
||||
})
|
||||
|
||||
return
|
||||
// Adjust parser size
|
||||
if sess.SshParser != nil {
|
||||
sess.SshParser.Resize(window.Width, window.Height)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// Notify connection is ready
|
||||
chs.ErrChan <- nil
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("unsupported protocol: %s", sess.Protocol)
|
||||
}
|
||||
|
||||
@@ -135,12 +135,12 @@ func connectSsh(ctx *gin.Context, sess *gsession.Session, asset *model.Asset, ac
|
||||
// HandleTerm handles terminal sessions
|
||||
func HandleTerm(sess *gsession.Session) (err error) {
|
||||
defer func() {
|
||||
logger.L().Debug("defer HandleSsh", zap.String("sessionId", sess.SessionId))
|
||||
logger.L().Debug("defer HandleTerm", zap.String("sessionId", sess.SessionId))
|
||||
sess.SshParser.Close(sess.Prompt)
|
||||
sess.Status = model.SESSIONSTATUS_OFFLINE
|
||||
sess.ClosedAt = lo.ToPtr(time.Now())
|
||||
if err = gsession.UpsertSession(sess); err != nil {
|
||||
logger.L().Error("offline ssh session failed", zap.String("sessionId", sess.SessionId), zap.Error(err))
|
||||
logger.L().Error("offline session failed", zap.String("sessionId", sess.SessionId), zap.Error(err))
|
||||
return
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -3,12 +3,10 @@ package connect
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/nicksnyder/go-i18n/v2/i18n"
|
||||
@@ -33,16 +31,8 @@ var (
|
||||
},
|
||||
}
|
||||
byteClearAll = []byte("\x15\r")
|
||||
byteClearCur = []byte("\b\x1b[J")
|
||||
byteDel = []byte{'\x7f'}
|
||||
byteR = []byte{'\r'}
|
||||
byteN = []byte{'\n'}
|
||||
byteT = []byte{'\t'}
|
||||
byteS = []byte{' '}
|
||||
byteRN = append(byteR, byteN...)
|
||||
|
||||
reRedis = regexp.MustCompile(`("[^"]*"|'[^']*'|\S+)`)
|
||||
border = lipgloss.RoundedBorder()
|
||||
wsWriteMutex = &sync.Mutex{}
|
||||
)
|
||||
|
||||
// WriteToMonitors sends data to all monitoring sessions
|
||||
@@ -117,6 +107,19 @@ func OfflineSession(ctx *gin.Context, sessionId string, closer string) {
|
||||
// HandleError handles errors from sessions
|
||||
func HandleError(ctx *gin.Context, sess *gsession.Session, err error, ws *websocket.Conn, chs *gsession.SessionChans) {
|
||||
defer func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.L().Error("Recovered from panic in HandleError",
|
||||
zap.Any("panic", r),
|
||||
zap.String("sessionId", func() string {
|
||||
if sess != nil {
|
||||
return sess.SessionId
|
||||
}
|
||||
return ""
|
||||
}()))
|
||||
}
|
||||
}()
|
||||
|
||||
if sess == nil || sess.Chans == nil {
|
||||
return
|
||||
}
|
||||
@@ -124,7 +127,12 @@ func HandleError(ctx *gin.Context, sess *gsession.Session, err error, ws *websoc
|
||||
if chs != nil {
|
||||
ch = chs.AwayChan
|
||||
}
|
||||
sess.Once.Do(func() { close(ch) })
|
||||
|
||||
sess.Once.Do(func() {
|
||||
logger.L().Debug("Closing AwayChan from HandleError",
|
||||
zap.String("sessionId", sess.SessionId))
|
||||
close(ch)
|
||||
})
|
||||
}()
|
||||
|
||||
if err == nil {
|
||||
@@ -150,19 +158,24 @@ func WriteErrMsg(sess *gsession.Session, msg string) {
|
||||
}
|
||||
|
||||
// Write writes data to the session output
|
||||
func Write(sess *gsession.Session) (err error) {
|
||||
// skipRecording: If true, it will skip recording to avoid duplicate recordings of manually recorded content
|
||||
func Write(sess *gsession.Session, skipRecording ...bool) (err error) {
|
||||
chs := sess.Chans
|
||||
out := chs.OutBuf.Bytes()
|
||||
|
||||
if sess.SessionType == model.SESSIONTYPE_WEB && sess.Ws != nil {
|
||||
if len(out) > 0 || sess.IsGuacd() {
|
||||
wsWriteMutex.Lock()
|
||||
defer wsWriteMutex.Unlock()
|
||||
err = sess.Ws.WriteMessage(websocket.TextMessage, out)
|
||||
}
|
||||
} else if sess.SessionType == model.SESSIONTYPE_CLIENT && len(out) > 0 {
|
||||
_, err = sess.CliRw.Write(out)
|
||||
}
|
||||
|
||||
if sess.SshRecoder != nil && len(out) > 0 && !sess.IsGuacd() {
|
||||
// Only write to recording if skipRecording is not specified or explicitly set to false
|
||||
shouldSkip := len(skipRecording) > 0 && skipRecording[0]
|
||||
if sess.SshRecoder != nil && len(out) > 0 && !sess.IsGuacd() && !shouldSkip {
|
||||
sess.SshRecoder.Write(out)
|
||||
}
|
||||
|
||||
|
||||
@@ -57,6 +57,7 @@ type Parser struct {
|
||||
Input []byte
|
||||
Output []byte
|
||||
SessionId string
|
||||
Protocol string
|
||||
Cmds []*model.Command
|
||||
isPrompt bool
|
||||
prompt string
|
||||
@@ -82,19 +83,49 @@ func (p *Parser) AddInput(bs []byte) (cmd string, forbidden bool) {
|
||||
p.lastCmd = ""
|
||||
p.lastRes = ""
|
||||
}
|
||||
|
||||
// Track command input, directly use curCmd field
|
||||
if len(bs) > 0 {
|
||||
if bytes.Equal(bs, []byte("\r")) {
|
||||
// Command ends, keep curCmd unchanged
|
||||
} else if len(bs) == 1 && (bs[0] == '\b' || bs[0] == 127) { // Backspace key
|
||||
if len(p.curCmd) > 0 {
|
||||
p.curCmd = p.curCmd[:len(p.curCmd)-1]
|
||||
}
|
||||
} else {
|
||||
// Check if all characters are printable
|
||||
input := string(bs)
|
||||
allPrintable := true
|
||||
for _, ch := range input {
|
||||
if ch < 32 || ch > 126 {
|
||||
allPrintable = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if allPrintable {
|
||||
p.curCmd += input
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
p.Input = append(p.Input, bs...)
|
||||
if !bytes.HasSuffix(p.Input, []byte("\r")) {
|
||||
return
|
||||
}
|
||||
|
||||
p.isPrompt = true
|
||||
p.curCmd = p.getCmdLocked()
|
||||
// Save current command
|
||||
currentCmd := strings.TrimSpace(p.curCmd)
|
||||
// Reset command buffer
|
||||
p.curCmd = ""
|
||||
p.resetLocked()
|
||||
|
||||
filter := ""
|
||||
if filter, forbidden = p.IsForbidden(p.curCmd); forbidden {
|
||||
if filter, forbidden = p.IsForbidden(currentCmd); forbidden {
|
||||
cmd = filter
|
||||
return
|
||||
}
|
||||
p.lastCmd = p.curCmd
|
||||
p.lastCmd = currentCmd
|
||||
return
|
||||
}
|
||||
|
||||
@@ -117,7 +148,7 @@ func (p *Parser) IsForbidden(cmd string) (string, bool) {
|
||||
}
|
||||
|
||||
func (p *Parser) WriteDb() {
|
||||
if p.lastCmd == "" {
|
||||
if p.lastCmd == "" || strings.TrimSpace(p.lastCmd) == "" {
|
||||
return
|
||||
}
|
||||
m := &model.SessionCmd{
|
||||
@@ -165,6 +196,12 @@ func (p *Parser) GetCmd() string {
|
||||
}
|
||||
|
||||
func (p *Parser) getCmdLocked() string {
|
||||
// If the current command being built is not empty, return it
|
||||
if p.curCmd != "" {
|
||||
return p.curCmd
|
||||
}
|
||||
|
||||
// Otherwise extract from output
|
||||
s := p.getOutputLocked()
|
||||
// TODO: some promot may change with its dir
|
||||
return strings.TrimPrefix(s, p.prompt)
|
||||
@@ -209,7 +246,20 @@ func (p *Parser) getOutputLocked() string {
|
||||
|
||||
p.lastRes = ""
|
||||
if ln > 1 {
|
||||
p.lastRes = strings.Join(res[:ln-1], "\n")
|
||||
// Process result based on protocol type
|
||||
if strings.HasPrefix(p.Protocol, "ssh") {
|
||||
// For SSH sessions, keep the original logic, retain all lines
|
||||
p.lastRes = strings.Join(res[:ln-1], "\n")
|
||||
} else {
|
||||
// For non-SSH sessions (like Redis, MySQL), remove first and last line
|
||||
startIdx := 1
|
||||
endIdx := ln - 1
|
||||
|
||||
// Ensure there are enough lines to process
|
||||
if ln > 2 {
|
||||
p.lastRes = strings.Join(res[startIdx:endIdx], "\n")
|
||||
}
|
||||
}
|
||||
}
|
||||
p.curRes = res[ln-1]
|
||||
return p.curRes
|
||||
|
||||
Reference in New Issue
Block a user