Skip to content

G
gost
Commit 445889c4 authored by ginuerzh's avatar ginuerzh
Browse files
Options

add more test cases for socks5

parent 72d7850c
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 1168 additions and 417 deletions
.travis.yml
View file @ 445889c4
......@@ -5,8 +5,8 @@ go:
install: true
script:
- env GO111MODULE=on go test -race -v -coverprofile=coverage.txt -covermode=atomic
- cd cmd/gost && env GO111MODULE=on go build
- go test -race -v -coverprofile=coverage.txt -covermode=atomic
- cd cmd/gost && go build
after_success:
- bash <(curl -s https://codecov.io/bash)
bypass_test.go
View file @ 445889c4
......@@ -2,6 +2,7 @@ package gost
import (
"bytes"
"fmt"
"io"
"testing"
"time"
......@@ -158,10 +159,13 @@ var bypassContainTests = []struct {
func TestBypassContains(t *testing.T) {
for i, tc := range bypassContainTests {
bp := NewBypassPatterns(tc.reversed, tc.patterns...)
if bp.Contains(tc.addr) != tc.bypassed {
t.Errorf("#%d test failed: %v, %s", i, tc.patterns, tc.addr)
}
tc := tc
t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) {
bp := NewBypassPatterns(tc.reversed, tc.patterns...)
if bp.Contains(tc.addr) != tc.bypassed {
t.Errorf("#%d test failed: %v, %s", i, tc.patterns, tc.addr)
}
})
}
}
......@@ -244,6 +248,9 @@ func TestByapssReload(t *testing.T) {
}
if tc.stopped {
bp.Stop()
if bp.Period() >= 0 {
t.Errorf("period of the stopped reloader should be minus value")
}
}
if bp.Stopped() != tc.stopped {
t.Errorf("#%d test failed: stopped value should be %v, got %v",
......
cmd/gost/.config/dns.txt
View file @ 445889c4
# resolver timeout, default 30s.
timeout 10s
# resolver cache TTL, default 60s, minus value means that cache is disabled.
ttl 300s
# resolver cache TTL,
# minus value means that cache is disabled,
# default to the TTL in DNS server response.
# ttl 300s
# period for live reloading
reload 10s
# ip[:port] [protocol] [hostname]
https://1.0.0.1/dns-query
1.1.1.1:853 tls cloudflare-dns.com
https://1.0.0.1/dns-query https
8.8.8.8
8.8.8.8 tcp
1.1.1.1 udp
......
cmd/gost/cfg.go
View file @ 445889c4
......@@ -10,7 +10,6 @@ import (
"net/url"
"os"
"strings"
"time"
"github.com/ginuerzh/gost"
)
......@@ -196,8 +195,6 @@ func parseResolver(cfg string) gost.Resolver {
if cfg == "" {
return nil
}
timeout := 30 * time.Second
ttl := 60 * time.Second
var nss []gost.NameServer
f, err := os.Open(cfg)
......@@ -237,11 +234,11 @@ func parseResolver(cfg string) gost.Resolver {
}
}
}
return gost.NewResolver(timeout, ttl, nss...)
return gost.NewResolver(0, nss...)
}
defer f.Close()
resolver := gost.NewResolver(timeout, ttl)
resolver := gost.NewResolver(0)
resolver.Reload(f)
go gost.PeriodReload(resolver, cfg)
......
common_test.go
View file @ 445889c4
package gost
import (
"bufio"
"bytes"
"crypto/tls"
"errors"
"fmt"
"io"
"io/ioutil"
"net"
"net/http"
"net/url"
"sync"
"time"
)
......@@ -37,6 +41,110 @@ var (
})
)
// proxyConn obtains a connection to the proxy server.
func proxyConn(client *Client, server *Server) (net.Conn, error) {
conn, err := client.Dial(server.Addr().String())
if err != nil {
return nil, err
}
cc, err := client.Handshake(conn, AddrHandshakeOption(server.Addr().String()))
if err != nil {
conn.Close()
return nil, err
}
return cc, nil
}
// httpRoundtrip does a HTTP request-response roundtrip, and checks the data received.
func httpRoundtrip(conn net.Conn, targetURL string, data []byte) (err error) {
req, err := http.NewRequest(
http.MethodGet,
targetURL,
bytes.NewReader(data),
)
if err != nil {
return
}
if err = req.Write(conn); err != nil {
return
}
resp, err := http.ReadResponse(bufio.NewReader(conn), req)
if err != nil {
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return errors.New(resp.Status)
}
recv, err := ioutil.ReadAll(resp.Body)
if err != nil {
return
}
if !bytes.Equal(data, recv) {
return fmt.Errorf("data not equal")
}
return
}
func udpRoundtrip(client *Client, server *Server, host string, data []byte) (err error) {
conn, err := proxyConn(client, server)
if err != nil {
return
}
defer conn.Close()
conn, err = client.Connect(conn, host)
if err != nil {
return
}
conn.SetDeadline(time.Now().Add(3 * time.Second))
defer conn.SetDeadline(time.Time{})
if _, err = conn.Write(data); err != nil {
return
}
recv := make([]byte, len(data))
if _, err = conn.Read(recv); err != nil {
return
}
if !bytes.Equal(data, recv) {
return fmt.Errorf("data not equal")
}
return
}
func proxyRoundtrip(client *Client, server *Server, targetURL string, data []byte) (err error) {
conn, err := proxyConn(client, server)
if err != nil {
return err
}
defer conn.Close()
u, err := url.Parse(targetURL)
if err != nil {
return
}
conn, err = client.Connect(conn, u.Host)
if err != nil {
return
}
conn.SetDeadline(time.Now().Add(500 * time.Millisecond))
defer conn.SetDeadline(time.Time{})
return httpRoundtrip(conn, targetURL, data)
}
type udpRequest struct {
Body io.Reader
RemoteAddr string
......@@ -55,11 +163,12 @@ type udpHandlerFunc func(w io.Writer, r *udpRequest)
// udpTestServer is a UDP server for test.
type udpTestServer struct {
ln net.PacketConn
handler udpHandlerFunc
wg sync.WaitGroup
mu sync.Mutex // guards closed and conns
closed bool
ln net.PacketConn
handler udpHandlerFunc
wg sync.WaitGroup
mu sync.Mutex // guards closed and conns
closed bool
exitChan chan struct{}
}
func newUDPTestServer(handler udpHandlerFunc) *udpTestServer {
......@@ -68,9 +177,13 @@ func newUDPTestServer(handler udpHandlerFunc) *udpTestServer {
if err != nil {
panic(fmt.Sprintf("udptest: failed to listen on a port: %v", err))
}
ln.SetReadBuffer(1024 * 1024)
ln.SetWriteBuffer(1024 * 1024)
return &udpTestServer{
ln: ln,
handler: handler,
ln: ln,
handler: handler,
exitChan: make(chan struct{}),
}
}
......@@ -83,7 +196,7 @@ func (s *udpTestServer) serve() {
data := make([]byte, 1024)
n, raddr, err := s.ln.ReadFrom(data)
if err != nil {
return
break
}
if s.handler != nil {
s.wg.Add(1)
......@@ -101,6 +214,9 @@ func (s *udpTestServer) serve() {
}()
}
}
// signal the listener has been exited.
close(s.exitChan)
}
func (s *udpTestServer) Addr() string {
......@@ -119,6 +235,8 @@ func (s *udpTestServer) Close() error {
s.closed = true
s.mu.Unlock()
<-s.exitChan
s.wg.Wait()
return err
......
forward.go
View file @ 445889c4
......@@ -662,10 +662,6 @@ func TCPRemoteForwardListener(addr string, chain *Chain) (Listener, error) {
go ln.listenLoop()
// if err = <-ln.errChan; err != nil {
// ln.Close()
// }
return ln, err
}
......@@ -680,17 +676,10 @@ func (l *tcpRemoteForwardListener) isChainValid() bool {
func (l *tcpRemoteForwardListener) listenLoop() {
var tempDelay time.Duration
// var once sync.Once
for {
conn, err := l.accept()
// once.Do(func() {
// l.errChan <- err
// log.Log("once.Do error:", err)
// close(l.errChan)
// })
select {
case <-l.closed:
if conn != nil {
......
forward_test.go
View file @ 445889c4
package gost
import (
"bytes"
"crypto/rand"
"fmt"
"net/http/httptest"
"net/url"
"testing"
"time"
)
func tcpDirectForwardRoundtrip(targetURL string, data []byte) error {
......@@ -122,37 +119,6 @@ func BenchmarkTCPDirectForwardParallel(b *testing.B) {
})
}
func udpRoundtrip(client *Client, server *Server, host string, data []byte) (err error) {
conn, err := proxyConn(client, server)
if err != nil {
return
}
defer conn.Close()
conn.SetDeadline(time.Now().Add(1 * time.Second))
defer conn.SetDeadline(time.Time{})
conn, err = client.Connect(conn, host)
if err != nil {
return
}
if _, err = conn.Write(data); err != nil {
return
}
recv := make([]byte, len(data))
if _, err = conn.Read(recv); err != nil {
return
}
if !bytes.Equal(data, recv) {
return fmt.Errorf("data not equal")
}
return
}
func udpDirectForwardRoundtrip(host string, data []byte) error {
ln, err := UDPDirectForwardListener("localhost:0", 0)
if err != nil {
......
gost.go
View file @ 445889c4
......@@ -7,8 +7,10 @@ import (
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"errors"
"io"
"math/big"
"net"
"sync"
"time"
......@@ -137,3 +139,48 @@ func (rw *readWriter) Read(p []byte) (n int, err error) {
func (rw *readWriter) Write(p []byte) (n int, err error) {
return rw.w.Write(p)
}
var (
nopClientConn = &nopConn{}
)
// a nop connection implements net.Conn,
// it does nothing.
type nopConn struct{}
func (c *nopConn) Read(b []byte) (n int, err error) {
return 0, &net.OpError{Op: "read", Net: "nop", Source: nil, Addr: nil, Err: errors.New("read not supported")}
}
func (c *nopConn) Write(b []byte) (n int, err error) {
return 0, &net.OpError{Op: "write", Net: "nop", Source: nil, Addr: nil, Err: errors.New("write not supported")}
}
func (c *nopConn) Close() error {
return nil
}
func (c *nopConn) LocalAddr() net.Addr {
return nil
}
func (c *nopConn) RemoteAddr() net.Addr {
return nil
}
func (c *nopConn) SetDeadline(t time.Time) error {
return &net.OpError{Op: "set", Net: "nop", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}
func (c *nopConn) SetReadDeadline(t time.Time) error {
return &net.OpError{Op: "set", Net: "nop", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}
func (c *nopConn) SetWriteDeadline(t time.Time) error {
return &net.OpError{Op: "set", Net: "nop", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}
// Accepter represents a network endpoint that can accept connection from peer.
type Accepter interface {
Accept() (net.Conn, error)
}
handler_test.go
View file @ 445889c4
......@@ -18,7 +18,6 @@ func autoHTTPProxyRoundtrip(targetURL string, data []byte, clientInfo *url.Useri
Connector: HTTPConnector(clientInfo),
Transporter: TCPTransporter(),
}
server := &Server{
Listener: ln,
Handler: AutoHandler(
......@@ -111,7 +110,7 @@ func TestAutoSOCKS5Proxy(t *testing.T) {
}
}
func autoSOCKS4ProxyRoundtrip(targetURL string, data []byte) error {
func autoSOCKS4ProxyRoundtrip(targetURL string, data []byte, options ...HandlerOption) error {
ln, err := TCPListener("")
if err != nil {
return err
......@@ -124,7 +123,7 @@ func autoSOCKS4ProxyRoundtrip(targetURL string, data []byte) error {
server := &Server{
Listener: ln,
Handler: AutoHandler(),
Handler: AutoHandler(options...),
}
go server.Run()
defer server.Close()
......@@ -139,14 +138,17 @@ func TestAutoSOCKS4Proxy(t *testing.T) {
sendData := make([]byte, 128)
rand.Read(sendData)
err := autoSOCKS4ProxyRoundtrip(httpSrv.URL, sendData)
// t.Logf("#%d %v", i, err)
if err != nil {
if err := autoSOCKS4ProxyRoundtrip(httpSrv.URL, sendData); err != nil {
t.Errorf("got error: %v", err)
}
if err := autoSOCKS4ProxyRoundtrip(httpSrv.URL, sendData,
UsersHandlerOption(url.UserPassword("admin", "123456"))); err == nil {
t.Errorf("authentication required auto handler for SOCKS4 should failed")
}
}
func autoSocks4aProxyRoundtrip(targetURL string, data []byte) error {
func autoSocks4aProxyRoundtrip(targetURL string, data []byte, options ...HandlerOption) error {
ln, err := TCPListener("")
if err != nil {
return err
......@@ -159,7 +161,7 @@ func autoSocks4aProxyRoundtrip(targetURL string, data []byte) error {
server := &Server{
Listener: ln,
Handler: AutoHandler(),
Handler: AutoHandler(options...),
}
go server.Run()
......@@ -175,11 +177,14 @@ func TestAutoSOCKS4AProxy(t *testing.T) {
sendData := make([]byte, 128)
rand.Read(sendData)
err := autoSocks4aProxyRoundtrip(httpSrv.URL, sendData)
// t.Logf("#%d %v", i, err)
if err != nil {
if err := autoSocks4aProxyRoundtrip(httpSrv.URL, sendData); err != nil {
t.Errorf("got error: %v", err)
}
if err := autoSocks4aProxyRoundtrip(httpSrv.URL, sendData,
UsersHandlerOption(url.UserPassword("admin", "123456"))); err == nil {
t.Errorf("authentication required auto handler for SOCKS4A should failed")
}
}
func autoSSProxyRoundtrip(targetURL string, data []byte, clientInfo *url.Userinfo, serverInfo *url.Userinfo) error {
......
hosts_test.go
View file @ 445889c4
......@@ -28,7 +28,8 @@ var hostsLookupTests = []struct {
func TestHostsLookup(t *testing.T) {
for i, tc := range hostsLookupTests {
hosts := NewHosts(tc.hosts...)
hosts := NewHosts()
hosts.AddHost(tc.hosts...)
ip := hosts.Lookup(tc.host)
if !ip.Equal(tc.ip) {
t.Errorf("#%d test failed: lookup should be %s, got %s", i, tc.ip, ip)
......@@ -61,6 +62,11 @@ var HostsReloadTests = []struct {
host: "example.com",
ip: nil,
},
{
r: bytes.NewBufferString("#reload 10s\ninvalid.ip.addr example.com"),
period: 0,
ip: nil,
},
{
r: bytes.NewBufferString("reload 10s\n192.168.1.1"),
period: 10 * time.Second,
......@@ -112,6 +118,9 @@ func TestHostsReload(t *testing.T) {
}
if tc.stopped {
hosts.Stop()
if hosts.Period() >= 0 {
t.Errorf("period of the stopped reloader should be minus value")
}
}
if hosts.Stopped() != tc.stopped {
t.Errorf("#%d test failed: stopped value should be %v, got %v",
......
http2.go
View file @ 445889c4
......@@ -569,12 +569,13 @@ func HTTP2Listener(addr string, config *tls.Config) (Listener, error) {
}
l.server = server
ln, err := tls.Listen("tcp", addr, config)
ln, err := net.Listen("tcp", addr)
if err != nil {
return nil, err
}
l.addr = ln.Addr()
ln = tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, config)
go func() {
err := server.Serve(ln)
if err != nil {
......@@ -875,42 +876,11 @@ func (c *http2ServerConn) SetWriteDeadline(t time.Time) error {
// a dummy HTTP2 client conn used by HTTP2 client connector
type http2ClientConn struct {
nopConn
addr string
client *http.Client
}
func (c *http2ClientConn) Read(b []byte) (n int, err error) {
return 0, &net.OpError{Op: "read", Net: "http2", Source: nil, Addr: nil, Err: errors.New("read not supported")}
}
func (c *http2ClientConn) Write(b []byte) (n int, err error) {
return 0, &net.OpError{Op: "write", Net: "http2", Source: nil, Addr: nil, Err: errors.New("write not supported")}
}
func (c *http2ClientConn) Close() error {
return nil
}
func (c *http2ClientConn) LocalAddr() net.Addr {
return nil
}
func (c *http2ClientConn) RemoteAddr() net.Addr {
return nil
}
func (c *http2ClientConn) SetDeadline(t time.Time) error {
return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}
func (c *http2ClientConn) SetReadDeadline(t time.Time) error {
return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}
func (c *http2ClientConn) SetWriteDeadline(t time.Time) error {
return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}
type flushWriter struct {
w io.Writer
}
......
http_test.go
View file @ 445889c4
package gost
import (
"bufio"
"bytes"
"crypto/rand"
"errors"
"fmt"
"io/ioutil"
"net"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"time"
)
// proxyConn obtains a connection to the proxy server.
func proxyConn(client *Client, server *Server) (net.Conn, error) {
conn, err := client.Dial(server.Addr().String())
if err != nil {
return nil, err
}
cc, err := client.Handshake(conn, AddrHandshakeOption(server.Addr().String()))
if err != nil {
conn.Close()
return nil, err
}
return cc, nil
}
// httpRoundtrip does a HTTP request-response roundtrip, and checks the data received.
func httpRoundtrip(conn net.Conn, targetURL string, data []byte) (err error) {
req, err := http.NewRequest(
http.MethodGet,
targetURL,
bytes.NewReader(data),
)
if err != nil {
return
}
if err = req.Write(conn); err != nil {
return
}
resp, err := http.ReadResponse(bufio.NewReader(conn), req)
if err != nil {
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return errors.New(resp.Status)
}
recv, err := ioutil.ReadAll(resp.Body)
if err != nil {
return
}
if !bytes.Equal(data, recv) {
return fmt.Errorf("data not equal")
}
return
}
func proxyRoundtrip(client *Client, server *Server, targetURL string, data []byte) (err error) {
conn, err := proxyConn(client, server)
if err != nil {
return err
}
defer conn.Close()
u, err := url.Parse(targetURL)
if err != nil {
return
}
conn, err = client.Connect(conn, u.Host)
if err != nil {
return
}
conn.SetDeadline(time.Now().Add(500 * time.Millisecond))
defer conn.SetDeadline(time.Time{})
return httpRoundtrip(conn, targetURL, data)
}
var httpProxyTests = []struct {
cliUser *url.Userinfo
srvUsers []*url.Userinfo
......
kcp.go
View file @ 445889c4
......@@ -123,6 +123,11 @@ func (tr *kcpTransporter) Dial(addr string, options ...DialOption) (conn net.Con
defer tr.sessionMutex.Unlock()
session, ok := tr.sessions[addr]
if session != nil && session.session != nil && session.session.IsClosed() {
session.Close()
delete(tr.sessions, addr) // session is dead
ok = false
}
if !ok {
timeout := opts.Timeout
if timeout <= 0 {
......
mux.go
View file @ 445889c4
......@@ -45,10 +45,16 @@ func (session *muxSession) Accept() (net.Conn, error) {
}
func (session *muxSession) Close() error {
if session.session == nil {
return nil
}
return session.session.Close()
}
func (session *muxSession) IsClosed() bool {
if session.session == nil {
return true
}
return session.session.IsClosed()
}
......
obfs.go
View file @ 445889c4
......@@ -331,7 +331,7 @@ func Obfs4Listener(addr string) (Listener, error) {
}
l := &obfs4Listener{
addr: addr,
Listener: ln,
Listener: tcpKeepAliveListener{ln.(*net.TCPListener)},
}
return l, nil
}
......
quic_test.go
View file @ 445889c4
......@@ -2,6 +2,7 @@ package gost
import (
"crypto/rand"
"crypto/sha256"
"fmt"
"net/http/httptest"
"net/url"
......@@ -405,3 +406,57 @@ func TestQUICForwardTunnel(t *testing.T) {
t.Error(err)
}
}
func httpOverCipherQUICRoundtrip(targetURL string, data []byte,
clientInfo *url.Userinfo, serverInfo []*url.Userinfo) error {
sum := sha256.Sum256([]byte("12345678"))
cfg := &QUICConfig{
Key: sum[:],
}
ln, err := QUICListener("localhost:0", cfg)
if err != nil {
return err
}
client := &Client{
Connector: HTTPConnector(clientInfo),
Transporter: QUICTransporter(cfg),
}
server := &Server{
Listener: ln,
Handler: HTTPHandler(
UsersHandlerOption(serverInfo...),
),
}
go server.Run()
defer server.Close()
return proxyRoundtrip(client, server, targetURL, data)
}
func TestHTTPOverCipherQUIC(t *testing.T) {
httpSrv := httptest.NewServer(httpTestHandler)
defer httpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
for i, tc := range httpProxyTests {
err := httpOverCipherQUICRoundtrip(httpSrv.URL, sendData, tc.cliUser, tc.srvUsers)
if err == nil {
if tc.errStr != "" {
t.Errorf("#%d should failed with error %s", i, tc.errStr)
}
} else {
if tc.errStr == "" {
t.Errorf("#%d got error %v", i, err)
}
if err.Error() != tc.errStr {
t.Errorf("#%d got error %v, want %v", i, err, tc.errStr)
}
}
}
}
resolver.go
View file @ 445889c4
......@@ -23,8 +23,6 @@ import (
var (
// DefaultResolverTimeout is the default timeout for name resolution.
DefaultResolverTimeout = 5 * time.Second
// DefaultResolverTTL is the default cache TTL for name resolution.
DefaultResolverTTL = 1 * time.Hour
)
// Resolver is a name resolver for domain name.
......@@ -53,13 +51,18 @@ type NameServer struct {
// Init initializes the name server.
func (ns *NameServer) Init() error {
timeout := ns.Timeout
if timeout <= 0 {
timeout = DefaultResolverTimeout
}
switch strings.ToLower(ns.Protocol) {
case "tcp":
ns.exchanger = &dnsExchanger{
endpoint: ns.Addr,
client: &dns.Client{
Net: "tcp",
Timeout: ns.Timeout,
Timeout: timeout,
},
}
case "tls":
......@@ -74,7 +77,7 @@ func (ns *NameServer) Init() error {
endpoint: ns.Addr,
client: &dns.Client{
Net: "tcp-tls",
Timeout: ns.Timeout,
Timeout: timeout,
TLSConfig: cfg,
},
}
......@@ -95,7 +98,7 @@ func (ns *NameServer) Init() error {
endpoint: u,
client: &http.Client{
Transport: transport,
Timeout: ns.Timeout,
Timeout: timeout,
},
}
case "udp":
......@@ -105,7 +108,7 @@ func (ns *NameServer) Init() error {
endpoint: ns.Addr,
client: &dns.Client{
Net: "udp",
Timeout: ns.Timeout,
Timeout: timeout,
},
}
}
......@@ -125,15 +128,9 @@ func (ns NameServer) String() string {
return fmt.Sprintf("%s/%s", addr, prot)
}
type resolverCacheItem struct {
IPs []net.IP
ts int64
}
type resolver struct {
Servers []NameServer
mCache *sync.Map
Timeout time.Duration
TTL time.Duration
period time.Duration
domain string
......@@ -142,22 +139,14 @@ type resolver struct {
}
// NewResolver create a new Resolver with the given name servers and resolution timeout.
func NewResolver(timeout, ttl time.Duration, servers ...NameServer) ReloadResolver {
r := newResolver(timeout, ttl, servers...)
if r.Timeout <= 0 {
r.Timeout = DefaultResolverTimeout
}
if r.TTL == 0 {
r.TTL = DefaultResolverTTL
}
func NewResolver(ttl time.Duration, servers ...NameServer) ReloadResolver {
r := newResolver(ttl, servers...)
return r
}
func newResolver(timeout, ttl time.Duration, servers ...NameServer) *resolver {
func newResolver(ttl time.Duration, servers ...NameServer) *resolver {
return &resolver{
Servers: servers,
Timeout: timeout,
TTL: ttl,
mCache: &sync.Map{},
stopped: make(chan struct{}),
......@@ -204,25 +193,25 @@ func (r *resolver) Resolve(host string) (ips []net.IP, err error) {
}
for _, ns := range servers {
ips, err = r.resolve(ns.exchanger, host)
ips, ttl, err = r.resolve(ns.exchanger, host)
if err != nil {
log.Logf("[resolver] %s via %s : %s", host, ns, err)
continue
}
if Debug {
log.Logf("[resolver] %s via %s %v", host, ns, ips)
log.Logf("[resolver] %s via %s %v(ttl: %v)", host, ns, ips, ttl)
}
if len(ips) > 0 {
break
}
}
r.storeCache(host, ips)
r.storeCache(host, ips, ttl)
return
}
func (*resolver) resolve(ex Exchanger, host string) (ips []net.IP, err error) {
func (*resolver) resolve(ex Exchanger, host string) (ips []net.IP, ttl time.Duration, err error) {
if ex == nil {
return
}
......@@ -236,11 +225,18 @@ func (*resolver) resolve(ex Exchanger, host string) (ips []net.IP, err error) {
for _, ans := range mr.Answer {
if ar, _ := ans.(*dns.A); ar != nil {
ips = append(ips, ar.A)
ttl = time.Duration(ar.Header().Ttl) * time.Second
}
}
return
}
type resolverCacheItem struct {
IPs []net.IP
ts int64
ttl time.Duration
}
func (r *resolver) loadCache(name string, ttl time.Duration) []net.IP {
if ttl < 0 {
return nil
......@@ -248,6 +244,10 @@ func (r *resolver) loadCache(name string, ttl time.Duration) []net.IP {
if v, ok := r.mCache.Load(name); ok {
item, _ := v.(*resolverCacheItem)
if ttl == 0 {
ttl = item.ttl
}
if item == nil || time.Since(time.Unix(item.ts, 0)) > ttl {
return nil
}
......@@ -257,13 +257,14 @@ func (r *resolver) loadCache(name string, ttl time.Duration) []net.IP {
return nil
}
func (r *resolver) storeCache(name string, ips []net.IP) {
func (r *resolver) storeCache(name string, ips []net.IP, ttl time.Duration) {
if name == "" || len(ips) == 0 {
return
}
r.mCache.Store(name, &resolverCacheItem{
IPs: ips,
ts: time.Now().Unix(),
ttl: ttl,
})
}
......@@ -343,10 +344,10 @@ func (r *resolver) Reload(rd io.Reader) error {
ns.Hostname = ss[2]
}
ns.Timeout = timeout
if timeout <= 0 {
ns.Timeout = DefaultResolverTimeout
if strings.HasPrefix(ns.Addr, "https") {
ns.Protocol = "https"
}
ns.Timeout = timeout
if err := ns.Init(); err == nil {
nss = append(nss, ns)
......@@ -359,7 +360,6 @@ func (r *resolver) Reload(rd io.Reader) error {
}
r.mux.Lock()
r.Timeout = timeout
r.TTL = ttl
r.domain = domain
r.period = period
......@@ -408,9 +408,9 @@ func (r *resolver) String() string {
defer r.mux.RUnlock()
b := &bytes.Buffer{}
fmt.Fprintf(b, "Timeout %v\n", r.Timeout)
fmt.Fprintf(b, "TTL %v\n", r.TTL)
fmt.Fprintf(b, "Reload %v\n", r.period)
fmt.Fprintf(b, "Domain %v\n", r.domain)
for i := range r.Servers {
fmt.Fprintln(b, r.Servers[i])
}
......
resolver_test.go
View file @ 445889c4
......@@ -46,7 +46,7 @@ func TestDNSResolver(t *testing.T) {
t.Error(err)
}
t.Log(ns)
r := NewResolver(0, 0, ns)
r := NewResolver(0, ns)
err := dnsResolverRoundtrip(t, r, tc.host)
if err != nil {
if tc.pass {
......@@ -103,13 +103,12 @@ var resolverReloadTests = []struct {
{
r: bytes.NewBufferString("1.1.1.1"),
ns: &NameServer{
Addr: "1.1.1.1",
Timeout: DefaultResolverTimeout,
Addr: "1.1.1.1",
},
stopped: true,
},
{
r: bytes.NewBufferString("timeout 10s\nsearch\nnameserver \nnameserver 1.1.1.1 udp"),
r: bytes.NewBufferString("\n# comment\ntimeout 10s\nsearch\nnameserver \nnameserver 1.1.1.1 udp"),
ns: &NameServer{
Protocol: "udp",
Addr: "1.1.1.1",
......@@ -123,7 +122,6 @@ var resolverReloadTests = []struct {
ns: &NameServer{
Addr: "1.1.1.1",
Protocol: "tcp",
Timeout: DefaultResolverTimeout,
},
stopped: true,
},
......@@ -133,7 +131,6 @@ var resolverReloadTests = []struct {
Addr: "1.1.1.1:853",
Protocol: "tls",
Hostname: "cloudflare-dns.com",
Timeout: DefaultResolverTimeout,
},
stopped: true,
},
......@@ -142,7 +139,6 @@ var resolverReloadTests = []struct {
ns: &NameServer{
Addr: "1.1.1.1:853",
Protocol: "tls",
Timeout: DefaultResolverTimeout,
},
stopped: true,
},
......@@ -151,11 +147,10 @@ var resolverReloadTests = []struct {
stopped: true,
},
{
r: bytes.NewBufferString("https://1.0.0.1/dns-query https"),
r: bytes.NewBufferString("https://1.0.0.1/dns-query"),
ns: &NameServer{
Addr: "https://1.0.0.1/dns-query",
Protocol: "https",
Timeout: DefaultResolverTimeout,
},
stopped: true,
},
......@@ -164,15 +159,11 @@ var resolverReloadTests = []struct {
func TestResolverReload(t *testing.T) {
for i, tc := range resolverReloadTests {
t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) {
r := newResolver(0, 0)
r := newResolver(0)
if err := r.Reload(tc.r); err != nil {
t.Error(err)
}
t.Log(r.String())
if r.Timeout != tc.timeout {
t.Errorf("timeout value should be %v, got %v",
tc.timeout, r.Timeout)
}
if r.TTL != tc.ttl {
t.Errorf("ttl value should be %v, got %v",
tc.ttl, r.TTL)
......@@ -198,6 +189,9 @@ func TestResolverReload(t *testing.T) {
if tc.stopped {
r.Stop()
if r.Period() >= 0 {
t.Errorf("period of the stopped reloader should be minus value")
}
}
if r.Stopped() != tc.stopped {
t.Errorf("stopped value should be %v, got %v",
......
socks.go
View file @ 445889c4
This diff is collapsed.
socks_test.go
View file @ 445889c4
package gost
import (
"bytes"
"crypto/rand"
"fmt"
"net"
"net/http/httptest"
"net/url"
"testing"
"time"
)
var socks5ProxyTests = []struct {
......@@ -407,6 +410,128 @@ func TestSOCKS5Bind(t *testing.T) {
}
}
func socks5MuxBindRoundtrip(t *testing.T, targetURL string, data []byte) (err error) {
ln, err := TCPListener("")
if err != nil {
return
}
l, err := net.Listen("tcp", "")
if err != nil {
return err
}
bindAddr := l.Addr().String()
l.Close()
client := &Client{
Connector: Socks5MuxBindConnector(),
Transporter: SOCKS5MuxBindTransporter(bindAddr),
}
server := &Server{
Handler: SOCKS5Handler(UsersHandlerOption(url.UserPassword("admin", "123456"))),
Listener: ln,
}
go server.Run()
defer server.Close()
return muxBindRoundtrip(client, server, bindAddr, targetURL, data)
}
func muxBindRoundtrip(client *Client, server *Server, bindAddr, targetURL string, data []byte) (err error) {
cn, err := client.Dial(server.Addr().String())
if err != nil {
return err
}
conn, err := client.Handshake(cn,
AddrHandshakeOption(server.Addr().String()),
UserHandshakeOption(url.UserPassword("admin", "123456")),
)
if err != nil {
cn.Close()
return err
}
defer conn.Close()
cc, err := net.Dial("tcp", bindAddr)
if err != nil {
return
}
defer cc.Close()
conn, err = client.Connect(conn, "")
if err != nil {
return
}
u, err := url.Parse(targetURL)
if err != nil {
return
}
hc, err := net.Dial("tcp", u.Host)
if err != nil {
return
}
defer hc.Close()
go transport(hc, conn)
return httpRoundtrip(cc, targetURL, data)
}
func TestSOCKS5MuxBind(t *testing.T) {
httpSrv := httptest.NewServer(httpTestHandler)
defer httpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
if err := socks5MuxBindRoundtrip(t, httpSrv.URL, sendData); err != nil {
t.Errorf("got error: %v", err)
}
}
func BenchmarkSOCKS5MuxBind(b *testing.B) {
httpSrv := httptest.NewServer(httpTestHandler)
defer httpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
ln, err := TCPListener("")
if err != nil {
b.Error(err)
}
l, err := net.Listen("tcp", "")
if err != nil {
b.Error(err)
}
bindAddr := l.Addr().String()
l.Close()
client := &Client{
Connector: Socks5MuxBindConnector(),
Transporter: SOCKS5MuxBindTransporter(bindAddr),
}
server := &Server{
Handler: SOCKS5Handler(UsersHandlerOption(url.UserPassword("admin", "123456"))),
Listener: ln,
}
go server.Run()
defer server.Close()
for i := 0; i < b.N; i++ {
if err := muxBindRoundtrip(client, server, bindAddr, httpSrv.URL, sendData); err != nil {
b.Error(err)
}
}
}
func socks5UDPRoundtrip(t *testing.T, host string, data []byte) (err error) {
ln, err := TCPListener("")
if err != nil {
......@@ -440,3 +565,226 @@ func TestSOCKS5UDP(t *testing.T) {
t.Errorf("got error: %v", err)
}
}
// TODO: fix a probability of timeout.
func BenchmarkSOCKS5UDP(b *testing.B) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
ln, err := TCPListener("")
if err != nil {
b.Error(err)
}
client := &Client{
Connector: SOCKS5UDPConnector(url.UserPassword("admin", "123456")),
Transporter: TCPTransporter(),
}
server := &Server{
Handler: SOCKS5Handler(UsersHandlerOption(url.UserPassword("admin", "123456"))),
Listener: ln,
}
go server.Run()
defer server.Close()
for i := 0; i < b.N; i++ {
if err := udpRoundtrip(client, server, udpSrv.Addr(), sendData); err != nil {
b.Error(err)
}
}
}
func BenchmarkSOCKS5UDPSingleConn(b *testing.B) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
ln, err := TCPListener("")
if err != nil {
b.Error(err)
}
client := &Client{
Connector: SOCKS5UDPConnector(url.UserPassword("admin", "123456")),
Transporter: TCPTransporter(),
}
server := &Server{
Handler: SOCKS5Handler(UsersHandlerOption(url.UserPassword("admin", "123456"))),
Listener: ln,
}
go server.Run()
defer server.Close()
conn, err := proxyConn(client, server)
if err != nil {
b.Error(err)
}
defer conn.Close()
conn, err = client.Connect(conn, udpSrv.Addr())
if err != nil {
b.Error(err)
}
roundtrip := func(conn net.Conn, data []byte) error {
conn.SetDeadline(time.Now().Add(1 * time.Second))
defer conn.SetDeadline(time.Time{})
if _, err = conn.Write(data); err != nil {
return err
}
recv := make([]byte, len(data))
if _, err = conn.Read(recv); err != nil {
return err
}
if !bytes.Equal(data, recv) {
return fmt.Errorf("data not equal")
}
return nil
}
for i := 0; i < b.N; i++ {
if err := roundtrip(conn, sendData); err != nil {
b.Error(err)
}
}
}
func socks5UDPTunRoundtrip(t *testing.T, host string, data []byte) (err error) {
ln, err := TCPListener("")
if err != nil {
return
}
client := &Client{
Connector: SOCKS5UDPTunConnector(url.UserPassword("admin", "123456")),
Transporter: TCPTransporter(),
}
server := &Server{
Handler: SOCKS5Handler(UsersHandlerOption(url.UserPassword("admin", "123456"))),
Listener: ln,
}
go server.Run()
defer server.Close()
return udpRoundtrip(client, server, host, data)
}
func TestSOCKS5UDPTun(t *testing.T) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
if err := socks5UDPTunRoundtrip(t, udpSrv.Addr(), sendData); err != nil {
t.Errorf("got error: %v", err)
}
}
func BenchmarkSOCKS5UDPTun(b *testing.B) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
ln, err := TCPListener("")
if err != nil {
b.Error(err)
}
client := &Client{
Connector: SOCKS5UDPTunConnector(url.UserPassword("admin", "123456")),
Transporter: TCPTransporter(),
}
server := &Server{
Handler: SOCKS5Handler(UsersHandlerOption(url.UserPassword("admin", "123456"))),
Listener: ln,
}
go server.Run()
defer server.Close()
for i := 0; i < b.N; i++ {
if err := udpRoundtrip(client, server, udpSrv.Addr(), sendData); err != nil {
b.Error(err)
}
}
}
func BenchmarkSOCKS5UDPTunSingleConn(b *testing.B) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
ln, err := TCPListener("")
if err != nil {
b.Error(err)
}
client := &Client{
Connector: SOCKS5UDPTunConnector(url.UserPassword("admin", "123456")),
Transporter: TCPTransporter(),
}
server := &Server{
Handler: SOCKS5Handler(UsersHandlerOption(url.UserPassword("admin", "123456"))),
Listener: ln,
}
go server.Run()
defer server.Close()
conn, err := proxyConn(client, server)
if err != nil {
b.Error(err)
}
defer conn.Close()
conn, err = client.Connect(conn, udpSrv.Addr())
if err != nil {
b.Error(err)
}
roundtrip := func(conn net.Conn, data []byte) error {
conn.SetDeadline(time.Now().Add(1 * time.Second))
defer conn.SetDeadline(time.Time{})
if _, err = conn.Write(data); err != nil {
return err
}
recv := make([]byte, len(data))
if _, err = conn.Read(recv); err != nil {
return err
}
if !bytes.Equal(data, recv) {
return fmt.Errorf("data not equal")
}
return nil
}
for i := 0; i < b.N; i++ {
if err := roundtrip(conn, sendData); err != nil {
b.Error(err)
}
}
}
ss_test.go
View file @ 445889c4
......@@ -2,6 +2,7 @@ package gost
import (
"crypto/rand"
"fmt"
"net/http/httptest"
"net/url"
"testing"
......@@ -148,20 +149,23 @@ func TestSSProxy(t *testing.T) {
rand.Read(sendData)
for i, tc := range ssTests {
err := ssProxyRoundtrip(httpSrv.URL, sendData,
tc.clientCipher,
tc.serverCipher,
)
if err == nil {
if !tc.pass {
t.Errorf("#%d should failed", i)
tc := tc
t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) {
err := ssProxyRoundtrip(httpSrv.URL, sendData,
tc.clientCipher,
tc.serverCipher,
)
if err == nil {
if !tc.pass {
t.Errorf("#%d should failed", i)
}
} else {
// t.Logf("#%d %v", i, err)
if tc.pass {
t.Errorf("#%d got error: %v", i, err)
}
}
} else {
// t.Logf("#%d %v", i, err)
if tc.pass {
t.Errorf("#%d got error: %v", i, err)
}
}
})
}
}
......@@ -317,7 +321,7 @@ func shadowUDPRoundtrip(t *testing.T, host string, data []byte) error {
return udpRoundtrip(client, server, host, data)
}
func TestShadowUDP(t *testing.T) {
func _TestShadowUDP(t *testing.T) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
......
tls.go
View file @ 445889c4
......@@ -58,21 +58,20 @@ func (tr *mtlsTransporter) Dial(addr string, options ...DialOption) (conn net.Co
option(opts)
}
timeout := opts.Timeout
if timeout <= 0 {
timeout = DialTimeout
}
tr.sessionMutex.Lock()
defer tr.sessionMutex.Unlock()
session, ok := tr.sessions[addr]
if session != nil && session.session != nil && session.session.IsClosed() {
session.Close()
if session != nil && session.IsClosed() {
delete(tr.sessions, addr)
ok = false
ok = false // session is dead
}
if !ok {
timeout := opts.Timeout
if timeout <= 0 {
timeout = DialTimeout
}
if opts.Chain == nil {
conn, err = net.DialTimeout("tcp", addr, timeout)
} else {
......@@ -159,10 +158,12 @@ func TLSListener(addr string, config *tls.Config) (Listener, error) {
if config == nil {
config = DefaultTLSConfig
}
ln, err := tls.Listen("tcp", addr, config)
ln, err := net.Listen("tcp", addr)
if err != nil {
return nil, err
}
ln = tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, config)
return &tlsListener{ln}, nil
}
......@@ -177,13 +178,13 @@ func MTLSListener(addr string, config *tls.Config) (Listener, error) {
if config == nil {
config = DefaultTLSConfig
}
ln, err := tls.Listen("tcp", addr, config)
ln, err := net.Listen("tcp", addr)
if err != nil {
return nil, err
}
l := &mtlsListener{
ln: ln,
ln: tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, config),
connChan: make(chan net.Conn, 1024),
errChan: make(chan error, 1),
}
......
ws.go
View file @ 445889c4
......@@ -28,92 +28,6 @@ type WSOptions struct {
UserAgent string
}
type websocketConn struct {
conn *websocket.Conn
rb []byte
}
func websocketClientConn(url string, conn net.Conn, tlsConfig *tls.Config, options *WSOptions) (net.Conn, error) {
if options == nil {
options = &WSOptions{}
}
timeout := options.HandshakeTimeout
if timeout <= 0 {
timeout = HandshakeTimeout
}
dialer := websocket.Dialer{
ReadBufferSize: options.ReadBufferSize,
WriteBufferSize: options.WriteBufferSize,
TLSClientConfig: tlsConfig,
HandshakeTimeout: timeout,
EnableCompression: options.EnableCompression,
NetDial: func(net, addr string) (net.Conn, error) {
return conn, nil
},
}
header := http.Header{}
header.Set("User-Agent", DefaultUserAgent)
if options.UserAgent != "" {
header.Set("User-Agent", options.UserAgent)
}
c, resp, err := dialer.Dial(url, header)
if err != nil {
return nil, err
}
resp.Body.Close()
return &websocketConn{conn: c}, nil
}
func websocketServerConn(conn *websocket.Conn) net.Conn {
// conn.EnableWriteCompression(true)
return &websocketConn{
conn: conn,
}
}
func (c *websocketConn) Read(b []byte) (n int, err error) {
if len(c.rb) == 0 {
_, c.rb, err = c.conn.ReadMessage()
}
n = copy(b, c.rb)
c.rb = c.rb[n:]
return
}
func (c *websocketConn) Write(b []byte) (n int, err error) {
err = c.conn.WriteMessage(websocket.BinaryMessage, b)
n = len(b)
return
}
func (c *websocketConn) Close() error {
return c.conn.Close()
}
func (c *websocketConn) LocalAddr() net.Addr {
return c.conn.LocalAddr()
}
func (c *websocketConn) RemoteAddr() net.Addr {
return c.conn.RemoteAddr()
}
func (c *websocketConn) SetDeadline(t time.Time) error {
if err := c.SetReadDeadline(t); err != nil {
return err
}
return c.SetWriteDeadline(t)
}
func (c *websocketConn) SetReadDeadline(t time.Time) error {
return c.conn.SetReadDeadline(t)
}
func (c *websocketConn) SetWriteDeadline(t time.Time) error {
return c.conn.SetWriteDeadline(t)
}
type wsTransporter struct {
tcpTransporter
options *WSOptions
......@@ -160,21 +74,20 @@ func (tr *mwsTransporter) Dial(addr string, options ...DialOption) (conn net.Con
option(opts)
}
timeout := opts.Timeout
if timeout <= 0 {
timeout = DialTimeout
}
tr.sessionMutex.Lock()
defer tr.sessionMutex.Unlock()
session, ok := tr.sessions[addr]
if session != nil && session.session != nil && session.session.IsClosed() {
session.Close()
if session != nil && session.IsClosed() {
delete(tr.sessions, addr)
ok = false
}
if !ok {
timeout := opts.Timeout
if timeout <= 0 {
timeout = DialTimeout
}
if opts.Chain == nil {
conn, err = net.DialTimeout("tcp", addr, timeout)
} else {
......@@ -302,21 +215,20 @@ func (tr *mwssTransporter) Dial(addr string, options ...DialOption) (conn net.Co
option(opts)
}
timeout := opts.Timeout
if timeout <= 0 {
timeout = DialTimeout
}
tr.sessionMutex.Lock()
defer tr.sessionMutex.Unlock()
session, ok := tr.sessions[addr]
if session != nil && session.session != nil && session.session.IsClosed() {
session.Close()
if session != nil && session.IsClosed() {
delete(tr.sessions, addr)
ok = false
}
if !ok {
timeout := opts.Timeout
if timeout <= 0 {
timeout = DialTimeout
}
if opts.Chain == nil {
conn, err = net.DialTimeout("tcp", addr, timeout)
} else {
......@@ -428,7 +340,11 @@ func WSListener(addr string, options *WSOptions) (Listener, error) {
mux := http.NewServeMux()
mux.Handle("/ws", http.HandlerFunc(l.upgrade))
l.srv = &http.Server{Addr: addr, Handler: mux}
l.srv = &http.Server{
Addr: addr,
Handler: mux,
ReadHeaderTimeout: 30 * time.Second,
}
ln, err := net.ListenTCP("tcp", tcpAddr)
if err != nil {
......@@ -517,7 +433,11 @@ func MWSListener(addr string, options *WSOptions) (Listener, error) {
mux := http.NewServeMux()
mux.Handle("/ws", http.HandlerFunc(l.upgrade))
l.srv = &http.Server{Addr: addr, Handler: mux}
l.srv = &http.Server{
Addr: addr,
Handler: mux,
ReadHeaderTimeout: 30 * time.Second,
}
ln, err := net.ListenTCP("tcp", tcpAddr)
if err != nil {
......@@ -634,9 +554,10 @@ func WSSListener(addr string, tlsConfig *tls.Config, options *WSOptions) (Listen
mux := http.NewServeMux()
mux.Handle("/ws", http.HandlerFunc(l.upgrade))
l.srv = &http.Server{
Addr: addr,
TLSConfig: tlsConfig,
Handler: mux,
Addr: addr,
TLSConfig: tlsConfig,
Handler: mux,
ReadHeaderTimeout: 30 * time.Second,
}
ln, err := net.ListenTCP("tcp", tcpAddr)
......@@ -694,9 +615,10 @@ func MWSSListener(addr string, tlsConfig *tls.Config, options *WSOptions) (Liste
mux := http.NewServeMux()
mux.Handle("/ws", http.HandlerFunc(l.upgrade))
l.srv = &http.Server{
Addr: addr,
TLSConfig: tlsConfig,
Handler: mux,
Addr: addr,
TLSConfig: tlsConfig,
Handler: mux,
ReadHeaderTimeout: 30 * time.Second,
}
ln, err := net.ListenTCP("tcp", tcpAddr)
......@@ -737,3 +659,89 @@ func generateChallengeKey() (string, error) {
}
return base64.StdEncoding.EncodeToString(p), nil
}
type websocketConn struct {
conn *websocket.Conn
rb []byte
}
func websocketClientConn(url string, conn net.Conn, tlsConfig *tls.Config, options *WSOptions) (net.Conn, error) {
if options == nil {
options = &WSOptions{}
}
timeout := options.HandshakeTimeout
if timeout <= 0 {
timeout = HandshakeTimeout
}
dialer := websocket.Dialer{
ReadBufferSize: options.ReadBufferSize,
WriteBufferSize: options.WriteBufferSize,
TLSClientConfig: tlsConfig,
HandshakeTimeout: timeout,
EnableCompression: options.EnableCompression,
NetDial: func(net, addr string) (net.Conn, error) {
return conn, nil
},
}
header := http.Header{}
header.Set("User-Agent", DefaultUserAgent)
if options.UserAgent != "" {
header.Set("User-Agent", options.UserAgent)
}
c, resp, err := dialer.Dial(url, header)
if err != nil {
return nil, err
}
resp.Body.Close()
return &websocketConn{conn: c}, nil
}
func websocketServerConn(conn *websocket.Conn) net.Conn {
// conn.EnableWriteCompression(true)
return &websocketConn{
conn: conn,
}
}
func (c *websocketConn) Read(b []byte) (n int, err error) {
if len(c.rb) == 0 {
_, c.rb, err = c.conn.ReadMessage()
}
n = copy(b, c.rb)
c.rb = c.rb[n:]
return
}
func (c *websocketConn) Write(b []byte) (n int, err error) {
err = c.conn.WriteMessage(websocket.BinaryMessage, b)
n = len(b)
return
}
func (c *websocketConn) Close() error {
return c.conn.Close()
}
func (c *websocketConn) LocalAddr() net.Addr {
return c.conn.LocalAddr()
}
func (c *websocketConn) RemoteAddr() net.Addr {
return c.conn.RemoteAddr()
}
func (c *websocketConn) SetDeadline(t time.Time) error {
if err := c.SetReadDeadline(t); err != nil {
return err
}
return c.SetWriteDeadline(t)
}
func (c *websocketConn) SetReadDeadline(t time.Time) error {
return c.conn.SetReadDeadline(t)
}
func (c *websocketConn) SetWriteDeadline(t time.Time) error {
return c.conn.SetWriteDeadline(t)
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment