Split out proxying into proxyConnection().

Author: sbruens

cmd/outline-ss-server/metrics.go

@@ -324,12 +324,12 @@ func (m *outlineMetrics) AddClosedTCPConnection(clientInfo ipinfo.IPInfo, client
 	}
 }
 
-func (m *outlineMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, clientProxyBytes, proxyTargetBytes int) {
+func (m *outlineMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, data metrics.ProxyMetrics) {
 	m.udpPacketsFromClientPerLocation.WithLabelValues(clientInfo.CountryCode.String(), asnLabel(clientInfo.ASN), status).Inc()
-	addIfNonZero(int64(clientProxyBytes), m.dataBytes, "c>p", "udp", accessKey)
-	addIfNonZero(int64(clientProxyBytes), m.dataBytesPerLocation, "c>p", "udp", clientInfo.CountryCode.String(), asnLabel(clientInfo.ASN))
-	addIfNonZero(int64(proxyTargetBytes), m.dataBytes, "p>t", "udp", accessKey)
-	addIfNonZero(int64(proxyTargetBytes), m.dataBytesPerLocation, "p>t", "udp", clientInfo.CountryCode.String(), asnLabel(clientInfo.ASN))
+	addIfNonZero(data.ClientProxy, m.dataBytes, "c>p", "udp", accessKey)
+	addIfNonZero(data.ClientProxy, m.dataBytesPerLocation, "c>p", "udp", clientInfo.CountryCode.String(), asnLabel(clientInfo.ASN))
+	addIfNonZero(data.ProxyTarget, m.dataBytes, "p>t", "udp", accessKey)
+	addIfNonZero(data.ProxyTarget, m.dataBytesPerLocation, "p>t", "udp", clientInfo.CountryCode.String(), asnLabel(clientInfo.ASN))
 }
 
 func (m *outlineMetrics) AddUDPPacketFromTarget(clientInfo ipinfo.IPInfo, accessKey, status string, targetProxyBytes, proxyClientBytes int) {

cmd/outline-ss-server/metrics_test.go

@@ -52,23 +52,29 @@ func init() {
 
 func TestMethodsDontPanic(t *testing.T) {
 	ssMetrics := newPrometheusOutlineMetrics(nil, prometheus.NewPedanticRegistry())
-	proxyMetrics := metrics.ProxyMetrics{
+	tcpProxyMetrics := metrics.ProxyMetrics{
 		ClientProxy: 1,
 		ProxyTarget: 2,
 		TargetProxy: 3,
 		ProxyClient: 4,
 	}
+	udpProxyMetrics := metrics.ProxyMetrics{
+		ClientProxy: 10,
+		ProxyTarget: 20,
+		TargetProxy: 30,
+		ProxyClient: 40,
+	}
 	ipInfo := ipinfo.IPInfo{CountryCode: "US", ASN: 100}
 	ssMetrics.SetBuildInfo("0.0.0-test")
 	ssMetrics.SetNumAccessKeys(20, 2)
 	ssMetrics.AddOpenTCPConnection(ipInfo)
 	ssMetrics.AddAuthenticatedTCPConnection(fakeAddr("127.0.0.1:9"), "0")
-	ssMetrics.AddClosedTCPConnection(ipInfo, fakeAddr("127.0.0.1:9"), "1", "OK", proxyMetrics, 10*time.Millisecond)
-	ssMetrics.AddUDPPacketFromClient(ipInfo, "2", "OK", 10, 20)
+	ssMetrics.AddClosedTCPConnection(ipInfo, fakeAddr("127.0.0.1:9"), "1", "OK", tcpProxyMetrics, 10*time.Millisecond)
+	ssMetrics.AddUDPPacketFromClient(ipInfo, "2", "OK", udpProxyMetrics)
 	ssMetrics.AddUDPPacketFromTarget(ipInfo, "3", "OK", 10, 20)
 	ssMetrics.AddUDPNatEntry(fakeAddr("127.0.0.1:9"), "key-1")
 	ssMetrics.RemoveUDPNatEntry(fakeAddr("127.0.0.1:9"), "key-1")
-	ssMetrics.AddTCPProbe("ERR_CIPHER", "eof", 443, proxyMetrics.ClientProxy)
+	ssMetrics.AddTCPProbe("ERR_CIPHER", "eof", 443, tcpProxyMetrics.ClientProxy)
 	ssMetrics.AddTCPCipherSearch(true, 10*time.Millisecond)
 	ssMetrics.AddUDPCipherSearch(true, 10*time.Millisecond)
 }
@@ -174,14 +180,19 @@ func BenchmarkProbe(b *testing.B) {
 
 func BenchmarkClientUDP(b *testing.B) {
 	ssMetrics := newPrometheusOutlineMetrics(nil, prometheus.NewRegistry())
+	proxyMetrics := metrics.ProxyMetrics{
+		ClientProxy: 1000,
+		ProxyTarget: 2000,
+		TargetProxy: 3000,
+		ProxyClient: 4000,
+	}
 	clientInfo := ipinfo.IPInfo{CountryCode: "ZZ", ASN: 100}
 	accessKey := "key 1"
 	status := "OK"
-	size := 1000
 	timeToCipher := time.Microsecond
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		ssMetrics.AddUDPPacketFromClient(clientInfo, accessKey, status, size, size)
+		ssMetrics.AddUDPPacketFromClient(clientInfo, accessKey, status, proxyMetrics)
 		ssMetrics.AddUDPCipherSearch(true, timeToCipher)
 	}
 }

internal/integration_test/integration_test.go

@@ -262,8 +262,8 @@ var _ service.UDPMetrics = (*fakeUDPMetrics)(nil)
 func (m *fakeUDPMetrics) GetIPInfo(ip net.IP) (ipinfo.IPInfo, error) {
 	return ipinfo.IPInfo{CountryCode: "QQ"}, nil
 }
-func (m *fakeUDPMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, clientProxyBytes, proxyTargetBytes int) {
-	m.up = append(m.up, udpRecord{clientInfo, accessKey, status, clientProxyBytes, proxyTargetBytes})
+func (m *fakeUDPMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, data metrics.ProxyMetrics) {
+	m.up = append(m.up, udpRecord{clientInfo, accessKey, status, int(data.ClientProxy), int(data.ProxyTarget)})
 }
 func (m *fakeUDPMetrics) AddUDPPacketFromTarget(clientInfo ipinfo.IPInfo, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
 	m.down = append(m.down, udpRecord{clientInfo, accessKey, status, targetProxyBytes, proxyClientBytes})

service/udp.go

@@ -15,6 +15,7 @@
 package service
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
@@ -26,6 +27,7 @@ import (
 	"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"
 )
@@ -35,7 +37,7 @@ type UDPMetrics interface {
 	ipinfo.IPInfoMap
 
 	// UDP metrics
-	AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, clientProxyBytes, proxyTargetBytes int)
+	AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, data metrics.ProxyMetrics)
 	AddUDPPacketFromTarget(clientInfo ipinfo.IPInfo, accessKey, status string, targetProxyBytes, proxyClientBytes int)
 	AddUDPNatEntry(clientAddr net.Addr, accessKey string)
 	RemoveUDPNatEntry(clientAddr net.Addr, accessKey string)
@@ -119,15 +121,15 @@ func (h *packetHandler) Handle(clientConn net.PacketConn) {
 
 	for {
 		status := "OK"
-		keyID, clientInfo, clientProxyBytes, proxyTargetBytes, connErr := h.handleConnection(clientConn)
+		keyID, clientInfo, proxyMetrics, connErr := h.handleConnection(context.TODO(), clientConn)
 		if connErr != nil {
 			if errors.Is(connErr.Cause, net.ErrClosed) {
 				break
 			}
 			logger.Debugf("UDP Error: %v: %v", connErr.Message, connErr.Cause)
 			status = connErr.Status
 		}
-		h.m.AddUDPPacketFromClient(clientInfo, keyID, status, clientProxyBytes, proxyTargetBytes)
+		h.m.AddUDPPacketFromClient(clientInfo, keyID, status, proxyMetrics)
 	}
 }
 
@@ -157,52 +159,59 @@ func (h *packetHandler) authenticate(clientConn net.PacketConn) (net.Addr, *Ciph
 	return clientAddr, cipherEntry, textData, clientProxyBytes, nil
 }
 
-func (h *packetHandler) handleConnection(clientConn net.PacketConn) (string, ipinfo.IPInfo, int, int, *onet.ConnectionError) {
+func (h *packetHandler) proxyConnection(ctx context.Context, clientAddr net.Addr, tgtAddr net.Addr, clientConn net.PacketConn, cipherEntry CipherEntry, payload []byte, proxyMetrics *metrics.ProxyMetrics) (ipinfo.IPInfo, *onet.ConnectionError) {
+	tgtConn := h.nm.Get(clientAddr.String())
+	if tgtConn == nil {
+		clientInfo, locErr := ipinfo.GetIPInfoFromAddr(h.m, clientAddr)
+		if locErr != nil {
+			logger.Warningf("Failed client info lookup: %v", locErr)
+		}
+		debugUDPAddr(clientAddr, "Got info \"%#v\"", clientInfo)
+
+		udpConn, err := net.ListenPacket("udp", "")
+		if err != nil {
+			return ipinfo.IPInfo{}, nil
+		}
+		tgtConn = h.nm.Add(clientAddr, clientConn, cipherEntry.CryptoKey, udpConn, clientInfo, cipherEntry.ID)
+	}
+
+	proxyTargetBytes, err := tgtConn.WriteTo(payload, tgtAddr)
+	proxyMetrics.ProxyTarget += int64(proxyTargetBytes)
+	if err != nil {
+		return tgtConn.clientInfo, onet.NewConnectionError("ERR_WRITE", "Failed to write to target", err)
+	}
+	return tgtConn.clientInfo, nil
+}
+
+func (h *packetHandler) handleConnection(ctx context.Context, clientConn net.PacketConn) (string, ipinfo.IPInfo, metrics.ProxyMetrics, *onet.ConnectionError) {
 	defer func() {
 		if r := recover(); r != nil {
 			logger.Errorf("Panic in UDP loop: %v. Continuing to listen.", r)
 			debug.PrintStack()
 		}
 	}()
 
+	var proxyMetrics metrics.ProxyMetrics
 	clientAddr, cipherEntry, textData, clientProxyBytes, authErr := h.authenticate(clientConn)
+	proxyMetrics.ClientProxy += int64(clientProxyBytes)
 	if authErr != nil {
-		return "", ipinfo.IPInfo{}, clientProxyBytes, 0, authErr
+		return "", ipinfo.IPInfo{}, proxyMetrics, authErr
 	}
 
-	targetConn := h.nm.Get(clientAddr.String())
-	if targetConn == nil {
-		udpConn, err := net.ListenPacket("udp", "")
-		if err != nil {
-			return "", ipinfo.IPInfo{}, clientProxyBytes, 0, onet.NewConnectionError("ERR_CREATE_SOCKET", "Failed to create UDP socket", err)
-		}
-
-		clientInfo, locErr := ipinfo.GetIPInfoFromAddr(h.m, clientAddr)
-		if locErr != nil {
-			logger.Warningf("Failed client info lookup: %v", locErr)
-		}
-		debugUDPAddr(clientAddr, "Got info \"%#v\"", clientInfo)
-
-		targetConn = h.nm.Add(clientAddr, clientConn, cipherEntry.CryptoKey, udpConn, clientInfo, cipherEntry.ID)
-	}
-
-	payload, tgtUDPAddr, onetErr := h.validatePacket(textData)
+	payload, tgtAddr, onetErr := h.getProxyRequest(textData)
 	if onetErr != nil {
-		return targetConn.keyID, targetConn.clientInfo, clientProxyBytes, 0, onetErr
+		return cipherEntry.ID, ipinfo.IPInfo{}, proxyMetrics, onetErr
 	}
+	debugUDPAddr(clientAddr, "Proxy exit %s", tgtAddr.String())
 
-	debugUDPAddr(targetConn.clientAddr, "Proxy exit %v", targetConn.LocalAddr())
-	proxyTargetBytes, err := targetConn.WriteTo(payload, tgtUDPAddr) // accept only UDPAddr despite the signature
-	if err != nil {
-		return targetConn.keyID, targetConn.clientInfo, clientProxyBytes, proxyTargetBytes, onet.NewConnectionError("ERR_WRITE", "Failed to write to target", err)
-	}
-	return targetConn.keyID, targetConn.clientInfo, clientProxyBytes, proxyTargetBytes, nil
+	clientInfo, err := h.proxyConnection(ctx, clientAddr, tgtAddr, clientConn, *cipherEntry, payload, &proxyMetrics)
+	return cipherEntry.ID, clientInfo, proxyMetrics, err
 }
 
 // Given the decrypted contents of a UDP packet, return
 // the payload and the destination address, or an error if
 // this packet cannot or should not be forwarded.
-func (h *packetHandler) validatePacket(textData []byte) ([]byte, *net.UDPAddr, *onet.ConnectionError) {
+func (h *packetHandler) getProxyRequest(textData []byte) ([]byte, *net.UDPAddr, *onet.ConnectionError) {
 	tgtAddr := socks.SplitAddr(textData)
 	if tgtAddr == nil {
 		return nil, nil, onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", nil)
@@ -227,9 +236,7 @@ func isDNS(addr net.Addr) bool {
 
 type natconn struct {
 	net.PacketConn
-	cryptoKey  *shadowsocks.EncryptionKey
-	keyID      string
-	clientAddr net.Addr
+	cryptoKey *shadowsocks.EncryptionKey
 	// We store the client information in the NAT map to avoid recomputing it
 	// for every downstream packet in a UDP-based connection.
 	clientInfo ipinfo.IPInfo
@@ -310,12 +317,9 @@ func (m *natmap) Get(key string) *natconn {
 	return m.keyConn[key]
 }
 
-func (m *natmap) set(clientAddr net.Addr, pc net.PacketConn, cryptoKey *shadowsocks.EncryptionKey, keyID string, clientInfo ipinfo.IPInfo) *natconn {
+func (m *natmap) set(clientAddr net.Addr, pc net.PacketConn, clientInfo ipinfo.IPInfo) *natconn {
 	entry := &natconn{
 		PacketConn:     pc,
-		cryptoKey:      cryptoKey,
-		keyID:          keyID,
-		clientAddr:     clientAddr,
 		clientInfo:     clientInfo,
 		defaultTimeout: m.timeout,
 	}
@@ -340,12 +344,12 @@ func (m *natmap) del(key string) net.PacketConn {
 }
 
 func (m *natmap) Add(clientAddr net.Addr, clientConn net.PacketConn, cryptoKey *shadowsocks.EncryptionKey, targetConn net.PacketConn, clientInfo ipinfo.IPInfo, keyID string) *natconn {
-	entry := m.set(clientAddr, targetConn, cryptoKey, keyID, clientInfo)
+	entry := m.set(clientAddr, targetConn, clientInfo)
 
 	m.metrics.AddUDPNatEntry(clientAddr, keyID)
 	m.running.Add(1)
 	go func() {
-		timedCopy(clientAddr, clientConn, entry, keyID, m.metrics)
+		timedCopy(clientAddr, clientConn, entry, cryptoKey, keyID, m.metrics)
 		m.metrics.RemoveUDPNatEntry(clientAddr, keyID)
 		if pc := m.del(clientAddr.String()); pc != nil {
 			pc.Close()
@@ -375,13 +379,13 @@ var maxAddrLen int = len(socks.ParseAddr("[2001:db8::1]:12345"))
 
 // copy from target to client until read timeout
 func timedCopy(clientAddr net.Addr, clientConn net.PacketConn, targetConn *natconn,
-	keyID string, sm UDPMetrics) {
+	cryptoKey *shadowsocks.EncryptionKey, keyID string, sm UDPMetrics) {
 	// pkt is used for in-place encryption of downstream UDP packets, with the layout
 	// [padding?][salt][address][body][tag][extra]
 	// Padding is only used if the address is IPv4.
 	pkt := make([]byte, serverUDPBufferSize)
 
-	saltSize := targetConn.cryptoKey.SaltSize()
+	saltSize := cryptoKey.SaltSize()
 	// Leave enough room at the beginning of the packet for a max-length header (i.e. IPv6).
 	bodyStart := saltSize + maxAddrLen
 
@@ -425,7 +429,7 @@ func timedCopy(clientAddr net.Addr, clientConn net.PacketConn, targetConn *natco
 			//           [            packBuf             ]
 			//           [          buf           ]
 			packBuf := pkt[saltStart:]
-			buf, err := shadowsocks.Pack(packBuf, plaintextBuf, targetConn.cryptoKey) // Encrypt in-place
+			buf, err := shadowsocks.Pack(packBuf, plaintextBuf, cryptoKey) // Encrypt in-place
 			if err != nil {
 				return onet.NewConnectionError("ERR_PACK", "Failed to pack data to client", err)
 			}
@@ -456,7 +460,7 @@ var _ UDPMetrics = (*NoOpUDPMetrics)(nil)
 func (m *NoOpUDPMetrics) GetIPInfo(net.IP) (ipinfo.IPInfo, error) {
 	return ipinfo.IPInfo{}, nil
 }
-func (m *NoOpUDPMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, clientProxyBytes, proxyTargetBytes int) {
+func (m *NoOpUDPMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, data metrics.ProxyMetrics) {
 }
 func (m *NoOpUDPMetrics) AddUDPPacketFromTarget(clientInfo ipinfo.IPInfo, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
 }

service/udp_test.go

@@ -26,6 +26,7 @@ import (
 	"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"
 	"github.com/stretchr/testify/assert"
@@ -93,9 +94,9 @@ func (conn *fakePacketConn) Close() error {
 }
 
 type udpReport struct {
-	clientInfo                         ipinfo.IPInfo
-	accessKey, status                  string
-	clientProxyBytes, proxyTargetBytes int
+	clientInfo        ipinfo.IPInfo
+	accessKey, status string
+	data              metrics.ProxyMetrics
 }
 
 // Stub metrics implementation for testing NAT behaviors.
@@ -109,8 +110,8 @@ var _ UDPMetrics = (*natTestMetrics)(nil)
 func (m *natTestMetrics) GetIPInfo(net.IP) (ipinfo.IPInfo, error) {
 	return ipinfo.IPInfo{}, nil
 }
-func (m *natTestMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, clientProxyBytes, proxyTargetBytes int) {
-	m.upstreamPackets = append(m.upstreamPackets, udpReport{clientInfo, accessKey, status, clientProxyBytes, proxyTargetBytes})
+func (m *natTestMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, data metrics.ProxyMetrics) {
+	m.upstreamPackets = append(m.upstreamPackets, udpReport{clientInfo, accessKey, status, data})
 }
 func (m *natTestMetrics) AddUDPPacketFromTarget(clientInfo ipinfo.IPInfo, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
 }
@@ -178,7 +179,7 @@ func TestIPFilter(t *testing.T) {
 
 		assert.Equal(t, 2, len(metrics.upstreamPackets), "Expected 2 reports, not %v", metrics.upstreamPackets)
 		for _, report := range metrics.upstreamPackets {
-			assert.Greater(t, report.proxyTargetBytes, 0, "Expected nonzero bytes to be sent for allowed packet")
+			assert.Greater(t, int(report.data.ProxyTarget), 0, "Expected nonzero bytes to be sent for allowed packet")
 		}
 	})
 
@@ -187,7 +188,7 @@ func TestIPFilter(t *testing.T) {
 
 		assert.Equal(t, 2, len(metrics.upstreamPackets), "Expected 2 reports, not %v", metrics.upstreamPackets)
 		for _, report := range metrics.upstreamPackets {
-			assert.Equal(t, 0, report.proxyTargetBytes, "No bytes should be sent due to a disallowed packet")
+			assert.EqualValues(t, 0, report.data.ProxyTarget, "No bytes should be sent due to a disallowed packet")
 		}
 	})
 }
@@ -197,13 +198,13 @@ func TestNATEntries(t *testing.T) {
 	payloads := [][]byte{[]byte("payload1"), []byte("payload2")}
 
 	t.Run("Valid cipher", func(t *testing.T) {
-		metrics := sendToDiscardWithValidCipher(payloads, onet.RequirePublicIP)
+		metrics := sendToDiscardWithValidCipher(payloads, allowAll)
 
 		assert.Equal(t, 1, metrics.natEntriesAdded, "Expected 1 NAT entry, not %d", metrics.natEntriesAdded)
 	})
 
 	t.Run("Invalid cipher", func(t *testing.T) {
-		metrics := sendToDiscardWithInValidCipher(payloads, onet.RequirePublicIP)
+		metrics := sendToDiscardWithInValidCipher(payloads, allowAll)
 
 		assert.Equal(t, 0, metrics.natEntriesAdded, "Unexpected NAT entry on rejected packet")
 	})
@@ -221,8 +222,8 @@ func TestUpstreamMetrics(t *testing.T) {
 
 	assert.Equal(t, N, len(metrics.upstreamPackets), "Expected %d reports, not %v", N, metrics.upstreamPackets)
 	for i, report := range metrics.upstreamPackets {
-		assert.Equal(t, i+1, report.proxyTargetBytes, "Expected %d payload bytes, not %d", i+1, report.proxyTargetBytes)
-		assert.Greater(t, report.clientProxyBytes, report.proxyTargetBytes, "Expected nonzero input overhead (%d > %d)", report.clientProxyBytes, report.proxyTargetBytes)
+		assert.EqualValues(t, i+1, report.data.ProxyTarget, "Expected %d payload bytes, not %d", i+1, report.data.ProxyTarget)
+		assert.Greater(t, report.data.ClientProxy, report.data.ProxyTarget, "Expected nonzero input overhead (%d > %d)", report.data.ClientProxy, report.data.ProxyTarget)
 		assert.Equal(t, "id-0", report.accessKey, "Unexpected access key name: %s", report.accessKey)
 		assert.Equal(t, "OK", report.status, "Wrong status: %s", report.status)
 	}