diff --git a/handler/handler.go b/handler/handler.go index 58f493f..8822188 100644 --- a/handler/handler.go +++ b/handler/handler.go @@ -17,6 +17,9 @@ import ( "go.uber.org/zap" ) +// BufferSize for coybuffer and websocket +const BufferSize = 256 * 1024 + var ( websocketUpstream uint = 1 upstreamWebsocket uint = 2 @@ -46,8 +49,8 @@ func New( logger *zap.Logger) (*Handler, error) { upgrader := websocket.Upgrader{ - ReadBufferSize: 1024, - WriteBufferSize: 1024, + ReadBufferSize: BufferSize, + WriteBufferSize: BufferSize, HandshakeTimeout: handshakeTimeout, CheckOrigin: func(r *http.Request) bool { return true @@ -175,6 +178,7 @@ func (h *Handler) Proxy(wg *sync.WaitGroup) func(w http.ResponseWriter, r *http. // websocket -> server go func() { defer func() { doneCh <- true }() + b := make([]byte, BufferSize) for { mt, r, err := conn.NextReader() if websocket.IsCloseError(err, @@ -201,7 +205,7 @@ func (h *Handler) Proxy(wg *sync.WaitGroup) func(w http.ResponseWriter, r *http. if h.dumpTCP > 0 { r = io.TeeReader(r, dr) } - n, err := io.Copy(s, r) + n, err := io.CopyBuffer(s, r, b) if err != nil { if !goClose { logger.Warn("Reading from websocket", zap.Error(err)) @@ -219,8 +223,8 @@ func (h *Handler) Proxy(wg *sync.WaitGroup) func(w http.ResponseWriter, r *http. // server -> websocket go func() { defer func() { doneCh <- true }() + b := make([]byte, BufferSize) for { - b := make([]byte, 64*1024) n, err := s.Read(b) if err != nil { if !goClose && err != io.EOF { @@ -233,11 +237,10 @@ func (h *Handler) Proxy(wg *sync.WaitGroup) func(w http.ResponseWriter, r *http. return } - b = b[:n] if h.dumpTCP > 1 { ds.Write(b) } - if err := conn.WriteMessage(websocket.BinaryMessage, b); err != nil { + if err := conn.WriteMessage(websocket.BinaryMessage, b[:n]); err != nil { if !goClose { logger.Warn("WriteMessage", zap.Error(err)) hasError = true