fix: some websocket error optimization (#4)

fix: websocket panic 
fix: websocket error disable
fix: s tools deploy
fix: remove debug dependency
chore: add debug pprof

.gitignore

@@ -26,4 +26,6 @@ dist/
 *.toml
 .config
 .secret
-node_modules
+node_modules
+bootstrap
+.s/

cmd/client.go

@@ -1,11 +1,13 @@
 package main
 
 import (
-	"github.com/DVKunion/SeaMoon/pkg/client"
-	"github.com/DVKunion/SeaMoon/pkg/consts"
+	"os"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
-	"os"
+
+	"github.com/DVKunion/SeaMoon/pkg/client"
+	"github.com/DVKunion/SeaMoon/pkg/consts"
 )
 
 var (

go.mod

@@ -4,19 +4,16 @@ go 1.18
 
 require (
 	github.com/BurntSushi/toml v0.3.1
-	github.com/elazarl/goproxy v0.0.0-20220529153421-8ea89ba92021
 	github.com/gin-gonic/gin v1.8.1
 	github.com/google/martian/v3 v3.3.2
 	github.com/gorilla/websocket v1.4.2
 	github.com/sirupsen/logrus v1.9.0
 	github.com/spf13/cobra v1.5.0
-	github.com/stretchr/testify v1.7.1
 	github.com/tg123/go-htpasswd v1.2.0
 )
 
 require (
 	github.com/GehirnInc/crypt v0.0.0-20200316065508-bb7000b8a962 // indirect
-	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
 	github.com/go-playground/locales v0.14.0 // indirect
 	github.com/go-playground/universal-translator v0.18.0 // indirect
@@ -29,7 +26,6 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.1 // indirect
-	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/ugorji/go/codec v1.2.7 // indirect
 	golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97 // indirect
@@ -38,5 +34,4 @@ require (
 	golang.org/x/text v0.3.6 // indirect
 	google.golang.org/protobuf v1.28.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )

go.sum

@@ -11,10 +11,6 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/elazarl/goproxy v0.0.0-20220529153421-8ea89ba92021 h1:EbF0UihnxWRcIMOwoVtqnAylsqcjzqpSvMdjF2Ud4rA=
-github.com/elazarl/goproxy v0.0.0-20220529153421-8ea89ba92021/go.mod h1:Ro8st/ElPeALwNFlcTpWmkr6IoMFfkjXAvTHpevnDsM=
-github.com/elazarl/goproxy/ext v0.0.0-20190711103511-473e67f1d7d2 h1:dWB6v3RcOy03t/bUadywsbyrQwCqZeNIEX6M1OtSZOM=
-github.com/elazarl/goproxy/ext v0.0.0-20190711103511-473e67f1d7d2/go.mod h1:gNh8nYJoAm43RfaxurUnxr+N1PwuFV3ZMl/efxlIlY8=
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.9-0.20210217033140-668b12f5399d/go.mod h1:cXg6YxExXjJnVBQHBLXeUAgxn2UodCpnH306RInaBQk=
@@ -88,7 +84,6 @@ github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsK
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
-github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
 github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
 github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
 github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=

pkg/client/client.go

@@ -3,15 +3,19 @@ package client
 import (
 	"bufio"
 	"context"
-	"github.com/DVKunion/SeaMoon/pkg/consts"
-	"github.com/DVKunion/SeaMoon/static"
-	"github.com/gin-gonic/gin"
-	log "github.com/sirupsen/logrus"
 	"html/template"
 	"io"
 	"net"
 	"net/http"
 	"os"
+
+	"github.com/gin-gonic/gin"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/DVKunion/SeaMoon/pkg/consts"
+	"github.com/DVKunion/SeaMoon/static"
+
+	_ "net/http/pprof"
 )
 
 type Client struct {
@@ -69,6 +73,11 @@ func Controller(sg *SigGroup, verbose bool, debug bool) {
 		ctx.HTML(200, "index.html", Config())
 	})
 
+	// pprof
+	if debug {
+		server.GET("/debug/pprof/*any", gin.WrapH(http.DefaultServeMux))
+	}
+
 	// controller set
 	server.POST("/", func(ctx *gin.Context) {
 		if err := ctx.ShouldBindJSON(Config()); err != nil {

pkg/client/socks.go

@@ -3,14 +3,17 @@ package client
 import (
 	"bufio"
 	"context"
-	"github.com/DVKunion/SeaMoon/pkg/consts"
-	"github.com/DVKunion/SeaMoon/pkg/server"
-	"github.com/DVKunion/SeaMoon/pkg/utils"
-	"github.com/gorilla/websocket"
-	log "github.com/sirupsen/logrus"
 	"net"
 	"net/http"
 	"strings"
+	"sync"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/DVKunion/SeaMoon/pkg/consts"
+	"github.com/DVKunion/SeaMoon/pkg/server"
+	"github.com/DVKunion/SeaMoon/pkg/utils"
 )
 
 type bufferedConn struct {
@@ -70,6 +73,7 @@ func NewSocks5Client(ctx context.Context, server net.Listener, proxyAddr string)
 	}()
 	go func() {
 		for {
+			lock := &sync.Mutex{}
 			conn, err := server.Accept()
 			if err == nil {
 				log.Debugf(consts.SOCKS5_ACCEPT_START, conn.RemoteAddr())
@@ -79,7 +83,7 @@ func NewSocks5Client(ctx context.Context, server net.Listener, proxyAddr string)
 					log.Errorf(consts.CLIENT_PROTOCOL_UNSUPPORT_ERROR, err)
 					return
 				}
-				go Socks5Handler(&bufferedConn{conn, br}, proxyAddr)
+				go Socks5Handler(&bufferedConn{conn, br}, proxyAddr, lock)
 			} else {
 				if closeFlag {
 					// except close
@@ -93,7 +97,7 @@ func NewSocks5Client(ctx context.Context, server net.Listener, proxyAddr string)
 	<-ctx.Done()
 }
 
-func Socks5Handler(conn net.Conn, raddr string) {
+func Socks5Handler(conn net.Conn, raddr string, lock *sync.Mutex) {
 	// select method
 	_, err := utils.ReadMethods(conn)
 	if err != nil {
@@ -119,7 +123,7 @@ func Socks5Handler(conn net.Conn, raddr string) {
 	}
 	switch request.Cmd {
 	case utils.CmdConnect:
-		handleConnect(conn, request, raddr)
+		handleConnect(conn, request, raddr, lock)
 		break
 	case utils.CmdBind:
 		log.Error("not support cmd bind")
@@ -132,7 +136,7 @@ func Socks5Handler(conn net.Conn, raddr string) {
 	}
 }
 
-func handleConnect(conn net.Conn, req *utils.Request, rAddr string) {
+func handleConnect(conn net.Conn, req *utils.Request, rAddr string, lock *sync.Mutex) {
 
 	log.Infof(consts.SOCKS5_CONNECT_SERVER, req.Addr, conn.RemoteAddr())
 
@@ -147,7 +151,7 @@ func handleConnect(conn net.Conn, req *utils.Request, rAddr string) {
 		return
 	}
 
-	newConn := server.NewWebsocketServer(wsConn)
+	newConn := server.NewWebsocketServer(wsConn, lock)
 
 	defer newConn.Close()
 

pkg/consts/log.go

@@ -43,7 +43,7 @@ const (
 	SOCKS5_LISTEN_START    string = "[Socks5] Client Start Listen At: %s"
 	SOCKS5_LISTEN_STOP     string = "[Socks5] Client Stop Listen"
 	SOCKS5_ACCEPT_START    string = "[Socks5] Client Accept Conn From: %s"
-	SOCKS5_CONNECT_SERVER  string = "[Socks5] Server Connect %s For %s"
+	SOCKS5_CONNECT_SERVER  string = "[Socks5] Server Connect %s From %s"
 	SOCKS5_CONNECT_ESTAB   string = "[Socks5] Connect Tunnel Established %s <-> %s"
 	SOCKS5_CONNECT_DIS     string = "[Socks5] Connect Tunnel Disconnected %s >-< %s"
 	SOCKS5_BIND_SERVER     string = "[Socks5] Bind For %s"

pkg/server/socks.go

@@ -2,15 +2,18 @@ package server
 
 import (
 	"errors"
-	"github.com/DVKunion/SeaMoon/pkg/consts"
-	"github.com/DVKunion/SeaMoon/pkg/utils"
-	"github.com/gorilla/websocket"
-	log "github.com/sirupsen/logrus"
 	"net"
 	"net/http"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/DVKunion/SeaMoon/pkg/consts"
+	"github.com/DVKunion/SeaMoon/pkg/utils"
 )
 
 type SocksServer struct {
@@ -32,16 +35,21 @@ func (s *SocksServer) Verification(w http.ResponseWriter) (bool, error) {
 }
 
 func (s *SocksServer) Serve(w http.ResponseWriter, r *http.Request) {
+	lock := &sync.Mutex{}
 	// socks upgrade websocket
-	var upGrader = websocket.Upgrader{
+	upGrader := websocket.Upgrader{
 		CheckOrigin: func(r *http.Request) bool {
 			return true
 		},
 	}
 
-	var conn, _ = upGrader.Upgrade(w, r, nil)
+	conn, err := upGrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Errorf("websocket upgrade error: %s", err)
+		return
+	}
 
-	wsConn := NewWebsocketServer(conn)
+	wsConn := NewWebsocketServer(conn, lock)
 	defer wsConn.Close()
 
 	switch s.request.Cmd {
@@ -52,7 +60,7 @@ func (s *SocksServer) Serve(w http.ResponseWriter, r *http.Request) {
 	case utils.CmdUDPOverTCP:
 		s.handleUDPOverTCP()
 	default:
-		conn.WriteMessage(websocket.BinaryMessage, []byte("UnSupport Command"))
+		log.Errorf("UnSupport Command")
 	}
 }
 

pkg/server/websocket.go

@@ -1,21 +1,26 @@
 package server
 
 import (
-	"github.com/gorilla/websocket"
 	"io"
 	"net"
+	"sync"
+
+	"github.com/gorilla/websocket"
 )
 
 type WebsocketServer struct {
 	net.Conn
-	wConn       *websocket.Conn
+	wConn     *websocket.Conn
+	writeLock *sync.Mutex
+
 	messageType int
 }
 
-func NewWebsocketServer(wConn *websocket.Conn) net.Conn {
+func NewWebsocketServer(wConn *websocket.Conn, lock *sync.Mutex) net.Conn {
 	return &WebsocketServer{
 		wConn:       wConn,
 		messageType: websocket.BinaryMessage,
+		writeLock:   lock,
 	}
 }
 
@@ -25,22 +30,27 @@ func (ws *WebsocketServer) RemoteAddr() net.Addr {
 
 func (ws *WebsocketServer) Close() error {
 	// ws need send close message first to avoid err : close 1006 (abnormal closure): unexpected EOF
-	// todo: panic - concurrent write to websocket connection
-	err := ws.wConn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "close"))
+	err := ws.write(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "close"))
 	if err != nil {
 		return err
 	}
 	return ws.wConn.Close()
 }
 
 func (ws *WebsocketServer) Write(b []byte) (n int, err error) {
-	err = ws.wConn.WriteMessage(ws.messageType, b)
+	err = ws.write(ws.messageType, b)
 	if err != nil {
 		return 0, err
 	}
 	return len(b), nil
 }
 
+func (ws *WebsocketServer) write(messageType int, data []byte) error {
+	ws.writeLock.TryLock()
+	defer ws.writeLock.Unlock()
+	return ws.wConn.WriteMessage(messageType, data)
+}
+
 func (ws *WebsocketServer) Read(b []byte) (n int, err error) {
 	mt, message, err := ws.wConn.ReadMessage()
 	if err != nil {

pkg/utils/socks.go

@@ -12,6 +12,8 @@ import (
 	"net"
 	"strconv"
 	"sync"
+
+	"github.com/gorilla/websocket"
 )
 
 // buffer pools
@@ -23,31 +25,25 @@ var (
 	} // small buff pool
 	LPool = sync.Pool{
 		New: func() interface{} {
-			return make([]byte, 64*1024+262)
+			return make([]byte, 32768)
 		},
 	} // large buff pool for udp
 )
 
 // Transport rw1 and rw2
 func Transport(rw1, rw2 io.ReadWriter) error {
-	errc := make(chan error, 1)
+	errC := make(chan error, 1)
 	go func() {
-		b := LPool.Get().([]byte)
-		defer LPool.Put(b)
-
-		_, err := io.CopyBuffer(rw1, rw2, b)
-		errc <- err
+		_, err := io.CopyBuffer(rw1, rw2, nil)
+		errC <- err
 	}()
 
 	go func() {
-		b := LPool.Get().([]byte)
-		defer LPool.Put(b)
-
-		_, err := io.CopyBuffer(rw2, rw1, b)
-		errc <- err
+		_, err := io.CopyBuffer(rw2, rw1, nil)
+		errC <- err
 	}()
 
-	if err := <-errc; err != nil && err != io.EOF {
+	if err := <-errC; err != nil && err != io.EOF && !websocket.IsUnexpectedCloseError(err) {
 		return err
 	}