Update On Fri May 31 20:31:22 CEST 2024

echo/internal/cli/config.go

@@ -11,6 +11,7 @@ import (
 	"github.com/Ehco1996/ehco/internal/relay/conf"
 	"github.com/Ehco1996/ehco/internal/tls"
 	"github.com/Ehco1996/ehco/internal/web"
+	"github.com/Ehco1996/ehco/pkg/buffer"
 	"github.com/Ehco1996/ehco/pkg/log"
 	"github.com/Ehco1996/ehco/pkg/xray"
 	"github.com/getsentry/sentry-go"
@@ -77,6 +78,12 @@ func initLogger(cfg *config.Config) error {
 	return nil
 }
 
+func initGlobalBufferPool() {
+	if BufferSize > 0 {
+		buffer.ReplaceBufferPool(BufferSize)
+	}
+}
+
 func InitConfigAndComponents() (*config.Config, error) {
 	cfg, err := loadConfig()
 	if err != nil {
@@ -88,7 +95,7 @@ func InitConfigAndComponents() (*config.Config, error) {
 	if err := initSentry(); err != nil {
 		return nil, err
 	}
-
+	initGlobalBufferPool()
 	return cfg, nil
 }
 

echo/internal/cli/flags.go

@@ -15,6 +15,7 @@ var (
 	SystemFilePath       = "/etc/systemd/system/ehco.service"
 	LogLevel             string
 	ConfigReloadInterval int
+	BufferSize           int
 )
 
 var RootFlags = []cli.Flag{
@@ -91,4 +92,10 @@ var RootFlags = []cli.Flag{
 		Destination: &ConfigReloadInterval,
 		DefaultText: "60",
 	},
+	&cli.IntFlag{
+		Name:        "buffer_size",
+		Usage:       "set buffer size to when transport data default 20 * 1024(20KB)",
+		EnvVars:     []string{"EHCO_BUFFER_SIZE"},
+		Destination: &BufferSize,
+	},
 }

echo/internal/conn/relay_conn.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/Ehco1996/ehco/internal/metrics"
+	"github.com/Ehco1996/ehco/pkg/buffer"
 	"github.com/Ehco1996/ehco/pkg/bytes"
 	"go.uber.org/zap"
 )
@@ -113,10 +114,12 @@ func connectionName(conn net.Conn) string {
 }
 
 func copyConn(conn1, conn2 *innerConn) error {
+	buf := buffer.BufferPool.Get()
+	defer buffer.BufferPool.Put(buf)
 	errCH := make(chan error, 1)
 	// copy conn1 to conn2,read from conn1 and write to conn2
 	go func() {
-		_, err := io.Copy(conn2, conn1)
+		_, err := io.CopyBuffer(conn2, conn1, buf)
 		_ = conn2.CloseWrite() // all data is written to conn2 now, so close the write side of conn2 to send eof
 		errCH <- err
 	}()

echo/internal/conn/ws_conn.go

@@ -1,4 +1,4 @@
-package transporter
+package conn
 
 import (
 	"fmt"
@@ -9,6 +9,7 @@ import (
 	"github.com/Ehco1996/ehco/pkg/buffer"
 	"github.com/gobwas/ws"
 	"github.com/gobwas/ws/wsutil"
+	"go.uber.org/zap"
 )
 
 type wsConn struct {
@@ -17,7 +18,7 @@ type wsConn struct {
 	buf      []byte
 }
 
-func newWsConn(conn net.Conn, isServer bool) *wsConn {
+func NewWSConn(conn net.Conn, isServer bool) *wsConn {
 	return &wsConn{conn: conn, isServer: isServer, buf: buffer.BufferPool.Get()}
 }
 
@@ -27,6 +28,7 @@ func (c *wsConn) Read(b []byte) (n int, err error) {
 		return 0, err
 	}
 	if header.Length > int64(cap(c.buf)) {
+		zap.S().Warnf("ws payload size:%d is larger than buffer size:%d", header.Length, cap(c.buf))
 		c.buf = make([]byte, header.Length)
 	}
 	payload := c.buf[:header.Length]
@@ -38,7 +40,7 @@ func (c *wsConn) Read(b []byte) (n int, err error) {
 		ws.Cipher(payload, header.Mask, 0)
 	}
 	if len(payload) > len(b) {
-		return 0, fmt.Errorf("buffer too small to transport ws msg")
+		return 0, fmt.Errorf("buffer size:%d too small to transport ws payload size:%d", len(b), len(payload))
 	}
 	copy(b, payload)
 	return len(payload), nil

echo/internal/conn/ws_conn_test.go

@@ -1,4 +1,4 @@
-package transporter
+package conn
 
 import (
 	"context"
@@ -23,7 +23,7 @@ func TestClientConn_ReadWrite(t *testing.T) {
 		}
 		go func() {
 			defer conn.Close()
-			wsc := newWsConn(conn, true)
+			wsc := NewWSConn(conn, true)
 
 			buf := make([]byte, 1024)
 			for {
@@ -53,7 +53,7 @@ func TestClientConn_ReadWrite(t *testing.T) {
 	}
 	defer conn.Close()
 
-	wsClientConn := newWsConn(conn, false)
+	wsClientConn := NewWSConn(conn, false)
 	for i := 0; i < 3; i++ {
 		// test write
 		n, err := wsClientConn.Write(data)

echo/internal/transporter/ws.go

@@ -9,6 +9,7 @@ import (
 	"github.com/gobwas/ws"
 	"github.com/labstack/echo/v4"
 
+	"github.com/Ehco1996/ehco/internal/conn"
 	"github.com/Ehco1996/ehco/internal/constant"
 	"github.com/Ehco1996/ehco/internal/metrics"
 	"github.com/Ehco1996/ehco/internal/web"
@@ -43,7 +44,7 @@ func (s *WsClient) TCPHandShake(remote *lb.Node) (net.Conn, error) {
 	latency := time.Since(t1)
 	metrics.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(latency.Milliseconds()))
 	remote.HandShakeDuration = latency
-	c := newWsConn(wsc, false)
+	c := conn.NewWSConn(wsc, false)
 	return c, nil
 }
 
@@ -91,8 +92,8 @@ func (s *WsServer) HandleRequest(w http.ResponseWriter, req *http.Request) {
 	if err != nil {
 		return
 	}
-	c := newWsConn(wsc, true)
-	if err := s.RelayTCPConn(c, s.relayer.TCPHandShake); err != nil {
+
+	if err := s.RelayTCPConn(conn.NewWSConn(wsc, true), s.relayer.TCPHandShake); err != nil {
 		s.l.Errorf("RelayTCPConn error: %s", err.Error())
 	}
 }

echo/internal/transporter/wss_mux.go

@@ -1,3 +1,6 @@
+// NOTE CAN NOT use real ws frame to transport smux frame
+// err: accept stream err: buffer size:8 too small to transport ws payload size:45
+// so this transport just use ws protocol to handshake and then use smux protocol to transport
 package transporter
 
 import (
@@ -101,12 +104,12 @@ func (s *MwssServer) ListenAndServe() error {
 }
 
 func (s *MwssServer) HandleRequest(w http.ResponseWriter, r *http.Request) {
-	conn, _, _, err := ws.UpgradeHTTP(r, w)
+	c, _, _, err := ws.UpgradeHTTP(r, w)
 	if err != nil {
 		s.l.Error(err)
 		return
 	}
-	s.mux(conn)
+	s.mux(c)
 }
 
 func (s *MwssServer) Close() error {