fix websocket config

conn.go

@@ -6,10 +6,6 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
-	"github.com/ginuerzh/gosocks4"
-	"github.com/ginuerzh/gosocks5"
-	"github.com/golang/glog"
-	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
 	"net"
 	"net/http"
 	"net/http/httputil"
@@ -18,6 +14,11 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/ginuerzh/gosocks4"
+	"github.com/ginuerzh/gosocks5"
+	"github.com/golang/glog"
+	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
 )
 
 type ProxyConn struct {
@@ -57,20 +58,40 @@ func (c *ProxyConn) handshake() error {
 
 	switch c.Node.Transport {
 	case "ws": // websocket connection
+		rbuf, _ := strconv.Atoi(c.Node.Get("rbuf"))
+		wbuf, _ := strconv.Atoi(c.Node.Get("wbuf"))
+		comp := c.Node.getBool("compression")
+		opt := WSOptions{
+			ReadBufferSize:    rbuf,
+			WriteBufferSize:   wbuf,
+			HandshakeTimeout:  DialTimeout,
+			EnableCompression: comp,
+		}
 		u := url.URL{Scheme: "ws", Host: c.Node.Addr, Path: "/ws"}
-		conn, err := WebsocketClientConn(u.String(), c.conn, nil)
+		conn, err := WebsocketClientConn(u.String(), c.conn, &opt)
 		if err != nil {
 			return err
 		}
 		c.conn = conn
 	case "wss": // websocket security
 		tlsUsed = true
-		u := url.URL{Scheme: "wss", Host: c.Node.Addr, Path: "/ws"}
-		config := &tls.Config{
+
+		rbuf, _ := strconv.Atoi(c.Node.Get("rbuf"))
+		wbuf, _ := strconv.Atoi(c.Node.Get("wbuf"))
+		comp := c.Node.getBool("compression")
+		opt := WSOptions{
+			ReadBufferSize:    rbuf,
+			WriteBufferSize:   wbuf,
+			HandshakeTimeout:  DialTimeout,
+			EnableCompression: comp,
+			TLSConfig: &tls.Config{
 				InsecureSkipVerify: c.Node.insecureSkipVerify(),
 				ServerName:         c.Node.serverName,
+			},
 		}
-		conn, err := WebsocketClientConn(u.String(), c.conn, config)
+
+		u := url.URL{Scheme: "wss", Host: c.Node.Addr, Path: "/ws"}
+		conn, err := WebsocketClientConn(u.String(), c.conn, &opt)
 		if err != nil {
 			return err
 		}

ws.go

@@ -2,12 +2,14 @@ package gost
 
 import (
 	"crypto/tls"
-	"github.com/golang/glog"
-	"gopkg.in/gorilla/websocket.v1"
 	"net"
 	"net/http"
 	"net/http/httputil"
+	"strconv"
 	"time"
+
+	"github.com/golang/glog"
+	"gopkg.in/gorilla/websocket.v1"
 )
 
 type WebsocketServer struct {
@@ -18,14 +20,18 @@ type WebsocketServer struct {
 }
 
 func NewWebsocketServer(base *ProxyServer) *WebsocketServer {
+	rbuf, _ := strconv.Atoi(base.Node.Get("rbuf"))
+	wbuf, _ := strconv.Atoi(base.Node.Get("wbuf"))
+	comp := base.Node.getBool("compression")
+
 	return &WebsocketServer{
 		Addr: base.Node.Addr,
 		Base: base,
 		upgrader: websocket.Upgrader{
-			ReadBufferSize:    1024,
-			WriteBufferSize:   1024,
+			ReadBufferSize:    rbuf,
+			WriteBufferSize:   wbuf,
 			CheckOrigin:       func(r *http.Request) bool { return true },
-			EnableCompression: true,
+			EnableCompression: comp,
 		},
 	}
 }
@@ -68,18 +74,29 @@ func (s *WebsocketServer) ListenAndServeTLS(config *tls.Config) error {
 	return server.ListenAndServeTLS("", "")
 }
 
+type WSOptions struct {
+	ReadBufferSize    int
+	WriteBufferSize   int
+	HandshakeTimeout  time.Duration
+	EnableCompression bool
+	TLSConfig         *tls.Config
+}
+
 type WebsocketConn struct {
 	conn *websocket.Conn
 	rb   []byte
 }
 
-func WebsocketClientConn(url string, conn net.Conn, config *tls.Config) (*WebsocketConn, error) {
+func WebsocketClientConn(url string, conn net.Conn, options *WSOptions) (*WebsocketConn, error) {
+	if options == nil {
+		options = &WSOptions{}
+	}
 	dialer := websocket.Dialer{
-		ReadBufferSize:    1024,
-		WriteBufferSize:   1024,
-		TLSClientConfig:   config,
-		HandshakeTimeout:  DialTimeout,
-		EnableCompression: true,
+		ReadBufferSize:    options.ReadBufferSize,
+		WriteBufferSize:   options.WriteBufferSize,
+		TLSClientConfig:   options.TLSConfig,
+		HandshakeTimeout:  options.HandshakeTimeout,
+		EnableCompression: options.EnableCompression,
 		NetDial: func(net, addr string) (net.Conn, error) {
 			return conn, nil
 		},