cleanup websocket code

transport/internet/websocket/dialer.go

@@ -84,12 +84,8 @@ func wsDial(ctx context.Context, dest v2net.Destination) (net.Conn, error) {
 		}
 		return nil, err
 	}
-	return func() net.Conn {
+
-		connv2ray := &wsconn{
+	return &wsconn{
 		wsc: conn,
-			connClosing: false,
+	}, nil
-		}
-		connv2ray.setup()
-		return connv2ray
-	}(), nil
 }

transport/internet/websocket/hub.go

@@ -120,9 +120,7 @@ func (wsl *WSListener) converttovws(w http.ResponseWriter, r *http.Request) (*ws
 		return nil, err
 	}
 
-	wrapedConn := &wsconn{wsc: conn, connClosing: false}
+	return &wsconn{wsc: conn}, nil
-	wrapedConn.setup()
-	return wrapedConn, nil
 }
 
 func (v *WSListener) Accept() (internet.Connection, error) {

transport/internet/websocket/wsconn.go

@@ -1,179 +1,87 @@
 package websocket
 
 import (
-	"bufio"
 	"io"
 	"net"
 	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
-	"v2ray.com/core/app/log"
 	"v2ray.com/core/common/errors"
 )
 
 type wsconn struct {
 	wsc    *websocket.Conn
-	readBuffer  *bufio.Reader
+	reader io.Reader
-	connClosing bool
+	rlock  sync.Mutex
-	rlock       *sync.Mutex
+	wlock  sync.Mutex
-	wlock       *sync.Mutex
 }
 
-func (ws *wsconn) Read(b []byte) (n int, err error) {
+func (c *wsconn) Read(b []byte) (int, error) {
-	ws.rlock.Lock()
+	c.rlock.Lock()
-	n, err = ws.read(b)
+	defer c.rlock.Unlock()
-	ws.rlock.Unlock()
-	return n, err
 
-}
+	for {
-
+		reader, err := c.getReader()
-func (ws *wsconn) read(b []byte) (n int, err error) {
-	if ws.connClosing {
-		return 0, io.EOF
-	}
-
-	n, err = ws.readNext(b)
-	return n, err
-}
-
-func (ws *wsconn) getNewReadBuffer() error {
-	_, r, err := ws.wsc.NextReader()
-	if err != nil {
-		log.Warning("WebSocket|Connection: Failed to get reader.", err)
-		ws.connClosing = true
-		ws.Close()
-		return err
-	}
-	ws.readBuffer = bufio.NewReader(r)
-	return nil
-}
-
-func (ws *wsconn) readNext(b []byte) (n int, err error) {
-	if ws.readBuffer == nil {
-		err = ws.getNewReadBuffer()
 		if err != nil {
 			return 0, err
 		}
-	}
-
-	n, err = ws.readBuffer.Read(b)
-
-	if err == nil {
-		return n, err
-	}
 
+		nBytes, err := reader.Read(b)
 		if errors.Cause(err) == io.EOF {
-		ws.readBuffer = nil
+			continue
-		if n == 0 {
-			return ws.readNext(b)
 		}
-		return n, nil
+		return nBytes, err
 	}
-	return n, err
-
 }
 
-func (ws *wsconn) Write(b []byte) (n int, err error) {
+func (c *wsconn) getReader() (io.Reader, error) {
-	ws.wlock.Lock()
+	if c.reader != nil {
-	if ws.connClosing {
+		return c.reader, nil
-		return 0, io.ErrClosedPipe
 	}
 
-	n, err = ws.write(b)
+	_, reader, err := c.wsc.NextReader()
-	ws.wlock.Unlock()
-	return n, err
-}
-
-func (ws *wsconn) write(b []byte) (n int, err error) {
-	wr, err := ws.wsc.NextWriter(websocket.BinaryMessage)
 	if err != nil {
-		log.Warning("WebSocket|Connection: Failed to get writer.", err)
+		return nil, err
-		ws.connClosing = true
+	}
-		ws.Close()
+	c.reader = reader
+	return reader, nil
+}
+
+func (c *wsconn) Write(b []byte) (int, error) {
+	c.wlock.Lock()
+	defer c.wlock.Unlock()
+
+	if err := c.wsc.WriteMessage(websocket.BinaryMessage, b); err != nil {
 		return 0, err
 	}
-	n, err = wr.Write(b)
+	return len(b), nil
-	if err != nil {
-		return 0, err
-	}
-	err = wr.Close()
-	if err != nil {
-		return 0, err
-	}
-	return n, err
 }
 
-func (ws *wsconn) Close() error {
+func (c *wsconn) Close() error {
-	ws.connClosing = true
+	c.wsc.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second*5))
-	ws.wlock.Lock()
+	return c.wsc.Close()
-	ws.wsc.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add((time.Second * 5)))
+}
-	ws.wlock.Unlock()
+
-	err := ws.wsc.Close()
+func (c *wsconn) LocalAddr() net.Addr {
+	return c.wsc.LocalAddr()
+}
+
+func (c *wsconn) RemoteAddr() net.Addr {
+	return c.wsc.RemoteAddr()
+}
+
+func (c *wsconn) SetDeadline(t time.Time) error {
+	if err := c.SetReadDeadline(t); err != nil {
 		return err
 	}
-func (ws *wsconn) LocalAddr() net.Addr {
+	return c.SetWriteDeadline(t)
-	return ws.wsc.LocalAddr()
-}
-func (ws *wsconn) RemoteAddr() net.Addr {
-	return ws.wsc.RemoteAddr()
-}
-func (ws *wsconn) SetDeadline(t time.Time) error {
-	if err := ws.SetReadDeadline(t); err != nil {
-		return err
-	}
-	return ws.SetWriteDeadline(t)
-}
-func (ws *wsconn) SetReadDeadline(t time.Time) error {
-	return ws.wsc.SetReadDeadline(t)
-}
-func (ws *wsconn) SetWriteDeadline(t time.Time) error {
-	return ws.wsc.SetWriteDeadline(t)
 }
 
-func (ws *wsconn) setup() {
+func (c *wsconn) SetReadDeadline(t time.Time) error {
-	ws.connClosing = false
+	return c.wsc.SetReadDeadline(t)
-
-	/*
-		https://godoc.org/github.com/gorilla/websocket#Conn.NextReader
-		https://godoc.org/github.com/gorilla/websocket#Conn.NextWriter
-
-		Both Read and write access are both exclusive.
-		And in both case it will need a lock.
-
-	*/
-	ws.rlock = &sync.Mutex{}
-	ws.wlock = &sync.Mutex{}
-
-	ws.pingPong()
 }
 
-func (ws *wsconn) pingPong() {
+func (c *wsconn) SetWriteDeadline(t time.Time) error {
-	pongRcv := make(chan int, 1)
+	return c.wsc.SetWriteDeadline(t)
-	ws.wsc.SetPongHandler(func(data string) error {
-		pongRcv <- 0
-		return nil
-	})
-
-	go func() {
-		for !ws.connClosing {
-			ws.wlock.Lock()
-			ws.wsc.WriteMessage(websocket.PingMessage, nil)
-			ws.wlock.Unlock()
-			tick := time.After(time.Second * 3)
-
-			select {
-			case <-pongRcv:
-			case <-tick:
-				if !ws.connClosing {
-					log.Debug("WS:Closing as ping is not responded~" + ws.wsc.UnderlyingConn().LocalAddr().String() + "-" + ws.wsc.UnderlyingConn().RemoteAddr().String())
-				}
-				ws.Close()
-			}
-			<-time.After(time.Second * 27)
-		}
-
-		return
-	}()
-
 }