Keep active ciphers at the front

Author: fortuna

server.go

@@ -54,7 +54,7 @@ func init() {
 type SSPort struct {
 	tcpService shadowsocks.TCPService
 	udpService shadowsocks.UDPService
-	keys       map[string]shadowaead.Cipher
+	cipherList shadowsocks.CipherList
 }
 
 type SSServer struct {
@@ -73,10 +73,10 @@ func (s *SSServer) startPort(portNum int) error {
 		return fmt.Errorf("Failed to start UDP on port %v: %v", portNum, err)
 	}
 	logger.Infof("Listening TCP and UDP on port %v", portNum)
-	port := &SSPort{keys: make(map[string]shadowaead.Cipher)}
+	port := &SSPort{cipherList: shadowsocks.NewCipherList()}
 	// TODO: Register initial data metrics at zero.
-	port.tcpService = shadowsocks.NewTCPService(listener, &port.keys, s.m)
-	port.udpService = shadowsocks.NewUDPService(packetConn, s.natTimeout, &port.keys, s.m)
+	port.tcpService = shadowsocks.NewTCPService(listener, &port.cipherList, s.m)
+	port.udpService = shadowsocks.NewUDPService(packetConn, s.natTimeout, &port.cipherList, s.m)
 	s.ports[portNum] = port
 	go port.udpService.Start()
 	go port.tcpService.Start()
@@ -108,13 +108,13 @@ func (s *SSServer) loadConfig(filename string) error {
 	}
 
 	portChanges := make(map[int]int)
-	portKeys := make(map[int]map[string]shadowaead.Cipher)
+	portCiphers := make(map[int]shadowsocks.CipherList)
 	for _, keyConfig := range config.Keys {
 		portChanges[keyConfig.Port] = 1
-		keys, ok := portKeys[keyConfig.Port]
+		cipherList, ok := portCiphers[keyConfig.Port]
 		if !ok {
-			keys = make(map[string]shadowaead.Cipher)
-			portKeys[keyConfig.Port] = keys
+			cipherList = shadowsocks.NewCipherList()
+			portCiphers[keyConfig.Port] = cipherList
 		}
 		cipher, err := core.PickCipher(keyConfig.Cipher, nil, keyConfig.Secret)
 		if err != nil {
@@ -127,7 +127,7 @@ func (s *SSServer) loadConfig(filename string) error {
 		if !ok {
 			return fmt.Errorf("Only AEAD ciphers are supported. Found %v", keyConfig.Cipher)
 		}
-		keys[keyConfig.ID] = aead
+		cipherList.PushBack(keyConfig.ID, aead)
 	}
 	for port := range s.ports {
 		portChanges[port] = portChanges[port] - 1
@@ -143,11 +143,11 @@ func (s *SSServer) loadConfig(filename string) error {
 			}
 		}
 	}
-	for portNum, keys := range portKeys {
-		s.ports[portNum].keys = keys
+	for portNum, cipherList := range portCiphers {
+		s.ports[portNum].cipherList = cipherList
 	}
 	logger.Infof("Loaded %v access keys", len(config.Keys))
-	s.m.SetNumAccessKeys(len(config.Keys), len(portKeys))
+	s.m.SetNumAccessKeys(len(config.Keys), len(portCiphers))
 	return nil
 }
 

shadowsocks/cipher_cache.go

@@ -0,0 +1,69 @@
+// Copyright 2018 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 shadowsocks
+
+import (
+	"container/list"
+	"sync"
+
+	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
+)
+
+// CipherEntry holds a Cipher with an identifier.
+type CipherEntry struct {
+	ID     string
+	Cipher shadowaead.Cipher
+}
+
+// CipherList is a list of CipherEntry elements that allows for thread-safe snapshotting and
+// moving to front.
+type CipherList interface {
+	PushBack(id string, cipher shadowaead.Cipher) *list.Element
+	SafeSnapshot() []*list.Element
+	SafeMoveToFront(e *list.Element)
+}
+
+type cipherList struct {
+	CipherList
+	list *list.List
+	mu   sync.RWMutex
+}
+
+// NewCipherList creates an empty CipherList
+func NewCipherList() CipherList {
+	return &cipherList{list: list.New()}
+}
+
+func (cl *cipherList) PushBack(id string, cipher shadowaead.Cipher) *list.Element {
+	return cl.list.PushBack(&CipherEntry{ID: id, Cipher: cipher})
+}
+
+func (cl *cipherList) SafeSnapshot() []*list.Element {
+	cl.mu.RLock()
+	defer cl.mu.RUnlock()
+	cipherArray := make([]*list.Element, cl.list.Len())
+	i := 0
+	for e := cl.list.Front(); e != nil; e = e.Next() {
+		cipherArray[i] = e
+		i++
+	}
+	return cipherArray
+}
+
+func (cl *cipherList) SafeMoveToFront(e *list.Element) {
+	cl.mu.Lock()
+	defer cl.mu.Unlock()
+	cl.list.MoveToFront(e)
+}

shadowsocks/cipher_list.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package testing
+package shadowsocks
 
 import (
 	"fmt"
@@ -21,16 +21,16 @@ import (
 	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
 )
 
-func MakeTestCiphers(numCiphers int) (map[string]shadowaead.Cipher, error) {
-	cipherList := make(map[string]shadowaead.Cipher)
+func MakeTestCiphers(numCiphers int) (CipherList, error) {
+	cipherList := NewCipherList()
 	for i := 0; i < numCiphers; i++ {
 		cipherID := fmt.Sprintf("id-%v", i)
 		secret := fmt.Sprintf("secret-%v", i)
 		cipher, err := core.PickCipher("chacha20-ietf-poly1305", nil, secret)
 		if err != nil {
 			return nil, fmt.Errorf("Failed to create cipher %v: %v", i, err)
 		}
-		cipherList[cipherID] = cipher.(shadowaead.Cipher)
+		cipherList.PushBack(cipherID, cipher.(shadowaead.Cipher))
 	}
 	return cipherList, nil
 }

shadowsocks/cipher_map.go

@@ -26,7 +26,6 @@ import (
 	"github.com/Jigsaw-Code/outline-ss-server/metrics"
 	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 
-	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
 	"github.com/shadowsocks/go-shadowsocks2/socks"
 )
 
@@ -48,7 +47,7 @@ func ensureBytes(reader io.Reader, buf []byte, bytesNeeded int) ([]byte, error)
 	return buf, err
 }
 
-func findAccessKey(clientConn onet.DuplexConn, cipherMap map[string]shadowaead.Cipher) (string, onet.DuplexConn, error) {
+func findAccessKey(clientConn onet.DuplexConn, cipherList CipherList) (string, onet.DuplexConn, error) {
 	// This must have enough space to hold the salt + 2 bytes chunk length + AEAD tag (Oeverhead) for any cipher
 	replayBytes := make([]byte, 0, 32+2+16)
 	// Constant of zeroes to use as the start chunk count. This must be as big as the max NonceSize() across all ciphers.
@@ -58,11 +57,10 @@ func findAccessKey(clientConn onet.DuplexConn, cipherMap map[string]shadowaead.C
 	var err error
 
 	// Try each cipher until we find one that authenticates successfully. This assumes that all ciphers are AEAD.
-	// We shuffle the cipher map so that every connection has the same expected time.
-	// TODO: Reorder list to try previously successful ciphers first for the client IP.
+	// We snapshot the list because it may be modified while we use it.
 	// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
-	for _, entry := range shuffleCipherMap(cipherMap) {
-		id, cipher := entry.id, entry.cipher
+	for _, entry := range cipherList.SafeSnapshot() {
+		id, cipher := entry.Value.(*CipherEntry).ID, entry.Value.(*CipherEntry).Cipher
 		replayBytes, err = ensureBytes(clientConn, replayBytes, cipher.SaltSize())
 		if err != nil {
 			if logger.IsEnabledFor(logging.DEBUG) {
@@ -90,6 +88,8 @@ func findAccessKey(clientConn onet.DuplexConn, cipherMap map[string]shadowaead.C
 		if logger.IsEnabledFor(logging.DEBUG) {
 			logger.Debugf("Selected TCP cipher %v", id)
 		}
+		// Move the active cipher to the front, so that the search is quicker next time.
+		cipherList.SafeMoveToFront(entry)
 		ssr := NewShadowsocksReader(io.MultiReader(bytes.NewReader(replayBytes), clientConn), cipher)
 		ssw := NewShadowsocksWriter(clientConn, cipher)
 		return id, onet.WrapConn(clientConn, ssr, ssw).(onet.DuplexConn), nil
@@ -99,13 +99,13 @@ func findAccessKey(clientConn onet.DuplexConn, cipherMap map[string]shadowaead.C
 
 type tcpService struct {
 	listener  *net.TCPListener
-	ciphers   *map[string]shadowaead.Cipher
+	ciphers   *CipherList
 	m         metrics.ShadowsocksMetrics
 	isRunning bool
 }
 
 // NewTCPService creates a TCPService
-func NewTCPService(listener *net.TCPListener, ciphers *map[string]shadowaead.Cipher, m metrics.ShadowsocksMetrics) TCPService {
+func NewTCPService(listener *net.TCPListener, ciphers *CipherList, m metrics.ShadowsocksMetrics) TCPService {
 	return &tcpService{listener: listener, ciphers: ciphers, m: m}
 }
 

shadowsocks/testing/ciphers.go renamed to shadowsocks/cipher_testing.go

@@ -18,7 +18,6 @@ import (
 	"net"
 	"testing"
 
-	sstest "github.com/Jigsaw-Code/outline-ss-server/shadowsocks/testing"
 	logging "github.com/op/go-logging"
 )
 
@@ -32,11 +31,11 @@ func BenchmarkTCPFindCipher(b *testing.B) {
 		b.Fatalf("ListenTCP failed: %v", err)
 	}
 
-	cipherList, err := sstest.MakeTestCiphers(100)
+	cipherList, err := MakeTestCiphers(100)
 	if err != nil {
 		b.Fatal(err)
 	}
-	testPayload := sstest.MakeTestPayload(50)
+	testPayload := MakeTestPayload(50)
 	for n := 0; n < b.N; n++ {
 		go func() {
 			conn, err := net.Dial("tcp", listener.Addr().String())