Browse Source

cleanup websocket code

Darien Raymond 8 years ago
parent
commit
3be6c72211

+ 6 - 10
transport/internet/websocket/dialer.go

@@ -6,8 +6,8 @@ import (
 	"net"
 
 	"github.com/gorilla/websocket"
-	"v2ray.com/core/common"
 	"v2ray.com/core/app/log"
+	"v2ray.com/core/common"
 	v2net "v2ray.com/core/common/net"
 	"v2ray.com/core/transport/internet"
 	"v2ray.com/core/transport/internet/internal"
@@ -24,12 +24,9 @@ func Dial(ctx context.Context, dest v2net.Destination) (internet.Connection, err
 	wsSettings := internet.TransportSettingsFromContext(ctx).(*Config)
 
 	id := internal.NewConnectionID(src, dest)
-	var conn *wsconn
+	var conn net.Conn
 	if dest.Network == v2net.Network_TCP && wsSettings.IsConnectionReuse() {
-		connt := globalCache.Get(id)
-		if connt != nil {
-			conn = connt.(*wsconn)
-		}
+		conn = globalCache.Get(id)
 	}
 	if conn == nil {
 		var err error
@@ -46,7 +43,7 @@ func init() {
 	common.Must(internet.RegisterTransportDialer(internet.TransportProtocol_WebSocket, Dial))
 }
 
-func wsDial(ctx context.Context, dest v2net.Destination) (*wsconn, error) {
+func wsDial(ctx context.Context, dest v2net.Destination) (net.Conn, error) {
 	src := internet.DialerSourceFromContext(ctx)
 	wsSettings := internet.TransportSettingsFromContext(ctx).(*Config)
 
@@ -87,13 +84,12 @@ func wsDial(ctx context.Context, dest v2net.Destination) (*wsconn, error) {
 		}
 		return nil, err
 	}
-	return func() internet.Connection {
+	return func() net.Conn {
 		connv2ray := &wsconn{
 			wsc:         conn,
 			connClosing: false,
-			config:      wsSettings,
 		}
 		connv2ray.setup()
 		return connv2ray
-	}().(*wsconn), nil
+	}(), nil
 }

+ 11 - 13
transport/internet/websocket/hub.go

@@ -8,15 +8,14 @@ import (
 	"sync"
 	"time"
 
+	"github.com/gorilla/websocket"
+	"v2ray.com/core/app/log"
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/errors"
-	"v2ray.com/core/app/log"
 	v2net "v2ray.com/core/common/net"
 	"v2ray.com/core/transport/internet"
-	v2tls "v2ray.com/core/transport/internet/tls"
-
-	"github.com/gorilla/websocket"
 	"v2ray.com/core/transport/internet/internal"
+	v2tls "v2ray.com/core/transport/internet/tls"
 )
 
 var (
@@ -68,7 +67,7 @@ func ListenWS(address v2net.Address, port v2net.Port, options internet.ListenOpt
 
 func (wsl *WSListener) listenws(address v2net.Address, port v2net.Port) error {
 	http.HandleFunc("/"+wsl.config.Path, func(w http.ResponseWriter, r *http.Request) {
-		con, err := wsl.converttovws(w, r)
+		conn, err := wsl.converttovws(w, r)
 		if err != nil {
 			log.Warning("WebSocket|Listener: Failed to convert connection: ", err)
 			return
@@ -76,17 +75,17 @@ func (wsl *WSListener) listenws(address v2net.Address, port v2net.Port) error {
 
 		select {
 		case wsl.awaitingConns <- &ConnectionWithError{
-			conn: con,
+			conn: conn,
 		}:
 		default:
-			if con != nil {
-				con.Close()
+			if conn != nil {
+				conn.Close()
 			}
 		}
 		return
 	})
 
-	errchan := make(chan error)
+	errchan := make(chan error, 1)
 
 	listenerfunc := func() error {
 		ol, err := net.Listen("tcp", address.String()+":"+strconv.Itoa(int(port.Value())))
@@ -114,14 +113,13 @@ func (wsl *WSListener) listenws(address v2net.Address, port v2net.Port) error {
 	go func() {
 		err := listenerfunc()
 		errchan <- err
-		return
 	}()
 
 	var err error
 	select {
 	case err = <-errchan:
-	case <-time.After(time.Second * 2):
-		//Should this listen fail after 2 sec, it could gone untracked.
+	case <-time.After(time.Second * 5):
+		//Should this listen fail after 5 sec, it could gone untracked.
 	}
 
 	if err != nil {
@@ -158,7 +156,7 @@ func (v *WSListener) Accept() (internet.Connection, error) {
 			if connErr.err != nil {
 				return nil, connErr.err
 			}
-			return internal.NewConnection(internal.ConnectionID{}, connErr.conn.(*wsconn), v, internal.ReuseConnection(v.config.IsConnectionReuse())), nil
+			return internal.NewConnection(internal.ConnectionID{}, connErr.conn, v, internal.ReuseConnection(v.config.IsConnectionReuse())), nil
 		case <-time.After(time.Second * 2):
 		}
 	}

+ 3 - 13
transport/internet/websocket/wsconn.go

@@ -8,18 +8,16 @@ import (
 	"time"
 
 	"github.com/gorilla/websocket"
-	"v2ray.com/core/common/errors"
 	"v2ray.com/core/app/log"
+	"v2ray.com/core/common/errors"
 )
 
 type wsconn struct {
 	wsc         *websocket.Conn
 	readBuffer  *bufio.Reader
 	connClosing bool
-	reusable    bool
 	rlock       *sync.Mutex
 	wlock       *sync.Mutex
-	config      *Config
 }
 
 func (ws *wsconn) Read(b []byte) (n int, err error) {
@@ -42,7 +40,7 @@ func (ws *wsconn) read(b []byte) (n int, err error) {
 func (ws *wsconn) getNewReadBuffer() error {
 	_, r, err := ws.wsc.NextReader()
 	if err != nil {
-		log.Warning("WS transport: ws connection NewFrameReader return ", err)
+		log.Warning("WebSocket|Connection: Failed to get reader.", err)
 		ws.connClosing = true
 		ws.Close()
 		return err
@@ -90,7 +88,7 @@ func (ws *wsconn) Write(b []byte) (n int, err error) {
 func (ws *wsconn) write(b []byte) (n int, err error) {
 	wr, err := ws.wsc.NextWriter(websocket.BinaryMessage)
 	if err != nil {
-		log.Warning("WS transport: ws connection NewFrameReader return ", err)
+		log.Warning("WebSocket|Connection: Failed to get writer.", err)
 		ws.connClosing = true
 		ws.Close()
 		return 0, err
@@ -150,14 +148,6 @@ func (ws *wsconn) setup() {
 	ws.pingPong()
 }
 
-func (ws *wsconn) Reusable() bool {
-	return ws.config.IsConnectionReuse() && ws.reusable && !ws.connClosing
-}
-
-func (ws *wsconn) SetReusable(reusable bool) {
-	ws.reusable = reusable
-}
-
 func (ws *wsconn) pingPong() {
 	pongRcv := make(chan int, 1)
 	ws.wsc.SetPongHandler(func(data string) error {