refactor: share Shadowsocks key search between TCP and UDP

service/shadowsocks.go

@@ -0,0 +1,49 @@
+// Copyright 2024 The Outline Authors
+//
+// 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
+//
+//     http://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 service
+
+import (
+	"errors"
+	"net/netip"
+	"time"
+
+	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
+)
+
+type DebugLoggerFunc func(tag string, template string, val interface{})
+
+// findShadowsocksAccessKey implements a trial decryption search. This assumes that all ciphers are AEAD.
+func findShadowsocksAccessKey(clientIP netip.Addr, bufferSize int, src []byte, cipherList CipherList, logDebug DebugLoggerFunc) (*CipherEntry, []byte, time.Duration, error) {
+	// We snapshot the list because it may be modified while we use it.
+	ciphers := cipherList.SnapshotForClientIP(clientIP)
+
+	unpackStart := time.Now()
+	// To hold the decrypted chunk length.
+	chunkLenBuf := make([]byte, bufferSize)
+	for ci, elt := range ciphers {
+		entry := elt.Value.(*CipherEntry)
+		buf, err := shadowsocks.Unpack(chunkLenBuf, src, entry.CryptoKey)
+		if err != nil {
+			logDebug(entry.ID, "Failed to unpack: %v", err)
+			continue
+		}
+		logDebug(entry.ID, "Found cipher at index %d", ci)
+
+		// Move the active cipher to the front, so that the search is quicker next time.
+		cipherList.MarkUsedByClientIP(elt, clientIP)
+		return entry, buf, time.Since(unpackStart), nil
+	}
+	return nil, nil, time.Since(unpackStart), errors.New("could not find valid cipher")
+}

service/shadowsocks_test.go

@@ -0,0 +1,152 @@
+// Copyright 2024 The Outline Authors
+//
+// 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
+//
+//     http://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 service
+
+import (
+	"io"
+	"net"
+	"net/netip"
+	"testing"
+
+	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
+	"github.com/stretchr/testify/require"
+)
+
+// Simulates receiving invalid TCP connection attempts on a server with 100 ciphers.
+func BenchmarkTCPFindCipherFail(b *testing.B) {
+	b.StopTimer()
+	b.ResetTimer()
+
+	clientIP := netip.MustParseAddr("127.0.0.1")
+	cipherList, err := MakeTestCiphers(makeTestSecrets(100))
+	if err != nil {
+		b.Fatal(err)
+	}
+	testPayload := makeTestPayload(50)
+	for n := 0; n < b.N; n++ {
+		b.StartTimer()
+		findShadowsocksAccessKey(clientIP, 2, testPayload, cipherList, func(tag string, template string, val interface{}) {})
+		b.StopTimer()
+	}
+}
+
+// Simulates receiving valid TCP connection attempts from 100 different users,
+// each with their own cipher and their own IP address.
+func BenchmarkTCPFindCipherRepeat(b *testing.B) {
+	b.StopTimer()
+	b.ResetTimer()
+
+	const numCiphers = 100 // Must be <256
+	cipherList, err := MakeTestCiphers(makeTestSecrets(numCiphers))
+	if err != nil {
+		b.Fatal(err)
+	}
+	cipherEntries := [numCiphers]*CipherEntry{}
+	snapshot := cipherList.SnapshotForClientIP(netip.Addr{})
+	for cipherNumber, element := range snapshot {
+		cipherEntries[cipherNumber] = element.Value.(*CipherEntry)
+	}
+	testPayload := makeTestPayload(50)
+	for n := 0; n < b.N; n++ {
+		cipherNumber := byte(n % numCiphers)
+		reader, writer := io.Pipe()
+		clientIP := netip.AddrFrom4([4]byte{192, 0, 2, cipherNumber})
+		addr := netip.AddrPortFrom(clientIP, 54321)
+		c := conn{clientAddr: net.TCPAddrFromAddrPort(addr), reader: reader, writer: writer}
+		cipher := cipherEntries[cipherNumber].CryptoKey
+		go shadowsocks.NewWriter(writer, cipher).Write(makeTestPayload(50))
+		b.StartTimer()
+		_, _, _, err := findShadowsocksAccessKey(clientIP, 2, testPayload, cipherList, nil)
+		b.StopTimer()
+		if err != nil {
+			b.Error(err)
+		}
+		c.Close()
+	}
+}
+
+// Simulates receiving invalid UDP packets on a server with 100 ciphers.
+func BenchmarkUDPUnpackFail(b *testing.B) {
+	cipherList, err := MakeTestCiphers(makeTestSecrets(100))
+	if err != nil {
+		b.Fatal(err)
+	}
+	testPayload := makeTestPayload(50)
+	testIP := netip.MustParseAddr("192.0.2.1")
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		findShadowsocksAccessKey(testIP, serverUDPBufferSize, testPayload, cipherList, nil)
+	}
+}
+
+// Simulates receiving valid UDP packets from 100 different users, each with
+// their own cipher and IP address.
+func BenchmarkUDPUnpackRepeat(b *testing.B) {
+	const numCiphers = 100 // Must be <256
+	cipherList, err := MakeTestCiphers(makeTestSecrets(numCiphers))
+	if err != nil {
+		b.Fatal(err)
+	}
+	packets := [numCiphers][]byte{}
+	ips := [numCiphers]netip.Addr{}
+	snapshot := cipherList.SnapshotForClientIP(netip.Addr{})
+	for i, element := range snapshot {
+		packets[i] = make([]byte, 0, serverUDPBufferSize)
+		plaintext := makeTestPayload(50)
+		packets[i], err = shadowsocks.Pack(make([]byte, serverUDPBufferSize), plaintext, element.Value.(*CipherEntry).CryptoKey)
+		if err != nil {
+			b.Error(err)
+		}
+		ips[i] = netip.AddrFrom4([4]byte{192, 0, 2, byte(i)})
+	}
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		cipherNumber := n % numCiphers
+		ip := ips[cipherNumber]
+		packet := packets[cipherNumber]
+		_, _, _, err := findShadowsocksAccessKey(ip, serverUDPBufferSize, packet, cipherList, nil)
+		if err != nil {
+			b.Error(err)
+		}
+	}
+}
+
+// Simulates receiving valid UDP packets from 100 different IP addresses,
+// all using the same cipher.
+func BenchmarkUDPUnpackSharedKey(b *testing.B) {
+	cipherList, err := MakeTestCiphers(makeTestSecrets(1)) // One widely shared key
+	if err != nil {
+		b.Fatal(err)
+	}
+	plaintext := makeTestPayload(50)
+	snapshot := cipherList.SnapshotForClientIP(netip.Addr{})
+	cryptoKey := snapshot[0].Value.(*CipherEntry).CryptoKey
+	packet, err := shadowsocks.Pack(make([]byte, serverUDPBufferSize), plaintext, cryptoKey)
+	require.Nil(b, err)
+
+	const numIPs = 100 // Must be <256
+	ips := [numIPs]netip.Addr{}
+	for i := 0; i < numIPs; i++ {
+		ips[i] = netip.AddrFrom4([4]byte{192, 0, 2, byte(i)})
+	}
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		ip := ips[n%numIPs]
+		_, _, _, err := findShadowsocksAccessKey(ip, serverUDPBufferSize, packet, cipherList, nil)
+		if err != nil {
+			b.Error(err)
+		}
+	}
+}

service/tcp.go

@@ -16,7 +16,6 @@ package service
 
 import (
 	"bytes"
-	"container/list"
 	"context"
 	"errors"
 	"fmt"
@@ -71,52 +70,6 @@ func debugTCP(cipherID, template string, val interface{}) {
 	}
 }
 
-// bytesForKeyFinding is the number of bytes to read for finding the AccessKey.
-// Is must satisfy provided >= bytesForKeyFinding >= required for every cipher in the list.
-// provided = saltSize + 2 + 2 * cipher.TagSize, the minimum number of bytes we will see in a valid connection
-// required = saltSize + 2 + cipher.TagSize, the number of bytes needed to authenticate the connection.
-const bytesForKeyFinding = 50
-
-func findAccessKey(clientReader io.Reader, clientIP netip.Addr, cipherList CipherList) (*CipherEntry, io.Reader, []byte, time.Duration, error) {
-	// We snapshot the list because it may be modified while we use it.
-	ciphers := cipherList.SnapshotForClientIP(clientIP)
-	firstBytes := make([]byte, bytesForKeyFinding)
-	if n, err := io.ReadFull(clientReader, firstBytes); err != nil {
-		return nil, clientReader, nil, 0, fmt.Errorf("reading header failed after %d bytes: %w", n, err)
-	}
-
-	findStartTime := time.Now()
-	entry, elt := findEntry(firstBytes, ciphers)
-	timeToCipher := time.Since(findStartTime)
-	if entry == nil {
-		// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
-		return nil, clientReader, nil, timeToCipher, fmt.Errorf("could not find valid TCP cipher")
-	}
-
-	// Move the active cipher to the front, so that the search is quicker next time.
-	cipherList.MarkUsedByClientIP(elt, clientIP)
-	salt := firstBytes[:entry.CryptoKey.SaltSize()]
-	return entry, io.MultiReader(bytes.NewReader(firstBytes), clientReader), salt, timeToCipher, nil
-}
-
-// Implements a trial decryption search.  This assumes that all ciphers are AEAD.
-func findEntry(firstBytes []byte, ciphers []*list.Element) (*CipherEntry, *list.Element) {
-	// To hold the decrypted chunk length.
-	chunkLenBuf := [2]byte{}
-	for ci, elt := range ciphers {
-		entry := elt.Value.(*CipherEntry)
-		cryptoKey := entry.CryptoKey
-		_, err := shadowsocks.Unpack(chunkLenBuf[:0], firstBytes[:cryptoKey.SaltSize()+2+cryptoKey.TagSize()], cryptoKey)
-		if err != nil {
-			debugTCP(entry.ID, "Failed to decrypt length: %v", err)
-			continue
-		}
-		debugTCP(entry.ID, "Found cipher at index %d", ci)
-		return entry, elt
-	}
-	return nil, nil
-}
-
 type StreamAuthenticateFunc func(clientConn transport.StreamConn) (string, transport.StreamConn, *onet.ConnectionError)
 
 // ShadowsocksTCPMetrics is used to report Shadowsocks metrics on TCP connections.
@@ -125,23 +78,37 @@ type ShadowsocksTCPMetrics interface {
 	AddTCPCipherSearch(accessKeyFound bool, timeToCipher time.Duration)
 }
 
+// bytesForKeyFinding is the number of bytes to read for finding the AccessKey.
+// Is must satisfy provided >= bytesForKeyFinding >= required for every cipher in the list.
+// provided = saltSize + 2 + 2 * cipher.TagSize, the minimum number of bytes we will see in a valid connection
+// required = saltSize + 2 + cipher.TagSize, the number of bytes needed to authenticate the connection.
+const bytesForKeyFinding = 50
+
 // NewShadowsocksStreamAuthenticator creates a stream authenticator that uses Shadowsocks.
 // TODO(fortuna): Offer alternative transports.
 func NewShadowsocksStreamAuthenticator(ciphers CipherList, replayCache *ReplayCache, metrics ShadowsocksTCPMetrics) StreamAuthenticateFunc {
 	return func(clientConn transport.StreamConn) (string, transport.StreamConn, *onet.ConnectionError) {
+		firstBytes := make([]byte, bytesForKeyFinding)
+		if n, err := io.ReadFull(clientConn, firstBytes); err != nil {
+			metrics.AddTCPCipherSearch(false, 0)
+			return "", clientConn, onet.NewConnectionError("ERR_CIPHER", fmt.Sprintf("Reading header failed after %d bytes", n), err)
+		}
+
 		// Find the cipher and acess key id.
-		cipherEntry, clientReader, clientSalt, timeToCipher, keyErr := findAccessKey(clientConn, remoteIP(clientConn), ciphers)
+		bufferSize := 2
+		cipherEntry, _, timeToCipher, keyErr := findShadowsocksAccessKey(remoteIP(clientConn), bufferSize, firstBytes, ciphers, debugTCP)
 		metrics.AddTCPCipherSearch(keyErr == nil, timeToCipher)
 		if keyErr != nil {
-			const status = "ERR_CIPHER"
-			return "", nil, onet.NewConnectionError(status, "Failed to find a valid cipher", keyErr)
+			// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
+			return "", clientConn, onet.NewConnectionError("ERR_CIPHER", "Failed to find a valid cipher", keyErr)
 		}
 		var id string
 		if cipherEntry != nil {
 			id = cipherEntry.ID
 		}
 
 		// Check if the connection is a replay.
+		clientSalt := firstBytes[:cipherEntry.CryptoKey.SaltSize()]
 		isServerSalt := cipherEntry.SaltGenerator.IsServerSalt(clientSalt)
 		// Only check the cache if findAccessKey succeeded and the salt is unrecognized.
 		if isServerSalt || !replayCache.Add(cipherEntry.ID, clientSalt) {
@@ -154,6 +121,7 @@ func NewShadowsocksStreamAuthenticator(ciphers CipherList, replayCache *ReplayCa
 			return id, nil, onet.NewConnectionError(status, "Replay detected", nil)
 		}
 
+		clientReader := io.MultiReader(bytes.NewReader(firstBytes), clientConn)
 		ssr := shadowsocks.NewReader(clientReader, cipherEntry.CryptoKey)
 		ssw := shadowsocks.NewWriter(clientConn, cipherEntry.CryptoKey)
 		ssw.SetSaltGenerator(cipherEntry.SaltGenerator)

service/tcp_test.go

@@ -74,39 +74,6 @@ func startDiscardServer(t testing.TB) (*net.TCPListener, *sync.WaitGroup) {
 	return listener, &running
 }
 
-// Simulates receiving invalid TCP connection attempts on a server with 100 ciphers.
-func BenchmarkTCPFindCipherFail(b *testing.B) {
-	b.StopTimer()
-	b.ResetTimer()
-
-	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
-	if err != nil {
-		b.Fatalf("ListenTCP failed: %v", err)
-	}
-
-	cipherList, err := MakeTestCiphers(makeTestSecrets(100))
-	if err != nil {
-		b.Fatal(err)
-	}
-	testPayload := makeTestPayload(50)
-	for n := 0; n < b.N; n++ {
-		go func() {
-			conn, err := net.Dial("tcp", listener.Addr().String())
-			require.NoErrorf(b, err, "Failed to dial %v: %v", listener.Addr(), err)
-			conn.Write(testPayload)
-			conn.Close()
-		}()
-		clientConn, err := listener.AcceptTCP()
-		if err != nil {
-			b.Fatalf("AcceptTCP failed: %v", err)
-		}
-		clientIP := clientConn.RemoteAddr().(*net.TCPAddr).AddrPort().Addr()
-		b.StartTimer()
-		findAccessKey(clientConn, clientIP, cipherList)
-		b.StopTimer()
-	}
-}
-
 func TestCompatibleCiphers(t *testing.T) {
 	for _, cipherName := range [](string){shadowsocks.CHACHA20IETFPOLY1305, shadowsocks.AES256GCM, shadowsocks.AES192GCM, shadowsocks.AES128GCM} {
 		cryptoKey, _ := shadowsocks.NewEncryptionKey(cipherName, "dummy secret")
@@ -180,40 +147,6 @@ func (c *conn) CloseWrite() error {
 	return nil
 }
 
-// Simulates receiving valid TCP connection attempts from 100 different users,
-// each with their own cipher and their own IP address.
-func BenchmarkTCPFindCipherRepeat(b *testing.B) {
-	b.StopTimer()
-	b.ResetTimer()
-
-	const numCiphers = 100 // Must be <256
-	cipherList, err := MakeTestCiphers(makeTestSecrets(numCiphers))
-	if err != nil {
-		b.Fatal(err)
-	}
-	cipherEntries := [numCiphers]*CipherEntry{}
-	snapshot := cipherList.SnapshotForClientIP(netip.Addr{})
-	for cipherNumber, element := range snapshot {
-		cipherEntries[cipherNumber] = element.Value.(*CipherEntry)
-	}
-	for n := 0; n < b.N; n++ {
-		cipherNumber := byte(n % numCiphers)
-		reader, writer := io.Pipe()
-		clientIP := netip.AddrFrom4([4]byte{192, 0, 2, cipherNumber})
-		addr := netip.AddrPortFrom(clientIP, 54321)
-		c := conn{clientAddr: net.TCPAddrFromAddrPort(addr), reader: reader, writer: writer}
-		cipher := cipherEntries[cipherNumber].CryptoKey
-		go shadowsocks.NewWriter(writer, cipher).Write(makeTestPayload(50))
-		b.StartTimer()
-		_, _, _, _, err := findAccessKey(&c, clientIP, cipherList)
-		b.StopTimer()
-		if err != nil {
-			b.Error(err)
-		}
-		c.Close()
-	}
-}
-
 // Stub metrics implementation for testing replay defense.
 type probeTestMetrics struct {
 	mu          sync.Mutex