feat: 更换websocket库

2026-02-04 05:31:44 +08:00 · 2025-11-23 01:49:36 +08:00
parent 378993f7c5
commit c60bba465d
5 changed files with 53 additions and 65 deletions
--- a/cmd/ace/wire_gen.go
+++ b/cmd/ace/wire_gen.go
@@ -143,7 +143,7 @@ func initWeb() (*app.Web, error) {
 	supervisorApp := supervisor.NewApp(locale)
 	loader := bootstrap.NewLoader(codeserverApp, dockerApp, fail2banApp, frpApp, giteaApp, memcachedApp, minioApp, mysqlApp, nginxApp, php74App, php80App, php81App, php82App, php83App, php84App, phpmyadminApp, podmanApp, postgresqlApp, pureftpdApp, redisApp, rsyncApp, s3fsApp, supervisorApp)
 	http := route.NewHttp(koanf, userService, userTokenService, dashboardService, taskService, websiteService, databaseService, databaseServerService, databaseUserService, backupService, certService, certDNSService, certAccountService, appService, cronService, processService, safeService, firewallService, sshService, containerService, containerComposeService, containerNetworkService, containerImageService, containerVolumeService, fileService, monitorService, settingService, systemctlService, toolboxSystemService, toolboxBenchmarkService, loader)
-	wsService := service.NewWsService(locale, koanf, sshRepo)
+	wsService := service.NewWsService(locale, koanf, logger, sshRepo)
 	ws := route.NewWs(wsService)
 	mux, err := bootstrap.NewRouter(locale, middlewares, http, ws)
 	if err != nil {
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.92.0
 	github.com/bddjr/hlfhr v1.4.0
 	github.com/beevik/ntp v1.5.0
+	github.com/coder/websocket v1.8.12
 	github.com/creack/pty v1.1.24
 	github.com/expr-lang/expr v1.17.6
 	github.com/go-chi/chi/v5 v5.2.3
@@ -20,7 +21,6 @@ require (
 	github.com/google/wire v0.7.0
 	github.com/gookit/color v1.6.0
 	github.com/gookit/validate v1.5.6
-	github.com/gorilla/websocket v1.5.3
 	github.com/hashicorp/go-version v1.7.0
 	github.com/jlaffaye/ftp v0.2.0
 	github.com/knadh/koanf/parsers/yaml v1.1.0
@@ -40,7 +40,7 @@ require (
 	github.com/libdns/westcn v1.0.2
 	github.com/libtnb/chix v1.3.2
 	github.com/libtnb/gormstore v1.1.1
-	github.com/libtnb/sessions v1.2.1
+	github.com/libtnb/sessions v1.2.2-0.20251122173530-a4002b1c459d
 	github.com/libtnb/utils v1.2.1
 	github.com/mholt/acmez/v3 v3.1.4
 	github.com/moby/moby/api v1.52.0
--- a/go.sum
+++ b/go.sum
@@ -76,6 +76,8 @@ github.com/boombuler/barcode v1.1.0 h1:ChaYjBR63fr4LFyGn8E8nt7dBSt3MiU3zMOZqFvVk
 github.com/boombuler/barcode v1.1.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
+github.com/coder/websocket v1.8.12 h1:5bUXkEPPIbewrnkU8LTCLVaxi4N4J8ahufH2vlo4NAo=
+github.com/coder/websocket v1.8.12/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs=
 github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
 github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M=
 github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE=
@@ -183,8 +185,6 @@ github.com/gookit/validate v1.5.6 h1:D6vbSZzreuKYpeeXm5FDDEJy3K5E4lcWsQE4saSMZbU
 github.com/gookit/validate v1.5.6/go.mod h1:WYEHndRNepIIkM+6CtgEX9MQ9ToIQRhXxmz5oLHF/fc=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
-github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
-github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
@@ -280,8 +280,8 @@ github.com/libtnb/gormstore v1.1.1 h1:FG/3P4PuWM6/vB4weVJ31meiSaoeXns1NQlP66quKe
 github.com/libtnb/gormstore v1.1.1/go.mod h1:8A5QzeZxi1MpSmjUVsHTDAL6KnU84feIXMutFLPawwA=
 github.com/libtnb/securecookie v1.2.0 h1:2uc0PBDm0foeSTrcZ9QTX1IEjf6kFEwfgEYSIXQSKrA=
 github.com/libtnb/securecookie v1.2.0/go.mod h1:ja+wNGnQzYqcqXQnJWu6icsaWi5JEBwNEMJ2ReTVDxA=
-github.com/libtnb/sessions v1.2.1 h1:O9gkEIeZuqyaxopXrUJcGxlNxmNfRBI8BOK43yLJXDI=
-github.com/libtnb/sessions v1.2.1/go.mod h1:45Bn9d6PseDINLIM1QaJrlCMbzSZ0NWpDbWkdrKJKw0=
+github.com/libtnb/sessions v1.2.2-0.20251122173530-a4002b1c459d h1:PIS6RcMg03UlAkLuif8go4G5fv1x6xFZBK7koBwNd4c=
+github.com/libtnb/sessions v1.2.2-0.20251122173530-a4002b1c459d/go.mod h1:qw+FWtBtrPDYCf6MfX0Lk5EhTArpvT72z5Ei4RUMTRg=
 github.com/libtnb/utils v1.2.1 h1:LJmReRREnpqfHyy9PZtNgBh3ZaIGct81b8ZaAsolMkM=
 github.com/libtnb/utils v1.2.1/go.mod h1:o6LEDeC42PXI21uLWdWJWTVYvR9BtAZfzzTGJVQoQiU=
 github.com/magiconair/properties v1.8.1/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
--- a/internal/service/ws.go
+++ b/internal/service/ws.go
@@ -3,10 +3,11 @@ package service
 import (
 	"bufio"
 	"context"
+	"log/slog"
 	"net/http"
 	"sync"

-	"github.com/gorilla/websocket"
+	"github.com/coder/websocket"
 	"github.com/knadh/koanf/v2"
 	"github.com/leonelquinteros/gotext"
 	stdssh "golang.org/x/crypto/ssh"
@@ -20,13 +21,15 @@ import (
 type WsService struct {
 	t       *gotext.Locale
 	conf    *koanf.Koanf
+	log     *slog.Logger
 	sshRepo biz.SSHRepo
 }

-func NewWsService(t *gotext.Locale, conf *koanf.Koanf, ssh biz.SSHRepo) *WsService {
+func NewWsService(t *gotext.Locale, conf *koanf.Koanf, log *slog.Logger, ssh biz.SSHRepo) *WsService {
 	return &WsService{
 		t:       t,
 		conf:    conf,
+		log:     log,
 		sshRepo: ssh,
 	}
 }
@@ -45,32 +48,28 @@ func (s *WsService) Session(w http.ResponseWriter, r *http.Request) {

 	ws, err := s.upgrade(w, r)
 	if err != nil {
-		ErrorSystem(w)
+		s.log.Warn("[Websocket] upgrade session ws error", slog.Any("error", err))
 		return
 	}
-	defer func(ws *websocket.Conn) {
-		_ = ws.Close()
-	}(ws)
+	defer func(ws *websocket.Conn) { _ = ws.CloseNow() }(ws)

 	client, err := ssh.NewSSHClient(info.Config)
 	if err != nil {
-		_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, err.Error()))
+		_ = ws.Close(websocket.StatusNormalClosure, err.Error())
 		return
 	}
-	defer func(client *stdssh.Client) {
-		_ = client.Close()
-	}(client)
-
-	turn, err := ssh.NewTurn(ws, client)
-	if err != nil {
-		_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, err.Error()))
-		return
-	}
-	defer func(turn *ssh.Turn) {
-		_ = turn.Close()
-	}(turn)
+	defer func(client *stdssh.Client) { _ = client.Close() }(client)

 	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	turn, err := ssh.NewTurn(ctx, ws, client)
+	if err != nil {
+		_ = ws.Close(websocket.StatusNormalClosure, err.Error())
+		return
+	}
+	defer func(turn *ssh.Turn) { _ = turn.Close() }(turn)
+
 	wg := sync.WaitGroup{}
 	wg.Add(2)

@@ -84,31 +83,29 @@ func (s *WsService) Session(w http.ResponseWriter, r *http.Request) {
 	}()

 	wg.Wait()
-	cancel()
 }

 func (s *WsService) Exec(w http.ResponseWriter, r *http.Request) {
 	ws, err := s.upgrade(w, r)
 	if err != nil {
-		ErrorSystem(w)
+		s.log.Warn("[Websocket] upgrade exec ws error", slog.Any("error", err))
 		return
 	}
-	defer func(ws *websocket.Conn) {
-		_ = ws.Close()
-	}(ws)
+	defer func(ws *websocket.Conn) { _ = ws.CloseNow() }(ws)

 	// 第一条消息是命令
-	_, cmd, err := ws.ReadMessage()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	_, cmd, err := ws.Read(ctx)
 	if err != nil {
-		_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, s.t.Get("failed to read command: %v", err)))
+		_ = ws.Close(websocket.StatusNormalClosure, s.t.Get("failed to read command: %v", err))
 		return
 	}

-	ctx, cancel := context.WithCancel(context.Background())
 	out, err := shell.ExecfWithPipe(ctx, string(cmd))
 	if err != nil {
-		_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, s.t.Get("failed to run command: %v", err)))
-		cancel()
+		_ = ws.Close(websocket.StatusNormalClosure, s.t.Get("failed to run command: %v", err))
 		return
 	}

@@ -116,38 +113,34 @@ func (s *WsService) Exec(w http.ResponseWriter, r *http.Request) {
 		scanner := bufio.NewScanner(out)
 		for scanner.Scan() {
 			line := scanner.Text()
-			_ = ws.WriteMessage(websocket.TextMessage, []byte(line))
+			_ = ws.Write(ctx, websocket.MessageText, []byte(line))
 		}
 		if err = scanner.Err(); err != nil {
-			_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, s.t.Get("failed to read command output: %v", err)))
+			_ = ws.Close(websocket.StatusNormalClosure, s.t.Get("failed to read command output: %v", err))
 		}
 	}()

-	s.readLoop(ws)
-	cancel()
+	s.readLoop(ctx, ws)
 }

 func (s *WsService) upgrade(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
-	upGrader := websocket.Upgrader{
-		ReadBufferSize:  4096,
-		WriteBufferSize: 4096,
+	opts := &websocket.AcceptOptions{
+		CompressionMode: websocket.CompressionContextTakeover,
 	}

 	// debug 模式下不校验 origin，方便 vite 代理调试
 	if s.conf.Bool("app.debug") {
-		upGrader.CheckOrigin = func(r *http.Request) bool {
-			return true
-		}
+		opts.InsecureSkipVerify = true
 	}

-	return upGrader.Upgrade(w, r, nil)
+	return websocket.Accept(w, r, opts)
 }

 // readLoop 阻塞直到客户端关闭连接
-func (s *WsService) readLoop(c *websocket.Conn) {
+func (s *WsService) readLoop(ctx context.Context, c *websocket.Conn) {
 	for {
-		if _, _, err := c.NextReader(); err != nil {
-			_ = c.Close()
+		if _, _, err := c.Read(ctx); err != nil {
+			_ = c.CloseNow()
 			break
 		}
 	}
--- a/pkg/ssh/turn.go
+++ b/pkg/ssh/turn.go
@@ -7,7 +7,7 @@ import (
 	"fmt"
 	"io"

-	"github.com/gorilla/websocket"
+	"github.com/coder/websocket"
 	"golang.org/x/crypto/ssh"
 )

@@ -18,12 +18,13 @@ type MessageResize struct {
 }

 type Turn struct {
+	ctx     context.Context
 	stdin   io.WriteCloser
 	session *ssh.Session
 	ws      *websocket.Conn
 }

-func NewTurn(ws *websocket.Conn, client *ssh.Client) (*Turn, error) {
+func NewTurn(ctx context.Context, ws *websocket.Conn, client *ssh.Client) (*Turn, error) {
 	sess, err := client.NewSession()
 	if err != nil {
 		return nil, err
@@ -34,7 +35,7 @@ func NewTurn(ws *websocket.Conn, client *ssh.Client) (*Turn, error) {
 		return nil, err
 	}

-	turn := &Turn{stdin: stdin, session: sess, ws: ws}
+	turn := &Turn{ctx: ctx, stdin: stdin, session: sess, ws: ws}
 	sess.Stdout = turn
 	sess.Stderr = turn

@@ -54,33 +55,27 @@ func NewTurn(ws *websocket.Conn, client *ssh.Client) (*Turn, error) {
 }

 func (t *Turn) Write(p []byte) (n int, err error) {
-	writer, err := t.ws.NextWriter(websocket.TextMessage)
-	if err != nil {
+	if err = t.ws.Write(t.ctx, websocket.MessageText, p); err != nil {
 		return 0, err
 	}
-	defer func(writer io.WriteCloser) {
-		_ = writer.Close()
-	}(writer)
-
-	return writer.Write(p)
+	return len(p), nil
 }

 func (t *Turn) Close() error {
 	if t.session != nil {
 		_ = t.session.Close()
 	}
-
-	return t.ws.Close()
+	return t.ws.CloseNow()
 }

-func (t *Turn) Handle(context context.Context) error {
+func (t *Turn) Handle(ctx context.Context) error {
 	var resize MessageResize
 	for {
 		select {
-		case <-context.Done():
+		case <-ctx.Done():
 			return errors.New("ssh context done exit")
 		default:
-			_, data, err := t.ws.ReadMessage()
+			_, data, err := t.ws.Read(ctx)
 			if err != nil {
 				// 通常是客户端关闭连接
 				return fmt.Errorf("reading ws message err: %v", err)