Pārlūkot izejas kodu

implement write multi buffer in websocket

Darien Raymond 8 gadi atpakaļ
vecāks
revīzija
69b51bffbb
1 mainītis faili ar 24 papildinājumiem un 2 dzēšanām
  1. 24 2
      transport/internet/websocket/connection.go

+ 24 - 2
transport/internet/websocket/connection.go

@@ -6,13 +6,15 @@ import (
 	"time"
 
 	"github.com/gorilla/websocket"
+	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/errors"
 )
 
 // connection is a wrapper for net.Conn over WebSocket connection.
 type connection struct {
-	wsc    *websocket.Conn
-	reader io.Reader
+	wsc         *websocket.Conn
+	reader      io.Reader
+	writeBuffer []byte
 }
 
 // Read implements net.Conn.Read()
@@ -52,6 +54,26 @@ func (c *connection) Write(b []byte) (int, error) {
 	return len(b), nil
 }
 
+func (c *connection) WriteMultiBuffer(mb buf.MultiBuffer) (int, error) {
+	defer mb.Release()
+
+	if c.writeBuffer == nil {
+		c.writeBuffer = make([]byte, 4096)
+	}
+	totalBytes := 0
+	for !mb.IsEmpty() {
+		nBytes, err := mb.Read(c.writeBuffer)
+		if err != nil {
+			return totalBytes, err
+		}
+		totalBytes += nBytes
+		if _, err := c.Write(c.writeBuffer[:nBytes]); err != nil {
+			return totalBytes, err
+		}
+	}
+	return totalBytes, nil
+}
+
 func (c *connection) Close() error {
 	c.wsc.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second*5))
 	return c.wsc.Close()