dashboard/controller/websocket.go
2025-05-21 19:14:12 +08:00

97 lines
2.0 KiB
Go

package controller
import (
"context"
"dashboard/utils"
"net/http"
"os/exec"
"github.com/creack/pty"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"golang.org/x/sync/errgroup"
)
func TerminalHandle() func(c *gin.Context) error {
upgrader := websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool {
return true // 允许跨域
},
}
return func(c *gin.Context) error {
log, _ := utils.GetLogFromContext(c)
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
log.Sugar().Errorf("WebSocket upgrade failed: %v", err)
return err
}
defer conn.Close()
cmd := exec.Command("sh")
ptyFile, err := pty.Start(cmd)
if err != nil {
log.Sugar().Errorf("Failed to start pty: %v", err)
return err
}
defer ptyFile.Close()
wg, ctx := errgroup.WithContext(context.Background())
wg.Go(func() error {
defer log.Sugar().Errorf("task read over")
for {
select {
case <-ctx.Done():
log.Sugar().Errorf("task canceld read")
return ctx.Err()
default:
}
_, msg, err := conn.ReadMessage()
if err != nil {
log.Sugar().Errorf("WebSocket read error: %v", err)
return err
}
_, err = ptyFile.Write(msg)
if err != nil {
log.Sugar().Errorf("PTY write error: %v", err)
return err
}
}
})
wg.Go(func() error {
defer log.Sugar().Errorf("task write over")
buf := make([]byte, 1024)
for {
select {
case <-ctx.Done():
log.Sugar().Errorf("task canceld write")
return ctx.Err()
default:
}
n, err := ptyFile.Read(buf)
if err != nil {
log.Sugar().Errorf("PTY read error: %v", err)
return err
}
err = conn.WriteMessage(websocket.BinaryMessage, buf[:n])
if err != nil {
log.Sugar().Errorf("WebSocket write error: %v", err)
return err
}
}
})
if err := wg.Wait(); err != nil {
log.Sugar().Errorf("WebSocket error: %v", err)
}
return nil
}
}