1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 10:08:15 -05:00

implement write multi buffer in websocket

This commit is contained in:
Darien Raymond 2017-04-18 20:56:29 +02:00
parent 33ba08be77
commit 69b51bffbb
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169

View File

@ -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()