refactor: pass a dialer to TCP serving (#150)

Author: fortuna

internal/integration_test/integration_test.go

@@ -26,7 +26,6 @@ import (
 	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
 	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
-	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 	"github.com/Jigsaw-Code/outline-ss-server/service"
 	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
 	sstest "github.com/Jigsaw-Code/outline-ss-server/shadowsocks"
@@ -41,7 +40,7 @@ func init() {
 	logging.SetLevel(logging.INFO, "")
 }
 
-func allowAll(ip net.IP) *onet.ConnectionError {
+func allowAll(ip net.IP) error {
 	// Allow access to localhost so that we can run integration tests with
 	// an actual destination server.
 	return nil
@@ -114,7 +113,7 @@ func TestTCPEcho(t *testing.T) {
 	replayCache := service.NewReplayCache(5)
 	const testTimeout = 200 * time.Millisecond
 	handler := service.NewTCPHandler(proxyListener.Addr().(*net.TCPAddr).Port, cipherList, &replayCache, &service.NoOpTCPMetrics{}, testTimeout)
-	handler.SetTargetIPValidator(allowAll)
+	handler.SetTargetDialer(&transport.TCPStreamDialer{})
 	done := make(chan struct{})
 	go func() {
 		service.StreamServe(func() (transport.StreamConn, error) { return proxyListener.AcceptTCP() }, handler.Handle)
@@ -362,7 +361,7 @@ func BenchmarkTCPThroughput(b *testing.B) {
 	}
 	const testTimeout = 200 * time.Millisecond
 	handler := service.NewTCPHandler(proxyListener.Addr().(*net.TCPAddr).Port, cipherList, nil, &service.NoOpTCPMetrics{}, testTimeout)
-	handler.SetTargetIPValidator(allowAll)
+	handler.SetTargetDialer(&transport.TCPStreamDialer{})
 	done := make(chan struct{})
 	go func() {
 		service.StreamServe(service.WrapStreamListener(proxyListener.AcceptTCP), handler.Handle)
@@ -424,7 +423,7 @@ func BenchmarkTCPMultiplexing(b *testing.B) {
 	replayCache := service.NewReplayCache(service.MaxCapacity)
 	const testTimeout = 200 * time.Millisecond
 	handler := service.NewTCPHandler(proxyListener.Addr().(*net.TCPAddr).Port, cipherList, &replayCache, &service.NoOpTCPMetrics{}, testTimeout)
-	handler.SetTargetIPValidator(allowAll)
+	handler.SetTargetDialer(&transport.TCPStreamDialer{})
 	done := make(chan struct{})
 	go func() {
 		service.StreamServe(service.WrapStreamListener(proxyListener.AcceptTCP), handler.Handle)

net/error.go

@@ -24,3 +24,23 @@ type ConnectionError struct {
 func NewConnectionError(status, message string, cause error) *ConnectionError {
 	return &ConnectionError{Status: status, Message: message, Cause: cause}
 }
+
+func (e *ConnectionError) Error() string {
+	if e == nil {
+		return "<nil>"
+	}
+	msg := e.Message
+	if len(e.Status) > 0 {
+		msg += " [" + e.Status + "]"
+	}
+	if e.Cause != nil {
+		msg += ": " + e.Cause.Error()
+	}
+	return msg
+}
+
+func (e *ConnectionError) Unwrap() error {
+	return e.Cause
+}
+
+var _ error = (*ConnectionError)(nil)

net/error_test.go

@@ -0,0 +1,49 @@
+// Copyright 2019 Jigsaw Operations LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestConnectionErrorUnwrapCause(t *testing.T) {
+	cause := errors.New("cause")
+	err := &ConnectionError{Cause: cause}
+	require.Equal(t, cause, err.Unwrap())
+	require.ErrorIs(t, err, cause)
+}
+
+func TestConnectionErrorString(t *testing.T) {
+	require.Equal(t, "example message", (&ConnectionError{Message: "example message"}).Error())
+	require.Equal(t, "example message [ERR_EXAMPLE]", (&ConnectionError{Message: "example message", Status: "ERR_EXAMPLE"}).Error())
+
+	cause := errors.New("cause")
+	err := &ConnectionError{Status: "ERR_EXAMPLE", Message: "example message", Cause: cause}
+	require.Equal(t, "example message [ERR_EXAMPLE]: cause", err.Error())
+}
+
+func TestConnectionErrorFromUnwrap(t *testing.T) {
+	connErr := &ConnectionError{Message: "connection error"}
+	topErr := fmt.Errorf("top error: %w", connErr)
+	require.NotEqual(t, topErr, connErr)
+	require.ErrorIs(t, topErr, connErr)
+	var unwrapped *ConnectionError
+	require.True(t, errors.As(topErr, &unwrapped))
+	require.Equal(t, connErr, unwrapped)
+}

net/private_net.go

@@ -48,11 +48,11 @@ func IsPrivateAddress(ip net.IP) bool {
 }
 
 // TargetIPValidator is a type alias for checking if an IP is allowed.
-type TargetIPValidator = func(net.IP) *ConnectionError
+type TargetIPValidator = func(net.IP) error
 
 // RequirePublicIP returns an error if the destination IP is not a
 // standard public IP.
-func RequirePublicIP(ip net.IP) *ConnectionError {
+func RequirePublicIP(ip net.IP) error {
 	if !ip.IsGlobalUnicast() {
 		return NewConnectionError("ERR_ADDRESS_INVALID", fmt.Sprintf("Address is not global unicast: %s", ip.String()), nil)
 	}

net/private_net_test.go

@@ -15,8 +15,11 @@
 package net
 
 import (
+	"errors"
 	"net"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 var privateAddressTests = []struct {
@@ -46,39 +49,50 @@ func TestIsLanAddress(t *testing.T) {
 }
 
 func TestRequirePublicIP(t *testing.T) {
-	if err := RequirePublicIP(net.ParseIP("8.8.8.8")); err != nil {
-		t.Error(err)
-	}
+	var err error
+
+	assert.Nil(t, RequirePublicIP(net.ParseIP("8.8.8.8")))
 
 	if err := RequirePublicIP(net.ParseIP("2001:4860:4860::8888")); err != nil {
 		t.Error(err)
 	}
 
-	err := RequirePublicIP(net.ParseIP("192.168.0.23"))
-	if err == nil {
-		t.Error("Expected error")
-	} else if err.Status != "ERR_ADDRESS_PRIVATE" {
-		t.Errorf("Wrong status %s", err.Status)
+	err = RequirePublicIP(net.ParseIP("192.168.0.23"))
+	if assert.NotNil(t, err) {
+		var connErr *ConnectionError
+		if assert.IsType(t, connErr, err) && assert.True(t, errors.As(err, &connErr)) {
+			assert.Equal(t, "ERR_ADDRESS_PRIVATE", connErr.Status)
+		}
 	}
 
 	err = RequirePublicIP(net.ParseIP("::1"))
-	if err == nil {
-		t.Error("Expected error")
-	} else if err.Status != "ERR_ADDRESS_INVALID" {
-		t.Errorf("Wrong status %s", err.Status)
+	if assert.NotNil(t, err) {
+		var connErr *ConnectionError
+		if assert.IsType(t, connErr, err) && assert.True(t, errors.As(err, &connErr)) {
+			assert.Equal(t, "ERR_ADDRESS_INVALID", connErr.Status)
+		}
 	}
 
 	err = RequirePublicIP(net.ParseIP("224.0.0.251"))
-	if err == nil {
-		t.Error("Expected error")
-	} else if err.Status != "ERR_ADDRESS_INVALID" {
-		t.Errorf("Wrong status %s", err.Status)
+	if assert.NotNil(t, err) {
+		var connErr *ConnectionError
+		if assert.IsType(t, connErr, err) && assert.True(t, errors.As(err, &connErr)) {
+			assert.Equal(t, "ERR_ADDRESS_INVALID", connErr.Status)
+		}
 	}
 
 	err = RequirePublicIP(net.ParseIP("ff02::fb"))
-	if err == nil {
-		t.Error("Expected error")
-	} else if err.Status != "ERR_ADDRESS_INVALID" {
-		t.Errorf("Wrong status %s", err.Status)
+	if assert.NotNil(t, err) {
+		var connErr *ConnectionError
+		if assert.IsType(t, connErr, err) && assert.True(t, errors.As(err, &connErr)) {
+			assert.Equal(t, "ERR_ADDRESS_INVALID", connErr.Status)
+		}
 	}
 }
+
+func TestRequirePublicIPInterface(t *testing.T) {
+	var err error
+	err = RequirePublicIP(net.ParseIP("8.8.8.8"))
+	assert.True(t, err == nil)
+	assert.Equal(t, nil, err)
+}

service/tcp.go

@@ -124,53 +124,53 @@ type tcpHandler struct {
 	m           TCPMetrics
 	readTimeout time.Duration
 	// `replayCache` is a pointer to SSServer.replayCache, to share the cache among all ports.
-	replayCache       *ReplayCache
-	targetIPValidator onet.TargetIPValidator
+	replayCache *ReplayCache
+	dialer      transport.StreamDialer
 }
 
 // NewTCPService creates a TCPService
 // `replayCache` is a pointer to SSServer.replayCache, to share the cache among all ports.
 func NewTCPHandler(port int, ciphers CipherList, replayCache *ReplayCache, m TCPMetrics, timeout time.Duration) TCPHandler {
 	return &tcpHandler{
-		port:              port,
-		ciphers:           ciphers,
-		m:                 m,
-		readTimeout:       timeout,
-		replayCache:       replayCache,
-		targetIPValidator: onet.RequirePublicIP,
+		port:        port,
+		ciphers:     ciphers,
+		m:           m,
+		readTimeout: timeout,
+		replayCache: replayCache,
+		dialer:      defaultDialer,
 	}
 }
 
+var defaultDialer = makeValidatingTCPStreamDialer(onet.RequirePublicIP)
+
+func makeValidatingTCPStreamDialer(targetIPValidator onet.TargetIPValidator) transport.StreamDialer {
+	return &transport.TCPStreamDialer{Dialer: net.Dialer{Control: func(network, address string, c syscall.RawConn) error {
+		ip, _, _ := net.SplitHostPort(address)
+		return targetIPValidator(net.ParseIP(ip))
+	}}}
+}
+
 // TCPService is a Shadowsocks TCP service that can be started and stopped.
 type TCPHandler interface {
 	Handle(ctx context.Context, conn transport.StreamConn)
-	// SetTargetIPValidator sets the function to be used to validate the target IP addresses.
-	SetTargetIPValidator(targetIPValidator onet.TargetIPValidator)
+	// SetTargetDialer sets the [transport.StreamDialer] to be used to connect to target addresses.
+	SetTargetDialer(dialer transport.StreamDialer)
 }
 
-func (s *tcpHandler) SetTargetIPValidator(targetIPValidator onet.TargetIPValidator) {
-	s.targetIPValidator = targetIPValidator
+func (s *tcpHandler) SetTargetDialer(dialer transport.StreamDialer) {
+	s.dialer = dialer
 }
 
-func dialTarget(tgtAddr socks.Addr, proxyMetrics *metrics.ProxyMetrics, targetIPValidator onet.TargetIPValidator) (transport.StreamConn, *onet.ConnectionError) {
-	var ipError *onet.ConnectionError
-	dialer := net.Dialer{Control: func(network, address string, c syscall.RawConn) error {
-		ip, _, _ := net.SplitHostPort(address)
-		ipError = targetIPValidator(net.ParseIP(ip))
-		if ipError != nil {
-			return errors.New(ipError.Message)
-		}
+func ensureConnectionError(err error, fallbackStatus string, fallbackMsg string) *onet.ConnectionError {
+	if err == nil {
 		return nil
-	}}
-	tgtConn, err := dialer.Dial("tcp", tgtAddr.String())
-	if ipError != nil {
-		return nil, ipError
-	} else if err != nil {
-		return nil, onet.NewConnectionError("ERR_CONNECT", "Failed to connect to target", err)
 	}
-	tgtTCPConn := tgtConn.(*net.TCPConn)
-	tgtTCPConn.SetKeepAlive(true)
-	return metrics.MeasureConn(tgtTCPConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy), nil
+	var connErr *onet.ConnectionError
+	if errors.As(err, &connErr) {
+		return connErr
+	} else {
+		return onet.NewConnectionError(fallbackStatus, fallbackMsg, err)
+	}
 }
 
 type StreamListener func() (transport.StreamConn, error)
@@ -226,7 +226,7 @@ func (h *tcpHandler) Handle(ctx context.Context, clientConn transport.StreamConn
 	measuredClientConn := metrics.MeasureConn(clientConn, &proxyMetrics.ProxyClient, &proxyMetrics.ClientProxy)
 	connStart := time.Now()
 
-	id, connError := h.handleConnection(h.port, measuredClientConn, &proxyMetrics)
+	id, connError := h.handleConnection(ctx, h.port, measuredClientConn, &proxyMetrics)
 
 	connDuration := time.Since(connStart)
 	status := "OK"
@@ -239,7 +239,7 @@ func (h *tcpHandler) Handle(ctx context.Context, clientConn transport.StreamConn
 	logger.Debugf("Done with status %v, duration %v", status, connDuration)
 }
 
-func (h *tcpHandler) handleConnection(listenerPort int, clientConn transport.StreamConn, proxyMetrics *metrics.ProxyMetrics) (string, *onet.ConnectionError) {
+func (h *tcpHandler) handleConnection(ctx context.Context, listenerPort int, clientConn transport.StreamConn, proxyMetrics *metrics.ProxyMetrics) (string, *onet.ConnectionError) {
 	// Set a deadline to receive the address to the target.
 	clientConn.SetReadDeadline(time.Now().Add(h.readTimeout))
 
@@ -275,18 +275,20 @@ func (h *tcpHandler) handleConnection(listenerPort int, clientConn transport.Str
 	// 3. Read target address and dial it.
 	ssr := shadowsocks.NewReader(clientReader, cipherEntry.CryptoKey)
 	tgtAddr, err := socks.ReadAddr(ssr)
+
 	// Clear the deadline for the target address
 	clientConn.SetReadDeadline(time.Time{})
 	if err != nil {
 		// Drain to prevent a close on cipher error.
 		io.Copy(io.Discard, clientConn)
 		return id, onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", err)
 	}
-	tgtConn, dialErr := dialTarget(tgtAddr, proxyMetrics, h.targetIPValidator)
+	tgtConn, dialErr := h.dialer.Dial(ctx, tgtAddr.String())
 	if dialErr != nil {
 		// We don't drain so dial errors and invalid addresses are communicated quickly.
-		return id, dialErr
+		return id, ensureConnectionError(dialErr, "ERR_CONNECT", "Failed to connect to target")
 	}
+	tgtConn = metrics.MeasureConn(tgtConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)
 	defer tgtConn.Close()
 
 	// 4. Bridge the client and target connections

service/tcp_test.go

@@ -28,7 +28,6 @@ import (
 	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
 	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
-	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
 	logging "github.com/op/go-logging"
 	"github.com/shadowsocks/go-shadowsocks2/socks"
@@ -39,7 +38,7 @@ func init() {
 	logging.SetLevel(logging.INFO, "")
 }
 
-func allowAll(ip net.IP) *onet.ConnectionError {
+func allowAll(ip net.IP) error {
 	// Allow access to localhost so that we can run integration tests with
 	// an actual destination server.
 	return nil
@@ -353,7 +352,7 @@ func TestProbeClientBytesBasicTruncated(t *testing.T) {
 	cipher := firstCipher(cipherList)
 	testMetrics := &probeTestMetrics{}
 	handler := NewTCPHandler(listener.Addr().(*net.TCPAddr).Port, cipherList, nil, testMetrics, 200*time.Millisecond)
-	handler.SetTargetIPValidator(allowAll)
+	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)
@@ -388,7 +387,7 @@ func TestProbeClientBytesBasicModified(t *testing.T) {
 	cipher := firstCipher(cipherList)
 	testMetrics := &probeTestMetrics{}
 	handler := NewTCPHandler(listener.Addr().(*net.TCPAddr).Port, cipherList, nil, testMetrics, 200*time.Millisecond)
-	handler.SetTargetIPValidator(allowAll)
+	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)
@@ -424,7 +423,7 @@ func TestProbeClientBytesCoalescedModified(t *testing.T) {
 	cipher := firstCipher(cipherList)
 	testMetrics := &probeTestMetrics{}
 	handler := NewTCPHandler(listener.Addr().(*net.TCPAddr).Port, cipherList, nil, testMetrics, 200*time.Millisecond)
-	handler.SetTargetIPValidator(allowAll)
+	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)

service/udp.go

@@ -230,7 +230,7 @@ func (h *packetHandler) validatePacket(textData []byte) ([]byte, *net.UDPAddr, *
 		return nil, nil, onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr), err)
 	}
 	if err := h.targetIPValidator(tgtUDPAddr.IP); err != nil {
-		return nil, nil, err
+		return nil, nil, ensureConnectionError(err, "ERR_ADDRESS_INVALID", "invalid address")
 	}
 
 	payload := textData[len(tgtAddr):]