🚲 ss2022: improve salt pool

- De-genericify SaltPool.
- Add a singly linked list to make evictions cheaper (O(n) -> O(1)).
- Improve the API.

Author: database64128

ss2022/saltpool.go

@@ -5,67 +5,104 @@ import (
 	"time"
 )
 
-// SaltPool stores salts for [retention, 2*retention) to protect against replay attacks
+// SaltPool stores salts for [ReplayWindowDuration] to protect against replay attacks
 // during the replay window.
-type SaltPool[T comparable] struct {
-	mu        sync.RWMutex
-	pool      map[T]time.Time
-	retention time.Duration
-	lastClean time.Time
+type SaltPool struct {
+	mu         sync.RWMutex
+	nodeBySalt map[[32]byte]*saltNode
+
+	// head is the oldest node.
+	head *saltNode
+	// tail is the newest node.
+	tail *saltNode
 }
 
-// clean removes expired salts from the pool,
-// if the amount of time since the last cleanup exceeds retention.
-func (p *SaltPool[T]) clean(now time.Time) {
-	if now.Sub(p.lastClean) > p.retention {
-		for salt, added := range p.pool {
-			if now.Sub(added) > p.retention {
-				delete(p.pool, salt)
-			}
-		}
-		p.lastClean = now
-	}
+type saltNode struct {
+	next      *saltNode
+	salt      [32]byte
+	expiresAt time.Time
 }
 
-// Check returns whether the given salt is valid (not in the pool).
-func (p *SaltPool[T]) Check(salt T) bool {
+// Contains returns whether the pool contains the given salt.
+func (p *SaltPool) Contains(salt [32]byte) bool {
 	p.mu.RLock()
-	_, ok := p.pool[salt]
+	_, ok := p.nodeBySalt[salt]
 	p.mu.RUnlock()
-	return !ok
+	return ok
 }
 
-// TryCheck is like Check, but it immediately returns true if the pool is contended.
-func (p *SaltPool[T]) TryCheck(salt T) bool {
+// TryContains is like Contains, but it immediately returns false if the pool is contended.
+func (p *SaltPool) TryContains(salt [32]byte) bool {
 	if p.mu.TryRLock() {
-		_, ok := p.pool[salt]
+		_, ok := p.nodeBySalt[salt]
 		p.mu.RUnlock()
-		return !ok
+		return ok
 	}
-	return true
+	return false
 }
 
-// Add cleans the pool, checks if the salt already exists in the pool,
-// and adds the salt to the pool if the salt is not already in the pool.
+// Add adds the salt to the pool if it is not already in the pool.
 // It returns true if the salt was added, false if it already exists.
-func (p *SaltPool[T]) Add(now time.Time, salt T) bool {
+func (p *SaltPool) Add(now time.Time, salt [32]byte) bool {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
-	p.clean(now)
-	if _, ok := p.pool[salt]; ok {
+	p.pruneExpired(now)
+	if _, ok := p.nodeBySalt[salt]; ok {
 		return false
 	}
-	p.pool[salt] = now
+	p.insert(now, salt)
 	return true
 }
 
-// NewSaltPool returns a new salt pool with the given retention as the minimum amount of time
-// for which an added salt is guaranteed to be kept in the pool.
-func NewSaltPool[T comparable](retention time.Duration) *SaltPool[T] {
-	return &SaltPool[T]{
-		pool:      make(map[T]time.Time),
-		retention: retention,
-		lastClean: time.Now(),
+// Clear removes all salts from the pool.
+func (p *SaltPool) Clear() {
+	p.mu.Lock()
+	clear(p.nodeBySalt)
+	p.head = nil
+	p.tail = nil
+	p.mu.Unlock()
+}
+
+// pruneExpired removes all expired salts from the pool.
+func (p *SaltPool) pruneExpired(now time.Time) {
+	node := p.head
+	if node == nil || node.expiresAt.After(now) {
+		return
+	}
+	for {
+		delete(p.nodeBySalt, node.salt)
+		node = node.next
+		if node == nil {
+			p.head = nil
+			p.tail = nil
+			return
+		}
+		if node.expiresAt.After(now) {
+			p.head = node
+			return
+		}
+	}
+}
+
+// insert adds the new salt to the pool.
+func (p *SaltPool) insert(now time.Time, salt [32]byte) {
+	node := &saltNode{
+		salt:      salt,
+		expiresAt: now.Add(ReplayWindowDuration),
+	}
+	p.nodeBySalt[salt] = node
+	if p.tail != nil {
+		p.tail.next = node
+	} else {
+		p.head = node
+	}
+	p.tail = node
+}
+
+// NewSaltPool returns a new salt pool.
+func NewSaltPool() *SaltPool {
+	return &SaltPool{
+		nodeBySalt: make(map[[32]byte]*saltNode),
 	}
 }

ss2022/saltpool_test.go

@@ -1,44 +1,89 @@
-package ss2022
+package ss2022_test
 
 import (
 	"crypto/rand"
 	"testing"
 	"time"
-)
 
-func TestSaltPoolAddDuplicateSalts(t *testing.T) {
-	const retention = 100 * time.Millisecond
-	var salt [32]byte
-	rand.Read(salt[:])
+	"github.com/database64128/shadowsocks-go/ss2022"
+)
 
-	pool := NewSaltPool[[32]byte](retention)
+func TestSaltPool(t *testing.T) {
+	pool := ss2022.NewSaltPool()
 	now := time.Now()
+	b := make([]byte, 64)
+	rand.Read(b)
+	salt0 := [32]byte(b)
+	salt1 := [32]byte(b[32:])
+
+	// Clear empty pool.
+	pool.Clear()
 
-	// Check fresh salt.
-	if !pool.Check(salt) {
-		t.Fatal("Denied fresh salt.")
+	// Check salt0 and salt1.
+	if pool.Contains(salt0) {
+		t.Fatal("pool.Contains(salt0) = true, want false")
+	}
+	if pool.TryContains(salt1) {
+		t.Fatal("pool.TryContains(salt1) = true, want false")
 	}
 
-	// Add fresh salt.
-	if !pool.Add(now, salt) {
-		t.Fatal("Denied fresh salt.")
+	// Add salt0.
+	if !pool.Add(now, salt0) {
+		t.Fatal("pool.Add(now, salt0) = false, want true")
+	}
+	if pool.Add(now, salt0) {
+		t.Fatal("pool.Add(now, salt0) = true, want false")
 	}
 
-	// Check the same salt again.
-	if pool.Check(salt) {
-		t.Fatal("Accepted duplicate salt.")
+	// Advance some time.
+	now = now.Add(ss2022.ReplayWindowDuration / 2)
+
+	// Add salt1.
+	if !pool.Add(now, salt1) {
+		t.Fatal("pool.Add(now, salt1) = false, want true")
+	}
+	if pool.Add(now, salt1) {
+		t.Fatal("pool.Add(now, salt1) = true, want false")
 	}
 
-	// Add the same salt again.
-	if pool.Add(now, salt) {
-		t.Fatal("Accepted duplicate salt.")
+	// Check salt0 and salt1.
+	if !pool.Contains(salt0) {
+		t.Fatal("pool.Contains(salt0) = false, want true")
+	}
+	if !pool.Contains(salt1) {
+		t.Fatal("pool.Contains(salt1) = false, want true")
 	}
 
-	// Advance time to let the salt expire.
-	now = now.Add(2 * retention)
+	// Advance some time to let salt0 expire.
+	now = now.Add(ss2022.ReplayWindowDuration / 2)
 
-	// Add the expired salt.
-	if !pool.Add(now, salt) {
-		t.Fatal("Denied expired salt.")
+	// Add salt0 and salt1.
+	if !pool.Add(now, salt0) {
+		t.Fatal("pool.Add(now, salt0) = false, want true")
+	}
+	if pool.Add(now, salt1) {
+		t.Fatal("pool.Add(now, salt1) = true, want false")
+	}
+
+	// Advance some time to let both expire.
+	now = now.Add(ss2022.ReplayWindowDuration)
+
+	// Add salt0 and salt1.
+	if !pool.Add(now, salt0) {
+		t.Fatal("pool.Add(now, salt0) = false, want true")
+	}
+	if !pool.Add(now, salt1) {
+		t.Fatal("pool.Add(now, salt1) = false, want true")
+	}
+
+	// Clear the pool.
+	pool.Clear()
+
+	// Check salt0 and salt1 again.
+	if pool.TryContains(salt0) {
+		t.Fatal("pool.TryContains(salt0) = true, want false")
+	}
+	if pool.TryContains(salt1) {
+		t.Fatal("pool.TryContains(salt1) = true, want false")
 	}
 }

ss2022/tcp.go

@@ -230,7 +230,7 @@ type StreamServerConfig struct {
 // NewStreamServer returns a new Shadowsocks 2022 stream server.
 func (c *StreamServerConfig) NewStreamServer() *StreamServer {
 	return &StreamServer{
-		saltPool:                   *NewSaltPool[[32]byte](ReplayWindowDuration),
+		saltPool:                   *NewSaltPool(),
 		readOnceOrFull:             readOnceOrFullFunc(c.AllowSegmentedFixedLengthHeader),
 		userCipherConfig:           c.UserCipherConfig,
 		identityCipherConfig:       c.IdentityCipherConfig,
@@ -246,7 +246,7 @@ func (c *StreamServerConfig) NewStreamServer() *StreamServer {
 // StreamServer implements [netio.StreamServer].
 type StreamServer struct {
 	CredStore
-	saltPool                   SaltPool[[32]byte]
+	saltPool                   SaltPool
 	readOnceOrFull             func(io.Reader, []byte) (int, error)
 	userCipherConfig           UserCipherConfig
 	identityCipherConfig       ServerIdentityCipherConfig
@@ -329,7 +329,7 @@ func (s *StreamServer) HandleStream(rawRW netio.Conn, logger *zap.Logger) (req n
 	extendedSalt := lengthExtendSalt(salt)
 
 	// Check but not add request salt to pool.
-	if !s.saltPool.TryCheck(extendedSalt) {
+	if s.saltPool.TryContains(extendedSalt) {
 		err = ErrRepeatedSalt
 		return
 	}