From 21ee1e2c6cbe39c606a1f1dd550ca2a4d2c2ae70 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Fri, 5 Mar 2021 16:39:09 +0300 Subject: [PATCH] Liniting for network --- mtglib/network/default.go | 8 +++++-- mtglib/network/network.go | 41 +++++++++++++++++++++++++++-------- mtglib/network/shadowsocks.go | 27 ++++++++++++++++++----- mtglib/network/socks5.go | 4 ++-- 4 files changed, 61 insertions(+), 19 deletions(-) diff --git a/mtglib/network/default.go b/mtglib/network/default.go index 67804a8..72b3da5 100644 --- a/mtglib/network/default.go +++ b/mtglib/network/default.go @@ -21,7 +21,7 @@ func (d *defaultDialer) Dial(network, address string) (net.Conn, error) { func (d *defaultDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) { switch network { - case "tcp", "tcp4", "tcp6": + case "tcp", "tcp4", "tcp6": // nolint: goconst default: return nil, fmt.Errorf("unsupported network %s", network) } @@ -35,21 +35,25 @@ func (d *defaultDialer) DialContext(ctx context.Context, network, address string if err := tcpConn.SetNoDelay(true); err != nil { conn.Close() + return nil, fmt.Errorf("cannot set TCP_NO_DELAY: %w", err) } if err := tcpConn.SetReadBuffer(d.bufferSize); err != nil { tcpConn.Close() + return nil, fmt.Errorf("cannot set read buffer size: %w", err) } if err := tcpConn.SetWriteBuffer(d.bufferSize); err != nil { tcpConn.Close() + return nil, fmt.Errorf("cannot set write buffer size: %w", err) } if err := tcpConn.SetKeepAlive(true); err != nil { tcpConn.Close() + return nil, fmt.Errorf("cannot enable keep-alive: %w", err) } @@ -61,7 +65,7 @@ func NewDefaultDialer(timeout time.Duration, bufferSize int) (Dialer, error) { case timeout < 0: return nil, fmt.Errorf("timeout %v should be positive number", timeout) case bufferSize < 0: - return nil, fmt.Errorf("buffer size %s should be positive number", bufferSize) + return nil, fmt.Errorf("buffer size %d should be positive number", bufferSize) } if timeout == 0 { diff --git a/mtglib/network/network.go b/mtglib/network/network.go index bc3d84e..ecbf0e7 100644 --- a/mtglib/network/network.go +++ b/mtglib/network/network.go @@ -6,6 +6,7 @@ import ( "math/rand" "net" "net/http" + "sync" "time" doh "github.com/babolivier/go-doh-client" @@ -50,26 +51,48 @@ func (d *Network) resolveIPs(network, address string) ([]string, error) { return []string{address}, nil } - var ips []string + ips := []string{} + wg := &sync.WaitGroup{} + mutex := &sync.Mutex{} switch network { case "tcp", "tcp4": - if recs, _, err := d.DNS.LookupA(address); err == nil { - for _, v := range recs { - ips = append(ips, v.IP4) + wg.Add(1) + + go func() { + defer wg.Done() + + if recs, _, err := d.DNS.LookupA(address); err == nil { + mutex.Lock() + defer mutex.Unlock() + + for _, v := range recs { + ips = append(ips, v.IP4) + } } - } + }() } switch network { case "tcp", "tcp6": - if recs, _, err := d.DNS.LookupAAAA(address); err == nil { - for _, v := range recs { - ips = append(ips, v.IP6) + wg.Add(1) + + go func() { + defer wg.Done() + + if recs, _, err := d.DNS.LookupAAAA(address); err == nil { + mutex.Lock() + defer mutex.Unlock() + + for _, v := range recs { + ips = append(ips, v.IP6) + } } - } + }() } + wg.Wait() + if len(ips) == 0 { return nil, fmt.Errorf("cannot find any ips for %s:%s", network, address) } diff --git a/mtglib/network/shadowsocks.go b/mtglib/network/shadowsocks.go index 5fff10a..ffad8aa 100644 --- a/mtglib/network/shadowsocks.go +++ b/mtglib/network/shadowsocks.go @@ -10,6 +10,7 @@ import ( "time" shadowsocks "github.com/shadowsocks/go-shadowsocks2/core" + "golang.org/x/net/proxy" ) type shadowsocksDialer struct { @@ -18,19 +19,23 @@ type shadowsocksDialer struct { cipher shadowsocks.StreamConnCipher } +func (s *shadowsocksDialer) Dial(network, address string) (net.Conn, error) { + return s.DialContext(context.Background(), network, address) +} + func (s *shadowsocksDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) { conn, err := s.Dialer.DialContext(ctx, network, address) if err != nil { - return nil, err + return nil, err // nolint: wrapcheck } return s.cipher.StreamConn(conn), nil } -func NewShadowsocksDialer(proxyUrl *url.URL, +func NewShadowsocksDialer(proxyURL *url.URL, timeout time.Duration, bufferSize int) (Dialer, error) { - username := proxyUrl.User.Username() + username := proxyURL.User.Username() decoded, err := base64.RawURLEncoding.DecodeString(username) if err != nil { @@ -47,14 +52,24 @@ func NewShadowsocksDialer(proxyUrl *url.URL, return nil, fmt.Errorf("cannot initialize shadowsocks cipher: %w", err) } + socks5URL := *proxyURL + socks5URL.Scheme = "socks5" + socks5URL.User = nil + dialer, err := NewDefaultDialer(timeout, bufferSize) if err != nil { return nil, fmt.Errorf("cannot initialize a base dialer: %w", err) - } - return &shadowsocksDialer{ + ssDialer := &shadowsocksDialer{ Dialer: dialer, cipher: cipher, - }, nil + } + + rv, err := proxy.FromURL(&socks5URL, ssDialer) + if err != nil { + return nil, fmt.Errorf("cannot initialize ss proxy dialer: %w", err) + } + + return rv.(Dialer), nil } diff --git a/mtglib/network/socks5.go b/mtglib/network/socks5.go index ddb1be3..8ed895c 100644 --- a/mtglib/network/socks5.go +++ b/mtglib/network/socks5.go @@ -8,13 +8,13 @@ import ( "golang.org/x/net/proxy" ) -func NewSocks5Dialer(proxyUrl *url.URL, timeout time.Duration, bufferSize int) (Dialer, error) { +func NewSocks5Dialer(proxyURL *url.URL, timeout time.Duration, bufferSize int) (Dialer, error) { dialer, err := NewDefaultDialer(timeout, bufferSize) if err != nil { return nil, fmt.Errorf("cannot initialize base dialer: %w", err) } - rv, err := proxy.FromURL(proxyUrl, dialer.(*defaultDialer)) + rv, err := proxy.FromURL(proxyURL, dialer.(*defaultDialer)) if err != nil { return nil, fmt.Errorf("cannot initialize socks5 proxy dialer: %w", err) }