diff --git a/conn.go b/conn.go index 3a6259a..f4cb36c 100644 --- a/conn.go +++ b/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{ - InsecureSkipVerify: c.Node.insecureSkipVerify(), - ServerName: c.Node.serverName, + + 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 } diff --git a/ws.go b/ws.go index da902e6..e92972a 100644 --- a/ws.go +++ b/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 },