Use CloseRead and CloseWrites

This commit is contained in:
9seconds
2021-12-01 10:37:31 +03:00
parent 7b1f86b75d
commit ffad717829
32 changed files with 320 additions and 99 deletions
+7 -5
View File
@@ -2,9 +2,10 @@ package network
import (
"context"
"net"
"sync/atomic"
"time"
"github.com/9seconds/mtg/v2/essentials"
)
const (
@@ -30,12 +31,12 @@ type circuitBreakerDialer struct {
resetFailuresTimeout time.Duration
}
func (c *circuitBreakerDialer) Dial(network, address string) (net.Conn, error) {
func (c *circuitBreakerDialer) Dial(network, address string) (essentials.Conn, error) {
return c.DialContext(context.Background(), network, address)
}
func (c *circuitBreakerDialer) DialContext(ctx context.Context,
network, address string) (net.Conn, error) {
network, address string) (essentials.Conn, error) {
switch atomic.LoadUint32(&c.state) {
case circuitBreakerStateClosed:
return c.doClosed(ctx, network, address)
@@ -47,7 +48,7 @@ func (c *circuitBreakerDialer) DialContext(ctx context.Context,
}
func (c *circuitBreakerDialer) doClosed(ctx context.Context,
network, address string) (net.Conn, error) {
network, address string) (essentials.Conn, error) {
conn, err := c.Dialer.DialContext(ctx, network, address)
select {
@@ -78,7 +79,8 @@ func (c *circuitBreakerDialer) doClosed(ctx context.Context,
return conn, err // nolint: wrapcheck
}
func (c *circuitBreakerDialer) doHalfOpened(ctx context.Context, network, address string) (net.Conn, error) {
func (c *circuitBreakerDialer) doHalfOpened(ctx context.Context,
network, address string) (essentials.Conn, error) {
if !atomic.CompareAndSwapUint32(&c.halfOpenAttempts, 0, 1) {
return nil, ErrCircuitBreakerOpened
}
+2 -2
View File
@@ -21,7 +21,7 @@ type CircuitBreakerTestSuite struct {
mutex sync.Mutex
ctx context.Context
ctxCancel context.CancelFunc
connMock *testlib.NetConnMock
connMock *testlib.EssentialsConnMock
baseDialerMock *DialerMock
}
@@ -29,7 +29,7 @@ func (suite *CircuitBreakerTestSuite) SetupTest() {
suite.mutex = sync.Mutex{}
suite.ctx, suite.ctxCancel = context.WithCancel(context.Background())
suite.baseDialerMock = &DialerMock{}
suite.connMock = &testlib.NetConnMock{}
suite.connMock = &testlib.EssentialsConnMock{}
suite.d = newCircuitBreakerDialer(suite.baseDialerMock,
3, 100*time.Millisecond, 50*time.Millisecond)
}
+5 -3
View File
@@ -5,17 +5,19 @@ import (
"fmt"
"net"
"time"
"github.com/9seconds/mtg/v2/essentials"
)
type defaultDialer struct {
net.Dialer
}
func (d *defaultDialer) Dial(network, address string) (net.Conn, error) {
func (d *defaultDialer) Dial(network, address string) (essentials.Conn, error) {
return d.DialContext(context.Background(), network, address)
}
func (d *defaultDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
func (d *defaultDialer) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) {
switch network {
case "tcp", "tcp4", "tcp6": // nolint: goconst
default:
@@ -34,7 +36,7 @@ func (d *defaultDialer) DialContext(ctx context.Context, network, address string
return nil, fmt.Errorf("cannot set socket options: %w", err)
}
return conn, nil
return conn.(essentials.Conn), nil
}
// NewDefaultDialer build a new dialer which dials bypassing proxies
+4 -3
View File
@@ -20,8 +20,9 @@ package network
import (
"context"
"errors"
"net"
"time"
"github.com/9seconds/mtg/v2/essentials"
)
const (
@@ -95,6 +96,6 @@ var (
// Dialer defines an interface which is required to bootstrap a network
// instance from.
type Dialer interface {
Dial(network, address string) (net.Conn, error)
DialContext(ctx context.Context, network, address string) (net.Conn, error)
Dial(network, address string) (essentials.Conn, error)
DialContext(ctx context.Context, network, address string) (essentials.Conn, error)
}
+5 -5
View File
@@ -2,8 +2,8 @@ package network
import (
"context"
"net"
"github.com/9seconds/mtg/v2/essentials"
"github.com/stretchr/testify/mock"
)
@@ -11,14 +11,14 @@ type DialerMock struct {
mock.Mock
}
func (d *DialerMock) Dial(network, address string) (net.Conn, error) {
func (d *DialerMock) Dial(network, address string) (essentials.Conn, error) {
args := d.Called(network, address)
return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck
return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck
}
func (d *DialerMock) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
func (d *DialerMock) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) {
args := d.Called(ctx, network, address)
return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck
return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck
}
+8 -5
View File
@@ -8,6 +8,7 @@ import (
"net/url"
"strings"
"github.com/9seconds/mtg/v2/essentials"
"github.com/9seconds/mtg/v2/network"
socks5 "github.com/armon/go-socks5"
"github.com/mccutchen/go-httpbin/httpbin"
@@ -18,16 +19,16 @@ type DialerMock struct {
mock.Mock
}
func (d *DialerMock) Dial(network, address string) (net.Conn, error) {
func (d *DialerMock) Dial(network, address string) (essentials.Conn, error) {
args := d.Called(network, address)
return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck
return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck
}
func (d *DialerMock) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
func (d *DialerMock) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) {
args := d.Called(ctx, network, address)
return args.Get(0).(net.Conn), args.Error(1) // nolint: wrapcheck
return args.Get(0).(essentials.Conn), args.Error(1) // nolint: wrapcheck
}
type HTTPServerTestSuite struct {
@@ -53,7 +54,9 @@ func (suite *HTTPServerTestSuite) MakeURL(path string) string {
func (suite *HTTPServerTestSuite) MakeHTTPClient(dialer network.Dialer) *http.Client {
return &http.Client{
Transport: &http.Transport{
DialContext: dialer.DialContext,
DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
return dialer.DialContext(ctx, network, address) // nolint: wrapcheck
},
},
}
}
+4 -3
View File
@@ -4,19 +4,20 @@ import (
"context"
"fmt"
"math/rand"
"net"
"net/url"
"github.com/9seconds/mtg/v2/essentials"
)
type loadBalancedSocks5Dialer struct {
dialers []Dialer
}
func (l loadBalancedSocks5Dialer) Dial(network, address string) (net.Conn, error) {
func (l loadBalancedSocks5Dialer) Dial(network, address string) (essentials.Conn, error) {
return l.DialContext(context.Background(), network, address)
}
func (l loadBalancedSocks5Dialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
func (l loadBalancedSocks5Dialer) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) {
length := len(l.dialers)
start := rand.Intn(length)
moved := false
+10 -6
View File
@@ -9,6 +9,7 @@ import (
"sync"
"time"
"github.com/9seconds/mtg/v2/essentials"
"github.com/9seconds/mtg/v2/mtglib"
)
@@ -30,11 +31,11 @@ type network struct {
dns *dnsResolver
}
func (n *network) Dial(protocol, address string) (net.Conn, error) {
func (n *network) Dial(protocol, address string) (essentials.Conn, error) {
return n.DialContext(context.Background(), protocol, address)
}
func (n *network) DialContext(ctx context.Context, protocol, address string) (net.Conn, error) {
func (n *network) DialContext(ctx context.Context, protocol, address string) (essentials.Conn, error) {
host, port, _ := net.SplitHostPort(address)
ips, err := n.dnsResolve(protocol, host)
@@ -46,7 +47,8 @@ func (n *network) DialContext(ctx context.Context, protocol, address string) (ne
ips[i], ips[j] = ips[j], ips[i]
})
var conn net.Conn
var conn essentials.Conn
for _, v := range ips {
conn, err = n.dialer.DialContext(ctx, protocol, net.JoinHostPort(v, port))
@@ -59,7 +61,7 @@ func (n *network) DialContext(ctx context.Context, protocol, address string) (ne
}
func (n *network) MakeHTTPClient(dialFunc func(ctx context.Context,
network, address string) (net.Conn, error)) *http.Client {
network, address string) (essentials.Conn, error)) *http.Client {
if dialFunc == nil {
dialFunc = n.DialContext
}
@@ -144,13 +146,15 @@ func NewNetwork(dialer Dialer,
func makeHTTPClient(userAgent string,
timeout time.Duration,
dialFunc func(ctx context.Context, network, address string) (net.Conn, error)) *http.Client {
dialFunc func(ctx context.Context, network, address string) (essentials.Conn, error)) *http.Client {
return &http.Client{
Timeout: timeout,
Transport: networkHTTPTransport{
userAgent: userAgent,
next: &http.Transport{
DialContext: dialFunc,
DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
return dialFunc(ctx, network, address)
},
},
},
}
+146 -5
View File
@@ -1,21 +1,162 @@
package network
import (
"context"
"fmt"
"io"
"net"
"net/url"
"golang.org/x/net/proxy"
"github.com/9seconds/mtg/v2/essentials"
"github.com/txthinking/socks5"
)
type socks5Dialer struct {
Dialer
username []byte
password []byte
proxyAddress string
}
func (s socks5Dialer) Dial(network, address string) (essentials.Conn, error) {
return s.DialContext(context.Background(), network, address)
}
func (s socks5Dialer) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) {
switch network {
case "tcp", "tcp4", "tcp6":
default:
return nil, fmt.Errorf("%s network type is not supported", network)
}
conn, err := s.Dialer.DialContext(ctx, network, s.proxyAddress)
if err != nil {
return nil, fmt.Errorf("cannot dial to the proxy: %w", err)
}
if err := s.handshake(conn); err != nil {
conn.Close()
return nil, fmt.Errorf("cannot perform a handshake: %w", err)
}
if err := s.connect(conn, address); err != nil {
conn.Close()
return nil, fmt.Errorf("cannot connect to a destination host %s: %w", address, err)
}
return conn, nil
}
func (s socks5Dialer) handshake(conn io.ReadWriter) error {
authMethod := socks5.MethodUsernamePassword
if len(s.username)+len(s.password) == 0 {
authMethod = socks5.MethodNone
}
if err := s.handshakeNegotiation(conn, authMethod); err != nil {
return fmt.Errorf("cannot perform negotiation: %w", err)
}
if authMethod == socks5.MethodNone {
return nil
}
if err := s.handshakeAuth(conn); err != nil {
return fmt.Errorf("cannot authenticate: %w", err)
}
return nil
}
func (s socks5Dialer) handshakeNegotiation(conn io.ReadWriter, authMethod byte) error {
request := socks5.NewNegotiationRequest([]byte{authMethod})
if _, err := request.WriteTo(conn); err != nil {
return fmt.Errorf("cannot send request: %w", err)
}
response, err := socks5.NewNegotiationReplyFrom(conn)
if err != nil {
return fmt.Errorf("cannot read response: %w", err)
}
if response.Method != authMethod {
return fmt.Errorf("%v is unsupported auth method", authMethod)
}
return nil
}
func (s socks5Dialer) handshakeAuth(conn io.ReadWriter) error {
request := socks5.NewUserPassNegotiationRequest(s.username, s.password)
if _, err := request.WriteTo(conn); err != nil {
return fmt.Errorf("cannot send a request: %w", err)
}
response, err := socks5.NewUserPassNegotiationReplyFrom(conn)
if err != nil {
return fmt.Errorf("cannot read a response: %w", err)
}
if response.Status != socks5.UserPassStatusSuccess {
return fmt.Errorf("authenticate has failed: %v", response.Status)
}
return nil
}
func (s socks5Dialer) connect(conn io.ReadWriter, address string) error {
addrType, host, port, err := socks5.ParseAddress(address)
if err != nil {
return fmt.Errorf("cannot parse address: %w", err)
}
if addrType == socks5.ATYPDomain {
host = host[1:]
}
request := socks5.NewRequest(socks5.CmdConnect, addrType, host, port)
if _, err := request.WriteTo(conn); err != nil {
return fmt.Errorf("cannot send a request: %w", err)
}
response, err := socks5.NewReplyFrom(conn)
if err != nil {
return fmt.Errorf("cannot read a response: %w", err)
}
if response.Rep != socks5.RepSuccess {
return fmt.Errorf("unsuccessful request: %v", response.Rep)
}
return nil
}
// NewSocks5Dialer build a new dialer from a given one (so, in theory
// you can chain here). Proxy parameters are passed with URI in a form of:
//
// socks5://[user:[password]]@host:port
func NewSocks5Dialer(baseDialer Dialer, proxyURL *url.URL) (Dialer, error) {
rv, err := proxy.FromURL(proxyURL, baseDialer)
if err != nil {
return nil, fmt.Errorf("cannot initialize socks5 proxy dialer: %w", err)
if _, _, err := net.SplitHostPort(proxyURL.Host); err != nil {
return nil, fmt.Errorf("incorrect url %s", proxyURL.Redacted())
}
return rv.(Dialer), nil
dialer := socks5Dialer{
Dialer: baseDialer,
proxyAddress: proxyURL.Host,
}
if proxyURL.User != nil {
password, isSet := proxyURL.User.Password()
if isSet {
dialer.username = []byte(proxyURL.User.Username())
dialer.password = []byte(password)
}
}
return dialer, nil
}
+1 -1
View File
@@ -55,7 +55,7 @@ func (suite *Socks5TestSuite) TestRequestOk() {
suite.Equal(http.StatusOK, resp.StatusCode)
}
func TestSocks5TestSuite(t *testing.T) {
func TestSocks5(t *testing.T) {
t.Parallel()
suite.Run(t, &Socks5TestSuite{})
}