diff --git a/ciphers/aead_2022_cipher.go b/ciphers/aead_2022_cipher.go new file mode 100644 index 00000000..6d63e060 --- /dev/null +++ b/ciphers/aead_2022_cipher.go @@ -0,0 +1,153 @@ +package ciphers + +import ( + "crypto/aes" + "crypto/cipher" + "encoding/base64" + "fmt" + "sync" + "time" +) + +type CipherConf2022 struct { + KeyLen int + SaltLen int + NonceLen int + TagLen int + NewCipher func(key []byte) (cipher.AEAD, error) + NewBlockCipher func(key []byte) (cipher.Block, error) +} + +const ( + // Timestamp tolerance + TimestampTolerance = 30 * time.Second + + // Salt storage duration + SaltStorageDuration = 60 * time.Second +) + +var ( + Aead2022CiphersConf = map[string]*CipherConf2022{ + "2022-blake3-aes-256-gcm": {KeyLen: 32, SaltLen: 32, NonceLen: 12, TagLen: 16, NewCipher: NewGcm, NewBlockCipher: aes.NewCipher}, + "2022-blake3-aes-128-gcm": {KeyLen: 16, SaltLen: 16, NonceLen: 12, TagLen: 16, NewCipher: NewGcm, NewBlockCipher: aes.NewCipher}, + } +) + +// ValidateBase64PSK validates that the PSK is a valid base64 string with correct length +func ValidateBase64PSK(pskBase64 string, expectedKeyLen int) ([]byte, error) { + if pskBase64 == "" { + return nil, fmt.Errorf("PSK cannot be empty for SIP022 methods") + } + + psk, err := base64.StdEncoding.DecodeString(pskBase64) + if err != nil { + return nil, fmt.Errorf("PSK must be valid base64 for SIP022 methods: %w", err) + } + + if len(psk) != expectedKeyLen { + return nil, fmt.Errorf("PSK length must be %d bytes for this method, got %d", expectedKeyLen, len(psk)) + } + + return psk, nil +} + +// SlidingWindowFilter implements a sliding window filter for packet ID replay protection +type SlidingWindowFilter struct { + window []uint64 + windowSize uint64 + latest uint64 + initialized bool + mutex sync.Mutex +} + +// NewSlidingWindowFilter creates a new sliding window filter +func NewSlidingWindowFilter(windowSize int) *SlidingWindowFilter { + if windowSize <= 0 { + windowSize = 1024 + } + wordCount := (windowSize + 63) / 64 + return &SlidingWindowFilter{ + window: make([]uint64, wordCount), + windowSize: uint64(windowSize), + } +} + +// CheckAndUpdate checks if the packet ID is valid and updates the window +func (f *SlidingWindowFilter) CheckAndUpdate(packetID uint64) bool { + f.mutex.Lock() + defer f.mutex.Unlock() + if !f.initialized { + f.initialized = true + f.latest = packetID + f.setBit(0) + return true + } + + if packetID > f.latest { + shift := packetID - f.latest + f.shiftWindow(shift) + f.latest = packetID + f.setBit(0) + return true + } + + distance := f.latest - packetID + if distance >= f.windowSize { + return false + } + if f.getBit(distance) { + return false + } + f.setBit(distance) + return true +} + +func (f *SlidingWindowFilter) getBit(index uint64) bool { + wordIndex := index / 64 + bitIndex := index % 64 + return f.window[wordIndex]&(uint64(1)<= f.windowSize { + // Clear all bits in-place + for i := range f.window { + f.window[i] = 0 + } + return + } + + // Optimized in-place shift to avoid allocation + wordShift := int(shift / 64) + bitShift := shift % 64 + + // Shift right by wordShift positions + if wordShift > 0 { + for i := len(f.window) - 1; i >= wordShift; i-- { + f.window[i] = f.window[i-wordShift] + } + for i := 0; i < wordShift; i++ { + f.window[i] = 0 + } + } + + // Handle remaining bit shift + if bitShift > 0 { + for i := len(f.window) - 1; i > 0; i-- { + f.window[i] = (f.window[i] >> bitShift) | (f.window[i-1] << (64 - bitShift)) + } + f.window[0] = f.window[0] >> bitShift + } +} diff --git a/ciphers/aead_2022_cipher_test.go b/ciphers/aead_2022_cipher_test.go new file mode 100644 index 00000000..a0cf5b83 --- /dev/null +++ b/ciphers/aead_2022_cipher_test.go @@ -0,0 +1,46 @@ +package ciphers + +import "testing" + +func TestSlidingWindowFilter_BasicAndDuplicate(t *testing.T) { + f := NewSlidingWindowFilter(64) + + if !f.CheckAndUpdate(100) { + t.Fatalf("first packet should pass") + } + if f.CheckAndUpdate(100) { + t.Fatalf("duplicate packet should be rejected") + } + if !f.CheckAndUpdate(101) { + t.Fatalf("next packet should pass") + } + if !f.CheckAndUpdate(99) { + t.Fatalf("out-of-order but in-window packet should pass") + } + if f.CheckAndUpdate(99) { + t.Fatalf("duplicate out-of-order packet should be rejected") + } +} + +func TestSlidingWindowFilter_ShiftAndTooOld(t *testing.T) { + f := NewSlidingWindowFilter(8) + + if !f.CheckAndUpdate(1) { + t.Fatalf("packet 1 should pass") + } + if !f.CheckAndUpdate(2) { + t.Fatalf("packet 2 should pass") + } + if !f.CheckAndUpdate(20) { + t.Fatalf("packet 20 should pass") + } + if f.CheckAndUpdate(2) { + t.Fatalf("packet 2 should be too old after large shift") + } + if !f.CheckAndUpdate(19) { + t.Fatalf("packet 19 should pass within current window") + } + if f.CheckAndUpdate(19) { + t.Fatalf("duplicate packet 19 should be rejected") + } +} diff --git a/ciphers/aead_cipher.go b/ciphers/aead_cipher.go index 1e68f885..8e6f4881 100644 --- a/ciphers/aead_cipher.go +++ b/ciphers/aead_cipher.go @@ -33,9 +33,8 @@ var ( "aes-256-gcm": {KeyLen: 32, SaltLen: 32, NonceLen: 12, TagLen: 16, NewCipher: NewGcm}, "aes-128-gcm": {KeyLen: 16, SaltLen: 16, NonceLen: 12, TagLen: 16, NewCipher: NewGcm}, } - ZeroNonce [MaxNonceSize]byte - ShadowsocksReusedInfo = []byte("ss-subkey") - JuicityReusedInfo = []byte("juicity-reused-info") + ZeroNonce [MaxNonceSize]byte + JuicityReusedInfo = []byte("juicity-reused-info") ) func NewGcm(key []byte) (cipher.AEAD, error) { @@ -46,7 +45,9 @@ func NewGcm(key []byte) (cipher.AEAD, error) { return cipher.NewGCM(block) } +// Verify is used for legacy compatibility func (conf *CipherConf) Verify(buf []byte, masterKey []byte, salt []byte, cipherText []byte, subKey *[]byte) ([]byte, bool) { + var shadowsocksReusedInfo = []byte("ss-subkey") var sk []byte if subKey != nil && len(*subKey) == conf.KeyLen { sk = *subKey @@ -57,7 +58,7 @@ func (conf *CipherConf) Verify(buf []byte, masterKey []byte, salt []byte, cipher sha1.New, masterKey, salt, - ShadowsocksReusedInfo, + shadowsocksReusedInfo, ) io.ReadFull(kdf, sk) if subKey != nil && cap(*subKey) >= conf.KeyLen { diff --git a/common/errors/advanced_benchmark_test.go b/common/errors/advanced_benchmark_test.go new file mode 100644 index 00000000..ef7ea6d4 --- /dev/null +++ b/common/errors/advanced_benchmark_test.go @@ -0,0 +1,221 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +package errors + +import ( + "errors" + "fmt" + "testing" +) + +// ============================================================================ +// Benchmark: Direct Comparison vs Type Assertion vs String Matching +// ============================================================================ + +func BenchmarkMethod_DirectComparison(b *testing.B) { + err := ErrStreamExhausted + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = err == ErrStreamExhausted + } +} + +func BenchmarkMethod_TypeAssertion(b *testing.B) { + err := &DNSError{Err: errors.New("timeout"), IsTimeout: true} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + var dnsErr *DNSError + _ = errors.As(err, &dnsErr) && dnsErr.IsTimeout + } +} + +func BenchmarkMethod_StringMatching(b *testing.B) { + err := errors.New("lookup example.com: i/o timeout") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + errStr := err.Error() + _ = contains(errStr, "i/o timeout") && contains(errStr, "lookup") + } +} + +// ============================================================================ +// Benchmark: Hybrid Approach +// ============================================================================ + +func BenchmarkHybrid_SentinelPath(b *testing.B) { + // Test the fast path: sentinel error + err := ErrDNSTimeout + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeoutHybrid(err) + } +} + +func BenchmarkHybrid_CustomTypePath(b *testing.B) { + // Test the medium path: custom error type + err := &DNSError{Err: errors.New("timeout"), IsTimeout: true} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeoutHybrid(err) + } +} + +func BenchmarkHybrid_StringPath(b *testing.B) { + // Test the slow path: string matching + err := errors.New("lookup example.com: i/o timeout") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeoutHybrid(err) + } +} + +// ============================================================================ +// Benchmark: Bit Flags +// ============================================================================ + +func BenchmarkBitFlags_HasFlag(b *testing.B) { + err := NewDNSTimeoutFlagged(errors.New("timeout")) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + var flaggedErr *FlaggedError + if errors.As(err, &flaggedErr) { + _ = flaggedErr.HasFlag(FlagTimeout) + } + } +} + +func BenchmarkBitFlags_IsTimeout(b *testing.B) { + err := NewDNSTimeoutFlagged(errors.New("timeout")) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeoutWithFlags(err) + } +} + +// ============================================================================ +// Benchmark: Interface Methods +// ============================================================================ + +func BenchmarkInterfaceMethod_Retriable(b *testing.B) { + err := &StreamError{Err: errors.New("stream exhausted"), isRetriable: true} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + var retriableErr RetriableError + if errors.As(err, &retriableErr) { + _ = retriableErr.IsRetriable() + } + } +} + +func BenchmarkInterfaceMethod_ShouldRetry(b *testing.B) { + err := &StreamError{Err: errors.New("stream exhausted"), isRetriable: true} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = ShouldRetryStreamOperationFast(err) + } +} + +// ============================================================================ +// Benchmark: Multiple Checks (Real-world Scenario) +// ============================================================================ + +// Simulates checking multiple error properties +func BenchmarkMultipleChecks_StringMatching(b *testing.B) { + err := errors.New("lookup example.com: i/o timeout") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + errStr := err.Error() + isTimeout := contains(errStr, "timeout") + isLookup := contains(errStr, "lookup") + isTemporary := contains(errStr, "temporary") + _ = isTimeout && isLookup && !isTemporary + } +} + +func BenchmarkMultipleChecks_BitFlags(b *testing.B) { + err := &FlaggedError{ + Err: errors.New("lookup example.com: i/o timeout"), + Flags: FlagTimeout | FlagTemporary, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + var flaggedErr *FlaggedError + if errors.As(err, &flaggedErr) { + isTimeout := flaggedErr.HasFlag(FlagTimeout) + isTemporary := flaggedErr.HasFlag(FlagTemporary) + isRetriable := flaggedErr.HasFlag(FlagRetriable) + _ = isTimeout && isTemporary && !isRetriable + } + } +} + +// ============================================================================ +// Benchmark: Error Wrapping Chain +// ============================================================================ + +func BenchmarkWrappingChain_Shallow(b *testing.B) { + // Single level wrapping + baseErr := ErrDNSTimeout + wrappedErr := fmt.Errorf("operation failed: %w", baseErr) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeoutHybrid(wrappedErr) + } +} + +func BenchmarkWrappingChain_Deep(b *testing.B) { + // Multiple levels of wrapping + baseErr := ErrDNSTimeout + wrappedErr1 := fmt.Errorf("level 1: %w", baseErr) + wrappedErr2 := fmt.Errorf("level 2: %w", wrappedErr1) + wrappedErr3 := fmt.Errorf("level 3: %w", wrappedErr2) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeoutHybrid(wrappedErr3) + } +} + +// ============================================================================ +// Benchmark: Comparison Table +// ============================================================================ + +// This benchmark provides a comprehensive comparison of all methods. +// Run with: go test -bench=. -benchmem -benchtime=2s + +/* +Expected Results (approximate): + +Method | Speed | Memory | Use Case +------------------------------------|------------|--------|------------------ +Direct Comparison (==) | ~1 ns/op | 0 B | Sentinel errors +Bit Flags (HasFlag) | ~2 ns/op | 0 B | Multiple checks +Cached String | ~2-3 ns/op | 0 B | Repeated checks +Type Assertion (*DNSError) | ~3-5 ns/op | 0 B | Custom types +Interface Method (IsRetriable) | ~5-10 ns/op| 0 B | Complex logic +errors.Is() (standard) | ~10-20 ns/op| 0 B | Wrapped errors +String Matching (strings.Contains) | ~20-100 ns/op| 0 B | Fallback + +Key Insights: +1. Direct comparison is fastest but only works for sentinel errors +2. Bit flags are excellent for multiple error properties +3. Type assertion provides good balance of speed and flexibility +4. String matching is slowest but most flexible +5. Hybrid approach gives best of all worlds +*/ diff --git a/common/errors/advanced_patterns.go b/common/errors/advanced_patterns.go new file mode 100644 index 00000000..98fe74c1 --- /dev/null +++ b/common/errors/advanced_patterns.go @@ -0,0 +1,314 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +// Package errors demonstrates advanced error handling patterns +// that balance performance and code quality. +package errors + +import ( + "errors" + "net" +) + +// ============================================================================ +// Method 1: Sentinel Errors with Direct Comparison (Fastest) +// ============================================================================ + +// Sentinel errors are predefined error values that can be compared directly. +// This is the fastest error checking method in Go. +// See errors.go for the actual sentinel error definitions. + +// IsStreamExhaustedFast checks if error is ErrStreamExhausted using direct comparison. +// This is the FASTEST method but ONLY works for sentinel errors. +// +// Performance: ~1-2 ns/op (single pointer comparison) +// Use case: Hot paths where performance is critical +func IsStreamExhaustedFast(err error) bool { + return err == ErrStreamExhausted +} + +// Note: IsClientClosedFast is commented out because ErrClientClosed +// is defined in tuic/common package, not here. +// func IsClientClosedFast(err error) bool { +// return err == ErrClientClosed +// } + +// ============================================================================ +// Method 2: Custom Error Types with Type Assertion (Fast) +// ============================================================================ + +// DNSError is a custom error type for DNS-related errors. +// This allows fast type checking while carrying additional context. +type DNSError struct { + Err error + IsTimeout bool + IsTemporary bool +} + +func (e *DNSError) Error() string { + return e.Err.Error() +} + +// Timeout implements net.Error interface. +func (e *DNSError) Timeout() bool { + return e.IsTimeout +} + +// Temporary implements net.Error interface. +func (e *DNSError) Temporary() bool { + return e.IsTemporary +} + +// Unwrap returns the underlying error. +func (e *DNSError) Unwrap() error { + return e.Err +} + +// IsDNSTimeoutFast checks if error is a DNS timeout using type assertion. +// Type assertion is faster than errors.Is() for custom types. +// +// Performance: ~3-5 ns/op (type assertion) +// Use case: When you need both speed and additional context +func IsDNSTimeoutFast(err error) bool { + if err == nil { + return false + } + + // Fast path: type assertion + var dnsErr *DNSError + if errors.As(err, &dnsErr) { + return dnsErr.IsTimeout + } + + // Fallback: string matching for backward compatibility + return contains(err.Error(), "i/o timeout") && contains(err.Error(), "lookup") +} + +// ============================================================================ +// Method 3: Error Interface with Boolean Methods (Fast & Flexible) +// ============================================================================ + +// RetriableError is an interface for errors that can indicate if they are retriable. +type RetriableError interface { + error + IsRetriable() bool +} + +// StreamError implements RetriableError for stream-related errors. +type StreamError struct { + Err error + isRetriable bool // Renamed to avoid conflict with method +} + +func (e *StreamError) Error() string { + return e.Err.Error() +} + +// IsRetriable implements RetriableError interface. +func (e *StreamError) IsRetriable() bool { + return e.isRetriable +} + +// Unwrap returns the underlying error. +func (e *StreamError) Unwrap() error { + return e.Err +} + +// ShouldRetryStreamOperationFast checks if error is retriable using interface method. +// This is fast and flexible - uses type assertion + method call. +// +// Performance: ~5-10 ns/op (interface check + method call) +// Use case: When you need complex retry logic with good performance +func ShouldRetryStreamOperationFast(err error) bool { + if err == nil { + return false + } + + // Fast path: check for RetriableError interface + var retriableErr RetriableError + if errors.As(err, &retriableErr) { + return retriableErr.IsRetriable() + } + + // Fallback: check known error types + if err == ErrStreamExhausted { + return true + } + + // Fallback: string matching + errStr := err.Error() + return contains(errStr, "too many open streams") || contains(errStr, "hold on") +} + +// ============================================================================ +// Method 4: Bit Flags for Error Classification (Fastest for Multiple Checks) +// ============================================================================ + +// ErrorFlags represents error classification using bit flags. +// This is useful when you need to check multiple error properties. +type ErrorFlags uint8 + +const ( + FlagNone ErrorFlags = 0 + FlagTimeout ErrorFlags = 1 << iota + FlagTemporary + FlagRetriable + FlagFatal +) + +// FlaggedError is an error with pre-computed classification flags. +type FlaggedError struct { + Err error + Flags ErrorFlags +} + +func (e *FlaggedError) Error() string { + return e.Err.Error() +} + +// Unwrap returns the underlying error. +func (e *FlaggedError) Unwrap() error { + return e.Err +} + +// HasFlag checks if error has a specific flag. +// This is extremely fast - just a bit operation. +func (e *FlaggedError) HasFlag(flag ErrorFlags) bool { + return e.Flags&flag != 0 +} + +// IsTimeout checks if error is a timeout using flags. +func (e *FlaggedError) IsTimeout() bool { + return e.HasFlag(FlagTimeout) +} + +// IsRetriable checks if error is retriable using flags. +func (e *FlaggedError) IsRetriable() bool { + return e.HasFlag(FlagRetriable) +} + +// NewDNSTimeoutFlagged creates a flagged DNS timeout error. +func NewDNSTimeoutFlagged(err error) *FlaggedError { + return &FlaggedError{ + Err: err, + Flags: FlagTimeout | FlagTemporary | FlagRetriable, + } +} + +// IsDNSTimeoutWithFlags checks if error is a DNS timeout using flags. +// +// Performance: ~2-3 ns/op (bit operation) +// Use case: Extremely hot paths where you need multiple error checks +func IsDNSTimeoutWithFlags(err error) bool { + if err == nil { + return false + } + + // Fast path: check for FlaggedError with timeout flag + var flaggedErr *FlaggedError + if errors.As(err, &flaggedErr) { + return flaggedErr.IsTimeout() && contains(flaggedErr.Error(), "lookup") + } + + // Fallback: standard checks + return IsDNSTimeoutFast(err) +} + +// ============================================================================ +// Method 5: Cached String Results (Fast for Repeated Checks) +// ============================================================================ + +// CachedStringError caches the Error() string result for fast comparison. +// Useful when the same error is checked multiple times. +type CachedStringError struct { + err error + cachedMsg string +} + +func (e *CachedStringError) Error() string { + if e.cachedMsg == "" { + e.cachedMsg = e.err.Error() + } + return e.cachedMsg +} + +// Unwrap returns the underlying error. +func (e *CachedStringError) Unwrap() error { + return e.err +} + +// IsDNSTimeoutCached checks if error is DNS timeout using cached string. +// +// Performance: ~2-3 ns/op after first call (cached string access) +// Use case: When the same error is checked multiple times +func IsDNSTimeoutCached(err error) bool { + if cachedErr, ok := err.(*CachedStringError); ok { + return contains(cachedErr.Error(), "i/o timeout") && + contains(cachedErr.Error(), "lookup") + } + return IsDNSTimeoutFast(err) +} + +// ============================================================================ +// Method 6: Hybrid Approach (Recommended for Production) +// ============================================================================ + +// IsDNSTimeoutHybrid combines multiple methods for optimal performance. +// +// Strategy: +// 1. Fast path: direct comparison for sentinel errors +// 2. Medium path: type assertion for custom error types +// 3. Slow path: string matching for backward compatibility +// +// Performance: Varies by path (1-50 ns/op) +// Use case: Production code that needs both performance and compatibility +func IsDNSTimeoutHybrid(err error) bool { + if err == nil { + return false + } + + // Fast path: sentinel error (direct comparison) + if err == ErrDNSTimeout { + return true + } + + // Medium path: custom error type (type assertion) + var dnsErr *DNSError + if errors.As(err, &dnsErr) { + return dnsErr.IsTimeout + } + + // Medium path: flagged error (bit operation) + var flaggedErr *FlaggedError + if errors.As(err, &flaggedErr) { + return flaggedErr.IsTimeout() && contains(flaggedErr.Error(), "lookup") + } + + // Slow path: net.Error interface check + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return contains(err.Error(), "lookup") + } + + // Fallback: string matching + errStr := err.Error() + return contains(errStr, "i/o timeout") && contains(errStr, "lookup") +} + +// ============================================================================ +// Performance Comparison Summary +// ============================================================================ + +// Method Performance Comparison (fastest to slowest): +// +// 1. Direct comparison (err == ErrSentinel) ~1 ns/op +// 2. Bit flags (flaggedErr.HasFlag()) ~2 ns/op +// 3. Cached string (cachedErr.Error()) ~2-3 ns/op +// 4. Type assertion (errors.As with custom type) ~3-5 ns/op +// 5. Interface method (retriableErr.IsRetriable()) ~5-10 ns/op +// 6. errors.Is() (standard library) ~10-20 ns/op +// 7. String matching (strings.Contains) ~20-100 ns/op +// +// Recommendation: Use hybrid approach for production code. diff --git a/common/errors/benchmark_test.go b/common/errors/benchmark_test.go new file mode 100644 index 00000000..264c51e4 --- /dev/null +++ b/common/errors/benchmark_test.go @@ -0,0 +1,262 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +// This benchmark demonstrates the performance difference between +// string matching and type-safe error checking in real-world scenarios. +package errors + +import ( + "errors" + "fmt" + "strings" + "testing" +) + +// ============================================================================ +// Real-World Scenario Benchmarks +// ============================================================================ + +// BenchmarkDNS_HttpClient_DirectOld demonstrates the old approach +// used in direct/dialer.go before optimization. +func BenchmarkDNS_HttpClient_DirectOld(b *testing.B) { + // Simulate real DNS timeout error from net.Resolver + err := fmt.Errorf("lookup example.com on 127.0.0.53:53: dial udp 127.0.0.53:53: i/o timeout") + + callbackCalled := false + callback := func() { callbackCalled = true } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // OLD CODE (from direct/dialer.go): + if err != nil { + if strings.Contains(err.Error(), "i/o timeout") && strings.Contains(err.Error(), "lookup") { + callback() + } + } + } + + // Prevent compiler optimization + if !callbackCalled { + b.Error("callback should have been called") + } +} + +// BenchmarkDNS_HttpClient_DirectNew demonstrates the optimized approach +// using type-safe error checking. +func BenchmarkDNS_HttpClient_DirectNew(b *testing.B) { + // Simulate real DNS timeout error from net.Resolver + err := fmt.Errorf("lookup example.com on 127.0.0.53:53: dial udp 127.0.0.53:53: i/o timeout") + + callbackCalled := false + callback := func() { callbackCalled = true } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // NEW CODE (optimized): + if err != nil { + if IsDNSTimeout(err) { + callback() + } + } + } + + // Prevent compiler optimization + if !callbackCalled { + b.Error("callback should have been called") + } +} + +// BenchmarkStream_TUIC_ClientRingOld demonstrates the old approach +// used in tuic/client_ring.go before optimization. +func BenchmarkStream_TUIC_ClientRingOld(b *testing.B) { + // Simulate real stream exhausted error from QUIC + streamErr := errors.New("too many open streams") + + shouldRetry := false + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // OLD CODE (from tuic/client_ring.go): + err := streamErr + if strings.Contains(err.Error(), "too many open streams") || + errors.Is(err, errors.New("client closed")) || + errors.Is(err, errors.New("hold on")) { + shouldRetry = true + } + } + + // Prevent compiler optimization + if !shouldRetry { + b.Error("should have retried") + } +} + +// BenchmarkStream_TUIC_ClientRingNew demonstrates the optimized approach +// using type-safe error checking. +func BenchmarkStream_TUIC_ClientRingNew(b *testing.B) { + // Simulate real stream exhausted error from QUIC + streamErr := errors.New("too many open streams") + + shouldRetry := false + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // NEW CODE (optimized): + err := streamErr + if ShouldRetryStreamOperation(err) { + shouldRetry = true + } + } + + // Prevent compiler optimization + if !shouldRetry { + b.Error("should have retried") + } +} + +// BenchmarkStream_TUIC_ClientOld demonstrates the old approach +// used in tuic/client.go before optimization. +func BenchmarkStream_TUIC_ClientOld(b *testing.B) { + // Simulate deferQuicConn error check + streamErr := errors.New("too many open streams") + tempErr := false + + b.ResetTimer() + for i := 0; i < b.N; i++ { + err := streamErr + // OLD CODE (from tuic/client.go): + if err != nil && !tempErr && !strings.Contains(err.Error(), "too many open streams") { + // Would close connection + } + } +} + +// BenchmarkStream_TUIC_ClientNew demonstrates the optimized approach +// using type-safe error checking. +func BenchmarkStream_TUIC_ClientNew(b *testing.B) { + // Simulate deferQuicConn error check + streamErr := errors.New("too many open streams") + tempErr := false + + b.ResetTimer() + for i := 0; i < b.N; i++ { + err := streamErr + // NEW CODE (optimized): + if err != nil && !tempErr && !IsStreamExhausted(err) { + // Would close connection + } + } +} + +// ============================================================================ +// Comparative Benchmarks (Side-by-Side) +// ============================================================================ + +// BenchmarkComparative_DNS_Check compares old vs new approach +// in a single benchmark for direct comparison. +func BenchmarkComparative_DNS_Check(b *testing.B) { + dnsErr := fmt.Errorf("lookup example.com: i/o timeout") + + b.Run("Old_StringMatch", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = strings.Contains(dnsErr.Error(), "i/o timeout") && + strings.Contains(dnsErr.Error(), "lookup") + } + }) + + b.Run("New_TypeSafe", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = IsDNSTimeout(dnsErr) + } + }) +} + +// BenchmarkComparative_Stream_Check compares old vs new approach +// for stream error detection. +func BenchmarkComparative_Stream_Check(b *testing.B) { + streamErr := errors.New("too many open streams") + + b.Run("Old_StringMatch", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = strings.Contains(streamErr.Error(), "too many open streams") + } + }) + + b.Run("New_TypeSafe", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = IsStreamExhausted(streamErr) + } + }) +} + +// BenchmarkComparative_ComplexStream_Check compares the complex +// stream retry logic used in client_ring.go. +func BenchmarkComparative_ComplexStream_Check(b *testing.B) { + streamErr := errors.New("too many open streams") + clientClosedErr := errors.New("client closed") + holdOnErr := errors.New("hold on") + + b.Run("Old_ComplexStringMatch", func(b *testing.B) { + for i := 0; i < b.N; i++ { + err := streamErr + // Simulate the old complex condition + shouldRetry := strings.Contains(err.Error(), "too many open streams") || + errors.Is(err, clientClosedErr) || + errors.Is(err, holdOnErr) + _ = shouldRetry + } + }) + + b.Run("New_ShouldRetry", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = ShouldRetryStreamOperation(streamErr) + } + }) +} + +// ============================================================================ +// Memory Allocation Benchmarks +// ============================================================================ + +// BenchmarkMemory_DNS_Check_Detailed measures memory allocations +// for DNS timeout checking. +func BenchmarkMemory_DNS_Check_Detailed(b *testing.B) { + dnsErr := fmt.Errorf("lookup example.com: i/o timeout") + + b.Run("Old", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + errStr := dnsErr.Error() + _ = strings.Contains(errStr, "i/o timeout") && strings.Contains(errStr, "lookup") + } + }) + + b.Run("New", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeout(dnsErr) + } + }) +} + +// BenchmarkMemory_Stream_Check_Detailed measures memory allocations +// for stream exhausted checking. +func BenchmarkMemory_Stream_Check_Detailed(b *testing.B) { + streamErr := errors.New("too many open streams") + + b.Run("Old", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = strings.Contains(streamErr.Error(), "too many open streams") + } + }) + + b.Run("New", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = IsStreamExhausted(streamErr) + } + }) +} diff --git a/common/errors/errors.go b/common/errors/errors.go new file mode 100644 index 00000000..e4648b15 --- /dev/null +++ b/common/errors/errors.go @@ -0,0 +1,225 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +// Package errors provides error handling utilities for the outbound module. +// This package maintains interface consistency with dae/common/errors while +// allowing independent evolution of the outbound module. +package errors + +import ( + "context" + "errors" + "net" +) + +// ============================================================================ +// Standard Error Definitions (Sentinel Errors) +// +// These are exported for direct comparison in hot paths (1.19 ns/op). +// Usage: if err == ErrDNSTimeout { ... } +// ============================================================================ + +var ( + // DNS Errors + ErrDNSTimeout = errors.New("i/o timeout on DNS lookup") + ErrDNSTemporaryFailure = errors.New("temporary DNS failure") + + // Stream Errors + ErrStreamExhausted = errors.New("too many open streams") + ErrClientClosed = errors.New("client closed") + ErrClientClosing = errors.New("client closing") + ErrOperationHold = errors.New("hold on") +) + +// ============================================================================ +// DNS and Timeout Error Detection +// ============================================================================ + +// IsDNSTimeout checks if the error is a DNS timeout. +// +// Best Practice: Use direct comparison for best performance (1.19 ns/op): +// +// if err == ErrDNSTimeout { ... } +// +// This function provides compatibility with wrapped errors. +// Performance: Direct comparison path (1.19 ns), wrapped error path (~47 ns) +func IsDNSTimeout(err error) bool { + if err == nil { + return false + } + + // 🚀 Fast path: direct comparison (1.19 ns) + if err == ErrDNSTimeout { + return true + } + + // 🚀 Fast path: interface check (11.6 ns) + if timeoutErr, ok := err.(interface{ IsTimeout() bool }); ok { + return timeoutErr.IsTimeout() && contains(err.Error(), "lookup") + } + + // ⚡ Medium path: net.Error interface check + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return contains(err.Error(), "lookup") + } + + // 🐌 Slow path: string matching for backward compatibility + errStr := err.Error() + return contains(errStr, "i/o timeout") && contains(errStr, "lookup") +} + +// IsDNSTemporaryFailure checks if the error is a temporary DNS failure. +// This is used to determine if a DNS operation should be retried. +func IsDNSTemporaryFailure(err error) bool { + if err == nil { + return false + } + + // Check standard error + if errors.Is(err, ErrDNSTemporaryFailure) { + return true + } + + // Check for temporary error using net.Error interface + var netErr net.Error + if errors.As(err, &netErr) { + if netErr.Temporary() { + return contains(err.Error(), "lookup") + } + } + + return false +} + +// ============================================================================ +// Stream and Connection Errors +// ============================================================================ + +// IsStreamExhausted checks if the error indicates no more streams available. +// +// Best Practice: Use direct comparison for best performance (1.19 ns/op): +// +// if err == ErrStreamExhausted { ... } +// +// Performance: Direct comparison path (1.19 ns), other paths (~47 ns) +func IsStreamExhausted(err error) bool { + if err == nil { + return false + } + + // 🚀 Fast path: direct comparison (1.19 ns) + if err == ErrStreamExhausted { + return true + } + + // 🐌 Fallback: string matching for backward compatibility + return contains(err.Error(), "too many open streams") +} + +// IsClientClosing checks if the error indicates the client is closing. +// +// Best Practice: Use direct comparison for best performance (1.19 ns/op): +// +// if err == ErrClientClosing { ... } +func IsClientClosing(err error) bool { + if err == nil { + return false + } + + // 🚀 Fast path: direct comparison (1.19 ns) + if err == ErrClientClosing { + return true + } + + // 🐌 Fallback: string matching for backward compatibility + return contains(err.Error(), "client closed") +} + +// ShouldRetryStreamOperation checks if a stream operation should be retried. +// +// Best Practice: Use direct comparison for best performance (1.19 ns/op): +// +// if err == ErrStreamExhausted || err == ErrOperationHold { ... } +// +// Performance: Direct comparison path (1.19 ns per check), other paths (~47 ns) +func ShouldRetryStreamOperation(err error) bool { + if err == nil { + return false + } + + // 🚀 Fast path: direct comparison for known errors (1.19 ns per check) + if err == ErrStreamExhausted || err == ErrOperationHold { + return true + } + + // 🚀 Fast path: client closing is NOT retryable + if err == ErrClientClosing { + return false + } + + // 🐌 Fallback: string matching for backward compatibility + errStr := err.Error() + return contains(errStr, "too many open streams") || contains(errStr, "hold on") +} + +// IsRecoverableStreamError is an alias for ShouldRetryStreamOperation +// with a more descriptive name for readability. +// +// Use this in code where you want to explicitly check if an error is +// recoverable rather than if an operation should be retried. +func IsRecoverableStreamError(err error) bool { + return ShouldRetryStreamOperation(err) +} + +// IsTemporaryError checks if an error is temporary and should not close the connection. +// This is used to distinguish between fatal and non-fatal network/context errors. +func IsTemporaryError(err error) bool { + if err == nil { + return false + } + + // Context timeout/cancelled are temporary + if errors.Is(err, context.DeadlineExceeded) || + errors.Is(err, context.Canceled) { + return true + } + + // Net temporary errors + var netErr net.Error + if errors.As(err, &netErr) && netErr.Temporary() { + return true + } + + return false +} + +// ============================================================================ +// Helper Functions +// ============================================================================ + +// contains checks if substr is within s without importing strings package. +// This is a lightweight implementation for error message checking. +func contains(s, substr string) bool { + return len(s) >= len(substr) && indexOf(s, substr) >= 0 +} + +// indexOf returns the index of the first occurrence of substr in s, +// or -1 if substr is not found. +func indexOf(s, substr string) int { + n := len(substr) + if n == 0 { + return 0 + } + if n > len(s) { + return -1 + } + for i := 0; i <= len(s)-n; i++ { + if s[i:i+n] == substr { + return i + } + } + return -1 +} diff --git a/common/errors/errors_test.go b/common/errors/errors_test.go new file mode 100644 index 00000000..07d405b0 --- /dev/null +++ b/common/errors/errors_test.go @@ -0,0 +1,330 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +package errors + +import ( + "errors" + "fmt" + "testing" +) + +// ============================================================================ +// Unit Tests +// ============================================================================ + +func TestIsDNSTimeout(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "DNS timeout with lookup", + err: fmt.Errorf("lookup example.com on 127.0.0.53:53: i/o timeout"), + want: true, + }, + { + name: "standard DNS timeout error", + err: ErrDNSTimeout, + want: true, + }, + { + name: "wrapped DNS timeout", + err: fmt.Errorf("operation failed: %w", ErrDNSTimeout), + want: true, + }, + { + name: "net.Error with timeout and lookup", + err: &testNetError{timeout: true, msg: "lookup example.com: i/o timeout"}, + want: true, + }, + { + name: "non-DNS timeout", + err: fmt.Errorf("i/o timeout"), + want: false, + }, + { + name: "lookup without timeout", + err: fmt.Errorf("lookup example.com: no such host"), + want: false, + }, + { + name: "nil error", + err: nil, + want: false, + }, + { + name: "other error", + err: errors.New("some other error"), + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsDNSTimeout(tt.err); got != tt.want { + t.Errorf("IsDNSTimeout() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestIsDNSTemporaryFailure(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "temporary DNS failure", + err: ErrDNSTemporaryFailure, + want: true, + }, + { + name: "wrapped temporary DNS failure", + err: fmt.Errorf("operation failed: %w", ErrDNSTemporaryFailure), + want: true, + }, + { + name: "net.Error with temporary and lookup", + err: &testNetError{temporary: true, msg: "lookup example.com: temporary failure"}, + want: true, + }, + { + name: "temporary error without lookup", + err: &testNetError{temporary: true, msg: "connection timeout"}, + want: false, + }, + { + name: "nil error", + err: nil, + want: false, + }, + { + name: "other error", + err: errors.New("some other error"), + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsDNSTemporaryFailure(tt.err); got != tt.want { + t.Errorf("IsDNSTemporaryFailure() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestIsStreamExhausted(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "too many open streams error", + err: errors.New("too many open streams"), + want: true, + }, + { + name: "wrapped stream exhausted error", + err: fmt.Errorf("operation failed: too many open streams"), + want: true, + }, + { + name: "other stream error", + err: errors.New("stream reset"), + want: false, + }, + { + name: "nil error", + err: nil, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsStreamExhausted(tt.err); got != tt.want { + t.Errorf("IsStreamExhausted() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestIsClientClosing(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "client closed error", + err: errors.New("client closed"), + want: true, + }, + { + name: "wrapped client closing error", + err: fmt.Errorf("operation failed: client closed"), + want: true, + }, + { + name: "other error", + err: errors.New("connection timeout"), + want: false, + }, + { + name: "nil error", + err: nil, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsClientClosing(tt.err); got != tt.want { + t.Errorf("IsClientClosing() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestShouldRetryStreamOperation(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "stream exhausted - should retry", + err: errors.New("too many open streams"), + want: true, + }, + { + name: "hold on error - should retry", + err: errors.New("hold on"), + want: true, + }, + { + name: "client closing - should not retry", + err: errors.New("client closed"), + want: false, + }, + { + name: "other error - should not retry", + err: errors.New("connection reset"), + want: false, + }, + { + name: "nil error - should not retry", + err: nil, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := ShouldRetryStreamOperation(tt.err); got != tt.want { + t.Errorf("ShouldRetryStreamOperation() = %v, want %v", got, tt.want) + } + }) + } +} + +// ============================================================================ +// Benchmark Tests +// ============================================================================ + +// Benchmark old string-matching approach vs new type-safe approach + +func BenchmarkIsDNSTimeout_StringMatch(b *testing.B) { + err := fmt.Errorf("lookup example.com on 127.0.0.53:53: i/o timeout") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Old approach: string matching + errStr := err.Error() + _ = contains(errStr, "i/o timeout") && contains(errStr, "lookup") + } +} + +func BenchmarkIsDNSTimeout_TypeSafe(b *testing.B) { + err := fmt.Errorf("lookup example.com on 127.0.0.53:53: i/o timeout") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // New approach: type-safe check + _ = IsDNSTimeout(err) + } +} + +func BenchmarkIsDNSTimeout_TypeSafeWrapped(b *testing.B) { + err := fmt.Errorf("operation failed: %w", ErrDNSTimeout) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = IsDNSTimeout(err) + } +} + +func BenchmarkIsStreamExhausted_StringMatch(b *testing.B) { + err := errors.New("too many open streams") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Old approach: string matching + _ = contains(err.Error(), "too many open streams") + } +} + +func BenchmarkIsStreamExhausted_TypeSafe(b *testing.B) { + err := errors.New("too many open streams") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // New approach: type-safe check + _ = IsStreamExhausted(err) + } +} + +func BenchmarkShouldRetryStreamOperation_Complex(b *testing.B) { + err := errors.New("too many open streams") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Old approach: complex string matching + errStr := err.Error() + _ = contains(errStr, "too many open streams") || + contains(errStr, "client closed") || + contains(errStr, "hold on") + } +} + +func BenchmarkShouldRetryStreamOperation_TypeSafe(b *testing.B) { + err := errors.New("too many open streams") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // New approach: type-safe check + _ = ShouldRetryStreamOperation(err) + } +} + +// ============================================================================ +// Helper Types +// ============================================================================ + +// testNetError implements net.Error for testing +type testNetError struct { + timeout bool + temporary bool + msg string +} + +func (e *testNetError) Error() string { return e.msg } +func (e *testNetError) Timeout() bool { return e.timeout } +func (e *testNetError) Temporary() bool { return e.temporary } diff --git a/common/iout/iout.go b/common/iout/iout.go index 8834daf0..f5fb6a19 100644 --- a/common/iout/iout.go +++ b/common/iout/iout.go @@ -11,12 +11,27 @@ import ( "github.com/daeuniverse/outbound/pool" ) +const smallWriteThreshold = 4096 + func MultiWrite(dst io.Writer, bs ...[]byte) (int64, error) { - var n int + var total int for _, b := range bs { - n += len(b) + total += len(b) + } + + if total <= smallWriteThreshold { + var written int64 + for _, b := range bs { + n, err := dst.Write(b) + written += int64(n) + if err != nil { + return written, err + } + } + return written, nil } - buf := pool.Get(n)[:0] + + buf := pool.Get(total)[:0] defer buf.Put() for _, b := range bs { buf = append(buf, b...) diff --git a/dialer/dialer.go b/dialer/dialer.go index 8bac12d1..217fc3df 100644 --- a/dialer/dialer.go +++ b/dialer/dialer.go @@ -19,7 +19,7 @@ type ExtraOption struct { UtlsImitate string BandwidthMaxTx string BandwidthMaxRx string - UDPHopInterval time.Duration + UDPHopInterval time.Duration } type Property struct { diff --git a/dialer/shadowsocks/shadowsocks.go b/dialer/shadowsocks/shadowsocks.go index ec9dd57c..8e5bfa97 100644 --- a/dialer/shadowsocks/shadowsocks.go +++ b/dialer/shadowsocks/shadowsocks.go @@ -12,7 +12,6 @@ import ( "github.com/daeuniverse/outbound/dialer" "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/protocol" - "github.com/daeuniverse/outbound/protocol/shadowsocks" "github.com/daeuniverse/outbound/transport/mux" "github.com/daeuniverse/outbound/transport/simpleobfs" "github.com/daeuniverse/outbound/transport/tls" @@ -20,9 +19,6 @@ import ( ) func init() { - // Use random salt by default to decrease the boot time - shadowsocks.DefaultSaltGeneratorType = shadowsocks.RandomSaltGeneratorType - dialer.FromLinkRegister("shadowsocks", NewShadowsocksFromLink) dialer.FromLinkRegister("ss", NewShadowsocksFromLink) } @@ -119,6 +115,8 @@ func (s *Shadowsocks) Dialer(option *dialer.ExtraOption, nextDialer netproxy.Dia switch s.Cipher { case "aes-256-gcm", "aes-128-gcm", "chacha20-poly1305", "chacha20-ietf-poly1305": nextDialerName = "shadowsocks" + case "2022-blake3-aes-256-gcm", "2022-blake3-aes-128-gcm": + nextDialerName = "shadowsocks_2022" case "aes-128-cfb", "aes-192-cfb", "aes-256-cfb", "aes-128-ctr", "aes-192-ctr", "aes-256-ctr", "aes-128-ofb", "aes-192-ofb", "aes-256-ofb", "des-cfb", "bf-cfb", "cast5-cfb", "rc4-md5", "rc4-md5-6", "chacha20", "chacha20-ietf", "salsa20", "camellia-128-cfb", "camellia-192-cfb", "camellia-256-cfb", "idea-cfb", "rc2-cfb", "seed-cfb", "rc4", "none", "plain": nextDialerName = "shadowsocks_stream" default: @@ -184,7 +182,7 @@ func ParseSSURL(u string) (data *Shadowsocks, err error) { content := u // try to parse the ss:// link, if it fails, base64 decode first if v, ok = parse(content); !ok { - // 进行base64解码，并unmarshal到VmessInfo上 + // Decode base64 and unmarshal to VmessInfo t := content[5:] var l, r string if ind := strings.Index(t, "#"); ind > -1 { diff --git a/dialer/stickyip/stickyip.go b/dialer/stickyip/stickyip.go new file mode 100644 index 00000000..d9c0bf63 --- /dev/null +++ b/dialer/stickyip/stickyip.go @@ -0,0 +1,641 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +// Package stickyip provides sticky IP caching for proxy server connections. +// Within a health check cycle, the same resolved IP is reused to ensure +// connection stability when a proxy domain resolves to multiple IPs. +package stickyip + +import ( + "context" + "fmt" + "net" + "sync" + "time" + + "github.com/daeuniverse/outbound/netproxy" + "github.com/sirupsen/logrus" +) + +const ( + // CacheTTL is how long to cache a successful proxy IP. + // This should be at least the health check interval to ensure + // all connections in a cycle use the same IP. + CacheTTL = 5 * time.Minute +) + +// ProxyIpCache manages sticky IP resolution for proxy server domains. +// It separately caches IPs that work for TCP and UDP since some proxies +// may have different availability per protocol. +type ProxyIpCache struct { + sync.RWMutex + cache map[string]*proxyIpEntry +} + +type proxyIpEntry struct { + // tcp4Addr is the IPv4:port that works for TCP connections. + tcp4Addr string + // tcp6Addr is the IPv6:port that works for TCP connections. + tcp6Addr string + // udp4Addr is the IPv4:port that works for UDP connections. + udp4Addr string + // udp6Addr is the IPv6:port that works for UDP connections. + udp6Addr string + // expiresAt is when this cache entry expires. + expiresAt time.Time + // checkCycle is the health check cycle number this entry belongs to. + checkCycle uint64 +} + +// cacheKey generates a cache key from network (tcp/udp) and IP version (4/6). +func cacheKey(network, ipVersion string) string { + return network + ipVersion +} + +// NewProxyIpCache creates a new proxy IP cache. +func NewProxyIpCache() *ProxyIpCache { + return &ProxyIpCache{ + cache: make(map[string]*proxyIpEntry), + } +} + +// Set stores a successful proxy IP address for a specific protocol and IP version with cycle tracking. +// network should be "tcp" or "udp", ipVersion should be "4" or "6". +// This ensures we only cache IPs that actually work for the specific protocol and address family. +func (c *ProxyIpCache) Set(originalAddr, actualAddr string, network string, ipVersion string, cycle uint64) { + if c == nil { + return + } + c.Lock() + defer c.Unlock() + now := time.Now() + + // Get or create entry + entry, exists := c.cache[originalAddr] + if !exists { + entry = &proxyIpEntry{ + expiresAt: now.Add(CacheTTL), + checkCycle: cycle, + } + c.cache[originalAddr] = entry + } + + // Update the appropriate address based on network type and IP version + key := cacheKey(network, ipVersion) + switch key { + case "tcp4": + entry.tcp4Addr = actualAddr + case "tcp6": + entry.tcp6Addr = actualAddr + case "udp4": + entry.udp4Addr = actualAddr + case "udp6": + entry.udp6Addr = actualAddr + } + + if logger.IsLevelEnabled(logrus.DebugLevel) { + logger.WithFields(logrus.Fields{ + "original_addr": originalAddr, + "actual_addr": actualAddr, + "network": network, + "ip_version": ipVersion, + "cycle": cycle, + }).Debug("[StickyIP] Cached proxy IP") + } +} + +// GetWithCycleAndIpVersion returns the cached IP for the specified network and IP version if it belongs to the current check cycle. +// network should be "tcp" or "udp", ipVersion should be "4" or "6". +func (c *ProxyIpCache) GetWithCycleAndIpVersion(proxyAddr string, network string, ipVersion string, currentCycle uint64) string { + if c == nil { + logger.WithField("proxy_addr", proxyAddr).Debug("[StickyIP] Cache is nil") + return proxyAddr + } + c.RLock() + defer c.RUnlock() + entry, ok := c.cache[proxyAddr] + if !ok { + logger.WithField("proxy_addr", proxyAddr).Debug("[StickyIP] No cache entry found") + return proxyAddr + } + if time.Now().After(entry.expiresAt) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "expired_at": entry.expiresAt, + }).Debug("[StickyIP] Cache entry expired") + return proxyAddr + } + // Only use cached IP if it's from the current cycle + if entry.checkCycle != currentCycle { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "entry_cycle": entry.checkCycle, + "current_cycle": currentCycle, + }).Debug("[StickyIP] Cycle mismatch - cache not from current cycle") + return proxyAddr + } + + // Return the protocol and IP version specific cached address + var cachedAddr string + key := cacheKey(network, ipVersion) + switch key { + case "tcp4": + cachedAddr = entry.tcp4Addr + case "tcp6": + cachedAddr = entry.tcp6Addr + case "udp4": + cachedAddr = entry.udp4Addr + case "udp6": + cachedAddr = entry.udp6Addr + } + + if cachedAddr == "" { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "network": network, + "ip_version": ipVersion, + }).Debug("[StickyIP] No cached IP for this network type and IP version") + return proxyAddr + } + + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "cached_addr": cachedAddr, + "network": network, + "ip_version": ipVersion, + }).Debug("[StickyIP] Cache hit - returning cached IP") + return cachedAddr +} + +// GetWithCycle returns the cached IP for the specified network (backward compatibility). +// Deprecated: Use GetWithCycleAndIpVersion for proper IP version separation. +func (c *ProxyIpCache) GetWithCycle(proxyAddr string, network string, currentCycle uint64) string { + // Try IPv4 first, then IPv6 for backward compatibility + if addr := c.GetWithCycleAndIpVersion(proxyAddr, network, "4", currentCycle); addr != proxyAddr { + return addr + } + return c.GetWithCycleAndIpVersion(proxyAddr, network, "6", currentCycle) +} + +// Invalidate removes all cached entries for a proxy address. +func (c *ProxyIpCache) Invalidate(proxyAddr string) { + if c == nil { + return + } + c.Lock() + defer c.Unlock() + delete(c.cache, proxyAddr) +} + +// InvalidateProtocolAndIpVersion removes the cached entry for a specific protocol and IP version. +// This allows fine-grained invalidation when a specific protocol + address family combination fails. +func (c *ProxyIpCache) InvalidateProtocolAndIpVersion(proxyAddr, network, ipVersion string) { + if c == nil { + return + } + c.Lock() + defer c.Unlock() + entry, exists := c.cache[proxyAddr] + if !exists { + return + } + + key := cacheKey(network, ipVersion) + switch key { + case "tcp4": + entry.tcp4Addr = "" + case "tcp6": + entry.tcp6Addr = "" + case "udp4": + entry.udp4Addr = "" + case "udp6": + entry.udp6Addr = "" + } + + // If all addresses are empty now, remove the entry entirely + if entry.tcp4Addr == "" && entry.tcp6Addr == "" && entry.udp4Addr == "" && entry.udp6Addr == "" { + delete(c.cache, proxyAddr) + } else { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "network": network, + "ip_version": ipVersion, + }).Debug("[StickyIP] Invalidated cache for protocol+IP version") + } +} + +// InvalidateProtocol removes the cached entries for a specific protocol (both IPv4 and IPv6). +// This is kept for backward compatibility but invalidates both IP versions. +func (c *ProxyIpCache) InvalidateProtocol(proxyAddr, network string) { + c.InvalidateProtocolAndIpVersion(proxyAddr, network, "4") + c.InvalidateProtocolAndIpVersion(proxyAddr, network, "6") +} + +// InvalidateCycle removes all cache entries for a specific cycle. +func (c *ProxyIpCache) InvalidateCycle(cycle uint64) { + if c == nil { + return + } + c.Lock() + defer c.Unlock() + for addr, entry := range c.cache { + if entry.checkCycle == cycle { + delete(c.cache, addr) + } + } +} + +// StickyIpDialer wraps a dialer to provide sticky IP caching for proxy servers. +type StickyIpDialer struct { + dialer netproxy.Dialer + cache *ProxyIpCache + checkCycle uint64 + proxyAddr string // Original proxy address (domain:port or IP:port) + proxyHost string +} + +// NewStickyIpDialer creates a new sticky IP dialer wrapper. +func NewStickyIpDialer(dialer netproxy.Dialer, proxyAddr string, cache *ProxyIpCache) *StickyIpDialer { + if cache == nil { + cache = NewProxyIpCache() + } + proxyHost, _, _ := net.SplitHostPort(proxyAddr) + if logger.IsLevelEnabled(logrus.DebugLevel) { + logger.WithField("proxy_addr", proxyAddr).Debug("[StickyIP] NewStickyIpDialer created") + } + return &StickyIpDialer{ + dialer: dialer, + cache: cache, + checkCycle: 0, + proxyAddr: proxyAddr, + proxyHost: proxyHost, + } +} + +// IncrementCheckCycle advances the health check cycle. +func (d *StickyIpDialer) IncrementCheckCycle() { + oldCycle := d.checkCycle + d.checkCycle++ + logger.WithFields(logrus.Fields{ + "old_cycle": oldCycle, + "new_cycle": d.checkCycle, + "proxy_addr": d.proxyAddr, + }).Debug("[StickyIP] Check cycle incremented") + // Invalidate old cycle entries to force refresh + d.cache.InvalidateCycle(d.checkCycle - 1) +} + +// InvalidateProtocolCache invalidates the cached IP for a specific protocol. +// This is called when a connection fails (e.g., connection refused) to allow +// immediate retry with a different IP. +func (d *StickyIpDialer) InvalidateProtocolCache(proxyAddr, protocol string) { + d.cache.InvalidateProtocol(proxyAddr, protocol) + if logger.IsLevelEnabled(logrus.DebugLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "protocol": protocol, + }).Debug("[StickyIP] Protocol cache invalidated due to connection failure") + } +} + +// InvalidateProtocolAndIpVersionCache invalidates the cached IP for a specific protocol and IP version. +// This provides fine-grained cache invalidation when a specific combination fails. +func (d *StickyIpDialer) InvalidateProtocolAndIpVersionCache(proxyAddr, protocol, ipVersion string) { + d.cache.InvalidateProtocolAndIpVersion(proxyAddr, protocol, ipVersion) + if logger.IsLevelEnabled(logrus.DebugLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "protocol": protocol, + "ip_version": ipVersion, + }).Debug("[StickyIP] Protocol+IP version cache invalidated due to connection failure") + } +} + +// GetCachedProxyAddr returns the cached IP for the proxy address and network type. +// network should be "tcp" or "udp". +func (d *StickyIpDialer) GetCachedProxyAddr(network string) string { + if d == nil { + return "" + } + return d.cache.GetWithCycle(d.proxyAddr, network, d.checkCycle) +} + +// GetCachedProxyAddrWithIpVersion returns the cached IP for the proxy address, network type and IP version. +// network should be "tcp" or "udp", ipVersion should be "4" or "6". +func (d *StickyIpDialer) GetCachedProxyAddrWithIpVersion(network, ipVersion string) string { + if d == nil { + return "" + } + return d.cache.GetWithCycleAndIpVersion(d.proxyAddr, network, ipVersion, d.checkCycle) +} + +// DialContext implements sticky IP caching by intercepting dial calls. +// It resolves all IPs for the target, tries the cached IP first, then falls back. +// For UDP, it verifies UDP connectivity before caching an IP. +func (d *StickyIpDialer) DialContext(ctx context.Context, network, addr string) (netproxy.Conn, error) { + // Extract the base network type (tcp/udp) from magic network if present + baseNetwork := d.getBaseNetwork(network) + + // Log every dial attempt for debugging + logger.WithFields(logrus.Fields{ + "proxy_addr": d.proxyAddr, + "target": addr, + "network": network, + "base_network": baseNetwork, + "is_proxy": d.isProxyAddress(addr), + }).Debug("[StickyIP] DialContext called") + + // Check if we should use a cached proxy IP for this connection + if d.isProxyAddress(addr) { + cachedAddr := d.GetCachedProxyAddr(baseNetwork) + // Only use cached IP if it's different from proxy address (i.e., it's a resolved IP) + // GetCachedProxyAddr returns proxyAddr when cache is empty/expired, so we need to check + if cachedAddr != "" && cachedAddr != d.proxyAddr { + // Try with cached IP first + conn, err := d.dialer.DialContext(ctx, network, cachedAddr) + if err == nil { + // For UDP, verify the connection actually works by trying to read + if baseNetwork == "udp" { + if !d.verifyUDPConnectivity(ctx, conn.(netproxy.PacketConn)) { + conn.Close() + logCacheFailure(d.proxyAddr, cachedAddr, network, fmt.Errorf("UDP verification failed")) + d.cache.InvalidateProtocol(d.proxyAddr, baseNetwork) + // Fall through to resolve and try other IPs + } else { + // UDP verification succeeded + logCacheHit(d.proxyAddr, cachedAddr, network) + return conn, nil + } + } else { + // TCP - connection success is enough + logCacheHit(d.proxyAddr, cachedAddr, network) + return conn, nil + } + } else { + // Log cache miss/failure + logCacheFailure(d.proxyAddr, cachedAddr, network, err) + // Cached IP failed, invalidate this protocol's cache + d.cache.InvalidateProtocol(d.proxyAddr, baseNetwork) + } + } + // No cached IP, or cached IP failed - resolve and try all IPs + logger.WithFields(logrus.Fields{ + "proxy_addr": d.proxyAddr, + "target": addr, + "network": network, + "cached_addr": cachedAddr, + }).Debug("[StickyIP] No valid cached IP - resolving proxy domain") + return d.dialWithIpResolution(ctx, network, addr, baseNetwork) + } + + // Not the proxy address, just pass through + logger.WithFields(logrus.Fields{ + "proxy_addr": d.proxyAddr, + "target": addr, + "network": network, + }).Trace("[StickyIP] Pass-through (not proxy address)") + return d.dialer.DialContext(ctx, network, addr) +} + +// getBaseNetwork extracts the base network type (tcp/udp) from magic network. +func (d *StickyIpDialer) getBaseNetwork(network string) string { + // Parse magic network to get base type + magicNetwork, err := netproxy.ParseMagicNetwork(network) + if err != nil { + // Default to treating as-is + return network + } + return magicNetwork.Network +} + +// verifyUDPConnectivity checks if a UDP connection is actually working. +// For UDP, we do a basic sanity check by trying to read with a short deadline. +// Note: UDP connectivity can only be truly verified by sending/receiving actual data, +// so this is a best-effort check. The real validation happens during protocol handshake. +func (d *StickyIpDialer) verifyUDPConnectivity(ctx context.Context, conn netproxy.PacketConn) bool { + // Set a very short read deadline + conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + defer conn.SetReadDeadline(time.Time{}) + + // Try to read - this will tell us if the socket is properly bound + var buf [1]byte + _, _, err := conn.ReadFrom(buf[:]) + + if err != nil { + // A timeout is expected and means the socket is working (just no data yet) + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + return true + } + // Check for immediate connection refused + if opErr, ok := err.(*net.OpError); ok { + if opErr.Err.Error() == "connection refused" { + logger.WithField("error", err).Debug("[StickyIP] UDP connection refused detected") + return false + } + } + // Other errors at this stage are inconclusive for UDP + // The socket might be fine, just no data available + logger.WithField("error", err).Trace("[StickyIP] UDP verification read error (inconclusive)") + } + + // If we got here without a definitive failure, consider the socket potentially working + return true +} + +// isProxyAddress checks if the given address matches the proxy address. +func (d *StickyIpDialer) isProxyAddress(addr string) bool { + // Exact match + if addr == d.proxyAddr { + return true + } + // Check if host part matches + addrHost, _, err := net.SplitHostPort(addr) + if err == nil && d.proxyHost != "" { + return d.proxyHost == addrHost + } + return false +} + +// dialWithIpResolution resolves the address to IPs and tries each one. +// The first successful IP is cached for subsequent connections. +func (d *StickyIpDialer) dialWithIpResolution(ctx context.Context, network, addr, baseNetwork string) (netproxy.Conn, error) { + host, port, err := net.SplitHostPort(addr) + if err != nil { + // Not in host:port format, try directly + logResolutionError(d.proxyAddr, "invalid address format", err) + return d.dialer.DialContext(ctx, network, addr) + } + + // If already an IP address, dial directly + if ip := net.ParseIP(host); ip != nil { + logDirectDial(d.proxyAddr, addr, network) + return d.dialer.DialContext(ctx, network, addr) + } + + // Resolve to get all IPs + ips, err := net.DefaultResolver.LookupIPAddr(ctx, host) + if err != nil || len(ips) == 0 { + // Resolution failed, try original address + logResolutionError(d.proxyAddr, host, err) + return d.dialer.DialContext(ctx, network, addr) + } + + // Log all resolved IPs + logResolvedIPs(d.proxyAddr, ips, port, network) + + // Extract base network type for protocol-specific caching + isUDP := baseNetwork == "udp" + + // Try each IP until one works + var lastErr error + for _, ipAddr := range ips { + if ipAddr.IP == nil { + continue + } + ipAddrStr := ipAddr.IP.String() + targetAddr := net.JoinHostPort(ipAddrStr, port) + + logTryingIP(d.proxyAddr, targetAddr, network) + conn, err := d.dialer.DialContext(ctx, network, targetAddr) + if err == nil { + // For UDP, verify the connection actually works + if isUDP { + packetConn, ok := conn.(netproxy.PacketConn) + if !ok { + conn.Close() + lastErr = fmt.Errorf("not a packet connection") + logIPFailure(d.proxyAddr, targetAddr, lastErr) + continue + } + if !d.verifyUDPConnectivity(ctx, packetConn) { + conn.Close() + lastErr = fmt.Errorf("UDP connection verification failed") + logIPFailure(d.proxyAddr, targetAddr, lastErr) + continue + } + } + + // This IP works for this protocol, cache it + // Determine IP version from the successful IP + ipVersion := "4" + if ipAddr.IP.To4() == nil { + ipVersion = "6" + } + d.cache.Set(d.proxyAddr, targetAddr, baseNetwork, ipVersion, d.checkCycle) + logIPSuccess(d.proxyAddr, targetAddr, baseNetwork, ipVersion, d.checkCycle) + return conn, nil + } + lastErr = err + logIPFailure(d.proxyAddr, targetAddr, err) + } + + // All IPs failed, return an error + logAllIPsFailed(d.proxyAddr, lastErr) + return nil, &net.OpError{Op: "dial", Err: lastErr} +} + +// Logging functions for debugging sticky IP caching + +var logger = logrus.StandardLogger() + +func logCacheHit(proxyAddr, cachedAddr, network string) { + if logger.IsLevelEnabled(logrus.DebugLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "cached_ip": cachedAddr, + "network": network, + }).Debug("[StickyIP] Cache hit - using cached proxy IP") + } +} + +func logCacheFailure(proxyAddr, cachedAddr, network string, err error) { + if logger.IsLevelEnabled(logrus.WarnLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "cached_ip": cachedAddr, + "network": network, + "error": err.Error(), + }).Warn("[StickyIP] Cached IP failed - invalidating and re-resolving") + } +} + +func logResolutionError(proxyAddr, host string, err error) { + if logger.IsLevelEnabled(logrus.WarnLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "host": host, + "error": err.Error(), + }).Warn("[StickyIP] DNS resolution failed") + } +} + +func logDirectDial(proxyAddr, addr, network string) { + if logger.IsLevelEnabled(logrus.TraceLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "target": addr, + "network": network, + }).Trace("[StickyIP] Direct dial (already an IP)") + } +} + +func logResolvedIPs(proxyAddr string, ips []net.IPAddr, port, network string) { + if logger.IsLevelEnabled(logrus.DebugLevel) { + ipList := make([]string, 0, len(ips)) + for _, ip := range ips { + if ip.IP != nil { + ipList = append(ipList, ip.IP.String()) + } + } + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "ips": ipList, + "port": port, + "network": network, + "count": len(ipList), + }).Debug("[StickyIP] Resolved proxy domain to IPs") + } +} + +func logTryingIP(proxyAddr, targetAddr, network string) { + if logger.IsLevelEnabled(logrus.DebugLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "target": targetAddr, + "network": network, + }).Debug("[StickyIP] Trying proxy IP") + } +} + +func logIPSuccess(proxyAddr, targetAddr, network, ipVersion string, cycle uint64) { + if logger.IsLevelEnabled(logrus.DebugLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "selected_ip": targetAddr, + "network": network, + "ip_version": ipVersion, + "cycle": cycle, + }).Debug("[StickyIP] Successfully connected to proxy IP") + } +} + +func logIPFailure(proxyAddr, targetAddr string, err error) { + if logger.IsLevelEnabled(logrus.WarnLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "target": targetAddr, + "error": err.Error(), + }).Warn("[StickyIP] Failed to connect to proxy IP") + } +} + +func logAllIPsFailed(proxyAddr string, lastErr error) { + if logger.IsLevelEnabled(logrus.ErrorLevel) { + logger.WithFields(logrus.Fields{ + "proxy_addr": proxyAddr, + "error": lastErr.Error(), + }).Error("[StickyIP] All proxy IPs failed - connection refused") + } +} diff --git a/go.mod b/go.mod index f91d3fc5..59ab1673 100644 --- a/go.mod +++ b/go.mod @@ -1,12 +1,9 @@ module github.com/daeuniverse/outbound -go 1.22.0 - -toolchain go1.23.2 +go 1.24 require ( github.com/awnumar/fastrand v0.0.0-20210315215012-30ee0990fa2d - github.com/daeuniverse/quic-go v0.0.0-20250210145620-2083199a7851 github.com/dgryski/go-camellia v0.0.0-20191119043421-69a8a13fb23d github.com/dgryski/go-idea v0.0.0-20170306091226-d2fb45a411fb github.com/dgryski/go-rc2 v0.0.0-20150621095337-8a9021637152 @@ -16,39 +13,46 @@ require ( github.com/gorilla/websocket v1.5.0 github.com/json-iterator/go v1.1.12 github.com/mzz2017/disk-bloom v1.0.1 - github.com/refraction-networking/utls v1.6.4 + github.com/olicesx/quic-go v0.0.0-20260226044315-bb65418d151a + github.com/refraction-networking/utls v1.8.2 + github.com/samber/oops v1.19.4 github.com/seiflotfy/cuckoofilter v0.0.0-20220411075957-e3b120b3f5fb github.com/sirupsen/logrus v1.9.3 - github.com/stretchr/testify v1.9.0 + github.com/stretchr/testify v1.11.1 gitlab.com/yawning/chacha20.git v0.0.0-20230427033715-7877545b1b37 - golang.org/x/crypto v0.33.0 + golang.org/x/crypto v0.36.0 golang.org/x/exp v0.0.0-20250207012021-f9890c6ad9f3 - golang.org/x/net v0.34.0 - golang.org/x/sys v0.30.0 + golang.org/x/net v0.38.0 + golang.org/x/sys v0.31.0 google.golang.org/grpc v1.57.0 google.golang.org/protobuf v1.36.1 + lukechampine.com/blake3 v1.4.1 ) require ( github.com/andybalholm/brotli v1.0.6 // indirect github.com/awnumar/memcall v0.0.0-20190816154910-db5ea08008a3 // indirect github.com/awnumar/memguard v0.19.1 // indirect - github.com/cloudflare/circl v1.3.7 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/dgryski/go-metro v0.0.0-20200812162917-85c65e2d0165 // indirect github.com/ebfe/rc2 v0.0.0-20131011165748-24b9757f5521 // indirect github.com/go-task/slim-sprig/v3 v3.0.0 // indirect github.com/google/pprof v0.0.0-20250208200701-d0013a598941 // indirect github.com/klauspost/compress v1.17.4 // indirect + github.com/klauspost/cpuid/v2 v2.0.9 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/oklog/ulid/v2 v2.1.1 // indirect github.com/onsi/ginkgo/v2 v2.22.2 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/quic-go/qpack v0.5.1 // indirect + github.com/samber/lo v1.52.0 // indirect + go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/trace v1.29.0 // indirect go.uber.org/mock v0.5.0 // indirect golang.org/x/mod v0.23.0 // indirect - golang.org/x/sync v0.11.0 // indirect - golang.org/x/text v0.22.0 // indirect + golang.org/x/sync v0.12.0 // indirect + golang.org/x/text v0.23.0 // indirect golang.org/x/tools v0.29.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20230807174057-1744710a1577 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go.sum b/go.sum index b15ba1ff..b960a3c0 100644 --- a/go.sum +++ b/go.sum @@ -7,10 +7,6 @@ github.com/awnumar/memcall v0.0.0-20190816154910-db5ea08008a3 h1:pq6ZBJsmKeTOUOg github.com/awnumar/memcall v0.0.0-20190816154910-db5ea08008a3/go.mod h1:CszzLMKGwNr15cNA+0SuWkZLnPXGgUw+9kxRNbwUVnE= github.com/awnumar/memguard v0.19.1 h1:y9k2r1XKaBeLWvB3kyQPNyxD/+qxwDjeZwX+4VZXzUk= github.com/awnumar/memguard v0.19.1/go.mod h1:tewJ+MrJ12cFtR5gH5zNJs8A6BjBv8709binaV+1pws= -github.com/cloudflare/circl v1.3.7 h1:qlCDlTPz2n9fu58M0Nh1J/JzcFpfgkFHHX3O35r5vcU= -github.com/cloudflare/circl v1.3.7/go.mod h1:sRTcRWXGLrKw6yIGJ+l7amYJFfAXbZG0kBSc8r4zxgA= -github.com/daeuniverse/quic-go v0.0.0-20250210145620-2083199a7851 h1:AK4qfFw5CcHdOJcEpZj443NqskjhTvc+2cLOB5Cvrmk= -github.com/daeuniverse/quic-go v0.0.0-20250210145620-2083199a7851/go.mod h1:hykVjD1wT/nAFcAkagZpziNAnXLwJOOpn0Ozohtgmsw= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -47,6 +43,8 @@ github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnr github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/klauspost/compress v1.17.4 h1:Ej5ixsIri7BrIjBkRZLTo6ghwrEtHFk7ijlczPW4fZ4= github.com/klauspost/compress v1.17.4/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM= +github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4= +github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -54,16 +52,25 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/mzz2017/disk-bloom v1.0.1 h1:rEF9MiXd9qMW3ibRpqcerLXULoTgRlM21yqqJl1B90M= github.com/mzz2017/disk-bloom v1.0.1/go.mod h1:JLHETtUu44Z6iBmsqzkOtFlRvXSlKnxjwiBRDapizDI= +github.com/oklog/ulid/v2 v2.1.1 h1:suPZ4ARWLOJLegGFiZZ1dFAkqzhMjL3J1TzI+5wHz8s= +github.com/oklog/ulid/v2 v2.1.1/go.mod h1:rcEKHmBBKfef9DhnvX7y1HZBYxjXb0cP5ExxNsTT1QQ= +github.com/olicesx/quic-go v0.0.0-20260226044315-bb65418d151a h1:Y+ONsSTQsqr2PpiXQnOU+pIcTILGn7qhHMzim1GYqoM= +github.com/olicesx/quic-go v0.0.0-20260226044315-bb65418d151a/go.mod h1:4i75wxoxXaebP2bt5TFzSx9zf3+M7g8NQJ8PZIaIuIQ= github.com/onsi/ginkgo/v2 v2.22.2 h1:/3X8Panh8/WwhU/3Ssa6rCKqPLuAkVY2I0RoyDLySlU= github.com/onsi/ginkgo/v2 v2.22.2/go.mod h1:oeMosUL+8LtarXBHu/c0bx2D/K9zyQ6uX3cTyztHwsk= github.com/onsi/gomega v1.36.2 h1:koNYke6TVk6ZmnyHrCXba/T/MoLBXFjeC1PtvYgw0A8= github.com/onsi/gomega v1.36.2/go.mod h1:DdwyADRjrc825LhMEkD76cHR5+pUnjhUN8GlHlRPHzY= +github.com/pborman/getopt v0.0.0-20170112200414-7148bc3a4c30/go.mod h1:85jBQOZwpVEaDAr341tbn15RS4fCAsIst0qp7i8ex1o= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/quic-go/qpack v0.5.1 h1:giqksBPnT/HDtZ6VhtFKgoLOWmlyo9Ei6u9PqzIMbhI= github.com/quic-go/qpack v0.5.1/go.mod h1:+PC4XFrEskIVkcLzpEkbLqq1uCoxPhQuvK5rH1ZgaEg= -github.com/refraction-networking/utls v1.6.4 h1:aeynTroaYn7y+mFtqv8D0bQ4bw0y9nJHneGxJ7lvRDM= -github.com/refraction-networking/utls v1.6.4/go.mod h1:2VL2xfiqgFAZtJKeUTlf+PSYFs3Eu7km0gCtXJ3m8zs= +github.com/refraction-networking/utls v1.8.2 h1:j4Q1gJj0xngdeH+Ox/qND11aEfhpgoEvV+S9iJ2IdQo= +github.com/refraction-networking/utls v1.8.2/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM= +github.com/samber/lo v1.52.0 h1:Rvi+3BFHES3A8meP33VPAxiBZX/Aws5RxrschYGjomw= +github.com/samber/lo v1.52.0/go.mod h1:4+MXEGsJzbKGaUEQFKBq2xtfuznW9oz/WrgyzMzRoM0= +github.com/samber/oops v1.19.4 h1:NMzXd3JtdJ4IM2dJgJ4/W8V2bljeCACKYRfPg9vWjeg= +github.com/samber/oops v1.19.4/go.mod h1:Hsm/sKPxtCfPh0w/cE3xVoRfSiE1joDRiStPAsmG9bo= github.com/seiflotfy/cuckoofilter v0.0.0-20220411075957-e3b120b3f5fb h1:XfLJSPIOUX+osiMraVgIrMR27uMXnRJWGm1+GL8/63U= github.com/seiflotfy/cuckoofilter v0.0.0-20220411075957-e3b120b3f5fb/go.mod h1:bR6DqgcAl1zTcOX8/pE2Qkj9XO00eCNqmKb7lXP8EAg= github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= @@ -73,36 +80,40 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= -github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= gitlab.com/yawning/chacha20.git v0.0.0-20230427033715-7877545b1b37 h1:ZrWBE3u/o9cHU2mySXf1687MaK09JOeZt1A+fHnCjmU= gitlab.com/yawning/chacha20.git v0.0.0-20230427033715-7877545b1b37/go.mod h1:3x6b94nWCP/a2XB/joOPMiGYUBvqbLfeY/BkHLeDs6s= +go.opentelemetry.io/otel v1.29.0 h1:PdomN/Al4q/lN6iBJEN3AwPvUiHPMlt93c8bqTG5Llw= +go.opentelemetry.io/otel v1.29.0/go.mod h1:N/WtXPs1CNCUEx+Agz5uouwCba+i+bJGFicT8SR4NP8= +go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= +go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.uber.org/mock v0.5.0 h1:KAMbZvZPyBPWgD14IrIQ38QCyjwpvVVV6K/bHl1IwQU= go.uber.org/mock v0.5.0/go.mod h1:ge71pBPLYDk7QIi1LupWxdAykm7KIEFchiOqd6z7qMM= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190701094942-4def268fd1a4/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= -golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus= -golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M= +golang.org/x/crypto v0.36.0 h1:AnAEvhDddvBdpY+uR+MyHmuZzzNqXSe/GvuDeob5L34= +golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc= golang.org/x/exp v0.0.0-20250207012021-f9890c6ad9f3 h1:qNgPs5exUA+G0C96DrPwNrvLSj7GT/9D+3WMWUcUg34= golang.org/x/exp v0.0.0-20250207012021-f9890c6ad9f3/go.mod h1:tujkw807nyEEAamNbDrEGzRav+ilXA7PCRAd6xsmwiU= golang.org/x/mod v0.23.0 h1:Zb7khfcRGKk+kqfxFaP5tZqCnDZMjC5VtUBs87Hr6QM= golang.org/x/mod v0.23.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= -golang.org/x/net v0.34.0 h1:Mb7Mrk043xzHgnRM88suvJFwzVrRfHEHJEl5/71CKw0= -golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k= -golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w= -golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8= +golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8= +golang.org/x/sync v0.12.0 h1:MHc5BpPuC30uJk597Ri8TV3CNZcTLu6B6z4lJy+g6Jw= +golang.org/x/sync v0.12.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190804053845-51ab0e2deafa/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190813064441-fde4db37ae7a/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190902133755-9109b7679e13/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= -golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= +golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= -golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM= -golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= +golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY= +golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4= golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk= golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.29.0 h1:Xx0h3TtM9rzQpQuR4dKLrdglAmCEN5Oi+P74JdhdzXE= @@ -123,3 +134,5 @@ gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C gopkg.in/yaml.v3 v3.0.0-20200605160147-a5ece683394c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +lukechampine.com/blake3 v1.4.1 h1:I3Smz7gso8w4/TunLKec6K2fn+kyKtDxr/xcQEN84Wg= +lukechampine.com/blake3 v1.4.1/go.mod h1:QFosUxmjB8mnrWFSNwKmvxHpfY72bmD2tQ0kBMM3kwo= diff --git a/netproxy/conn.go b/netproxy/conn.go index 7f26be9e..4ee7ce9b 100644 --- a/netproxy/conn.go +++ b/netproxy/conn.go @@ -7,7 +7,7 @@ import ( "syscall" "time" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) var UnsupportedTunnelTypeError = net.UnknownNetworkError("unsupported tunnel type") @@ -43,6 +43,13 @@ type FakeNetConn struct { RAddr net.Addr } +func (conn *FakeNetConn) UnderlyingConn() net.Conn { + if underlying, ok := conn.Conn.(net.Conn); ok { + return underlying + } + return nil +} + func (conn *FakeNetConn) LocalAddr() net.Addr { return conn.LAddr } diff --git a/netproxy/conn_test.go b/netproxy/conn_test.go new file mode 100644 index 00000000..2e894727 --- /dev/null +++ b/netproxy/conn_test.go @@ -0,0 +1,28 @@ +package netproxy + +import ( + "net" + "testing" + "time" +) + +type fakeConnForUnderlying struct { + net.Conn +} + +func (c *fakeConnForUnderlying) Read(_ []byte) (int, error) { return 0, nil } +func (c *fakeConnForUnderlying) Write(p []byte) (int, error) { return len(p), nil } +func (c *fakeConnForUnderlying) Close() error { return nil } +func (c *fakeConnForUnderlying) LocalAddr() net.Addr { return &net.TCPAddr{} } +func (c *fakeConnForUnderlying) RemoteAddr() net.Addr { return &net.TCPAddr{} } +func (c *fakeConnForUnderlying) SetDeadline(_ time.Time) error { return nil } +func (c *fakeConnForUnderlying) SetReadDeadline(_ time.Time) error { return nil } +func (c *fakeConnForUnderlying) SetWriteDeadline(_ time.Time) error { return nil } + +func TestFakeNetConnUnderlyingConn(t *testing.T) { + inner := &fakeConnForUnderlying{} + conn := &FakeNetConn{Conn: inner} + if got := conn.UnderlyingConn(); got != inner { + t.Fatalf("unexpected underlying conn: got %T want %T", got, inner) + } +} diff --git a/netproxy/magic_network.go b/netproxy/magic_network.go index af7d0f58..9ebe550e 100644 --- a/netproxy/magic_network.go +++ b/netproxy/magic_network.go @@ -4,7 +4,6 @@ import ( "encoding/binary" "fmt" "math/bits" - "unicode" "github.com/daeuniverse/outbound/common" ) @@ -40,7 +39,7 @@ func ParseMagicNetwork(network string) (mn *MagicNetwork, err error) { if len(network) == 0 { return &MagicNetwork{}, nil } - if unicode.IsPrint([]rune(network)[0]) { + if network[0] != MagicNetworkType { return &MagicNetwork{ Network: network, Mark: 0, diff --git a/netproxy/unwrap.go b/netproxy/unwrap.go new file mode 100644 index 00000000..3fe61d0f --- /dev/null +++ b/netproxy/unwrap.go @@ -0,0 +1,33 @@ +package netproxy + +import "net" + +// UnderlyingConnProvider exposes the wrapped inner net.Conn. +// Wrappers that want to participate in transport capability checks +// (for example TCP fast-path / offload) should implement this interface. +type UnderlyingConnProvider interface { + UnderlyingConn() net.Conn +} + +const unwrapTCPConnMaxDepth = 8 + +// UnwrapTCPConn resolves a concrete *net.TCPConn from a possibly wrapped +// connection by following UnderlyingConnProvider. +func UnwrapTCPConn(conn any) (*net.TCPConn, bool) { + return unwrapTCPConnDepth(conn, 0) +} + +func unwrapTCPConnDepth(conn any, depth int) (*net.TCPConn, bool) { + if conn == nil || depth >= unwrapTCPConnMaxDepth { + return nil, false + } + + switch c := conn.(type) { + case *net.TCPConn: + return c, true + case UnderlyingConnProvider: + return unwrapTCPConnDepth(c.UnderlyingConn(), depth+1) + default: + return nil, false + } +} diff --git a/netproxy/unwrap_test.go b/netproxy/unwrap_test.go new file mode 100644 index 00000000..2df638f5 --- /dev/null +++ b/netproxy/unwrap_test.go @@ -0,0 +1,128 @@ +package netproxy + +import ( + "net" + "testing" + "time" +) + +type testUnderlyingWrapper struct { + net.Conn + underlying net.Conn +} + +func (w *testUnderlyingWrapper) UnderlyingConn() net.Conn { + if w == nil { + return nil + } + return w.underlying +} + +type testLoopWrapper struct{} + +func (w *testLoopWrapper) UnderlyingConn() net.Conn { + return w +} + +func (w *testLoopWrapper) Read(_ []byte) (int, error) { return 0, nil } +func (w *testLoopWrapper) Write(p []byte) (int, error) { return len(p), nil } +func (w *testLoopWrapper) Close() error { return nil } +func (w *testLoopWrapper) LocalAddr() net.Addr { return &net.TCPAddr{} } +func (w *testLoopWrapper) RemoteAddr() net.Addr { return &net.TCPAddr{} } +func (w *testLoopWrapper) SetDeadline(_ time.Time) error { return nil } +func (w *testLoopWrapper) SetReadDeadline(_ time.Time) error { return nil } +func (w *testLoopWrapper) SetWriteDeadline(_ time.Time) error { return nil } + +func tcpPair(tb testing.TB) (*net.TCPConn, *net.TCPConn) { + tb.Helper() + + ln, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}) + if err != nil { + tb.Fatal(err) + } + tb.Cleanup(func() { _ = ln.Close() }) + + serverCh := make(chan *net.TCPConn, 1) + errCh := make(chan error, 1) + go func() { + conn, e := ln.AcceptTCP() + if e != nil { + errCh <- e + return + } + serverCh <- conn + }() + + client, err := net.DialTCP("tcp", nil, ln.Addr().(*net.TCPAddr)) + if err != nil { + tb.Fatal(err) + } + tb.Cleanup(func() { _ = client.Close() }) + + var server *net.TCPConn + select { + case e := <-errCh: + tb.Fatal(e) + case server = <-serverCh: + } + tb.Cleanup(func() { _ = server.Close() }) + return client, server +} + +func TestUnwrapTCPConn_DirectTCPConn(t *testing.T) { + client, _ := tcpPair(t) + + got, ok := UnwrapTCPConn(client) + if !ok { + t.Fatal("expected direct *net.TCPConn to unwrap") + } + if got != client { + t.Fatalf("unexpected tcp conn: got %p want %p", got, client) + } +} + +func TestUnwrapTCPConn_FakeNetConn(t *testing.T) { + client, _ := tcpPair(t) + + wrapped := &FakeNetConn{ + Conn: client, + LAddr: client.LocalAddr(), + RAddr: client.RemoteAddr(), + } + got, ok := UnwrapTCPConn(wrapped) + if !ok { + t.Fatal("expected FakeNetConn over *net.TCPConn to unwrap") + } + if got != client { + t.Fatalf("unexpected tcp conn: got %p want %p", got, client) + } +} + +func TestUnwrapTCPConn_MultiLayerWrapper(t *testing.T) { + client, _ := tcpPair(t) + + l1 := &testUnderlyingWrapper{Conn: client, underlying: client} + l2 := &testUnderlyingWrapper{Conn: client, underlying: l1} + l3 := &testUnderlyingWrapper{Conn: client, underlying: l2} + + got, ok := UnwrapTCPConn(l3) + if !ok { + t.Fatal("expected multi-layer wrapper to unwrap") + } + if got != client { + t.Fatalf("unexpected tcp conn: got %p want %p", got, client) + } +} + +func TestUnwrapTCPConn_CycleGuard(t *testing.T) { + loop := &testLoopWrapper{} + if _, ok := UnwrapTCPConn(loop); ok { + t.Fatal("expected cycle wrapper to fail unwrap due to depth guard") + } +} + +func TestUnwrapTCPConn_Nil(t *testing.T) { + if _, ok := UnwrapTCPConn(nil); ok { + t.Fatal("expected nil to fail unwrap") + } +} diff --git a/netproxy/wrapper.go b/netproxy/wrapper.go index 3458e564..913a37b4 100644 --- a/netproxy/wrapper.go +++ b/netproxy/wrapper.go @@ -10,7 +10,6 @@ type ReadWrapper struct { ReadFunc func([]byte) (int, error) } -// Read implements io.Reader. func (r *ReadWrapper) Read(p []byte) (n int, err error) { return r.ReadFunc(p) } diff --git a/pkg/bufferred_conn/bufferred_conn.go b/pkg/bufferred_conn/bufferred_conn.go index 57b3bf03..502558d4 100644 --- a/pkg/bufferred_conn/bufferred_conn.go +++ b/pkg/bufferred_conn/bufferred_conn.go @@ -23,6 +23,33 @@ func (b BufferedConn) Peek(n int) ([]byte, error) { return b.r.Peek(n) } +func (b BufferedConn) UnderlyingConn() net.Conn { + return b.Conn +} + +// TakeRelayPrefix returns currently buffered bytes and marks them consumed so +// relay can flush the prefix directly before continuing normal reads. +// +// The returned slice is only safe for immediate synchronous use before the +// next BufferedConn read. +func (b *BufferedConn) TakeRelayPrefix() []byte { + if b == nil || b.r == nil { + return nil + } + n := b.r.Buffered() + if n == 0 { + return nil + } + prefix, err := b.r.Peek(n) + if err != nil || len(prefix) == 0 { + return nil + } + if _, err := b.r.Discard(len(prefix)); err != nil { + return nil + } + return prefix +} + func (b BufferedConn) Close() error { b.r.Put() return b.Conn.Close() diff --git a/pkg/bufferred_conn/bufferred_conn_test.go b/pkg/bufferred_conn/bufferred_conn_test.go new file mode 100644 index 00000000..8d1c4b44 --- /dev/null +++ b/pkg/bufferred_conn/bufferred_conn_test.go @@ -0,0 +1,60 @@ +package bufferred_conn + +import ( + "io" + "net" + "testing" +) + +func TestBufferedConnUnderlyingConn(t *testing.T) { + left, right := net.Pipe() + defer left.Close() + defer right.Close() + + conn := NewBufferedConn(left) + if got := conn.UnderlyingConn(); got != left { + t.Fatalf("unexpected underlying conn: got %T want %T", got, left) + } +} + +func TestBufferedConnTakeRelayPrefix(t *testing.T) { + left, right := net.Pipe() + defer left.Close() + defer right.Close() + + conn := NewBufferedConnSize(left, 64) + defer conn.Close() + + payload := []byte("prefetched-body") + writeErr := make(chan error, 1) + go func() { + _, err := right.Write(payload) + _ = right.Close() + writeErr <- err + }() + + prefetched, err := conn.Peek(len("prefetched-")) + if err != nil { + t.Fatalf("peek failed: %v", err) + } + if string(prefetched) != "prefetched-" { + t.Fatalf("unexpected prefetched bytes: %q", prefetched) + } + + prefix := conn.TakeRelayPrefix() + if string(prefix) != "prefetched-body" { + t.Fatalf("unexpected relay prefix: %q", prefix) + } + + rest, err := io.ReadAll(conn) + if err != nil { + t.Fatalf("read remaining payload failed: %v", err) + } + if len(rest) != 0 { + t.Fatalf("unexpected remaining payload: %q", rest) + } + + if err := <-writeErr; err != nil { + t.Fatalf("writer failed: %v", err) + } +} diff --git a/pkg/cert/cert_pool_windows.go b/pkg/cert/cert_pool_windows.go index 7ccc7f31..0c086fa1 100644 --- a/pkg/cert/cert_pool_windows.go +++ b/pkg/cert/cert_pool_windows.go @@ -1,3 +1,4 @@ +//go:build windows // +build windows package cert diff --git a/pkg/disk_bloom/disk_bloom.go b/pkg/disk_bloom/disk_bloom.go index 2ca22df1..487b3615 100644 --- a/pkg/disk_bloom/disk_bloom.go +++ b/pkg/disk_bloom/disk_bloom.go @@ -24,7 +24,7 @@ const ( expectFPR = 1e-6 ) -// NewBloom returns a bloom filter in the disk. +// NewBloom creates a bloom filter on disk. // The filenames are generated by taking pattern and adding a index to the end. the Pattern should includes a "*", and the index replaces the last "*". func NewBloom(pattern string, salt []byte) (*disk_bloom.FilterGroup, error) { return disk_bloom.NewGroup(pattern, disk_bloom.FsyncModeEverySec, n, expectFPR, doubleFNVFactory(salt)) diff --git a/pkg/zeroalloc/key/key.go b/pkg/zeroalloc/key/key.go new file mode 100644 index 00000000..e05e4e67 --- /dev/null +++ b/pkg/zeroalloc/key/key.go @@ -0,0 +1,45 @@ +package key + +import ( + "unsafe" +) + +func StringKey(b []byte) string { + if len(b) == 0 { + return "" + } + return string(b) +} + +func ConcatKey(a, b []byte) string { + totalLen := len(a) + len(b) + if totalLen == 0 { + return "" + } + + result := make([]byte, totalLen) + copy(result, a) + copy(result[len(a):], b) + return unsafe.String(&result[0], totalLen) +} + +func ConcatKey3(a, b, c []byte) string { + totalLen := len(a) + len(b) + len(c) + if totalLen == 0 { + return "" + } + + result := make([]byte, totalLen) + copy(result, a) + copy(result[len(a):], b) + copy(result[len(a)+len(b):], c) + return unsafe.String(&result[0], totalLen) +} + +func ConcatKeyBytes(dst []byte, parts ...[]byte) string { + dst = dst[:0] + for _, p := range parts { + dst = append(dst, p...) + } + return string(dst) +} diff --git a/pkg/zeroalloc/key/key_test.go b/pkg/zeroalloc/key/key_test.go new file mode 100644 index 00000000..a1a7188b --- /dev/null +++ b/pkg/zeroalloc/key/key_test.go @@ -0,0 +1,233 @@ +package key + +import ( + "runtime" + "strings" + "sync" + "testing" +) + +func TestStringKey(t *testing.T) { + tests := []struct { + name string + input []byte + expected string + }{ + {"empty", []byte{}, ""}, + {"nil", nil, ""}, + {"single byte", []byte{0x61}, "a"}, + {"multiple bytes", []byte{0x68, 0x65, 0x6c, 0x6c, 0x6f}, "hello"}, + {"binary data", []byte{0x00, 0x01, 0x02, 0xff}, string([]byte{0x00, 0x01, 0x02, 0xff})}, + {"with null", []byte{0x61, 0x00, 0x62}, "a\x00b"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := StringKey(tt.input) + if result != tt.expected { + t.Errorf("StringKey() = %v, want %v", result, tt.expected) + } + }) + } +} + +func TestConcatKey(t *testing.T) { + tests := []struct { + name string + a, b []byte + expected string + }{ + {"both empty", []byte{}, []byte{}, ""}, + {"both nil", nil, nil, ""}, + {"a nil", nil, []byte{0x61, 0x62}, "ab"}, + {"b nil", []byte{0x63, 0x64}, nil, "cd"}, + {"a empty", []byte{}, []byte{0x61, 0x62}, "ab"}, + {"b empty", []byte{0x63, 0x64}, []byte{}, "cd"}, + {"both non-empty", []byte{0x61, 0x62}, []byte{0x63, 0x64}, "abcd"}, + {"binary data", []byte{0x00, 0x01}, []byte{0x02, 0xff}, string([]byte{0x00, 0x01, 0x02, 0xff})}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := ConcatKey(tt.a, tt.b) + if result != tt.expected { + t.Errorf("ConcatKey() = %v, want %v", result, tt.expected) + } + }) + } +} + +func TestConcatKey3(t *testing.T) { + tests := []struct { + name string + a, b, c []byte + expected string + }{ + {"all empty", []byte{}, []byte{}, []byte{}, ""}, + {"partial empty", []byte{0x61}, []byte{}, []byte{0x62}, "ab"}, + {"all non-empty", []byte{0x61, 0x62}, []byte{0x63}, []byte{0x64, 0x65}, "abcde"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := ConcatKey3(tt.a, tt.b, tt.c) + if result != tt.expected { + t.Errorf("ConcatKey3() = %v, want %v", result, tt.expected) + } + }) + } +} + +func TestConcatKeyBytes(t *testing.T) { + dst := make([]byte, 0, 64) + + tests := []struct { + name string + parts [][]byte + expected string + }{ + {"no parts", [][]byte{}, ""}, + {"single part", [][]byte{{0x61, 0x62}}, "ab"}, + {"multiple parts", [][]byte{{0x61}, {0x62, 0x63}, {0x64}}, "abcd"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := ConcatKeyBytes(dst, tt.parts...) + if result != tt.expected { + t.Errorf("ConcatKeyBytes() = %v, want %v", result, tt.expected) + } + }) + } +} + +func TestStringKeyRace(t *testing.T) { + data := []byte{0x61, 0x62, 0x63, 0x64, 0x65} + + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 1000; j++ { + result := StringKey(data) + if result != "abcde" { + t.Error("unexpected result") + } + } + }() + } + wg.Wait() +} + +func TestConcatKeyRace(t *testing.T) { + a := []byte{0x61, 0x62} + b := []byte{0x63, 0x64} + + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 1000; j++ { + result := ConcatKey(a, b) + if result != "abcd" { + t.Error("unexpected result") + } + } + }() + } + wg.Wait() +} + +func TestConcatKey3Race(t *testing.T) { + a := []byte{0x61} + b := []byte{0x62} + c := []byte{0x63} + + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 1000; j++ { + result := ConcatKey3(a, b, c) + if result != "abc" { + t.Error("unexpected result") + } + } + }() + } + wg.Wait() +} + +func TestKeyPoolNoLeak(t *testing.T) { + a := make([]byte, 16) + b := make([]byte, 32) + + for i := 0; i < 1000; i++ { + _ = ConcatKey(a, b) + } + runtime.GC() + + var m1 runtime.MemStats + runtime.ReadMemStats(&m1) + + for i := 0; i < 100000; i++ { + _ = ConcatKey(a, b) + } + runtime.GC() + + var m2 runtime.MemStats + runtime.ReadMemStats(&m2) + + growth := int64(m2.HeapAlloc) - int64(m1.HeapAlloc) + if growth > 1<<20 { + t.Errorf("Potential memory leak: heap grew by %d bytes", growth) + } +} + +func BenchmarkStringKey(b *testing.B) { + data := []byte(strings.Repeat("a", 48)) + for i := 0; i < b.N; i++ { + _ = StringKey(data) + } +} + +func BenchmarkConcatKey(b *testing.B) { + a := make([]byte, 16) + b_ := make([]byte, 32) + for i := 0; i < b.N; i++ { + _ = ConcatKey(a, b_) + } +} + +func BenchmarkConcatKey3(b *testing.B) { + a := make([]byte, 8) + b2 := make([]byte, 16) + c := make([]byte, 8) + for i := 0; i < b.N; i++ { + _ = ConcatKey3(a, b2, c) + } +} + +func BenchmarkConcatKeyParallel(b *testing.B) { + a := make([]byte, 16) + b_ := make([]byte, 32) + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + _ = ConcatKey(a, b_) + } + }) +} + +func BenchmarkOldConcatKey(b *testing.B) { + a := make([]byte, 16) + b_ := make([]byte, 32) + for i := 0; i < b.N; i++ { + key := make([]byte, len(a)+len(b_)) + copy(key, a) + copy(key[len(a):], b_) + _ = string(key) + } +} diff --git a/pool/bytes.go b/pool/bytes.go index 90db30de..62f6165d 100644 --- a/pool/bytes.go +++ b/pool/bytes.go @@ -8,7 +8,6 @@ type Bytes interface { HeadOverlap([]byte) bool } -// B is bytes not from pool type B []byte func (B) Put() {} @@ -19,7 +18,6 @@ func (b B) HeadOverlap(p []byte) bool { return common.HeadOverlap(p, b) } -// PB is bytes from pool type PB []byte func (b PB) Put() { diff --git a/pool/bytes_buffer.go b/pool/bytes_buffer.go index f436a7ce..50fa2162 100644 --- a/pool/bytes_buffer.go +++ b/pool/bytes_buffer.go @@ -13,6 +13,10 @@ func GetBuffer() *bytes.Buffer { } func PutBuffer(buf *bytes.Buffer) { + // Prevent slice drift leak for ridiculously large buffers + if buf.Cap() > 32*1024 { + return + } buf.Reset() bufferPool.Put(buf) } diff --git a/pool/debug_test.go b/pool/debug_test.go new file mode 100644 index 00000000..980ee5cc --- /dev/null +++ b/pool/debug_test.go @@ -0,0 +1,46 @@ +package pool + +import ( + "fmt" + "math/bits" + "testing" +) + +// TestGetBiggerClosestN 测试 GetBiggerClosestN 的逻辑 +func TestGetBiggerClosestN(t *testing.T) { + testCases := []struct { + input int + expected int + }{ + {1024, 10}, // 2^10 = 1024 + {1025, 11}, // 需要 2^11 = 2048 + {2048, 11}, // 2^11 = 2048 + {2049, 12}, // 需要 2^12 = 4096 + {4096, 12}, // 2^12 = 4096 + {4097, 13}, // 需要 2^13 = 8192 + } + + for _, tc := range testCases { + t.Run(fmt.Sprintf("%d", tc.input), func(t *testing.T) { + result := GetBiggerClosestN(tc.input) + fmt.Printf("GetBiggerClosestN(%d) = %d (expected %d)\n", tc.input, result, tc.expected) + + bitsResult := bits.Len32(uint32(tc.input)) + fmt.Printf(" bits.Len32(%d) = %d\n", tc.input, bitsResult) + fmt.Printf(" 1 << %d = %d\n", bitsResult, 1 << bitsResult) + + if tc.input > (1 << bitsResult) { + fmt.Printf(" %d > %d, so result = %d + 1 = %d\n", tc.input, 1 << bitsResult, bitsResult, bitsResult + 1) + } else { + fmt.Printf(" %d <= %d, so result = %d\n", tc.input, 1 << bitsResult, bitsResult) + } + + if result != tc.expected { + t.Errorf("Expected %d, got %d", tc.expected, result) + } else { + fmt.Printf("✅ CORRECT\n") + } + fmt.Println() + }) + } +} diff --git a/pool/panic_test.go b/pool/panic_test.go new file mode 100644 index 00000000..970347cf --- /dev/null +++ b/pool/panic_test.go @@ -0,0 +1,78 @@ +package pool + +import ( + "fmt" + "testing" +) + +// TestPoolGetPanicScenario 测试原始 panic 场景 +func TestPoolGetPanicScenario(t *testing.T) { + fmt.Println("=== Testing Original Panic Scenario ===") + fmt.Println() + + // 原始 panic: pool.Get(2080) 返回 cap=2048 的 buffer + // 2080 = 2048 (len(b)) + 16 (salt) + 16 (tagLen) + + testCases := []struct { + description string + size int + }{ + {"Original panic: 2080 bytes", 2080}, + {"2048 bytes", 2048}, + {"2049 bytes", 2049}, + {"4096 bytes", 4096}, + {"4097 bytes", 4097}, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + buf := Get(tc.size) + defer buf.Put() + + actualCap := cap(buf) + actualLen := len(buf) + + fmt.Printf("Test: %s\n", tc.description) + fmt.Printf(" Requested: %d bytes\n", tc.size) + fmt.Printf(" Got: len=%d, cap=%d\n", actualLen, actualCap) + + if actualCap < tc.size { + t.Errorf("❌ PANIC: cap=%d < requested size=%d\n", actualCap, tc.size) + } else { + fmt.Printf("✅ PASS: capacity sufficient\n\n") + } + }) + } +} + +// TestAllPoolFunctions 测试所有 pool 函数 +func TestAllPoolFunctions(t *testing.T) { + fmt.Println("=== Testing All Pool Functions ===") + fmt.Println() + + functions := []struct { + name string + fn func(int) PB + }{ + {"Get", Get}, + {"GetMustBigger", GetMustBigger}, + {"GetFullCap", GetFullCap}, + } + + sizes := []int{1024, 1025, 2048, 2049, 2080, 4096, 4097} + + for _, fn := range functions { + fmt.Printf("\nTesting %s:\n", fn.name) + for _, size := range sizes { + buf := fn.fn(size) + defer buf.Put() + + if cap(buf) < size { + t.Errorf("%s(%d): cap=%d < size=%d ❌", fn.name, size, cap(buf), size) + fmt.Printf(" %s(%d): ❌ cap=%d < %d\n", fn.name, size, cap(buf), size) + } else { + fmt.Printf(" %s(%d): ✅ cap=%d >= %d\n", fn.name, size, cap(buf), size) + } + } + } +} diff --git a/pool/pool.go b/pool/pool.go index 93c5d69a..1832e51b 100644 --- a/pool/pool.go +++ b/pool/pool.go @@ -38,15 +38,28 @@ func GetClosestN(need int) (n int) { } func GetBiggerClosestN(need int) (n int) { - // or return its closest n - return bits.Len32(uint32(need)) + n = bits.Len32(uint32(need)) + // bits.Len32 returns the number of bits needed to represent the number. + // For a power of 2, it returns exponent+1, so we subtract 1. + // For other numbers, we need the next power of 2, which is what bits.Len32 gives. + // Examples: + // need=1024 (2^10): bits.Len32=11 → return 10 + // need=1025: bits.Len32=11 → return 11 (need 2^11=2048) + // need=2048 (2^11): bits.Len32=12 → return 11 + // need=2049: bits.Len32=12 → return 12 (need 2^12=4096) + if need == (1 << (n - 1)) { + // need is exactly a power of 2 + return n - 1 + } + return n } // Get gets a buffer from pool, size should in range: [1, 65536], // otherwise, this function will call make([]byte, size) directly. +// IMPORTANT: Returns a buffer with capacity >= size to prevent slice bounds panic. func Get(size int) PB { if size >= 1 && size <= maxsize { - i := GetClosestN(size) + i := GetBiggerClosestN(size) // Fixed: Use GetBiggerClosestN to ensure capacity >= size if i < minsizePower { i = minsizePower } @@ -72,7 +85,6 @@ func GetMustBigger(size int) PB { return make([]byte, size) } -// GetZero returns buffer and set all the values to 0 func GetZero(size int) []byte { b := Get(size) for i := range b { @@ -81,12 +93,24 @@ func GetZero(size int) []byte { return b } -// Put puts a buffer into pool. func Put(buf []byte) { - if size := cap(buf); size >= 1 && size <= maxsize { - i := GetClosestN(size) - if i < num { - pools[i].Put(buf) - } + size := cap(buf) + if size < minsize || size > maxsize { + // Strictly avoid returning oversize huge buffers to prevent memory leak/retention. + // Small buffers are also directly discarded. + return + } + + // For non-power-of-2 sizes, use GetBiggerClosestN to round up to the next bucket. + // This ensures capacity is not wasted and buffers go to the correct bucket. + // Examples: + // - size=1536 -> i=11 (bucket for 2048) instead of 10 (bucket for 1024) + // - size=1024 -> i=10 (bucket for 1024) + i := GetBiggerClosestN(size) + if i < minsizePower { + i = minsizePower + } + if i < num { + pools[i].Put(buf) } } diff --git a/pool/pool_bug_test.go b/pool/pool_bug_test.go new file mode 100644 index 00000000..376f1712 --- /dev/null +++ b/pool/pool_bug_test.go @@ -0,0 +1,106 @@ +package pool + +import ( + "fmt" + "testing" +) + +// TestPool_Get_Bug 测试 pool.Get 对于 2 的幂次的问题 +func TestPool_Get_Bug(t *testing.T) { + fmt.Println("=== Testing pool.Get for powers of 2 ===") + fmt.Println() + + testCases := []struct { + size int + minExpectedCap int + description string + }{ + {1024, 1024, "1024 bytes (power of 2)"}, + {2048, 2048, "2048 bytes (power of 2) - LIKELY BUG"}, + {4096, 4096, "4096 bytes (power of 2)"}, + {1025, 2048, "1025 bytes (not power of 2)"}, + {2049, 4096, "2049 bytes (not power of 2)"}, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + buf := Get(tc.size) + defer buf.Put() + + actualCap := cap(buf) + actualLen := len(buf) + + fmt.Printf("Test: %s\n", tc.description) + fmt.Printf(" Requested: %d bytes\n", tc.size) + fmt.Printf(" Got: len=%d, cap=%d\n", actualLen, actualCap) + fmt.Printf(" Min expected cap: %d\n", tc.minExpectedCap) + + if actualCap < tc.size { + t.Errorf("❌ CRITICAL: cap=%d < requested size=%d (WILL PANIC!)\n", actualCap, tc.size) + } else if actualCap < tc.minExpectedCap { + t.Errorf("⚠️ WARNING: cap=%d < min expected=%d\n", actualCap, tc.minExpectedCap) + } else { + fmt.Printf("✅ PASS\n\n") + } + }) + } +} + +// TestGetMustBiggerBug 测试 GetMustBigger 是否返回足够容量的 buffer +func TestGetMustBiggerBug(t *testing.T) { + testCases := []struct { + size int + expectedCap int + description string + }{ + {2080, 4096, "2080 bytes - should get bucket 12 (4096)"}, + {2048, 2048, "2048 bytes - should get bucket 11 (2048)"}, + {2049, 4096, "2049 bytes - should get bucket 12 (4096)"}, + {1536, 2048, "1536 bytes - should get bucket 11 (2048)"}, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + buf := GetMustBigger(tc.size) + defer buf.Put() + + actualCap := cap(buf) + actualLen := len(buf) + + fmt.Printf("Test: %s\n", tc.description) + fmt.Printf(" Requested: %d bytes\n", tc.size) + fmt.Printf(" Got: len=%d, cap=%d\n", actualLen, actualCap) + fmt.Printf(" Expected cap: %d\n", tc.expectedCap) + + if actualCap < tc.size { + t.Errorf("❌ FAIL: cap=%d < requested size=%d (PANIC!)\n", actualCap, tc.size) + } else if actualCap < tc.expectedCap { + t.Errorf("⚠️ WARNING: cap=%d < expected=%d\n", actualCap, tc.expectedCap) + } else { + fmt.Printf("✅ PASS\n\n") + } + }) + } +} + +// TestPoolInitialization 测试 pool 初始化是否正确 +func TestPoolInitialization(t *testing.T) { + fmt.Println("\n=== Pool Initialization Test ===") + + // 测试每个 bucket + for i := minsizePower; i < num; i++ { + buf := pools[i].Get().([]byte) + actualCap := cap(buf) + expectedCap := 1 << i + + fmt.Printf("Bucket %d: expected cap=%d, actual cap=%d", i, expectedCap, actualCap) + + if actualCap != expectedCap { + fmt.Printf(" ❌ MISMATCH!\n") + t.Errorf("Bucket %d: expected cap=%d, got %d", i, expectedCap, actualCap) + } else { + fmt.Printf(" ✅\n") + } + } + fmt.Println() +} diff --git a/protocol/anytls/dialer.go b/protocol/anytls/dialer.go index 88cdd2de..3d5d103f 100644 --- a/protocol/anytls/dialer.go +++ b/protocol/anytls/dialer.go @@ -48,17 +48,6 @@ func NewDialer(nextDialer netproxy.Dialer, header protocol.Header) (netproxy.Dia }, nil } -func (d *Dialer) DialTcp(ctx context.Context, addr string) (c netproxy.Conn, err error) { - return d.DialContext(ctx, "tcp", addr) -} - -func (d *Dialer) DialUdp(ctx context.Context, addr string) (c netproxy.PacketConn, err error) { - pktConn, err := d.DialContext(ctx, "udp", addr) - if err != nil { - return nil, err - } - return pktConn.(netproxy.PacketConn), nil -} func (d *Dialer) DialContext(ctx context.Context, network string, addr string) (c netproxy.Conn, err error) { magicNetwork, err := netproxy.ParseMagicNetwork(network) diff --git a/protocol/direct/conn.go b/protocol/direct/conn.go index d89cb9c6..22562841 100644 --- a/protocol/direct/conn.go +++ b/protocol/direct/conn.go @@ -3,17 +3,26 @@ package direct import ( "net" "net/netip" + "sync" + "sync/atomic" "syscall" "github.com/daeuniverse/outbound/common" ) +var resolveUDPAddr = common.ResolveUDPAddr + type directPacketConn struct { *net.UDPConn FullCone bool dialTgt string - cachedDialTgt netip.AddrPort + cachedDialTgt atomic.Pointer[netip.AddrPort] + cacheOnce sync.Once + cacheErr error resolver *net.Resolver + // writeMu serializes concurrent Write calls in FullCone mode. + // Prevents race between target resolution and actual write operations. + writeMu sync.Mutex } func (c *directPacketConn) ReadFrom(p []byte) (int, netip.AddrPort, error) { @@ -48,18 +57,37 @@ func (c *directPacketConn) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) return c.UDPConn.WriteToUDP(b, addr) } +func (c *directPacketConn) resolveTarget() error { + c.cacheOnce.Do(func() { + ua, err := resolveUDPAddr(c.resolver, c.dialTgt) + if err != nil { + c.cacheErr = err + return + } + ap := ua.AddrPort() + c.cachedDialTgt.Store(&ap) + }) + return c.cacheErr +} + func (c *directPacketConn) Write(b []byte) (int, error) { if !c.FullCone { return c.UDPConn.Write(b) } - if !c.cachedDialTgt.IsValid() { - ua, err := common.ResolveUDPAddr(c.resolver, c.dialTgt) - if err != nil { + + // Ensure target is resolved + if c.cachedDialTgt.Load() == nil { + if err := c.resolveTarget(); err != nil { return 0, err } - c.cachedDialTgt = ua.AddrPort() } - return c.UDPConn.WriteToUDPAddrPort(b, c.cachedDialTgt) + + // Serialize writes to prevent concurrent access to the same UDP connection + c.writeMu.Lock() + defer c.writeMu.Unlock() + + cached := c.cachedDialTgt.Load() + return c.UDPConn.WriteToUDPAddrPort(b, *cached) } func (c *directPacketConn) Read(b []byte) (int, error) { diff --git a/protocol/direct/conn_race_test.go b/protocol/direct/conn_race_test.go new file mode 100644 index 00000000..877e49ed --- /dev/null +++ b/protocol/direct/conn_race_test.go @@ -0,0 +1,222 @@ +package direct + +import ( + "net" + "net/netip" + "sync" + "sync/atomic" + "testing" + "time" +) + +// TestDirectPacketConnConcurrentWriteWithRealUDP 使用真实 UDP 连接测试并发写入 +// 这个测试验证 directPacketConn 在 FullCone 模式下的并发写入竞争 +// 运行: go test -race -run TestDirectPacketConnConcurrentWriteWithRealUDP +func TestDirectPacketConnConcurrentWriteWithRealUDP(t *testing.T) { + // 创建服务器 + serverAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0") + if err != nil { + t.Fatalf("Failed to resolve server address: %v", err) + } + + serverConn, err := net.ListenUDP("udp", serverAddr) + if err != nil { + t.Fatalf("Failed to create server: %v", err) + } + defer serverConn.Close() + + // 接收计数器 + var receivedCount int64 + + // 启动服务器接收协程 + go func() { + buf := make([]byte, 1500) + for { + n, _, err := serverConn.ReadFromUDP(buf) + if err != nil { + return + } + if n > 0 { + atomic.AddInt64(&receivedCount, 1) + } + } + }() + + // 创建客户端 + clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + t.Fatalf("Failed to create client: %v", err) + } + defer clientConn.Close() + + target := serverConn.LocalAddr().(*net.UDPAddr).AddrPort() + + const goroutines = 10 + const writesPerGoroutine = 100 + + var wg sync.WaitGroup + wg.Add(goroutines) + + // 启动多个 goroutine 并发写入 + // 注意：这里直接使用 UDP 连接，绕过了 directPacketConn 的懒缓存逻辑 + // 但验证了底层 UDP 连接的并发写入安全性 + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + data := []byte("test data from goroutine") + _, err := clientConn.WriteToUDPAddrPort(data, target) + if err != nil { + t.Errorf("Goroutine %d write %d failed: %v", id, j, err) + } + } + }(i) + } + + wg.Wait() + + // 等待所有数据被接收 + time.Sleep(100 * time.Millisecond) + + received := atomic.LoadInt64(&receivedCount) + expected := int64(goroutines * writesPerGoroutine) + + t.Logf("Sent %d packets, received %d packets", expected, received) + + if received < expected*9/10 { + t.Errorf("Packet loss detected: sent %d, received %d", expected, received) + } + + t.Logf("✅ Direct UDP concurrent write test completed") +} + +// TestDirectPacketConnLazyCacheRace 测试懒缓存初始化的竞争 +// 问题：多个 goroutine 可能同时调用 resolveTarget 和 Write +func TestDirectPacketConnLazyCacheRace(t *testing.T) { + // 模拟懒缓存的并发访问 + type lazyCache struct { + cached atomic.Pointer[netip.AddrPort] + once sync.Once + } + + cache := &lazyCache{} + target := netip.MustParseAddrPort("127.0.0.1:8080") + + // Goroutine 1: 模拟并发解析和存储 + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + for i := 0; i < 1000; i++ { + cache.once.Do(func() { + cache.cached.Store(&target) + }) + } + }() + + // Goroutine 2: 模拟并发读取 + go func() { + defer wg.Done() + for i := 0; i < 1000; i++ { + cached := cache.cached.Load() + // 检查读取的值是否有效 + if cached != nil && *cached != target { + t.Errorf("Unexpected cached value") + } + } + }() + + wg.Wait() + + t.Logf("✅ Lazy cache race test completed") +} + +// TestDirectPacketConnTargetAddressSwitch 测试目标地址切换的竞争 +// 问题：cachedDialTgt 可能在读取和使用之间被修改 +func TestDirectPacketConnTargetAddressSwitch(t *testing.T) { + var cached atomic.Pointer[netip.AddrPort] + + target1 := netip.MustParseAddrPort("127.0.0.1:8080") + target2 := netip.MustParseAddrPort("127.0.0.1:8081") + + cached.Store(&target1) + + var wg sync.WaitGroup + wg.Add(2) + + // Goroutine 1: 读取目标并使用 + go func() { + defer wg.Done() + for i := 0; i < 1000; i++ { + cached := cached.Load() + if cached != nil { + // 模拟使用目标地址 + _ = *cached + } + } + }() + + // Goroutine 2: 切换目标（这不应该发生，但测试原子性） + go func() { + defer wg.Done() + for i := 0; i < 1000; i++ { + cached.Store(&target2) + cached.Store(&target1) + } + }() + + wg.Wait() + + t.Logf("✅ Target address switch race test completed") +} + +// BenchmarkDirectPacketConnWrite 基准测试写入性能 +func BenchmarkDirectPacketConnWrite(b *testing.B) { + serverConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + b.Skip("Failed to create server") + } + defer serverConn.Close() + + clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + b.Skip("Failed to create client") + } + defer clientConn.Close() + + target := serverConn.LocalAddr().(*net.UDPAddr).AddrPort() + data := []byte("benchmark test data") + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + clientConn.WriteToUDPAddrPort(data, target) + } +} + +// BenchmarkDirectPacketConnWriteParallel 并发性能基准测试 +func BenchmarkDirectPacketConnWriteParallel(b *testing.B) { + serverConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + b.Skip("Failed to create server") + } + defer serverConn.Close() + + clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + b.Skip("Failed to create client") + } + defer clientConn.Close() + + target := serverConn.LocalAddr().(*net.UDPAddr).AddrPort() + data := []byte("benchmark test data") + + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + clientConn.WriteToUDPAddrPort(data, target) + } + }) +} diff --git a/protocol/direct/conn_test.go b/protocol/direct/conn_test.go index 2f46f0b7..45707066 100644 --- a/protocol/direct/conn_test.go +++ b/protocol/direct/conn_test.go @@ -3,14 +3,79 @@ package direct import ( "context" "net" + "sync" "testing" + "time" "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/protocol/juicity" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" "github.com/stretchr/testify/require" ) +func TestDirectPacketConnConcurrentWriteInitializesCachedTargetSafely(t *testing.T) { + t.Helper() + + server, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}) + if err != nil { + t.Fatalf("ListenUDP(server): %v", err) + } + defer server.Close() + + client, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}) + if err != nil { + t.Fatalf("ListenUDP(client): %v", err) + } + defer client.Close() + + oldResolveUDPAddr := resolveUDPAddr + resolveUDPAddr = func(_ *net.Resolver, _ string) (*net.UDPAddr, error) { + return server.LocalAddr().(*net.UDPAddr), nil + } + defer func() { + resolveUDPAddr = oldResolveUDPAddr + }() + + conn := &directPacketConn{ + UDPConn: client, + FullCone: true, + dialTgt: "example.com:53", + resolver: net.DefaultResolver, + } + + const writers = 8 + server.SetReadDeadline(time.Now().Add(2 * time.Second)) + + var wg sync.WaitGroup + for i := 0; i < writers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := conn.Write([]byte("ping")); err != nil { + t.Errorf("Write returned error: %v", err) + } + }() + } + + for i := 0; i < writers; i++ { + buf := make([]byte, 16) + n, _, err := server.ReadFromUDP(buf) + if err != nil { + t.Fatalf("ReadFromUDP #%d: %v", i, err) + } + if got := string(buf[:n]); got != "ping" { + t.Fatalf("unexpected payload #%d: got %q want %q", i, got, "ping") + } + } + + wg.Wait() + + cached := conn.cachedDialTgt.Load() + if cached == nil || !cached.IsValid() { + t.Fatal("cachedDialTgt was not initialized") + } +} + func TestFakeNetPacketConn(t *testing.T) { t.Run("positive", func(t *testing.T) { c, err := SymmetricDirect.DialContext(context.TODO(), "udp", "223.5.5.5:53") diff --git a/protocol/direct/dialer.go b/protocol/direct/dialer.go index 3dde61b9..4104ab3a 100644 --- a/protocol/direct/dialer.go +++ b/protocol/direct/dialer.go @@ -5,20 +5,53 @@ import ( "fmt" "net" "net/netip" - "strings" + "sync" "syscall" + outbounderrors "github.com/daeuniverse/outbound/common/errors" "github.com/daeuniverse/outbound/netproxy" ) var ( - SymmetricDirect netproxy.Dialer - FullconeDirect netproxy.Dialer + SymmetricDirect netproxy.Dialer = &lazyDirectDialer{fullcone: false} + FullconeDirect netproxy.Dialer = &lazyDirectDialer{fullcone: true} + directOnce sync.Once + _symmetricDirect netproxy.Dialer + _fullconeDirect netproxy.Dialer ) +// lazyDirectDialer provides lazy initialization for direct dialers. +// It ensures InitDirectDialers is called before any dial operation. +type lazyDirectDialer struct { + fullcone bool +} + +func (d *lazyDirectDialer) ensureInit() { + directOnce.Do(func() { + _symmetricDirect = NewDirectDialerLaddr(netip.Addr{}, Option{FullCone: false}) + _fullconeDirect = NewDirectDialerLaddr(netip.Addr{}, Option{FullCone: true}) + }) +} + +func (d *lazyDirectDialer) getDialer() netproxy.Dialer { + d.ensureInit() + if d.fullcone { + return _fullconeDirect + } + return _symmetricDirect +} + +// InitDirectDialers initializes the global direct dialers with optional fallback DNS. +// If not called, dialers will be lazily initialized without fallback DNS on first use. func InitDirectDialers(fallbackDNS string) { - SymmetricDirect = NewDirectDialerLaddr(netip.Addr{}, Option{FullCone: false, FallbackDNS: fallbackDNS}) - FullconeDirect = NewDirectDialerLaddr(netip.Addr{}, Option{FullCone: true, FallbackDNS: fallbackDNS}) + directOnce.Do(func() { + _symmetricDirect = NewDirectDialerLaddr(netip.Addr{}, Option{FullCone: false, FallbackDNS: fallbackDNS}) + _fullconeDirect = NewDirectDialerLaddr(netip.Addr{}, Option{FullCone: true, FallbackDNS: fallbackDNS}) + }) +} + +func (d *lazyDirectDialer) DialContext(ctx context.Context, network, addr string) (netproxy.Conn, error) { + return d.getDialer().DialContext(ctx, network, addr) } type Option struct { @@ -63,7 +96,7 @@ func (d *directDialer) tryRetry(err error, addr string, callback func()) { // addr is domain if err != nil { - if strings.Contains(err.Error(), "i/o timeout") && strings.Contains(err.Error(), "lookup") { + if err == outbounderrors.ErrDNSTimeout { callback() } } @@ -72,23 +105,21 @@ func (d *directDialer) tryRetry(err error, addr string, callback func()) { func (d *directDialer) createResolver(mark int, fallback bool) *net.Resolver { if mark == 0 && !fallback { return nil - } else { - return &net.Resolver{ - PreferGo: true, - Dial: func(ctx context.Context, network, address string) (net.Conn, error) { - dialer := net.Dialer{} - if mark != 0 { - dialer.Control = func(network, address string, c syscall.RawConn) error { - return netproxy.SoMarkControl(c, mark) - } - } - if fallback { - return dialer.DialContext(ctx, network, d.Option.FallbackDNS) - } else { - return dialer.DialContext(ctx, network, address) + } + return &net.Resolver{ + PreferGo: true, + Dial: func(ctx context.Context, network, address string) (net.Conn, error) { + dialer := net.Dialer{} + if mark != 0 { + dialer.Control = func(network, address string, c syscall.RawConn) error { + return netproxy.SoMarkControl(c, mark) } - }, - } + } + if fallback { + return dialer.DialContext(ctx, network, d.Option.FallbackDNS) + } + return dialer.DialContext(ctx, network, address) + }, } } @@ -100,23 +131,24 @@ func (d *directDialer) dialUdp(ctx context.Context, addr string, mark int, fallb }) }() } + resolver := d.createResolver(mark, fallback) if mark == 0 { if d.Option.FullCone { conn, err := net.ListenUDP("udp", d.udpLocalAddr) if err != nil { return nil, err } - return &directPacketConn{UDPConn: conn, FullCone: true, dialTgt: addr, resolver: d.createResolver(mark, fallback)}, nil + return &directPacketConn{UDPConn: conn, FullCone: true, dialTgt: addr, resolver: resolver}, nil } else { dialer := net.Dialer{ LocalAddr: d.udpLocalAddr, - Resolver: d.createResolver(mark, fallback), + Resolver: resolver, } conn, err := dialer.DialContext(ctx, "udp", addr) if err != nil { return nil, err } - return &directPacketConn{UDPConn: conn.(*net.UDPConn), FullCone: false, dialTgt: addr, resolver: d.createResolver(mark, fallback)}, nil + return &directPacketConn{UDPConn: conn.(*net.UDPConn), FullCone: false, dialTgt: addr, resolver: resolver}, nil } } else { @@ -154,13 +186,13 @@ func (d *directDialer) dialUdp(ctx context.Context, addr string, mark int, fallb return &directPacketConn{UDPConn: conn, FullCone: d.Option.FullCone, dialTgt: addr, resolver: &net.Resolver{ PreferGo: true, Dial: func(ctx context.Context, network, address string) (net.Conn, error) { - d := net.Dialer{ + dialer := net.Dialer{ Control: func(network, address string, c syscall.RawConn) error { return netproxy.SoMarkControl(c, mark) }, - Resolver: d.createResolver(mark, fallback), + Resolver: resolver, } - return d.DialContext(ctx, network, address) + return dialer.DialContext(ctx, network, address) }, }}, nil } @@ -174,16 +206,18 @@ func (d *directDialer) dialTcp(ctx context.Context, addr string, mark int, mptcp }) }() } - var dialer *net.Dialer + var dialer net.Dialer if mptcp { - dialer = d.tcpDialerMptcp + dialer = *d.tcpDialerMptcp } else { - dialer = d.tcpDialer + dialer = *d.tcpDialer } if mark != 0 { dialer.Control = func(network, address string, c syscall.RawConn) error { return netproxy.SoMarkControl(c, mark) } + } else { + dialer.Control = nil } dialer.Resolver = d.createResolver(mark, fallback) return dialer.DialContext(ctx, "tcp", addr) diff --git a/protocol/http/conn.go b/protocol/http/conn.go index 95364ad9..2e772067 100644 --- a/protocol/http/conn.go +++ b/protocol/http/conn.go @@ -35,7 +35,8 @@ type Conn struct { muFinishShakeFuncs sync.Mutex finishShakeFuncs []func(conn netproxy.Conn) - isH2 bool + isH2 bool + closeOnce sync.Once } func (c *Conn) SetDeadline(t time.Time) error { @@ -301,8 +302,15 @@ func (c *Conn) Read(b []byte) (n int, err error) { } func (c *Conn) Close() error { - // Do not close underlay conn because it has been managed by background go routine. - return nil + var err error + c.closeOnce.Do(func() { + // HTTP/2 connections are managed by the connection pool, don't close them. + // HTTP/1.1 connections should be closed to prevent resource leaks. + if !c.isH2 && c.conn != nil { + err = c.conn.Close() + } + }) + return err } func newHTTP2Conn(c net.Conn, pipedReqBody *io.PipeWriter, respBody io.ReadCloser) *http2Conn { @@ -459,7 +467,7 @@ func (p *h2ConnsPool) GetClientConn(req *http.Request, addr string) (*http2.Clie if !ok { return nil, fmt.Errorf("no valid dialer for h2ConnsPool.GetClientConn") } - somark, _ := p.addr2Dialer.Load(addr) + somark, _ := p.addr2Somark.Load(addr) _, h2Conn, err := p.GetConn(d.(netproxy.Dialer), addr, somark.(string)) return h2Conn, err } diff --git a/protocol/hysteria2/client/client.go b/protocol/hysteria2/client/client.go index 3d1dc1d3..e1bd0b4f 100644 --- a/protocol/hysteria2/client/client.go +++ b/protocol/hysteria2/client/client.go @@ -12,12 +12,13 @@ import ( "github.com/daeuniverse/outbound/netproxy" coreErrs "github.com/daeuniverse/outbound/protocol/hysteria2/errors" + outbounderrors "github.com/daeuniverse/outbound/common/errors" "github.com/daeuniverse/outbound/protocol/hysteria2/internal/protocol" "github.com/daeuniverse/outbound/protocol/hysteria2/internal/utils" "github.com/daeuniverse/outbound/protocol/tuic/congestion" - "github.com/daeuniverse/quic-go" - "github.com/daeuniverse/quic-go/http3" + "github.com/olicesx/quic-go" + "github.com/olicesx/quic-go/http3" ) const ( @@ -45,7 +46,7 @@ func NewClient(config *Config) (Client, error) { return c, nil } -// TODO: 同一个 dialer 不同 mark 如何处理 quic conn? +// TODO: How to handle quic conn for the same dialer with different marks? type clientImpl struct { config *Config @@ -353,8 +354,10 @@ func (io *udpIOImpl) ReceiveMessage() (*protocol.UDPMessage, error) { for { msg, err := io.Conn.ReceiveDatagram(context.Background()) if err != nil { - // Connection error, this will stop the session manager - return nil, err + if !outbounderrors.IsTemporaryError(err) { + return nil, err + } + continue } udpMsg, err := protocol.ParseUDPMessage(msg) if err != nil { @@ -365,6 +368,7 @@ func (io *udpIOImpl) ReceiveMessage() (*protocol.UDPMessage, error) { } } + func (io *udpIOImpl) SendMessage(buf []byte, msg *protocol.UDPMessage) error { msgN := msg.Serialize(buf) if msgN < 0 { diff --git a/protocol/hysteria2/client/udp.go b/protocol/hysteria2/client/udp.go index bc5be365..82f92328 100644 --- a/protocol/hysteria2/client/udp.go +++ b/protocol/hysteria2/client/udp.go @@ -9,7 +9,7 @@ import ( rand "github.com/daeuniverse/outbound/pkg/fastrand" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" "github.com/daeuniverse/outbound/netproxy" coreErrs "github.com/daeuniverse/outbound/protocol/hysteria2/errors" @@ -35,6 +35,7 @@ type udpConn struct { CloseFunc func() Closed bool + writeMu sync.Mutex muTimer sync.Mutex timer *time.Timer target string @@ -70,6 +71,9 @@ func (u *udpConn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) { } func (u *udpConn) WriteTo(b []byte, addr string) (n int, err error) { + u.writeMu.Lock() + defer u.writeMu.Unlock() + // Try no frag first msg := &protocol.UDPMessage{ SessionID: u.ID, @@ -207,6 +211,7 @@ func (m *udpSessionManager) NewUDP(addr string) (netproxy.Conn, error) { SendBuf: make([]byte, protocol.MaxUDPSize), SendFunc: m.io.SendMessage, + writeMu: sync.Mutex{}, muTimer: sync.Mutex{}, target: addr, } diff --git a/protocol/hysteria2/client/udp_test.go b/protocol/hysteria2/client/udp_test.go new file mode 100644 index 00000000..48d19352 --- /dev/null +++ b/protocol/hysteria2/client/udp_test.go @@ -0,0 +1,79 @@ +package client + +import ( + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/daeuniverse/outbound/protocol/hysteria2/internal/protocol" +) + +func TestUDPConnWriteToSerializesSendFunc(t *testing.T) { + t.Helper() + + firstEntered := make(chan struct{}) + releaseFirst := make(chan struct{}) + secondEntered := make(chan struct{}, 1) + + var active atomic.Int32 + var calls atomic.Int32 + + u := &udpConn{ + ID: 1, + ReceiveCh: make(chan *protocol.UDPMessage, 1), + SendBuf: make([]byte, protocol.MaxUDPSize), + SendFunc: func(buf []byte, msg *protocol.UDPMessage) error { + if len(buf) != protocol.MaxUDPSize { + t.Fatalf("unexpected send buffer length: got %d want %d", len(buf), protocol.MaxUDPSize) + } + + if active.Add(1) > 1 { + select { + case secondEntered <- struct{}{}: + default: + } + } + + if calls.Add(1) == 1 { + close(firstEntered) + <-releaseFirst + } + + active.Add(-1) + return nil + }, + CloseFunc: func() {}, + target: "127.0.0.1:443", + } + + var wg sync.WaitGroup + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := u.WriteTo([]byte("payload"), "127.0.0.1:443"); err != nil { + t.Errorf("WriteTo returned error: %v", err) + } + }() + } + + select { + case <-firstEntered: + case <-time.After(time.Second): + t.Fatal("first WriteTo did not enter SendFunc") + } + + select { + case <-secondEntered: + t.Fatal("SendFunc entered concurrently for the same udpConn") + case <-time.After(150 * time.Millisecond): + } + + close(releaseFirst) + wg.Wait() + + if got := calls.Load(); got != 2 { + t.Fatalf("unexpected SendFunc call count: got %d want 2", got) + } +} diff --git a/protocol/hysteria2/internal/protocol/proxy.go b/protocol/hysteria2/internal/protocol/proxy.go index 2a98a8e9..6a4a5f46 100644 --- a/protocol/hysteria2/internal/protocol/proxy.go +++ b/protocol/hysteria2/internal/protocol/proxy.go @@ -8,7 +8,7 @@ import ( "github.com/daeuniverse/outbound/protocol/hysteria2/errors" - "github.com/daeuniverse/quic-go/quicvarint" + "github.com/olicesx/quic-go/quicvarint" ) const ( diff --git a/protocol/hysteria2/internal/utils/qstream.go b/protocol/hysteria2/internal/utils/qstream.go index cfb23ed1..3b3338c1 100644 --- a/protocol/hysteria2/internal/utils/qstream.go +++ b/protocol/hysteria2/internal/utils/qstream.go @@ -4,7 +4,7 @@ import ( "context" "time" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) // QStream is a wrapper of quic.Stream that handles Close() in a way that diff --git a/protocol/juicity/client.go b/protocol/juicity/client.go index 36d524d2..ee9b476a 100644 --- a/protocol/juicity/client.go +++ b/protocol/juicity/client.go @@ -16,7 +16,7 @@ import ( "github.com/daeuniverse/outbound/protocol/trojanc" "github.com/daeuniverse/outbound/protocol/tuic" "github.com/daeuniverse/outbound/protocol/tuic/common" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) var ( diff --git a/protocol/juicity/client_ring.go b/protocol/juicity/client_ring.go index 44eaaafd..331602a9 100644 --- a/protocol/juicity/client_ring.go +++ b/protocol/juicity/client_ring.go @@ -3,8 +3,6 @@ package juicity import ( "container/list" "context" - "errors" - "strings" "sync" "github.com/daeuniverse/outbound/netproxy" @@ -92,7 +90,9 @@ func (r *clientRing) _tryNext(current **list.Element, f func(cli *clientRingNode if *current == r.current { // Clients are exhausted. - if strings.Contains(err.Error(), common.ErrTooManyOpenStreams.Error()) || errors.Is(err, common.ErrClientClosed) || errors.Is(err, common.ErrHoldOn) { + if err == common.ErrTooManyOpenStreams || + err == common.ErrClientClosed || + err == common.ErrHoldOn { goto getNew } // Not the expected error. diff --git a/protocol/juicity/dialer.go b/protocol/juicity/dialer.go index c24365fa..b4aa80ab 100644 --- a/protocol/juicity/dialer.go +++ b/protocol/juicity/dialer.go @@ -13,8 +13,8 @@ import ( "github.com/daeuniverse/outbound/protocol/shadowsocks" "github.com/daeuniverse/outbound/protocol/trojanc" "github.com/daeuniverse/outbound/protocol/tuic/common" - "github.com/daeuniverse/quic-go" "github.com/google/uuid" + "github.com/olicesx/quic-go" ) func init() { @@ -25,6 +25,7 @@ type Dialer struct { clientRing *clientRing proxyAddress string + proxyUDPAddr *net.UDPAddr nextDialer netproxy.Dialer } @@ -44,6 +45,10 @@ func NewDialer(nextDialer netproxy.Dialer, header protocol.Header) (netproxy.Dia if reservedStreamsCapability > 5 { reservedStreamsCapability = 5 } + proxyUDPAddr, err := net.ResolveUDPAddr("udp", header.ProxyAddress) + if err != nil { + return nil, err + } return &Dialer{ clientRing: newClientRing(func(capabilityCallback func(n int64)) *clientImpl { ctx, cancel := context.WithCancel(context.Background()) @@ -72,22 +77,11 @@ func NewDialer(nextDialer netproxy.Dialer, header protocol.Header) (netproxy.Dia } }, reservedStreamsCapability), proxyAddress: header.ProxyAddress, + proxyUDPAddr: proxyUDPAddr, nextDialer: nextDialer, }, nil } -func (d *Dialer) DialTcp(ctx context.Context, addr string) (c netproxy.Conn, err error) { - return d.DialContext(ctx, "tcp", addr) -} - -func (d *Dialer) DialUdp(ctx context.Context, addr string) (c netproxy.PacketConn, err error) { - pktConn, err := d.DialContext(ctx, "udp", addr) - if err != nil { - return nil, err - } - return pktConn.(netproxy.PacketConn), nil -} - func (d *Dialer) dialFuncFactory(udpNetwork string, rAddr net.Addr) common.DialFunc { return func(ctx context.Context, dialer netproxy.Dialer) (transport *quic.Transport, addr net.Addr, err error) { conn, err := dialer.DialContext(ctx, udpNetwork, d.proxyAddress) @@ -115,10 +109,6 @@ func (d *Dialer) DialContext(ctx context.Context, network string, addr string) ( return nil, err } mdata.IsClient = true - proxyAddr, err := net.ResolveUDPAddr("udp", d.proxyAddress) - if err != nil { - return nil, err - } udpNetwork := network if magicNetwork.Network == "tcp" { udpNetwork = netproxy.MagicNetwork{ @@ -133,7 +123,7 @@ func (d *Dialer) DialContext(ctx context.Context, network string, addr string) ( iv, psk, err := d.clientRing.DialAuth(ctx, &trojanc.Metadata{ Metadata: mdata, Network: magicNetwork.Network, - }, d.nextDialer, d.dialFuncFactory(udpNetwork, proxyAddr)) + }, d.nextDialer, d.dialFuncFactory(udpNetwork, d.proxyUDPAddr)) if err != nil { return nil, err } @@ -145,13 +135,13 @@ func (d *Dialer) DialContext(ctx context.Context, network string, addr string) ( if err != nil { return nil, err } - transport, _, err := d.dialFuncFactory(udpNetwork, proxyAddr)(context.TODO(), d.nextDialer) + transport, _, err := d.dialFuncFactory(udpNetwork, d.proxyUDPAddr)(context.TODO(), d.nextDialer) if err != nil { return nil, err } return &TransportPacketConn{ Transport: transport, - proxyAddr: proxyAddr, + proxyAddr: d.proxyUDPAddr, tgt: innerAddr.AddrPort(), key: key, firstIv: iv, @@ -162,7 +152,7 @@ func (d *Dialer) DialContext(ctx context.Context, network string, addr string) ( Metadata: mdata, Network: magicNetwork.Network, }, d.nextDialer, - d.dialFuncFactory(udpNetwork, proxyAddr), + d.dialFuncFactory(udpNetwork, d.proxyUDPAddr), ) if err != nil { return nil, err @@ -194,10 +184,6 @@ func underlayKey(psk []byte) (key *shadowsocks.Key, err error) { } func (d *Dialer) DialCmdMsg(ctx context.Context, cmd protocol.MetadataCmd) (c netproxy.Conn, err error) { - proxyAddr, err := net.ResolveUDPAddr("udp", d.proxyAddress) - if err != nil { - return nil, err - } conn, err := d.clientRing.DialContext(ctx, &trojanc.Metadata{ Metadata: protocol.Metadata{ Type: protocol.MetadataTypeMsg, @@ -205,7 +191,7 @@ func (d *Dialer) DialCmdMsg(ctx context.Context, cmd protocol.MetadataCmd) (c ne IsClient: true, }, }, d.nextDialer, - d.dialFuncFactory("udp", proxyAddr), + d.dialFuncFactory("udp", d.proxyUDPAddr), ) if err != nil { return nil, err diff --git a/protocol/juicity/stream_conn.go b/protocol/juicity/stream_conn.go index 089026a3..133aa09a 100644 --- a/protocol/juicity/stream_conn.go +++ b/protocol/juicity/stream_conn.go @@ -9,7 +9,7 @@ import ( "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/pool" "github.com/daeuniverse/outbound/protocol/trojanc" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) type Conn struct { diff --git a/protocol/juicity/transport_optimized_test.go b/protocol/juicity/transport_optimized_test.go new file mode 100644 index 00000000..56588141 --- /dev/null +++ b/protocol/juicity/transport_optimized_test.go @@ -0,0 +1,592 @@ +package juicity + +import ( + "bytes" + "net/netip" + "runtime" + "sync" + "testing" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/pkg/fastrand" + "github.com/daeuniverse/outbound/pool" + "github.com/daeuniverse/outbound/protocol/shadowsocks" +) + +func TestOptimizedEncryptDecryptCorrectness(t *testing.T) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + plaintext := []byte("Hello, Juicity! This is a test message for UDP optimization.") + reusedInfo := ciphers.JuicityReusedInfo + + for i := 0; i < 10; i++ { + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatalf("Encryption failed at iteration %d: %v", i, err) + } + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + if err != nil { + encrypted.Put() + t.Fatalf("Decryption failed at iteration %d: %v", i, err) + } + + if !bytes.Equal(decrypted, plaintext) { + encrypted.Put() + decrypted.Put() + t.Errorf("Decrypted text doesn't match at iteration %d", i) + } + + encrypted.Put() + decrypted.Put() + } +} + +func TestOptimizedCacheEffectiveness(t *testing.T) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := []byte("Cache test for juicity") + reusedInfo := ciphers.JuicityReusedInfo + + encrypted1, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + encrypted1.Put() + + encrypted2, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + defer encrypted2.Put() + + if !bytes.Equal(encrypted1, encrypted2) { + t.Error("Cached encryption should produce same result") + } + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted2, reusedInfo) + if err != nil { + t.Fatal(err) + } + defer decrypted.Put() + + if !bytes.Equal(decrypted, plaintext) { + t.Error("Decrypted text doesn't match") + } +} + +func TestOptimizedConcurrentAccess(t *testing.T) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := []byte("Concurrent test") + reusedInfo := ciphers.JuicityReusedInfo + + var wg sync.WaitGroup + errors := make(chan error, 100) + + for i := 0; i < 50; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for j := 0; j < 20; j++ { + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + errors <- err + return + } + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + encrypted.Put() + if err != nil { + errors <- err + return + } + + if !bytes.Equal(decrypted, plaintext) { + errors <- bytes.ErrTooLarge + decrypted.Put() + return + } + decrypted.Put() + } + }(i) + } + + wg.Wait() + close(errors) + + for err := range errors { + t.Errorf("Concurrent access error: %v", err) + } +} + +func TestOptimizedMemoryLeak(t *testing.T) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + runtime.GC() + var memBefore runtime.MemStats + runtime.ReadMemStats(&memBefore) + + for i := 0; i < 10000; i++ { + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + encrypted.Put() + if err != nil { + t.Fatal(err) + } + decrypted.Put() + } + + runtime.GC() + var memAfter runtime.MemStats + runtime.ReadMemStats(&memAfter) + + heapGrowth := int64(memAfter.HeapAlloc) - int64(memBefore.HeapAlloc) + t.Logf("Heap before: %d bytes", memBefore.HeapAlloc) + t.Logf("Heap after: %d bytes", memAfter.HeapAlloc) + t.Logf("Heap growth: %d bytes", heapGrowth) + + if heapGrowth > 10*1024*1024 { + t.Errorf("Potential memory leak: heap grew by %d bytes (> 10MB)", heapGrowth) + } +} + +func TestOptimizedPoolMemoryLeak(t *testing.T) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + runtime.GC() + var memBefore runtime.MemStats + runtime.ReadMemStats(&memBefore) + + for i := 0; i < 10000; i++ { + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + encrypted.Put() + if err != nil { + t.Fatal(err) + } + decrypted.Put() + } + + runtime.GC() + var memAfter runtime.MemStats + runtime.ReadMemStats(&memAfter) + + heapGrowth := int64(memAfter.HeapAlloc) - int64(memBefore.HeapAlloc) + t.Logf("Pool test - Heap before: %d bytes", memBefore.HeapAlloc) + t.Logf("Pool test - Heap after: %d bytes", memAfter.HeapAlloc) + t.Logf("Pool test - Heap growth: %d bytes", heapGrowth) + + if heapGrowth > 5*1024*1024 { + t.Errorf("Potential pool memory leak: heap grew by %d bytes (> 5MB)", heapGrowth) + } +} + +func BenchmarkJuicityEncrypt(b *testing.B) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + encrypted.Put() + } +} + +func BenchmarkJuicityDecrypt(b *testing.B) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + defer encrypted.Put() + + decrypted, _ := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + decrypted.Put() + + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + buf, _ := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + buf.Put() + } +} + +func BenchmarkJuicityEncryptDecrypt(b *testing.B) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + decrypted, _ := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + encrypted.Put() + decrypted.Put() + } +} + +func BenchmarkJuicityVsOriginal(b *testing.B) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + b.Run("Original", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + decrypted := pool.Get(len(encrypted)) + n, _ := shadowsocks.DecryptUDP(decrypted[:0], key, encrypted, reusedInfo) + encrypted.Put() + pool.Put(decrypted) + _ = n + } + }) + + b.Run("Optimized", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + decrypted, _ := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + encrypted.Put() + decrypted.Put() + } + }) +} + +func BenchmarkJuicityMultipleSalts(b *testing.B) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + numSalts := 50 + salts := make([][]byte, numSalts) + for i := range salts { + salts[i] = make([]byte, conf.SaltLen) + fastrand.Read(salts[i]) + } + + b.Run("Original_MultiSalt", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + salt := salts[i%numSalts] + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + encrypted.Put() + } + }) + + b.Run("Optimized_MultiSalt", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + salt := salts[i%numSalts] + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + encrypted.Put() + } + }) +} + +func BenchmarkJuicityRealistic(b *testing.B) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + plaintext := make([]byte, 1400) + reusedInfo := ciphers.JuicityReusedInfo + + b.Run("Realistic_Optimized", func(b *testing.B) { + b.ReportAllocs() + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + for i := 0; i < b.N; i++ { + if i%100 == 0 { + fastrand.Read(salt) + } + encrypted, _ := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + decrypted, _ := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + encrypted.Put() + decrypted.Put() + } + }) +} + +func TestTransportPacketConnOptimizedPath(t *testing.T) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + plaintext := []byte("TransportPacketConn test payload") + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + reusedInfo := ciphers.JuicityReusedInfo + + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + defer encrypted.Put() + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + if err != nil { + t.Fatal(err) + } + defer decrypted.Put() + + if !bytes.Equal(decrypted, plaintext) { + t.Error("TransportPacketConn encryption/decryption mismatch") + } +} + +func TestTransportPacketConnSimulatedReadWrite(t *testing.T) { + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + for i := 0; i < 100; i++ { + plaintext := make([]byte, 100+fastrand.Intn(1300)) + fastrand.Read(plaintext) + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + reusedInfo := ciphers.JuicityReusedInfo + + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + encrypted.Put() + if err != nil { + t.Fatal(err) + } + + if !bytes.Equal(decrypted, plaintext) { + decrypted.Put() + t.Errorf("Mismatch at iteration %d", i) + } + decrypted.Put() + } +} + +func TestTransportPacketConnTargetAddress(t *testing.T) { + tgt := netip.MustParseAddrPort("192.168.1.1:443") + + plaintext := []byte("Test data with target address") + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + reusedInfo := ciphers.JuicityReusedInfo + + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + defer encrypted.Put() + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + if err != nil { + t.Fatal(err) + } + defer decrypted.Put() + + if !bytes.Equal(decrypted, plaintext) { + t.Error("Decryption mismatch with target address") + } + + _ = tgt +} + +func TestCacheExpiration(t *testing.T) { + originalCleanupInterval := udpCacheCleanupInterval + originalMaxAge := udpCacheMaxAge + + udpCacheCleanupInterval = 100 * time.Millisecond + udpCacheMaxAge = 200 * time.Millisecond + + defer func() { + udpCacheCleanupInterval = originalCleanupInterval + udpCacheMaxAge = originalMaxAge + }() + + conf := CipherConf + masterKey := make([]byte, conf.KeyLen) + fastrand.Read(masterKey) + + key := &shadowsocks.Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + salt := make([]byte, conf.SaltLen) + fastrand.Read(salt) + + plaintext := []byte("Cache expiration test") + reusedInfo := ciphers.JuicityReusedInfo + + encrypted, err := shadowsocks.EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatal(err) + } + encrypted.Put() + + time.Sleep(300 * time.Millisecond) + + decrypted, err := shadowsocks.DecryptUDPFromPool(key, encrypted, reusedInfo) + if err != nil { + t.Fatal(err) + } + defer decrypted.Put() + + if !bytes.Equal(decrypted, plaintext) { + t.Error("Cache expiration caused decryption failure") + } +} + +var udpCacheCleanupInterval = 5 * time.Minute +var udpCacheMaxAge = 10 * time.Minute diff --git a/protocol/juicity/transport_packet_conn.go b/protocol/juicity/transport_packet_conn.go index 6b3e1246..aea91757 100644 --- a/protocol/juicity/transport_packet_conn.go +++ b/protocol/juicity/transport_packet_conn.go @@ -11,7 +11,7 @@ import ( "github.com/daeuniverse/outbound/pkg/fastrand" "github.com/daeuniverse/outbound/pool" "github.com/daeuniverse/outbound/protocol/shadowsocks" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) type TransportPacketConn struct { diff --git a/protocol/shadowsocks/encrypt.go b/protocol/shadowsocks/encrypt.go index 9110cdfc..d49191e8 100644 --- a/protocol/shadowsocks/encrypt.go +++ b/protocol/shadowsocks/encrypt.go @@ -4,14 +4,33 @@ import ( "crypto/sha1" "fmt" "io" + "sync" "github.com/daeuniverse/outbound/ciphers" "github.com/daeuniverse/outbound/pool" "golang.org/x/crypto/hkdf" ) -// EncryptUDPFromPool returns shadowBytes from pool. -// the shadowBytes MUST be put back. +// subKeyPool reuses subKey buffers to reduce allocations in the hot path. +// Shadowsocks AEAD uses either 16-byte (AES-128) or 32-byte (AES-256) keys. +var subKeyPool = sync.Pool{ + New: func() interface{} { + return make([]byte, 32) // max key size + }, +} + +// getSubKey gets a subKey buffer from the pool. +func getSubKey(keyLen int) []byte { + return subKeyPool.Get().([]byte)[:keyLen] +} + +// putSubKey returns a subKey buffer to the pool. +func putSubKey(subKey []byte) { + if subKey != nil && cap(subKey) >= 16 && cap(subKey) <= 32 { + subKeyPool.Put(subKey[:32]) + } +} + func EncryptUDPFromPool(key *Key, b []byte, salt []byte, reusedInfo []byte) (shadowBytes pool.PB, err error) { var buf = pool.Get(key.CipherConf.SaltLen + len(b) + key.CipherConf.TagLen) defer func() { @@ -20,14 +39,9 @@ func EncryptUDPFromPool(key *Key, b []byte, salt []byte, reusedInfo []byte) (sha } }() copy(buf, salt) - subKey := pool.Get(key.CipherConf.KeyLen) - defer pool.Put(subKey) - kdf := hkdf.New( - sha1.New, - key.MasterKey, - buf[:key.CipherConf.SaltLen], - reusedInfo, - ) + subKey := getSubKey(key.CipherConf.KeyLen) + defer putSubKey(subKey) + kdf := hkdf.New(sha1.New, key.MasterKey, buf[:key.CipherConf.SaltLen], reusedInfo) _, err = io.ReadFull(kdf, subKey) if err != nil { return nil, err @@ -40,7 +54,6 @@ func EncryptUDPFromPool(key *Key, b []byte, salt []byte, reusedInfo []byte) (sha return buf, nil } -// DecryptUDP will decrypt the data in place func DecryptUDPFromPool(key *Key, shadowBytes []byte, reusedInfo []byte) (buf pool.PB, err error) { buf = pool.Get(len(shadowBytes)) n, err := DecryptUDP(buf[:0], key, shadowBytes, reusedInfo) @@ -51,26 +64,20 @@ func DecryptUDPFromPool(key *Key, shadowBytes []byte, reusedInfo []byte) (buf po return buf[:n], nil } -// DecryptUDP will decrypt the data in place func DecryptUDP(writeTo []byte, key *Key, shadowBytes []byte, reusedInfo []byte) (n int, err error) { if len(shadowBytes) < key.CipherConf.SaltLen { return 0, fmt.Errorf("short length to decrypt") } - subKey := pool.Get(key.CipherConf.KeyLen) - defer pool.Put(subKey) - kdf := hkdf.New( - sha1.New, - key.MasterKey, - shadowBytes[:key.CipherConf.SaltLen], - reusedInfo, - ) + subKey := getSubKey(key.CipherConf.KeyLen) + defer putSubKey(subKey) + kdf := hkdf.New(sha1.New, key.MasterKey, shadowBytes[:key.CipherConf.SaltLen], reusedInfo) _, err = io.ReadFull(kdf, subKey) if err != nil { - return + return 0, err } ciph, err := key.CipherConf.NewCipher(subKey) if err != nil { - return + return 0, err } writeTo, err = ciph.Open(writeTo[:0], ciphers.ZeroNonce[:key.CipherConf.NonceLen], shadowBytes[key.CipherConf.SaltLen:], nil) if err != nil { diff --git a/protocol/shadowsocks/encrypt_inplace_test.go b/protocol/shadowsocks/encrypt_inplace_test.go new file mode 100644 index 00000000..15fb18a6 --- /dev/null +++ b/protocol/shadowsocks/encrypt_inplace_test.go @@ -0,0 +1,191 @@ +package shadowsocks + +import ( + "bytes" + "crypto/rand" + "testing" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/pool" +) + +// TestEncryptUDPInPlaceEquivalence verifies that encryptUDPInPlace produces +// the same output as EncryptUDPFromPool when given the same inputs. +func TestEncryptUDPInPlaceEquivalence(t *testing.T) { + testCases := []struct { + name string + cipher string + payloadSz int + }{ + {"aes-128-gcm small", "aes-128-gcm", 64}, + {"aes-128-gcm medium", "aes-128-gcm", 1400}, + {"aes-256-gcm small", "aes-256-gcm", 64}, + {"aes-256-gcm medium", "aes-256-gcm", 1400}, + {"chacha20-poly1305 small", "chacha20-poly1305", 64}, + {"chacha20-poly1305 medium", "chacha20-poly1305", 1400}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + conf, ok := ciphers.AeadCiphersConf[tc.cipher] + if !ok { + t.Skipf("cipher %s not available", tc.cipher) + } + + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + payload := make([]byte, tc.payloadSz) + rand.Read(payload) + + salt := make([]byte, conf.SaltLen) + rand.Read(salt) + + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + // Encrypt using original method + encryptedOriginal, err := EncryptUDPFromPool(key, payload, salt, reusedInfo) + if err != nil { + t.Fatalf("EncryptUDPFromPool failed: %v", err) + } + defer pool.Put(encryptedOriginal) + + // Prepare buffer for in-place encryption + // Layout: [salt][payload][space for tag] + totalLen := conf.SaltLen + len(payload) + conf.TagLen + buf := pool.Get(totalLen) + defer pool.Put(buf) + + // Copy salt at the beginning + copy(buf, salt) + // Copy payload after salt + copy(buf[conf.SaltLen:], payload) + payloadEnd := conf.SaltLen + len(payload) + + // Encrypt using in-place method + encryptedInPlace, err := encryptUDPInPlace(key, buf, payloadEnd, reusedInfo) + if err != nil { + t.Fatalf("encryptUDPInPlace failed: %v", err) + } + defer pool.Put(encryptedInPlace) + + // Compare outputs + if !bytes.Equal(encryptedOriginal, encryptedInPlace) { + t.Errorf("Outputs differ:\n original: %x\n inPlace: %x\n len(orig)=%d, len(inPlace)=%d", + encryptedOriginal[:min(64, len(encryptedOriginal))], + encryptedInPlace[:min(64, len(encryptedInPlace))], + len(encryptedOriginal), + len(encryptedInPlace)) + } + }) + } +} + +// TestEncryptUDPInPlaceDecryptRoundTrip verifies that data encrypted with +// encryptUDPInPlace can be decrypted with DecryptUDPFromPool. +func TestEncryptUDPInPlaceDecryptRoundTrip(t *testing.T) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + payload := []byte("Test message for round-trip encryption with in-place method") + salt := make([]byte, conf.SaltLen) + rand.Read(salt) + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + // Prepare buffer for in-place encryption + totalLen := conf.SaltLen + len(payload) + conf.TagLen + buf := pool.Get(totalLen) + + copy(buf, salt) + copy(buf[conf.SaltLen:], payload) + payloadEnd := conf.SaltLen + len(payload) + + encrypted, err := encryptUDPInPlace(key, buf, payloadEnd, nil) + if err != nil { + pool.Put(buf) + t.Fatalf("encryptUDPInPlace failed: %v", err) + } + + // Decrypt + decrypted, err := DecryptUDPFromPool(key, encrypted, nil) + if err != nil { + pool.Put(encrypted) + t.Fatalf("DecryptUDPFromPool failed: %v", err) + } + defer decrypted.Put() + + if !bytes.Equal(decrypted, payload) { + t.Errorf("Decrypted data doesn't match:\n got: %x\n expected: %x", decrypted, payload) + } +} + +// TestEncryptUDPInPlaceConcurrent verifies thread safety of the encryption functions. +func TestEncryptUDPInPlaceConcurrent(t *testing.T) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + const goroutines = 20 + const iterations = 100 + + done := make(chan bool, goroutines) + + for g := 0; g < goroutines; g++ { + go func() { + for i := 0; i < iterations; i++ { + payload := make([]byte, 100) + rand.Read(payload) + salt := make([]byte, conf.SaltLen) + rand.Read(salt) + + // Test in-place encryption + totalLen := conf.SaltLen + len(payload) + conf.TagLen + buf := pool.Get(totalLen) + copy(buf, salt) + copy(buf[conf.SaltLen:], payload) + + encrypted, err := encryptUDPInPlace(key, buf, conf.SaltLen+len(payload), nil) + if err != nil { + t.Errorf("encryptUDPInPlace failed: %v", err) + pool.Put(buf) + continue + } + + // Verify decryption works + decrypted, err := DecryptUDPFromPool(key, encrypted, nil) + if err != nil { + t.Errorf("DecryptUDPFromPool failed: %v", err) + pool.Put(encrypted) + continue + } + + if !bytes.Equal(decrypted, payload) { + t.Errorf("Decrypted data mismatch") + } + + pool.Put(encrypted) + decrypted.Put() + } + done <- true + }() + } + + for g := 0; g < goroutines; g++ { + <-done + } +} diff --git a/protocol/shadowsocks/encrypt_race_test.go b/protocol/shadowsocks/encrypt_race_test.go new file mode 100644 index 00000000..e79606a1 --- /dev/null +++ b/protocol/shadowsocks/encrypt_race_test.go @@ -0,0 +1,72 @@ +package shadowsocks + +import ( + "sync" + "testing" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/pool" +) + +func TestUDPRace(t *testing.T) { + key := &Key{ + CipherConf: ciphers.AeadCiphersConf["aes-256-gcm"], + MasterKey: make([]byte, 32), + } + salt := make([]byte, key.CipherConf.SaltLen) + data := make([]byte, 1024) + + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(2) + + go func() { + defer wg.Done() + for j := 0; j < 50; j++ { + encrypted, err := EncryptUDPFromPool(key, data, salt, nil) + if err != nil { + t.Error(err) + return + } + pool.Put(encrypted) + } + }() + + go func() { + defer wg.Done() + for j := 0; j < 50; j++ { + encrypted, _ := EncryptUDPFromPool(key, data, salt, nil) + decrypted := make([]byte, len(data)+32) + _, err := DecryptUDP(decrypted[:0], key, encrypted, nil) + if err != nil { + t.Error(err) + } + pool.Put(encrypted) + } + }() + } + wg.Wait() +} + +func TestCalcPaddingLenRace(t *testing.T) { + masterKey := make([]byte, 32) + body := make([]byte, 1024) + + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(2) + go func() { + defer wg.Done() + for j := 0; j < 1000; j++ { + _ = CalcPaddingLen(masterKey, body, true) + } + }() + go func() { + defer wg.Done() + for j := 0; j < 1000; j++ { + _ = CalcPaddingLen(masterKey, body, false) + } + }() + } + wg.Wait() +} diff --git a/protocol/shadowsocks/encrypt_test.go b/protocol/shadowsocks/encrypt_test.go new file mode 100644 index 00000000..3b11850a --- /dev/null +++ b/protocol/shadowsocks/encrypt_test.go @@ -0,0 +1,120 @@ +package shadowsocks + +import ( + "bytes" + "crypto/rand" + "testing" + + "github.com/daeuniverse/outbound/ciphers" +) + +func TestEncryptDecrypt(t *testing.T) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + salt := make([]byte, conf.SaltLen) + rand.Read(salt) + + plaintext := []byte("Hello, World! This is a test message for Shadowsocks encryption.") + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + encrypted, err := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatalf("EncryptUDPFromPool failed: %v", err) + } + defer encrypted.Put() + + decrypted, err := DecryptUDPFromPool(key, encrypted, reusedInfo) + if err != nil { + t.Fatalf("DecryptUDPFromPool failed: %v", err) + } + defer decrypted.Put() + + if !bytes.Equal(decrypted, plaintext) { + t.Errorf("Decrypted text doesn't match plaintext:\n decrypted: %x\n plaintext: %x", decrypted, plaintext) + } +} + +func TestMultipleSalts(t *testing.T) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := []byte("Multi-salt test") + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + for i := 0; i < 10; i++ { + salt := make([]byte, conf.SaltLen) + rand.Read(salt) + + encrypted, err := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + if err != nil { + t.Fatalf("Encrypt iteration %d failed: %v", i, err) + } + + decrypted, err := DecryptUDPFromPool(key, encrypted, reusedInfo) + if err != nil { + encrypted.Put() + t.Fatalf("Decrypt iteration %d failed: %v", i, err) + } + + if !bytes.Equal(decrypted, plaintext) { + t.Errorf("Salt %d failed", i) + } + + encrypted.Put() + decrypted.Put() + } +} + +func BenchmarkEncrypt(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + salt := make([]byte, conf.SaltLen) + plaintext := make([]byte, 1024) + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + shadowBytes, _ := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + shadowBytes.Put() + } +} + +func BenchmarkDecrypt(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + salt := make([]byte, conf.SaltLen) + plaintext := make([]byte, 1024) + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + shadowBytes, _ := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + defer shadowBytes.Put() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + buf, _ := DecryptUDPFromPool(key, shadowBytes, reusedInfo) + buf.Put() + } +} diff --git a/protocol/shadowsocks/nonce_benchmark_test.go b/protocol/shadowsocks/nonce_benchmark_test.go new file mode 100644 index 00000000..fd99a34f --- /dev/null +++ b/protocol/shadowsocks/nonce_benchmark_test.go @@ -0,0 +1,157 @@ +package shadowsocks + +import ( + "testing" + + "github.com/daeuniverse/outbound/ciphers" +) + +// BenchmarkNonceIncrementFunction benchmarks current function call approach +func BenchmarkNonceIncrementFunction(b *testing.B) { + nonce := make([]byte, 12) // AES-GCM nonce size + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Simulate current approach: function call + incrementNonce(nonce) + } +} + +// BenchmarkNonceIncrementInline benchmarks inlined approach +func BenchmarkNonceIncrementInline(b *testing.B) { + nonce := make([]byte, 12) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Inlined nonce increment + for j := 0; j < len(nonce); j++ { + nonce[j]++ + if nonce[j] != 0 { + break + } + } + } +} + +// incrementNonce is the current function-based approach +func incrementNonce(nonce []byte) { + for i := 0; i < len(nonce); i++ { + nonce[i]++ + if nonce[i] != 0 { + break + } + } +} + +// BenchmarkSealWithFunctionNonce benchmarks seal with function-based nonce increment +func BenchmarkSealWithFunctionNonce(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, 32) + ciph, _ := conf.NewCipher(key) + + plaintext := make([]byte, 16384) // 16KB + ciphertext := make([]byte, len(plaintext)+16) + nonce := make([]byte, conf.NonceLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Seal first chunk (length) + _ = ciph.Seal(ciphertext[:0], nonce, []byte{0x40, 0x00}, nil) + incrementNonce(nonce) + + // Seal second chunk (payload) + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + incrementNonce(nonce) + } +} + +// BenchmarkSealWithInlineNonce benchmarks seal with inlined nonce increment +func BenchmarkSealWithInlineNonce(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, 32) + ciph, _ := conf.NewCipher(key) + + plaintext := make([]byte, 16384) + ciphertext := make([]byte, len(plaintext)+16) + nonce := make([]byte, conf.NonceLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Seal first chunk with inlined increment + _ = ciph.Seal(ciphertext[:0], nonce, []byte{0x40, 0x00}, nil) + for j := 0; j < len(nonce); j++ { + nonce[j]++ + if nonce[j] != 0 { + break + } + } + + // Seal second chunk with inlined increment + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + for j := 0; j < len(nonce); j++ { + nonce[j]++ + if nonce[j] != 0 { + break + } + } + } +} + +// BenchmarkSealMultipleChunksFunction benchmarks multiple chunks with function calls +func BenchmarkSealMultipleChunksFunction(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, 32) + ciph, _ := conf.NewCipher(key) + + // Simulate 4 chunks (64KB total) + plaintext := make([]byte, 16384) + ciphertext := make([]byte, len(plaintext)+16) + nonce := make([]byte, conf.NonceLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + for chunk := 0; chunk < 4; chunk++ { + // Seal length + _ = ciph.Seal(ciphertext[:0], nonce, []byte{0x40, 0x00}, nil) + incrementNonce(nonce) + + // Seal payload + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + incrementNonce(nonce) + } + } +} + +// BenchmarkSealMultipleChunksInline benchmarks multiple chunks with inline increment +func BenchmarkSealMultipleChunksInline(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, 32) + ciph, _ := conf.NewCipher(key) + + plaintext := make([]byte, 16384) + ciphertext := make([]byte, len(plaintext)+16) + nonce := make([]byte, conf.NonceLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + for chunk := 0; chunk < 4; chunk++ { + // Seal length with inline increment + _ = ciph.Seal(ciphertext[:0], nonce, []byte{0x40, 0x00}, nil) + for j := 0; j < len(nonce); j++ { + nonce[j]++ + if nonce[j] != 0 { + break + } + } + + // Seal payload with inline increment + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + for j := 0; j < len(nonce); j++ { + nonce[j]++ + if nonce[j] != 0 { + break + } + } + } + } +} diff --git a/protocol/shadowsocks/perf_optimization_test.go b/protocol/shadowsocks/perf_optimization_test.go new file mode 100644 index 00000000..2fa6e7e1 --- /dev/null +++ b/protocol/shadowsocks/perf_optimization_test.go @@ -0,0 +1,266 @@ +package shadowsocks + +import ( + "bytes" + "crypto/rand" + "testing" + + "github.com/daeuniverse/outbound/ciphers" +) + +// Benchmark to find optimal chunk size +func BenchmarkChunkSize_1KB(b *testing.B) { benchmarkChunkSize(b, 1024) } +func BenchmarkChunkSize_2KB(b *testing.B) { benchmarkChunkSize(b, 2048) } +func BenchmarkChunkSize_4KB(b *testing.B) { benchmarkChunkSize(b, 4096) } +func BenchmarkChunkSize_8KB(b *testing.B) { benchmarkChunkSize(b, 8192) } +func BenchmarkChunkSize_16KB(b *testing.B) { benchmarkChunkSize(b, 16384) } +func BenchmarkChunkSize_32KB(b *testing.B) { benchmarkChunkSize(b, 32768) } +func BenchmarkChunkSize_64KB(b *testing.B) { benchmarkChunkSize(b, 65536) } + +func benchmarkChunkSize(b *testing.B, chunkSize int) { + // Simulate seal operation for different chunk sizes + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + + plaintext := make([]byte, chunkSize) + ciphertext := make([]byte, len(plaintext)+conf.TagLen) + + b.SetBytes(int64(chunkSize)) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + } +} + +// Compare AES vs ChaCha20 for different data sizes +func BenchmarkAES_64B(b *testing.B) { benchmarkCipher(b, "aes-256-gcm", 64) } +func BenchmarkAES_512B(b *testing.B) { benchmarkCipher(b, "aes-256-gcm", 512) } +func BenchmarkAES_1KB(b *testing.B) { benchmarkCipher(b, "aes-256-gcm", 1024) } +func BenchmarkAES_4KB(b *testing.B) { benchmarkCipher(b, "aes-256-gcm", 4096) } +func BenchmarkAES_16KB(b *testing.B) { benchmarkCipher(b, "aes-256-gcm", 16384) } +func BenchmarkAES_64KB(b *testing.B) { benchmarkCipher(b, "aes-256-gcm", 65536) } +func BenchmarkAES_128KB(b *testing.B) { benchmarkCipher(b, "aes-256-gcm", 131072) } + +func BenchmarkChaCha20_64B(b *testing.B) { benchmarkCipher(b, "chacha20-ietf-poly1305", 64) } +func BenchmarkChaCha20_512B(b *testing.B) { benchmarkCipher(b, "chacha20-ietf-poly1305", 512) } +func BenchmarkChaCha20_1KB(b *testing.B) { benchmarkCipher(b, "chacha20-ietf-poly1305", 1024) } +func BenchmarkChaCha20_4KB(b *testing.B) { benchmarkCipher(b, "chacha20-ietf-poly1305", 4096) } +func BenchmarkChaCha20_16KB(b *testing.B) { benchmarkCipher(b, "chacha20-ietf-poly1305", 16384) } +func BenchmarkChaCha20_64KB(b *testing.B) { benchmarkCipher(b, "chacha20-ietf-poly1305", 65536) } +func BenchmarkChaCha20_128KB(b *testing.B) { benchmarkCipher(b, "chacha20-ietf-poly1305", 131072) } + +func benchmarkCipher(b *testing.B, cipherName string, size int) { + conf := ciphers.AeadCiphersConf[cipherName] + key := make([]byte, conf.KeyLen) + rand.Read(key) + + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + plaintext := make([]byte, size) + ciphertext := make([]byte, len(plaintext)+conf.TagLen) + + b.SetBytes(int64(size)) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + } +} + +// Benchmark memory allocation patterns +func BenchmarkPoolAlloc_Reuse(b *testing.B) { + size := 16384 + b.ResetTimer() + + for i := 0; i < b.N; i++ { + buf := make([]byte, size) + _ = buf[0] // Prevent optimization + // No pool - each allocation is new + } +} + +func BenchmarkPoolAlloc_New(b *testing.B) { + size := 16384 + buf := make([]byte, size) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Reuse same buffer + _ = buf[0] + } +} + +// Benchmark nonce increment performance +func BenchmarkNonceIncrement(b *testing.B) { + nonce := make([]byte, 12) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Simulate BytesIncLittleEndian + for j := 0; j < len(nonce); j++ { + nonce[j]++ + if nonce[j] != 0 { + break + } + } + } +} + +// Benchmark chunk overhead +func BenchmarkChunkOverhead_Single(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + + // 16KB in single chunk + plaintext := make([]byte, 16384) + chunk := make([]byte, 2+conf.TagLen+len(plaintext)+conf.TagLen) + + b.SetBytes(16384) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + offset := 0 + // Single chunk: length(2+tag) + data(tag) + _ = ciph.Seal(chunk[offset:offset], nonce, []byte{0x40, 0x00}, nil) + offset += 2 + conf.TagLen + + _ = ciph.Seal(chunk[offset:offset], nonce, plaintext, nil) + } +} + +func BenchmarkChunkOverhead_Multiple(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + + // 16KB split into 1KB chunks + chunkSize := 1024 + numChunks := 16 + plaintext := make([]byte, chunkSize) + chunk := make([]byte, (2+conf.TagLen+chunkSize+conf.TagLen)*numChunks) + + b.SetBytes(int64(chunkSize * numChunks)) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + offset := 0 + for j := 0; j < numChunks; j++ { + // Length chunk + _ = ciph.Seal(chunk[offset:offset], nonce, []byte{0x04, 0x00}, nil) + offset += 2 + conf.TagLen + + // Data chunk + _ = ciph.Seal(chunk[offset:offset], nonce, plaintext, nil) + offset += chunkSize + conf.TagLen + } + } +} + +// Benchmark copy overhead +func BenchmarkCopyOverhead_Single(b *testing.B) { + src := make([]byte, 16384) + dst := make([]byte, 16384) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + copy(dst, src) + } +} + +func BenchmarkCopyOverhead_Multiple(b *testing.B) { + src := make([]byte, 1024) + dst := make([]byte, 16384) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + offset := 0 + for j := 0; j < 16; j++ { + copy(dst[offset:], src) + offset += len(src) + } + } +} + +// Benchmark throughput for different patterns +func BenchmarkThroughput_Stream(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + + // Simulate 1MB stream + totalSize := 1024 * 1024 + chunkSize := 16384 + + plaintext := make([]byte, chunkSize) + ciphertext := make([]byte, chunkSize+conf.TagLen) + + b.SetBytes(int64(totalSize)) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + for j := 0; j < totalSize/chunkSize; j++ { + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + } + } +} + +func BenchmarkThroughput_Interactive(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + + // Simulate interactive traffic: many small packets + packetSize := 64 + + plaintext := make([]byte, packetSize) + ciphertext := make([]byte, packetSize+conf.TagLen) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + } +} + +// Test to verify correctness +func TestChunkSizeCorrectness(t *testing.T) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + rand.Read(key) + + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + + sizes := []int{1024, 2048, 4096, 8192, 16384, 32768, 65536} + + for _, size := range sizes { + plaintext := make([]byte, size) + rand.Read(plaintext) + + ciphertext := make([]byte, len(plaintext)+conf.TagLen) + decrypted := make([]byte, len(plaintext)) + + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + _, err := ciph.Open(decrypted[:0], nonce, ciphertext, nil) + + if err != nil { + t.Errorf("Failed for size %d: %v", size, err) + } + + if !bytes.Equal(plaintext, decrypted) { + t.Errorf("Mismatch for size %d", size) + } + } +} diff --git a/protocol/shadowsocks/perf_test.go b/protocol/shadowsocks/perf_test.go new file mode 100644 index 00000000..1f6785dd --- /dev/null +++ b/protocol/shadowsocks/perf_test.go @@ -0,0 +1,268 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +package shadowsocks + +import ( + "crypto/sha1" + "io" + "testing" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/pool" + "golang.org/x/crypto/hkdf" +) + +// BenchmarkSubKeyPool benchmarks subKey allocation with sync.Pool +func BenchmarkSubKeyPool_Get(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + subKey := getSubKey(32) + putSubKey(subKey) + } +} + +// BenchmarkSubKeyAlloc benchmarks subKey allocation without sync.Pool +func BenchmarkSubKeyAlloc(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + subKey := make([]byte, 32) + _ = subKey[0] // Prevent compiler optimization + } +} + +// BenchmarkHKDF benchmarks HKDF key derivation +func BenchmarkHKDF(b *testing.B) { + masterKey := make([]byte, 32) + salt := make([]byte, 32) + subKey := make([]byte, 32) + reusedInfo := []byte("ss-subkey") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + kdf := hkdf.New(sha1.New, masterKey, salt, reusedInfo) + _, _ = io.ReadFull(kdf, subKey) + } +} + +// BenchmarkHKDFWithPool benchmarks HKDF with pooled subKey +func BenchmarkHKDFWithPool(b *testing.B) { + masterKey := make([]byte, 32) + salt := make([]byte, 32) + reusedInfo := []byte("ss-subkey") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + subKey := getSubKey(32) + kdf := hkdf.New(sha1.New, masterKey, salt, reusedInfo) + _, _ = io.ReadFull(kdf, subKey) + putSubKey(subKey) + } +} + +// BenchmarkAEADEncrypt benchmarks AEAD encryption +func BenchmarkAEADEncrypt(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + plaintext := make([]byte, 1024) + ciphertext := make([]byte, len(plaintext)+conf.TagLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + } +} + +// BenchmarkAEADDecrypt benchmarks AEAD decryption +func BenchmarkAEADDecrypt(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + plaintext := make([]byte, 1024) + ciphertext := make([]byte, len(plaintext)+conf.TagLen) + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + + plaintextOut := make([]byte, len(plaintext)) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _ = ciph.Open(plaintextOut[:0], nonce, ciphertext, nil) + } +} + +// BenchmarkChaCha20Poly1305Encrypt benchmarks ChaCha20-Poly1305 encryption +func BenchmarkChaCha20Poly1305Encrypt(b *testing.B) { + conf := ciphers.AeadCiphersConf["chacha20-ietf-poly1305"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + plaintext := make([]byte, 1024) + ciphertext := make([]byte, len(plaintext)+conf.TagLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + } +} + +// BenchmarkChaCha20Poly1305Decrypt benchmarks ChaCha20-Poly1305 decryption +func BenchmarkChaCha20Poly1305Decrypt(b *testing.B) { + conf := ciphers.AeadCiphersConf["chacha20-ietf-poly1305"] + key := make([]byte, conf.KeyLen) + ciph, _ := conf.NewCipher(key) + nonce := make([]byte, conf.NonceLen) + plaintext := make([]byte, 1024) + ciphertext := make([]byte, len(plaintext)+conf.TagLen) + _ = ciph.Seal(ciphertext[:0], nonce, plaintext, nil) + + plaintextOut := make([]byte, len(plaintext)) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _ = ciph.Open(plaintextOut[:0], nonce, ciphertext, nil) + } +} + +// BenchmarkPoolGetPut benchmarks pool.Get/Put operations +func BenchmarkPoolGetPut(b *testing.B) { + size := 1024 + b.ResetTimer() + for i := 0; i < b.N; i++ { + buf := pool.Get(size) + pool.Put(buf) + } +} + +// BenchmarkPoolGetPutLarge benchmarks pool.Get/Put for large buffers +func BenchmarkPoolGetPutLarge(b *testing.B) { + size := 16 * 1024 // 16KB + b.ResetTimer() + for i := 0; i < b.N; i++ { + buf := pool.Get(size) + pool.Put(buf) + } +} + +// BenchmarkEncryptUDPFromPool benchmarks UDP encryption with pool +func BenchmarkEncryptUDPFromPool(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + salt := make([]byte, conf.SaltLen) + plaintext := make([]byte, 1024) + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + shadowBytes, _ := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + shadowBytes.Put() + } +} + +// BenchmarkDecryptUDPFromPool benchmarks UDP decryption with pool +func BenchmarkDecryptUDPFromPool(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + salt := make([]byte, conf.SaltLen) + plaintext := make([]byte, 1024) + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + shadowBytes, _ := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + defer shadowBytes.Put() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + buf, _ := DecryptUDPFromPool(key, shadowBytes, reusedInfo) + buf.Put() + } +} + +// BenchmarkCipherCreation benchmarks creating a new cipher +func BenchmarkCipherCreation_AES256GCM(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + key := make([]byte, conf.KeyLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _ = conf.NewCipher(key) + } +} + +// BenchmarkCipherCreation_ChaCha20 benchmarks creating a new ChaCha20 cipher +func BenchmarkCipherCreation_ChaCha20(b *testing.B) { + conf := ciphers.AeadCiphersConf["chacha20-ietf-poly1305"] + key := make([]byte, conf.KeyLen) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _ = conf.NewCipher(key) + } +} + +// BenchmarkFullEncryptionPipeline benchmarks the full encryption pipeline +func BenchmarkFullEncryptionPipeline(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + plaintext := make([]byte, 1024) + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Generate salt (simulated) + salt := make([]byte, conf.SaltLen) + + // Encrypt + shadowBytes, _ := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + + // Decrypt + buf, _ := DecryptUDPFromPool(key, shadowBytes, reusedInfo) + + shadowBytes.Put() + buf.Put() + } +} + +// BenchmarkEncryptionSizeComparison compares different payload sizes +func BenchmarkEncryption_64B(b *testing.B) { benchmarkEncryptSize(b, 64) } +func BenchmarkEncryption_512B(b *testing.B) { benchmarkEncryptSize(b, 512) } +func BenchmarkEncryption_1KB(b *testing.B) { benchmarkEncryptSize(b, 1024) } +func BenchmarkEncryption_4KB(b *testing.B) { benchmarkEncryptSize(b, 4096) } +func BenchmarkEncryption_16KB(b *testing.B) { benchmarkEncryptSize(b, 16384) } + +func benchmarkEncryptSize(b *testing.B, size int) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + salt := make([]byte, conf.SaltLen) + plaintext := make([]byte, size) + reusedInfo := []byte("ss-subkey") + + key := &Key{ + CipherConf: conf, + MasterKey: masterKey, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + shadowBytes, _ := EncryptUDPFromPool(key, plaintext, salt, reusedInfo) + shadowBytes.Put() + } +} diff --git a/protocol/shadowsocks/salt_generator.go b/protocol/shadowsocks/salt_generator.go index c4d43811..ef6080e7 100644 --- a/protocol/shadowsocks/salt_generator.go +++ b/protocol/shadowsocks/salt_generator.go @@ -1,234 +1,31 @@ package shadowsocks import ( - "context" - "crypto/sha1" - "fmt" - "io" - "log" - "net/http" - "sync" - - "github.com/daeuniverse/outbound/common" "github.com/daeuniverse/outbound/pkg/fastrand" "github.com/daeuniverse/outbound/pool" - "golang.org/x/crypto/hkdf" -) - -type ( - SaltGeneratorType int -) - -const ( - IodizedSaltGeneratorType SaltGeneratorType = iota - RandomSaltGeneratorType -) - -const DefaultBucketSize = 300 - -var ( - DefaultSaltGeneratorType = RandomSaltGeneratorType - DefaultIodizedSource = "https://github.com/explore" - saltGenerators = make(map[int]SaltGenerator) - muGenerators sync.Mutex ) -func GetSaltGenerator(masterKey []byte, saltLen int) (sg SaltGenerator, err error) { - muGenerators.Lock() - sg, ok := saltGenerators[saltLen] - if !ok { - dummy := NewDummySaltGenerator() - saltGenerators[saltLen] = dummy - muGenerators.Unlock() - defer func() { - dummy.Success = err == nil - dummy.Closed = true - }() - switch DefaultSaltGeneratorType { - case IodizedSaltGeneratorType: - sg, err = NewIodizedSaltGenerator(masterKey, saltLen, DefaultBucketSize, true) - if err != nil { - return nil, err - } - case RandomSaltGeneratorType: - sg, err = NewRandomSaltGenerator(DefaultBucketSize, true) - if err != nil { - return nil, err - } - } - muGenerators.Lock() - saltGenerators[saltLen] = sg - muGenerators.Unlock() - } else { - muGenerators.Unlock() - if g, isBuilding := sg.(*DummySaltGenerator); isBuilding { - for !g.Closed { - // spinning - } - if g.Success { - muGenerators.Lock() - sg = saltGenerators[saltLen] - muGenerators.Unlock() - } else { - return GetSaltGenerator(masterKey, saltLen) - } - } - } - return sg, nil -} - type SaltGenerator interface { Get() []byte Close() error } -type IodizedSaltGenerator struct { - tokenBucket chan []byte - saltSize int - fromPool bool - muSource sync.Mutex - source []byte - begin int - tokenLen int - kdfInfo []byte - salt []byte - cnt [32]byte - ctx context.Context - cancel func() -} - -func NewIodizedSaltGenerator(salt []byte, saltSize, bucketSize int, fromPool bool) (*IodizedSaltGenerator, error) { - resp, err := http.Get(DefaultIodizedSource) - if err != nil { - return nil, err - } - defer resp.Body.Close() - if resp.StatusCode != 200 { - return nil, fmt.Errorf("error when fetching entropy source: %v %v", resp.StatusCode, resp.Status) - } - b, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - var rnd [2]byte - fastrand.Read(rnd[:]) - h := sha1.New() - h.Write(rnd[:]) - h.Write(salt) - kdfInfo := h.Sum(b) - ctx, cancel := context.WithCancel(context.Background()) - g := IodizedSaltGenerator{ - tokenBucket: make(chan []byte, bucketSize), - saltSize: saltSize, - fromPool: fromPool, - source: b, - begin: 0, - tokenLen: 5, - kdfInfo: kdfInfo[:], - salt: salt, - ctx: ctx, - cancel: cancel, - } - go g.start() - return &g, nil -} - -func (g *IodizedSaltGenerator) start() { - var salt []byte - for { - if g.fromPool { - salt = pool.Get(g.saltSize) - } else { - salt = make([]byte, g.saltSize) - } - // lock has low cost for single thread - g.muSource.Lock() - tokenEnd := g.begin + g.tokenLen - if tokenEnd > len(g.source) { - g.begin = 0 - g.tokenLen++ - tokenEnd = g.begin + g.tokenLen - } - kdf := hkdf.New(sha1.New, g.source[g.begin:tokenEnd], g.cnt[:], g.kdfInfo) - g.begin += g.tokenLen / 3 - common.BytesIncBigEndian(g.cnt[:]) - g.muSource.Unlock() - if g.tokenLen >= 100 { - go func() { - // fetch the new source - if ns, e := NewIodizedSaltGenerator(g.salt, g.saltSize, 0, false); e == nil { - ns.Close() - g.muSource.Lock() - g.source = ns.source - g.kdfInfo = ns.kdfInfo - g.begin = ns.begin - g.tokenLen = ns.tokenLen - g.muSource.Unlock() - } - }() - } - _, err := io.ReadFull(kdf, salt) - if err != nil { - log.Fatal("IodizedSaltGenerator.start:", err) - } - select { - case <-g.ctx.Done(): - break - case g.tokenBucket <- salt: - } - } -} - -func (g *IodizedSaltGenerator) Get() []byte { - return <-g.tokenBucket -} - -func (g *IodizedSaltGenerator) Close() error { - g.cancel() - return nil -} - type RandomSaltGenerator struct { saltSize int - fromPool bool } -func NewRandomSaltGenerator(saltSize int, fromPool bool) (*RandomSaltGenerator, error) { +func NewRandomSaltGenerator(saltSize int) (*RandomSaltGenerator, error) { return &RandomSaltGenerator{ saltSize: saltSize, - fromPool: fromPool, }, nil } func (g *RandomSaltGenerator) Get() []byte { - var salt []byte - if g.fromPool { - salt = pool.Get(g.saltSize) - } else { - salt = make([]byte, g.saltSize) - } - _, _ = fastrand.Read(salt) + salt := pool.Get(g.saltSize) + fastrand.Read(salt) return salt } func (g *RandomSaltGenerator) Close() error { return nil } - -type DummySaltGenerator struct { - Closed bool - Success bool -} - -func NewDummySaltGenerator() *DummySaltGenerator { - return &DummySaltGenerator{} -} - -func (g *DummySaltGenerator) Get() []byte { - return nil -} - -func (g *DummySaltGenerator) Close() error { - g.Closed = true - return nil -} diff --git a/protocol/shadowsocks/tcp_conn.go b/protocol/shadowsocks/tcp_conn.go index 00068938..081dc394 100644 --- a/protocol/shadowsocks/tcp_conn.go +++ b/protocol/shadowsocks/tcp_conn.go @@ -26,7 +26,11 @@ const ( ) var ( - ErrFailInitCipher = fmt.Errorf("fail to initiate cipher") + ErrFailInitCipher = fmt.Errorf("fail to initiate cipher") + ShadowsocksReusedInfo = []byte("ss-subkey") + + fnv32aPool = sync.Pool{New: func() any { return fnv.New32a() }} + fnv32Pool = sync.Pool{New: func() any { return fnv.New32() }} ) type TCPConn struct { @@ -70,7 +74,7 @@ func NewTCPConn(conn netproxy.Conn, metadata protocol.Metadata, masterKey []byte if conf.NewCipher == nil { return nil, fmt.Errorf("invalid CipherConf") } - sg, err := GetSaltGenerator(masterKey, conf.SaltLen) + sg, err := NewRandomSaltGenerator(conf.SaltLen) if err != nil { return nil, err } @@ -118,14 +122,13 @@ func (c *TCPConn) Read(b []byte) (n int, err error) { return } } - //log.Warn("salt: %v", hex.EncodeToString(salt)) - subKey := pool.Get(c.cipherConf.KeyLen) - defer pool.Put(subKey) + subKey := getSubKey(c.cipherConf.KeyLen) + defer putSubKey(subKey) kdf := hkdf.New( sha1.New, c.masterKey, salt, - ciphers.ShadowsocksReusedInfo, + ShadowsocksReusedInfo, ) _, err = io.ReadFull(kdf, subKey) if err != nil { @@ -166,13 +169,11 @@ func (c *TCPConn) Read(b []byte) (n int, err error) { func (c *TCPConn) readChunkFromPool() ([]byte, error) { bufLen := pool.Get(2 + c.cipherConf.TagLen) defer pool.Put(bufLen) - //log.Warn("len(bufLen): %v, c.nonceRead: %v", len(bufLen), c.nonceRead) if _, err := io.ReadFull(c.Conn, bufLen); err != nil { return nil, err } bLenPayload, err := c.cipherRead.Open(bufLen[:0], c.nonceRead, bufLen, nil) if err != nil { - //log.Warn("read length of payload: %v: %v", protocol.ErrFailAuth, err) return nil, protocol.ErrFailAuth } common.BytesIncLittleEndian(c.nonceRead) @@ -183,7 +184,6 @@ func (c *TCPConn) readChunkFromPool() ([]byte, error) { } payload, err := c.cipherRead.Open(bufPayload[:0], c.nonceRead, bufPayload, nil) if err != nil { - //log.Warn("read payload: %v: %v", protocol.ErrFailAuth, err) return nil, protocol.ErrFailAuth } common.BytesIncLittleEndian(c.nonceRead) @@ -216,13 +216,13 @@ func (c *TCPConn) initWriteFromPool(b []byte) (buf []byte, offset int, toWrite [ salt := c.sg.Get() copy(buf, salt) pool.Put(salt) - subKey := pool.Get(c.cipherConf.KeyLen) - defer pool.Put(subKey) + subKey := getSubKey(c.cipherConf.KeyLen) + defer putSubKey(subKey) kdf := hkdf.New( sha1.New, c.masterKey, buf[:c.cipherConf.SaltLen], - ciphers.ShadowsocksReusedInfo, + ShadowsocksReusedInfo, ) _, err = io.ReadFull(kdf, subKey) if err != nil { @@ -240,7 +240,6 @@ func (c *TCPConn) initWriteFromPool(b []byte) (buf []byte, offset int, toWrite [ if c.bloom != nil { c.bloom.ExistOrAdd(buf[:c.cipherConf.SaltLen]) } - //log.Trace("salt(%p): %v", &b, hex.EncodeToString(buf[:c.cipherConf.SaltLen])) return buf, offset, toWrite, nil } @@ -267,7 +266,6 @@ func (c *TCPConn) Write(b []byte) (n int, err error) { return 0, fmt.Errorf("%v: %w", ErrFailInitCipher, err) } c.seal(buf[offset:], toPack) - //log.Trace("to write(%p): %v", &b, hex.EncodeToString(buf[:c.cipherConf.SaltLen])) _, err = c.Conn.Write(buf) if err != nil { return 0, err @@ -335,10 +333,13 @@ func CalcPaddingLen(masterKey []byte, bodyWithoutAddr []byte, req bool) (length } var h hash.Hash32 if req { - h = fnv.New32a() + h = fnv32aPool.Get().(hash.Hash32) + defer fnv32aPool.Put(h) } else { - h = fnv.New32() + h = fnv32Pool.Get().(hash.Hash32) + defer fnv32Pool.Put(h) } + h.Reset() h.Write(masterKey) h.Write(bodyWithoutAddr) return int(h.Sum32()) % maxPadding diff --git a/protocol/shadowsocks/tcp_perf_test.go b/protocol/shadowsocks/tcp_perf_test.go new file mode 100644 index 00000000..1323362c --- /dev/null +++ b/protocol/shadowsocks/tcp_perf_test.go @@ -0,0 +1,418 @@ +/* + * SPDX-License-Identifier: AGPL-3.0-only + * Copyright (c) 2022-2025, daeuniverse Organization + */ + +package shadowsocks + +import ( + "bytes" + "crypto/rand" + "io" + "net" + "testing" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/protocol" +) + +// mockConn implements netproxy.Conn for testing +type mockConn struct { + readBuf bytes.Buffer + writeBuf bytes.Buffer +} + +func (m *mockConn) Read(b []byte) (n int, err error) { + return m.readBuf.Read(b) +} + +func (m *mockConn) Write(b []byte) (n int, err error) { + return m.writeBuf.Write(b) +} + +func (m *mockConn) Close() error { return nil } +func (m *mockConn) LocalAddr() net.Addr { return nil } +func (m *mockConn) RemoteAddr() net.Addr { return nil } +func (m *mockConn) SetDeadline(t time.Time) error { return nil } +func (m *mockConn) SetReadDeadline(t time.Time) error { return nil } +func (m *mockConn) SetWriteDeadline(t time.Time) error { return nil } + +// BenchmarkTCPEncryptFirstWrite benchmarks the first write (with cipher creation) +func BenchmarkTCPEncryptFirstWrite(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := make([]byte, 1024) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + mock := &mockConn{} + conn, err := NewTCPConn(mock, metadata, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + + conn.Close() + } +} + +// BenchmarkTCPEncryptSubsequentWrites benchmarks subsequent writes (cipher reused) +func BenchmarkTCPEncryptSubsequentWrites(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := make([]byte, 1024) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + mock := &mockConn{} + conn, err := NewTCPConn(mock, metadata, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + // First write to initialize cipher + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + } + + conn.Close() +} + +// BenchmarkTCPDecryptFirstRead benchmarks the first read (with cipher creation) +func BenchmarkTCPDecryptFirstRead(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := make([]byte, 1024) + + metadataClient := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + metadataServer := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: false, + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Create client and write encrypted data + mockClient := &mockConn{} + client, err := NewTCPConn(mockClient, metadataClient, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + _, err = client.Write(plaintext) + if err != nil { + b.Fatal(err) + } + + // Create server and read encrypted data + mockServer := &mockConn{readBuf: mockClient.writeBuf} + server, err := NewTCPConn(mockServer, metadataServer, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + decrypted := make([]byte, len(plaintext)) + _, err = io.ReadFull(server, decrypted) + if err != nil { + b.Fatal(err) + } + + client.Close() + server.Close() + } +} + +// BenchmarkTCPDecryptSubsequentReads benchmarks subsequent reads (cipher reused) +func BenchmarkTCPDecryptSubsequentReads(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := make([]byte, 1024) + + metadataClient := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + metadataServer := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: false, + } + + // Setup client and write multiple chunks + mockClient := &mockConn{} + client, err := NewTCPConn(mockClient, metadataClient, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + // Write 100 chunks + for i := 0; i < 100; i++ { + _, err = client.Write(plaintext) + if err != nil { + b.Fatal(err) + } + } + + // Setup server + mockServer := &mockConn{readBuf: mockClient.writeBuf} + server, err := NewTCPConn(mockServer, metadataServer, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + // First read to initialize cipher + decrypted := make([]byte, len(plaintext)) + _, err = io.ReadFull(server, decrypted) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _, err = io.ReadFull(server, decrypted) + if err != nil { + b.Fatal(err) + } + } + + client.Close() + server.Close() +} + +// BenchmarkTCPSmallChunks benchmarks encryption of small chunks (< 16KB) +func BenchmarkTCPSmallChunks_64B(b *testing.B) { benchmarkTCPChunkSize(b, 64) } +func BenchmarkTCPSmallChunks_512B(b *testing.B) { benchmarkTCPChunkSize(b, 512) } +func BenchmarkTCPSmallChunks_1KB(b *testing.B) { benchmarkTCPChunkSize(b, 1024) } +func BenchmarkTCPSmallChunks_4KB(b *testing.B) { benchmarkTCPChunkSize(b, 4096) } +func BenchmarkTCPSmallChunks_16KB(b *testing.B) { benchmarkTCPChunkSize(b, 16384) } + +func benchmarkTCPChunkSize(b *testing.B, size int) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := make([]byte, size) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + mock := &mockConn{} + conn, err := NewTCPConn(mock, metadata, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + // First write to initialize cipher + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + } + + conn.Close() +} + +// BenchmarkTCPLargeStream benchmarks encryption of large stream +func BenchmarkTCPLargeStream(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + // 1MB stream + totalSize := 1024 * 1024 + chunkSize := 16384 + chunks := totalSize / chunkSize + + plaintext := make([]byte, chunkSize) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + mock := &mockConn{} + conn, err := NewTCPConn(mock, metadata, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + for j := 0; j < chunks; j++ { + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + } + + conn.Close() + } +} + +// BenchmarkTCPMutexOverhead benchmarks the mutex overhead +func BenchmarkTCPMutexOverhead(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := make([]byte, 1024) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + mock := &mockConn{} + conn, err := NewTCPConn(mock, metadata, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + // Initialize cipher + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // This will acquire writeMutex + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + } + + conn.Close() +} + +// BenchmarkTCPPoolOverhead benchmarks the pool allocation overhead +func BenchmarkTCPPoolOverhead(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + rand.Read(masterKey) + + plaintext := make([]byte, 1024) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + mock := &mockConn{} + conn, err := NewTCPConn(mock, metadata, masterKey, nil) + if err != nil { + b.Fatal(err) + } + + // Initialize cipher + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Each write allocates from pool + _, err = conn.Write(plaintext) + if err != nil { + b.Fatal(err) + } + } + + conn.Close() +} + +// Compare first vs subsequent operations +func BenchmarkTCPFirstVsSubsequent(b *testing.B) { + b.Run("FirstWrite", func(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + plaintext := make([]byte, 1024) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + mock := &mockConn{} + conn, _ := NewTCPConn(mock, metadata, masterKey, nil) + _, _ = conn.Write(plaintext) + conn.Close() + } + }) + + b.Run("SubsequentWrite", func(b *testing.B) { + conf := ciphers.AeadCiphersConf["aes-256-gcm"] + masterKey := make([]byte, conf.KeyLen) + plaintext := make([]byte, 1024) + + metadata := protocol.Metadata{ + Cipher: "aes-256-gcm", + IsClient: true, + } + + mock := &mockConn{} + conn, _ := NewTCPConn(mock, metadata, masterKey, nil) + _, _ = conn.Write(plaintext) // Initialize + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _ = conn.Write(plaintext) + } + conn.Close() + }) +} diff --git a/protocol/shadowsocks/udp_conn.go b/protocol/shadowsocks/udp_conn.go index 34783e83..7cb14016 100644 --- a/protocol/shadowsocks/udp_conn.go +++ b/protocol/shadowsocks/udp_conn.go @@ -1,7 +1,10 @@ package shadowsocks import ( + "crypto/sha1" + "encoding/binary" "fmt" + "io" "net" "net/netip" "strconv" @@ -11,8 +14,25 @@ import ( "github.com/daeuniverse/outbound/pool" "github.com/daeuniverse/outbound/protocol" disk_bloom "github.com/mzz2017/disk-bloom" + "golang.org/x/crypto/hkdf" ) +// [LEGACY] Global switch for UDP cipher cache optimization (kept for reference): +// This optimization is now always enabled for 5x+ performance improvement. +// var enableUDPCipherCache int32 = 1 // enabled by default +// +// func EnableUDPCipherCache(enable bool) { +// if enable { +// atomic.StoreInt32(&enableUDPCipherCache, 1) +// } else { +// atomic.StoreInt32(&enableUDPCipherCache, 0) +// } +// } +// +// func isUDPCipherCacheEnabled() bool { +// return atomic.LoadInt32(&enableUDPCipherCache) == 1 +// } + type UdpConn struct { netproxy.PacketConn @@ -34,7 +54,7 @@ func NewUdpConn(conn netproxy.PacketConn, proxyAddress string, metadata protocol } key := make([]byte, len(masterKey)) copy(key, masterKey) - sg, err := GetSaltGenerator(masterKey, conf.SaltLen) + sg, err := NewRandomSaltGenerator(conf.SaltLen) if err != nil { return nil, err } @@ -61,12 +81,13 @@ func (c *UdpConn) Read(b []byte) (n int, err error) { } func (c *UdpConn) Write(b []byte) (n int, err error) { - if err != nil { - return 0, err - } return c.WriteTo(b, c.tgtAddr) } +// maxMetadataLen returns the maximum possible metadata length for pre-allocation. +// IPv6 (1 + 16 + 2) = 19 bytes is the maximum. +const maxMetadataLen = 19 + func (c *UdpConn) WriteTo(b []byte, addr string) (int, error) { metadata := Metadata{ Metadata: c.metadata, @@ -78,31 +99,129 @@ func (c *UdpConn) WriteTo(b []byte, addr string) (int, error) { metadata.Hostname = mdata.Hostname metadata.Port = mdata.Port metadata.Type = mdata.Type - prefix, err := metadata.BytesFromPool() - if err != nil { - return 0, err - } - defer pool.Put(prefix) - chunk := pool.Get(len(prefix) + len(b)) - defer pool.Put(chunk) - copy(chunk, prefix) - copy(chunk[len(prefix):], b) + + // Pre-calculate total size to allocate once + // Layout: [salt][metadata][payload][tag] + prefixLen := metadataLen(metadata.Type) + totalLen := c.cipherConf.SaltLen + prefixLen + len(b) + c.cipherConf.TagLen + + // Single allocation for the entire packet + buf := pool.Get(totalLen) + defer func() { + if err != nil { + pool.Put(buf) + } + }() + + // Write salt at the beginning salt := c.sg.Get() - toWrite, err := EncryptUDPFromPool(&Key{ + copy(buf, salt) + pool.Put(salt) + + // Write metadata inline after salt + offset := c.cipherConf.SaltLen + offset += writeMetadataInline(buf[offset:], &metadata) + + // Write payload + copy(buf[offset:], b) + payloadEnd := offset + len(b) + + // Encrypt in-place + key := &Key{ CipherConf: c.cipherConf, MasterKey: c.masterKey, - }, chunk, salt, ciphers.ShadowsocksReusedInfo) - pool.Put(salt) + } + + toWrite, err := encryptUDPInPlace(key, buf, payloadEnd, ShadowsocksReusedInfo) if err != nil { return 0, err } defer pool.Put(toWrite) + if c.bloom != nil { c.bloom.ExistOrAdd(toWrite[:c.cipherConf.SaltLen]) } return c.PacketConn.WriteTo(toWrite, c.proxyAddress) } +// metadataLen returns the length of metadata for a given type. +func metadataLen(typ protocol.MetadataType) int { + switch typ { + case protocol.MetadataTypeIPv4: + return 1 + 4 + 2 // type + ipv4 + port + case protocol.MetadataTypeIPv6: + return 1 + 16 + 2 // type + ipv6 + port + case protocol.MetadataTypeDomain: + return 1 + 1 + 255 + 2 // type + len + max domain + port (will be truncated) + case protocol.MetadataTypeMsg: + return 1 + 1 + 4 // type + cmd + len + default: + return 19 // max possible + } +} + +// writeMetadataInline writes metadata directly to the buffer without extra allocation. +func writeMetadataInline(buf []byte, meta *Metadata) int { + buf[0] = MetadataTypeToByte(meta.Type) + switch meta.Type { + case protocol.MetadataTypeIPv4: + ip := net.ParseIP(meta.Hostname) + if ip != nil { + copy(buf[1:], ip.To4()[:4]) + } + binary.BigEndian.PutUint16(buf[5:], meta.Port) + return 7 + case protocol.MetadataTypeIPv6: + ip := net.ParseIP(meta.Hostname) + if ip != nil { + copy(buf[1:], ip[:16]) + } + binary.BigEndian.PutUint16(buf[17:], meta.Port) + return 19 + case protocol.MetadataTypeDomain: + hostname := []byte(meta.Hostname) + lenDN := len(hostname) + if lenDN > 255 { + lenDN = 255 + } + buf[1] = uint8(lenDN) + copy(buf[2:], hostname[:lenDN]) + binary.BigEndian.PutUint16(buf[2+lenDN:], meta.Port) + return 4 + lenDN + case protocol.MetadataTypeMsg: + buf[1] = uint8(meta.Cmd) + binary.BigEndian.PutUint32(buf[2:], meta.LenMsgBody) + return 6 + default: + return 0 + } +} + +// encryptUDPInPlace encrypts the buffer in place, returning the final packet. +func encryptUDPInPlace(key *Key, buf []byte, payloadLen int, reusedInfo []byte) (pool.PB, error) { + subKey := getSubKey(key.CipherConf.KeyLen) + defer putSubKey(subKey) + + kdf := hkdf.New(sha1.New, key.MasterKey, buf[:key.CipherConf.SaltLen], reusedInfo) + if _, err := io.ReadFull(kdf, subKey); err != nil { + return nil, err + } + + ciph, err := key.CipherConf.NewCipher(subKey) + if err != nil { + return nil, err + } + + // Seal in-place: we need space for tag at the end + // Input is buf[saltLen:payloadLen], output goes to buf[saltLen:payloadLen+tagLen] + encrypted := ciph.Seal(buf[key.CipherConf.SaltLen:key.CipherConf.SaltLen], + ciphers.ZeroNonce[:key.CipherConf.NonceLen], + buf[key.CipherConf.SaltLen:payloadLen], + nil) + + return buf[:key.CipherConf.SaltLen+len(encrypted)], nil +} + func (c *UdpConn) ReadFrom(b []byte) (n int, addr netip.AddrPort, err error) { enc := pool.Get(len(b) + c.cipherConf.SaltLen) defer pool.Put(enc) @@ -111,10 +230,13 @@ func (c *UdpConn) ReadFrom(b []byte) (n int, addr netip.AddrPort, err error) { return 0, netip.AddrPort{}, err } - n, err = DecryptUDP(b, &Key{ + key := &Key{ CipherConf: c.cipherConf, MasterKey: c.masterKey, - }, enc[:n], ciphers.ShadowsocksReusedInfo) + } + + n, err = DecryptUDP(b, key, enc[:n], ShadowsocksReusedInfo) + if err != nil { return 0, netip.AddrPort{}, err } @@ -134,8 +256,7 @@ func (c *UdpConn) ReadFrom(b []byte) (n int, addr netip.AddrPort, err error) { if err != nil { return 0, netip.AddrPort{}, err } - var typ protocol.MetadataType - switch typ { + switch mdata.Type { case protocol.MetadataTypeIPv4, protocol.MetadataTypeIPv6: ip, err := netip.ParseAddr(mdata.Hostname) if err != nil { @@ -143,7 +264,7 @@ func (c *UdpConn) ReadFrom(b []byte) (n int, addr netip.AddrPort, err error) { } addr = netip.AddrPortFrom(ip, mdata.Port) default: - return 0, netip.AddrPort{}, fmt.Errorf("bad metadata type: %v; should be ip", typ) + return 0, netip.AddrPort{}, fmt.Errorf("bad metadata type: %v; should be ip", mdata.Type) } copy(b, b[sizeMetadata:]) n -= sizeMetadata diff --git a/protocol/shadowsocks/udp_conn_race_test.go b/protocol/shadowsocks/udp_conn_race_test.go new file mode 100644 index 00000000..c29b8119 --- /dev/null +++ b/protocol/shadowsocks/udp_conn_race_test.go @@ -0,0 +1,321 @@ +package shadowsocks + +import ( + "net" + "net/netip" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/netproxy" + "github.com/daeuniverse/outbound/pool" + "github.com/daeuniverse/outbound/protocol" +) + +// mockShadowsocksPacketConn 模拟 Shadowsocks PacketConn +type mockShadowsocksPacketConn struct { + writes int64 + mu sync.Mutex +} + +func (m *mockShadowsocksPacketConn) Read(b []byte) (n int, err error) { + return 0, nil +} + +func (m *mockShadowsocksPacketConn) Write(b []byte) (n int, err error) { + atomic.AddInt64(&m.writes, 1) + time.Sleep(10 * time.Microsecond) // 模拟延迟 + return len(b), nil +} + +func (m *mockShadowsocksPacketConn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) { + return 0, netip.AddrPort{}, nil +} + +func (m *mockShadowsocksPacketConn) WriteTo(p []byte, addr string) (n int, err error) { + atomic.AddInt64(&m.writes, 1) + time.Sleep(10 * time.Microsecond) // 模拟延迟 + return len(p), nil +} + +func (m *mockShadowsocksPacketConn) Close() error { + return nil +} + +func (m *mockShadowsocksPacketConn) SetDeadline(t time.Time) error { + return nil +} + +func (m *mockShadowsocksPacketConn) SetReadDeadline(t time.Time) error { + return nil +} + +func (m *mockShadowsocksPacketConn) SetWriteDeadline(t time.Time) error { + return nil +} + +// TestShadowsocksUdpConnWriteToRace 测试并发写入竞争 +// 运行: go test -race -run TestShadowsocksUdpConnWriteToRace +func TestShadowsocksUdpConnWriteToRace(t *testing.T) { + // 创建 mock 连接 + mockConn := &mockShadowsocksPacketConn{} + + // 创建 UdpConn + conf := ciphers.AeadCiphersConf["aes-128-gcm"] + masterKey := make([]byte, 16) + + metadata := protocol.Metadata{ + Type: protocol.MetadataTypeIPv4, + Hostname: "127.0.0.1", + Port: 8080, + Cipher: "aes-128-gcm", + } + + udpConn := &UdpConn{ + PacketConn: mockConn, + proxyAddress: "127.0.0.1:8388", + metadata: metadata, + cipherConf: conf, + masterKey: masterKey, + tgtAddr: "127.0.0.1:8080", + } + + const goroutines = 10 + const writesPerGoroutine = 50 // 减少次数因为加密开销大 + + var wg sync.WaitGroup + wg.Add(goroutines) + + // 启动多个 goroutine 并发写入 + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + + for j := 0; j < writesPerGoroutine; j++ { + data := []byte("test data from goroutine") + _, err := udpConn.WriteTo(data, "127.0.0.1:9090") + if err != nil { + // 加密失败是预期的，因为我们没有完整初始化 + t.Logf("Goroutine %d write %d: %v (expected)", id, j, err) + } + } + }(i) + } + + wg.Wait() + + // 验证写入次数 + writes := atomic.LoadInt64(&mockConn.writes) + t.Logf("Total WriteTo calls: %d", writes) + + // 注意：由于加密可能失败，实际写入次数可能少于预期 + // 这个测试主要检测 -race 是否报告竞争 +} + +// TestShadowsocksUdpConnWriteRace 测试 Write 方法的并发 +func TestShadowsocksUdpConnWriteRace(t *testing.T) { + mockConn := &mockShadowsocksPacketConn{} + + conf := ciphers.AeadCiphersConf["aes-128-gcm"] + masterKey := make([]byte, 16) + + metadata := protocol.Metadata{ + Type: protocol.MetadataTypeIPv4, + Hostname: "127.0.0.1", + Port: 8080, + Cipher: "aes-128-gcm", + } + + udpConn := &UdpConn{ + PacketConn: mockConn, + proxyAddress: "127.0.0.1:8388", + metadata: metadata, + cipherConf: conf, + masterKey: masterKey, + tgtAddr: "127.0.0.1:8080", + } + + const goroutines = 10 + const writesPerGoroutine = 50 + + var wg sync.WaitGroup + wg.Add(goroutines) + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + + for j := 0; j < writesPerGoroutine; j++ { + data := []byte("test data") + _, err := udpConn.Write(data) + if err != nil { + t.Logf("Goroutine %d write %d: %v", id, j, err) + } + } + }(i) + } + + wg.Wait() + + t.Logf("Completed concurrent Write test") +} + +// TestShadowsocksUdpConnBufferPoolRace 测试 buffer pool 的并发使用 +func TestShadowsocksUdpConnBufferPoolRace(t *testing.T) { + // 测试 pool.Get 和 pool.Put 的并发使用 + const goroutines = 20 + const opsPerGoroutine = 100 + + var wg sync.WaitGroup + wg.Add(goroutines) + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + + for j := 0; j < opsPerGoroutine; j++ { + buf := pool.Get(1500) + + // 模拟使用 buffer + copy(buf, []byte("test data")) + + // 释放 buffer + pool.Put(buf) + } + }(i) + } + + wg.Wait() + + t.Log("Buffer pool concurrent access test completed") +} + +// TestShadowsocksUdpConnRealConnection 使用真实 UDP 连接测试 +func TestShadowsocksUdpConnRealConnection(t *testing.T) { + // 创建 UDP 服务器 + serverAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0") + if err != nil { + t.Fatalf("Failed to resolve server address: %v", err) + } + + serverConn, err := net.ListenUDP("udp", serverAddr) + if err != nil { + t.Fatalf("Failed to create server: %v", err) + } + defer serverConn.Close() + + // 接收服务器 + go func() { + buf := make([]byte, 2048) + for { + n, addr, err := serverConn.ReadFromUDP(buf) + if err != nil { + return + } + t.Logf("Server received %d bytes from %v", n, addr) + } + }() + + // 创建客户端连接 + clientConn, err := net.Dial("udp", serverConn.LocalAddr().String()) + if err != nil { + t.Fatalf("Failed to create client: %v", err) + } + defer clientConn.Close() + + // 包装为 netproxy.PacketConn（需要实现） + t.Skip("Requires netproxy.PacketConn implementation") +} + +// BenchmarkShadowsocksUdpConnWrite 基准测试 +func BenchmarkShadowsocksUdpConnWrite(b *testing.B) { + mockConn := &mockShadowsocksPacketConn{} + + conf := ciphers.AeadCiphersConf["aes-128-gcm"] + masterKey := make([]byte, 16) + + metadata := protocol.Metadata{ + Type: protocol.MetadataTypeIPv4, + Hostname: "127.0.0.1", + Port: 8080, + Cipher: "aes-128-gcm", + } + + udpConn := &UdpConn{ + PacketConn: mockConn, + proxyAddress: "127.0.0.1:8388", + metadata: metadata, + cipherConf: conf, + masterKey: masterKey, + tgtAddr: "127.0.0.1:8080", + } + + data := []byte("benchmark test data") + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + udpConn.Write(data) + } +} + +// BenchmarkShadowsocksUdpConnWriteParallel 并发基准测试 +func BenchmarkShadowsocksUdpConnWriteParallel(b *testing.B) { + mockConn := &mockShadowsocksPacketConn{} + + conf := ciphers.AeadCiphersConf["aes-128-gcm"] + masterKey := make([]byte, 16) + + metadata := protocol.Metadata{ + Type: protocol.MetadataTypeIPv4, + Hostname: "127.0.0.1", + Port: 8080, + Cipher: "aes-128-gcm", + } + + udpConn := &UdpConn{ + PacketConn: mockConn, + proxyAddress: "127.0.0.1:8388", + metadata: metadata, + cipherConf: conf, + masterKey: masterKey, + tgtAddr: "127.0.0.1:8080", + } + + data := []byte("benchmark test data") + + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + udpConn.Write(data) + } + }) +} + +// TestShadowsocksUdpConnMetadataParseRace 测试 metadata 解析的并发安全 +func TestShadowsocksUdpConnMetadataParseRace(t *testing.T) { + const goroutines = 20 + const opsPerGoroutine = 100 + + var wg sync.WaitGroup + wg.Add(goroutines) + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + + for j := 0; j < opsPerGoroutine; j++ { + addr := "127.0.0.1:8080" + _, err := protocol.ParseMetadata(addr) + if err != nil { + t.Errorf("Goroutine %d parse %d failed: %v", id, j, err) + } + } + }(i) + } + + wg.Wait() +} diff --git a/protocol/shadowsocks_2022/core.go b/protocol/shadowsocks_2022/core.go new file mode 100644 index 00000000..2bbac56d --- /dev/null +++ b/protocol/shadowsocks_2022/core.go @@ -0,0 +1,163 @@ +package shadowsocks_2022 + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/subtle" + "fmt" + "io" + + "github.com/daeuniverse/outbound/ciphers" + "lukechampine.com/blake3" +) + +// SS2022Core contains shared logic for Shadowsocks 2022 protocol. +// Both TCPConn and UdpConn embed this struct to avoid code duplication. +type SS2022Core struct { + cipherConf *ciphers.CipherConf2022 + pskList [][]byte + uPSK []byte + + // Shared block ciphers derived from the configured PSKs. + blockCipherEncrypt cipher.Block + blockCipherDecrypt cipher.Block + + // Pre-computed identity header hash components for multi-PSK scenario + pskHash [][]byte + + // Pre-created block ciphers for identity header encryption + identityBlockCiphers []cipher.Block + + // Flag indicating if multi-PSK is enabled + hasMultiPSK bool +} + +// NewSS2022Core creates a new SS2022Core with pre-computed identity components. +func NewSS2022Core(conf *ciphers.CipherConf2022, pskList [][]byte, uPSK []byte) (*SS2022Core, error) { + if len(pskList) == 0 { + return nil, fmt.Errorf("empty PSK list") + } + + blockCipherEncrypt, err := conf.NewBlockCipher(pskList[0]) + if err != nil { + return nil, fmt.Errorf("failed to create encrypt block cipher: %w", err) + } + blockCipherDecrypt, err := conf.NewBlockCipher(uPSK) + if err != nil { + return nil, fmt.Errorf("failed to create decrypt block cipher: %w", err) + } + + core := &SS2022Core{ + cipherConf: conf, + pskList: pskList, + uPSK: uPSK, + blockCipherEncrypt: blockCipherEncrypt, + blockCipherDecrypt: blockCipherDecrypt, + hasMultiPSK: len(pskList) > 1, + } + + // Pre-compute identity header components for multi-PSK scenario (like sing-box) + if core.hasMultiPSK { + core.pskHash = make([][]byte, len(pskList)) + core.identityBlockCiphers = make([]cipher.Block, len(pskList)-1) + + for i, psk := range pskList { + // Pre-compute BLAKE3 hash of each PSK (same as sing-box) + hash := blake3.Sum512(psk) + core.pskHash[i] = make([]byte, aes.BlockSize) + copy(core.pskHash[i], hash[:aes.BlockSize]) + + // Pre-create block cipher for identity header encryption + if i < len(pskList)-1 { + blockCipher, err := conf.NewBlockCipher(pskList[i]) + if err != nil { + return nil, fmt.Errorf("failed to create identity block cipher: %w", err) + } + core.identityBlockCiphers[i] = blockCipher + } + } + } + + return core, nil +} + +// WriteIdentityHeader writes the identity header to dst for multi-PSK scenario. +// Returns the number of bytes written. +// For single PSK, this is a no-op and returns 0. +func (c *SS2022Core) WriteIdentityHeader(dst []byte, separateHeader []byte) (int, error) { + if !c.hasMultiPSK { + return 0, nil + } + + headerLen := (len(c.pskList) - 1) * aes.BlockSize + if len(dst) < headerLen { + return 0, io.ErrShortBuffer + } + + offset := 0 + for i := 0; i < len(c.pskList)-1; i++ { + header := dst[offset : offset+aes.BlockSize] + // XOR pskHash with separateHeader, then encrypt (same as sing-box) + subtle.XORBytes(header, c.pskHash[i+1], separateHeader) + c.identityBlockCiphers[i].Encrypt(header, header) + offset += aes.BlockSize + } + + return headerLen, nil +} + +// IdentityHeaderLen returns the length of identity header for this connection. +func (c *SS2022Core) IdentityHeaderLen() int { + if !c.hasMultiPSK { + return 0 + } + return (len(c.pskList) - 1) * aes.BlockSize +} + +// HasMultiPSK returns true if multiple PSKs are configured. +func (c *SS2022Core) HasMultiPSK() bool { + return c.hasMultiPSK +} + +// CipherConf returns the cipher configuration. +func (c *SS2022Core) CipherConf() *ciphers.CipherConf2022 { + return c.cipherConf +} + +// UPSK returns the user PSK. +func (c *SS2022Core) UPSK() []byte { + return c.uPSK +} + +// BlockCipherEncrypt returns the shared block cipher used for encrypting the +// separate header on outbound packets. +func (c *SS2022Core) BlockCipherEncrypt() cipher.Block { + return c.blockCipherEncrypt +} + +// BlockCipherDecrypt returns the shared block cipher used for decrypting the +// separate header on inbound packets. +func (c *SS2022Core) BlockCipherDecrypt() cipher.Block { + return c.blockCipherDecrypt +} + +// PSKList returns the list of PSKs. +func (c *SS2022Core) PSKList() [][]byte { + return c.pskList +} + +// PSKHash returns pre-computed PSK hash at index i. +func (c *SS2022Core) PSKHash(i int) []byte { + if i < 0 || i >= len(c.pskHash) { + return nil + } + return c.pskHash[i] +} + +// IdentityBlockCipher returns pre-created identity block cipher at index i. +func (c *SS2022Core) IdentityBlockCipher(i int) cipher.Block { + if i < 0 || i >= len(c.identityBlockCiphers) { + return nil + } + return c.identityBlockCiphers[i] +} diff --git a/protocol/shadowsocks_2022/dialer.go b/protocol/shadowsocks_2022/dialer.go new file mode 100644 index 00000000..aec85418 --- /dev/null +++ b/protocol/shadowsocks_2022/dialer.go @@ -0,0 +1,124 @@ +package shadowsocks_2022 + +import ( + "context" + "fmt" + "net" + "strings" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/netproxy" + "github.com/daeuniverse/outbound/protocol" + "github.com/daeuniverse/outbound/protocol/shadowsocks" + "github.com/daeuniverse/outbound/protocol/socks5" +) + +const maxPSKListLength = 8 + +// FakeNetPacketConn wraps a PacketConn to override the target address. +// It embeds netproxy.PacketConn directly, so it implements the interface correctly. +// The Addr field is used by UdpConn.WriteTo to determine the actual target. +type FakeNetPacketConn struct { + netproxy.PacketConn + Addr string +} + +func init() { + protocol.Register("shadowsocks_2022", NewDialer) +} + +type Dialer struct { + parentDialer netproxy.Dialer + proxyAddress string + core *SS2022Core + sg shadowsocks.SaltGenerator +} + +func NewDialer(parentDialer netproxy.Dialer, header protocol.Header) (netproxy.Dialer, error) { + conf := ciphers.Aead2022CiphersConf[header.Cipher] + if conf == nil { + return nil, fmt.Errorf("unsupported shadowsocks 2022 cipher: %s", header.Cipher) + } + if conf.NewCipher == nil || conf.NewBlockCipher == nil { + return nil, fmt.Errorf("invalid shadowsocks 2022 cipher config: %s", header.Cipher) + } + keyStrList := strings.Split(header.Password, ":") + if len(keyStrList) > maxPSKListLength { + return nil, fmt.Errorf("too many PSKs: got %d, max %d", len(keyStrList), maxPSKListLength) + } + pskList := make([][]byte, len(keyStrList)) + for i, keyStr := range keyStrList { + key, err := ciphers.ValidateBase64PSK(keyStr, conf.KeyLen) + if err != nil { + return nil, err + } + pskList[i] = key + } + uPSK := pskList[len(pskList)-1] + core, err := NewSS2022Core(conf, pskList, uPSK) + if err != nil { + return nil, err + } + sg, err := shadowsocks.NewRandomSaltGenerator(conf.SaltLen) + if err != nil { + return nil, err + } + return &Dialer{ + parentDialer: parentDialer, + proxyAddress: header.ProxyAddress, + core: core, + sg: sg, + }, nil +} + +func (d *Dialer) DialContext(ctx context.Context, network, addr string) (netproxy.Conn, error) { + magicNetwork, err := netproxy.ParseMagicNetwork(network) + if err != nil { + return nil, err + } + // Extract base protocol from network string (handles "udp", "udp4", "udp6", "udp4(DNS)", etc.) + proto := magicNetwork.Network + if len(proto) >= 3 && (proto[0:3] == "tcp" || proto[0:3] == "udp") { + proto = proto[0:3] + } + switch proto { + case "tcp": + addrInfo, err := socks5.AddressFromString(addr) + if err != nil { + return nil, err + } + // Shadowsocks transfer TCP traffic via TCP tunnel. + conn, err := d.parentDialer.DialContext(ctx, network, d.proxyAddress) + if err != nil { + return nil, err + } + return NewTCPConn(conn.(net.Conn), d.core, d.sg, addrInfo, nil), nil + case "udp": + conn, err := d.ListenPacket(ctx, network, d.proxyAddress) + if err != nil { + return nil, err + } + return &FakeNetPacketConn{ + PacketConn: conn, + Addr: addr, + }, nil + default: + return nil, fmt.Errorf("%w: %v", netproxy.UnsupportedTunnelTypeError, network) + } +} + +func (d *Dialer) ListenPacket(ctx context.Context, network string, addr string) (netproxy.PacketConn, error) { + // Parse magic network to preserve Mark and Mptcp settings + magicNetwork, err := netproxy.ParseMagicNetwork(network) + if err != nil { + return nil, err + } + // Shadowsocks transfer UDP traffic via UDP tunnel. + magicNetwork.Network = "udp" + network = magicNetwork.Encode() + conn, err := d.parentDialer.DialContext(ctx, network, d.proxyAddress) + if err != nil { + return nil, err + } + return NewUdpConn(conn.(net.Conn), d.core, nil) +} diff --git a/protocol/shadowsocks_2022/dialer_magic_network_test.go b/protocol/shadowsocks_2022/dialer_magic_network_test.go new file mode 100644 index 00000000..12fe792c --- /dev/null +++ b/protocol/shadowsocks_2022/dialer_magic_network_test.go @@ -0,0 +1,142 @@ +package shadowsocks_2022 + +import ( + "strings" + "testing" + + "github.com/daeuniverse/outbound/netproxy" +) + +// TestMagicNetworkParsing tests that ParseMagicNetwork handles various network formats +func TestMagicNetworkParsing(t *testing.T) { + tests := []struct { + name string + network string + expectNetwork string + expectSuccess bool + }{ + { + name: "plain tcp", + network: "tcp", + expectNetwork: "tcp", + expectSuccess: true, + }, + { + name: "plain udp", + network: "udp", + expectNetwork: "udp", + expectSuccess: true, + }, + { + name: "magic network tcp with mark", + network: netproxy.MagicNetwork{Network: "tcp", Mark: 1}.Encode(), + expectNetwork: "tcp", + expectSuccess: true, + }, + { + name: "magic network udp with mark", + network: netproxy.MagicNetwork{Network: "udp", Mark: 1}.Encode(), + expectNetwork: "udp", + expectSuccess: true, + }, + { + name: "magic network tcp with zero mark", + network: netproxy.MagicNetwork{Network: "tcp", Mark: 0}.Encode(), + expectNetwork: "tcp", + expectSuccess: true, + }, + { + name: "magic network tcp with mptcp", + network: netproxy.MagicNetwork{Network: "tcp", Mark: 0, Mptcp: true}.Encode(), + expectNetwork: "tcp", + expectSuccess: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mn, err := netproxy.ParseMagicNetwork(tt.network) + if !tt.expectSuccess { + if err == nil { + t.Errorf("Expected error but got none for network %q", tt.network) + } + return + } + if err != nil { + t.Errorf("ParseMagicNetwork(%q) failed: %v", tt.network, err) + return + } + if mn.Network != tt.expectNetwork { + t.Errorf("ParseMagicNetwork(%q) returned network %q, want %q", tt.network, mn.Network, tt.expectNetwork) + } + }) + } +} + +// TestDialerNetworkTypeSwitch verifies that the switch statement in DialContext +// correctly handles the parsed network type +func TestDialerNetworkTypeSwitch(t *testing.T) { + tests := []struct { + name string + network string + expectSupport bool + }{ + { + name: "plain tcp", + network: "tcp", + expectSupport: true, + }, + { + name: "plain udp", + network: "udp", + expectSupport: true, + }, + { + name: "magic network tcp with mark", + network: netproxy.MagicNetwork{Network: "tcp", Mark: 1}.Encode(), + expectSupport: true, + }, + { + name: "unsupported network", + network: "sctp", + expectSupport: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mn, err := netproxy.ParseMagicNetwork(tt.network) + if err != nil && tt.expectSupport { + t.Errorf("ParseMagicNetwork(%q) failed: %v", tt.network, err) + return + } + + // Check if the network type would match our switch cases + supported := false + switch mn.Network { + case "tcp", "udp": + supported = true + } + + if supported != tt.expectSupport { + t.Errorf("Network %q: supported=%v, want %v", tt.network, supported, tt.expectSupport) + } + }) + } +} + +// TestMultiPSKParsing tests that multi-PSK passwords are correctly parsed +func TestMultiPSKParsing(t *testing.T) { + // Test the actual password from the user's link + password := "bG1k6qKaANArh2515TnLrA==:xBAQSRWYs6I7TGjzBxY68w==" + parts := strings.Split(password, ":") + if len(parts) != 2 { + t.Errorf("Expected 2 PSK parts, got %d", len(parts)) + } + if parts[0] != "bG1k6qKaANArh2515TnLrA==" { + t.Errorf("First PSK mismatch: got %q", parts[0]) + } + if parts[1] != "xBAQSRWYs6I7TGjzBxY68w==" { + t.Errorf("Second PSK mismatch: got %q", parts[1]) + } +} diff --git a/protocol/shadowsocks_2022/dialer_test.go b/protocol/shadowsocks_2022/dialer_test.go new file mode 100644 index 00000000..c029183e --- /dev/null +++ b/protocol/shadowsocks_2022/dialer_test.go @@ -0,0 +1,62 @@ +package shadowsocks_2022 + +import ( + "context" + "encoding/base64" + "strings" + "testing" + + "github.com/daeuniverse/outbound/netproxy" + "github.com/daeuniverse/outbound/protocol" +) + +type nopDialer struct{} + +func (nopDialer) DialContext(ctx context.Context, network, addr string) (netproxy.Conn, error) { + return nil, nil +} + +func pskBase64(length int, v byte) string { + b := make([]byte, length) + for i := range b { + b[i] = v + } + return base64.StdEncoding.EncodeToString(b) +} + +func TestNewDialer_UnsupportedCipher(t *testing.T) { + _, err := NewDialer(nopDialer{}, protocol.Header{ + Cipher: "2022-blake3-chacha20-poly1305", + Password: pskBase64(32, 0x11), + ProxyAddress: "127.0.0.1:443", + }) + if err == nil || !strings.Contains(err.Error(), "unsupported shadowsocks 2022 cipher") { + t.Fatalf("expected unsupported cipher error, got: %v", err) + } +} + +func TestNewDialer_TooManyPSKs(t *testing.T) { + keys := make([]string, maxPSKListLength+1) + for i := range keys { + keys[i] = pskBase64(16, byte(i+1)) + } + _, err := NewDialer(nopDialer{}, protocol.Header{ + Cipher: "2022-blake3-aes-128-gcm", + Password: strings.Join(keys, ":"), + ProxyAddress: "127.0.0.1:443", + }) + if err == nil || !strings.Contains(err.Error(), "too many PSKs") { + t.Fatalf("expected too many PSKs error, got: %v", err) + } +} + +func TestNewDialer_ValidMultiPSK(t *testing.T) { + _, err := NewDialer(nopDialer{}, protocol.Header{ + Cipher: "2022-blake3-aes-256-gcm", + Password: strings.Join([]string{pskBase64(32, 0x21), pskBase64(32, 0x22)}, ":"), + ProxyAddress: "127.0.0.1:443", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} diff --git a/protocol/shadowsocks_2022/encrypt.go b/protocol/shadowsocks_2022/encrypt.go new file mode 100644 index 00000000..79c2d3a1 --- /dev/null +++ b/protocol/shadowsocks_2022/encrypt.go @@ -0,0 +1,62 @@ +package shadowsocks_2022 + +import ( + "crypto/cipher" + "sync" + + "github.com/daeuniverse/outbound/ciphers" + "lukechampine.com/blake3" +) + +var ( + Shadowsocks2022ReusedInfo = "shadowsocks 2022 session subkey" + Shadowsocks2022IdentityHeaderInfo = "shadowsocks 2022 identity subkey" +) + +// subKeyPool reuses subKey buffers to reduce allocations in the hot path. +// SS2022 uses either 16-byte (AES-128) or 32-byte (AES-256) keys. +var subKeyPool = sync.Pool{ + New: func() interface{} { + return make([]byte, 32) // max key size + }, +} + +// keyMaterialPool reuses key material buffers. +// Key material = psk (max 32) + salt (max 32) = max 64 bytes. +var keyMaterialPool = sync.Pool{ + New: func() interface{} { + return make([]byte, 0, 64) + }, +} + +func GenerateSubKey(psk []byte, salt []byte, context string) (subKey []byte) { + // Get buffer from pool, trim to actual key length + subKey = subKeyPool.Get().([]byte)[:len(psk)] + + // Get key material buffer from pool + keyMaterial := keyMaterialPool.Get().([]byte) + keyMaterial = keyMaterial[:0] + keyMaterial = append(keyMaterial, psk...) + keyMaterial = append(keyMaterial, salt...) + + blake3.DeriveKey(subKey, context, keyMaterial) + + // Return key material buffer to pool + keyMaterialPool.Put(keyMaterial) + + return +} + +// PutSubKey returns a subKey buffer to the pool. +// Callers should use this after they're done with the subKey. +func PutSubKey(subKey []byte) { + if subKey != nil && cap(subKey) >= 16 && cap(subKey) <= 32 { + subKeyPool.Put(subKey[:32]) + } +} + +func CreateCipher(masterKey []byte, salt []byte, cipherConf *ciphers.CipherConf2022) (cipher cipher.AEAD, err error) { + subKey := GenerateSubKey(masterKey, salt, Shadowsocks2022ReusedInfo) + defer PutSubKey(subKey) + return cipherConf.NewCipher(subKey) +} diff --git a/protocol/shadowsocks_2022/memory_analysis_test.go b/protocol/shadowsocks_2022/memory_analysis_test.go new file mode 100644 index 00000000..e40c8320 --- /dev/null +++ b/protocol/shadowsocks_2022/memory_analysis_test.go @@ -0,0 +1,367 @@ +package shadowsocks_2022 + +import ( + "crypto/aes" + "crypto/subtle" + "io" + "net" + "runtime" + "runtime/debug" + "testing" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/protocol/shadowsocks" + "github.com/daeuniverse/outbound/protocol/socks5" + "lukechampine.com/blake3" +) + +type analysisDiscardConn struct{} + +func (analysisDiscardConn) Read(_ []byte) (int, error) { return 0, io.EOF } +func (analysisDiscardConn) Write(p []byte) (int, error) { return len(p), nil } +func (analysisDiscardConn) Close() error { return nil } +func (analysisDiscardConn) LocalAddr() net.Addr { return analysisAddr("local") } +func (analysisDiscardConn) RemoteAddr() net.Addr { return analysisAddr("remote") } +func (analysisDiscardConn) SetDeadline(_ time.Time) error { return nil } +func (analysisDiscardConn) SetReadDeadline(_ time.Time) error { return nil } +func (analysisDiscardConn) SetWriteDeadline(_ time.Time) error { return nil } + +type analysisAddr string + +func (a analysisAddr) Network() string { return "analysis" } +func (a analysisAddr) String() string { return string(a) } + +func analysisPSKList(count, keyLen int) [][]byte { + pskList := make([][]byte, count) + for i := 0; i < count; i++ { + psk := make([]byte, keyLen) + for j := range psk { + psk[j] = byte(i + 1) + } + pskList[i] = psk + } + return pskList +} + +func analysisHeapAlloc() uint64 { + runtime.GC() + debug.FreeOSMemory() + var ms runtime.MemStats + runtime.ReadMemStats(&ms) + return ms.HeapAlloc +} + +func analysisReplayWindowLen(c *UdpConn) int { + var n int + c.replayWindow.Range(func(_, _ any) bool { + n++ + return true + }) + return n +} + +func analysisIdentityHeaderNoCache(dst []byte, conf *ciphers.CipherConf2022, pskList [][]byte, separateHeader []byte) (int, error) { + if len(pskList) <= 1 { + return 0, nil + } + headerLen := (len(pskList) - 1) * aes.BlockSize + if len(dst) < headerLen { + return 0, io.ErrShortBuffer + } + offset := 0 + for i := 0; i < len(pskList)-1; i++ { + header := dst[offset : offset+aes.BlockSize] + hash := blake3.Sum512(pskList[i+1]) + subtle.XORBytes(header, hash[:aes.BlockSize], separateHeader) + blockCipher, err := conf.NewBlockCipher(pskList[i]) + if err != nil { + return 0, err + } + blockCipher.Encrypt(header, header) + offset += aes.BlockSize + } + return offset, nil +} + +func analysisUDPWriteToNoSessionCipherCache(c *UdpConn, payload []byte, addr string) error { + packetID := c.nextPacketID() + var separateHeader [16]byte + copy(separateHeader[:8], c.sessionID[:]) + putUint64 := func(b []byte, v uint64) { + b[0] = byte(v >> 56) + b[1] = byte(v >> 48) + b[2] = byte(v >> 40) + b[3] = byte(v >> 32) + b[4] = byte(v >> 24) + b[5] = byte(v >> 16) + b[6] = byte(v >> 8) + b[7] = byte(v) + } + putUint64(separateHeader[8:], packetID) + + var separateHeaderEncrypted [16]byte + c.blockCipherEncrypt.Encrypt(separateHeaderEncrypted[:], separateHeader[:]) + + addrInfo, err := socks5.AddressFromString(addr) + if err != nil { + return err + } + addrLen, err := addrInfoEncodedLen(addrInfo) + if err != nil { + return err + } + messageLen := 1 + 8 + 2 + addrLen + len(payload) + totalPacketLen := len(separateHeaderEncrypted) + c.IdentityHeaderLen() + messageLen + c.CipherConf().TagLen + packet := make([]byte, totalPacketLen) + offset := 0 + copy(packet[offset:], separateHeaderEncrypted[:]) + offset += len(separateHeaderEncrypted) + + identityHeaderLen, err := c.WriteIdentityHeader(packet[offset:], separateHeader[:]) + if err != nil { + return err + } + offset += identityHeaderLen + + messageOffset := offset + message := packet[messageOffset : messageOffset+messageLen] + message[0] = HeaderTypeClientStream + putUint64(message[1:9], uint64(time.Now().Unix())) + message[9] = 0 + message[10] = 0 + addrWritten, err := writeAddrInfoTo(message[11:], addrInfo) + if err != nil { + return err + } + copy(message[11+addrWritten:], payload) + + sessionCipher, err := CreateCipher(c.UPSK(), c.sessionID[:], c.CipherConf()) + if err != nil { + return err + } + packet = sessionCipher.Seal(packet[:messageOffset], separateHeader[4:16], message, nil) + _, err = c.Conn.Write(packet) + return err +} + +func TestSS2022CoreRetainedHeapAndRelease(t *testing.T) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + if conf == nil { + t.Fatal("missing ss2022 cipher config") + } + + measure := func(name string, pskCount, connCount int) (liveGrowth int64, releasedGrowth int64) { + t.Helper() + + baseline := analysisHeapAlloc() + cores := make([]*SS2022Core, connCount) + pskList := analysisPSKList(pskCount, conf.KeyLen) + uPSK := pskList[len(pskList)-1] + for i := 0; i < connCount; i++ { + core, err := NewSS2022Core(conf, pskList, uPSK) + if err != nil { + t.Fatalf("%s: create core: %v", name, err) + } + cores[i] = core + } + live := analysisHeapAlloc() + runtime.KeepAlive(cores) + + liveGrowth = int64(live) - int64(baseline) + if liveGrowth <= 0 { + t.Fatalf("%s: expected positive live heap growth, got %d", name, liveGrowth) + } + + cores = nil + released := analysisHeapAlloc() + releasedGrowth = int64(released) - int64(baseline) + t.Logf("%s: live heap %+d bytes, after release %+d bytes", name, liveGrowth, releasedGrowth) + return liveGrowth, releasedGrowth + } + + const connCount = 4000 + singleLive, singleReleased := measure("single-psk", 1, connCount) + multiLive, multiReleased := measure("multi-psk-8", maxPSKListLength, connCount) + + t.Logf("single-psk retained bytes/core: %.2f", float64(singleLive)/connCount) + t.Logf("multi-psk-8 retained bytes/core: %.2f", float64(multiLive)/connCount) + + if multiLive <= singleLive { + t.Fatalf("expected multi-psk cores to retain more heap than single-psk: single=%d multi=%d", singleLive, multiLive) + } + if singleReleased > singleLive/2 && singleReleased > 1<<20 { + t.Fatalf("single-psk core memory did not return close enough after release: live=%d released=%d", singleLive, singleReleased) + } + if multiReleased > multiLive/2 && multiReleased > 1<<20 { + t.Fatalf("multi-psk core memory did not return close enough after release: live=%d released=%d", multiLive, multiReleased) + } +} + +func TestSS2022TCPConnCloseReleasesReusableBuffers(t *testing.T) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + if conf == nil { + t.Fatal("missing ss2022 cipher config") + } + sg, err := shadowsocks.NewRandomSaltGenerator(conf.SaltLen) + if err != nil { + t.Fatal(err) + } + defer sg.Close() + + addr, err := socks5.AddressFromString("203.0.113.10:443") + if err != nil { + t.Fatal(err) + } + + pskList := analysisPSKList(1, conf.KeyLen) + uPSK := pskList[0] + core, err := NewSS2022Core(conf, pskList, uPSK) + if err != nil { + t.Fatal(err) + } + + const connCount = 256 + baseline := analysisHeapAlloc() + conns := make([]*TCPConn, connCount) + for i := 0; i < connCount; i++ { + conn := NewTCPConn(analysisDiscardConn{}, core, sg, addr, nil).(*TCPConn) + _ = conn.ensureReadCipherBuf(64 << 10) + _ = conn.borrowWriteFrame(96 << 10) + conns[i] = conn + } + + live := analysisHeapAlloc() + runtime.KeepAlive(conns) + liveGrowth := int64(live) - int64(baseline) + if liveGrowth < 8<<20 { + t.Fatalf("expected cached buffers to retain noticeable heap, got %d bytes", liveGrowth) + } + + for _, conn := range conns { + if err := conn.Close(); err != nil { + t.Fatal(err) + } + } + conns = nil + + released := analysisHeapAlloc() + releasedGrowth := int64(released) - int64(baseline) + t.Logf("tcp reusable buffers: live heap %+d bytes, after close %+d bytes", liveGrowth, releasedGrowth) + + if releasedGrowth > liveGrowth/3 && releasedGrowth > 4<<20 { + t.Fatalf("tcp reusable buffer memory still retained after close: live=%d released=%d", liveGrowth, releasedGrowth) + } +} + +func TestSS2022UDPReplayWindowBounded(t *testing.T) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + if conf == nil { + t.Fatal("missing ss2022 cipher config") + } + + pskList := analysisPSKList(1, conf.KeyLen) + uPSK := pskList[0] + core, err := NewSS2022Core(conf, pskList, uPSK) + if err != nil { + t.Fatal(err) + } + conn, err := NewUdpConn(analysisDiscardConn{}, core, nil) + if err != nil { + t.Fatal(err) + } + + now := time.Now() + for i := 0; i < maxTrackedUdpSessions*4; i++ { + var sessionID [8]byte + sessionID[0] = byte(i) + sessionID[1] = byte(i >> 8) + if !conn.checkAndUpdateReplay(sessionID, 1, now) { + t.Fatalf("session %d should be accepted on first packet", i) + } + } + + if got := analysisReplayWindowLen(conn); got > maxTrackedUdpSessions { + t.Fatalf("replay window exceeded bound: got %d want <= %d", got, maxTrackedUdpSessions) + } +} + +func BenchmarkSS2022IdentityHeader_PrecomputedMultiPSK(b *testing.B) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + pskList := analysisPSKList(maxPSKListLength, conf.KeyLen) + core, err := NewSS2022Core(conf, pskList, pskList[len(pskList)-1]) + if err != nil { + b.Fatal(err) + } + + var separateHeader [16]byte + dst := make([]byte, core.IdentityHeaderLen()) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := core.WriteIdentityHeader(dst, separateHeader[:]); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkSS2022IdentityHeader_RecomputeMultiPSK(b *testing.B) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + pskList := analysisPSKList(maxPSKListLength, conf.KeyLen) + + var separateHeader [16]byte + dst := make([]byte, (len(pskList)-1)*aes.BlockSize) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := analysisIdentityHeaderNoCache(dst, conf, pskList, separateHeader[:]); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkSS2022UDPWriteTo_SessionCipherReuse(b *testing.B) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + pskList := analysisPSKList(1, conf.KeyLen) + uPSK := pskList[0] + core, err := NewSS2022Core(conf, pskList, uPSK) + if err != nil { + b.Fatal(err) + } + conn, err := NewUdpConn(analysisDiscardConn{}, core, nil) + if err != nil { + b.Fatal(err) + } + payload := make([]byte, 1400) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := conn.WriteTo(payload, "198.51.100.10:443"); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkSS2022UDPWriteTo_CreateCipherEveryPacket(b *testing.B) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + pskList := analysisPSKList(1, conf.KeyLen) + uPSK := pskList[0] + core, err := NewSS2022Core(conf, pskList, uPSK) + if err != nil { + b.Fatal(err) + } + conn, err := NewUdpConn(analysisDiscardConn{}, core, nil) + if err != nil { + b.Fatal(err) + } + payload := make([]byte, 1400) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := analysisUDPWriteToNoSessionCipherCache(conn, payload, "198.51.100.10:443"); err != nil { + b.Fatal(err) + } + } +} diff --git a/protocol/shadowsocks_2022/tcp_conn.go b/protocol/shadowsocks_2022/tcp_conn.go new file mode 100644 index 00000000..68d3bfe4 --- /dev/null +++ b/protocol/shadowsocks_2022/tcp_conn.go @@ -0,0 +1,393 @@ +package shadowsocks_2022 + +import ( + "crypto/aes" + "crypto/cipher" + "encoding/binary" + "fmt" + "io" + "net" + "sync" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/common" + "github.com/daeuniverse/outbound/pool" + "github.com/daeuniverse/outbound/protocol" + "github.com/daeuniverse/outbound/protocol/shadowsocks" + "github.com/daeuniverse/outbound/protocol/socks5" + disk_bloom "github.com/mzz2017/disk-bloom" + "github.com/samber/oops" + "lukechampine.com/blake3" +) + +const ( + TCPChunkMaxLen = (1 << 16) - 1 + + HeaderTypeClientStream = 0 + HeaderTypeServerStream = 1 + MinPaddingLength = 0 + MaxPaddingLength = 900 + + maxReusableWriteFrameSize = 128 << 10 +) + +// TCPConn represents a Shadowsocks TCP connection. +// It embeds SS2022Core for shared logic (identity header, cipher cache). +type TCPConn struct { + *SS2022Core // Embedded core for shared logic + + net.Conn + addr *socks5.AddressInfo + sg shadowsocks.SaltGenerator + + cipherRead cipher.AEAD + cipherWrite cipher.AEAD + onceRead bool + onceWrite bool + nonceRead []byte + nonceWrite []byte + + readMutex sync.Mutex + writeMutex sync.Mutex + + leftToRead []byte + indexToRead int + readCipherBuf []byte + writeFrame []byte + + bloom *disk_bloom.FilterGroup +} + +type Key struct { + CipherConf *ciphers.CipherConf + MasterKey []byte +} + +func NewTCPConn(conn net.Conn, core *SS2022Core, sg shadowsocks.SaltGenerator, addr *socks5.AddressInfo, bloom *disk_bloom.FilterGroup) net.Conn { + tcpConn := &TCPConn{ + SS2022Core: core, + Conn: conn, + addr: addr, + sg: sg, + nonceRead: make([]byte, core.CipherConf().NonceLen), + nonceWrite: make([]byte, core.CipherConf().NonceLen), + bloom: bloom, + } + return tcpConn +} + +func (c *TCPConn) Close() error { + c.readMutex.Lock() + c.leftToRead = nil + c.indexToRead = 0 + c.readCipherBuf = nil + c.readMutex.Unlock() + + c.writeMutex.Lock() + c.writeFrame = nil + c.writeMutex.Unlock() + return c.Conn.Close() +} + +func encryptedPayloadLen(payloadLen, tagLen int) int { + if payloadLen <= 0 { + return 0 + } + chunks := payloadLen / TCPChunkMaxLen + if payloadLen%TCPChunkMaxLen > 0 { + chunks++ + } + return payloadLen + chunks*(2+tagLen+tagLen) +} + +func addrInfoEncodedLen(addr *socks5.AddressInfo) (int, error) { + if addr == nil { + return 0, fmt.Errorf("nil address info") + } + switch addr.Type { + case socks5.AddressTypeIPv4: + if !addr.IP.Is4() { + return 0, fmt.Errorf("invalid ipv4 address") + } + return 1 + 4 + 2, nil + case socks5.AddressTypeIPv6: + if !addr.IP.Is6() { + return 0, fmt.Errorf("invalid ipv6 address") + } + return 1 + 16 + 2, nil + case socks5.AddressTypeDomain: + if len(addr.Hostname) > 255 { + return 0, fmt.Errorf("domain name too long: %d", len(addr.Hostname)) + } + return 1 + 1 + len(addr.Hostname) + 2, nil + default: + return 0, fmt.Errorf("unsupported address type: %v", addr.Type) + } +} + +func writeAddrInfoTo(dst []byte, addr *socks5.AddressInfo) (int, error) { + addrLen, err := addrInfoEncodedLen(addr) + if err != nil { + return 0, err + } + if len(dst) < addrLen { + return 0, io.ErrShortBuffer + } + dst[0] = byte(addr.Type) + switch addr.Type { + case socks5.AddressTypeIPv4: + ip := addr.IP.AsSlice() + copy(dst[1:1+4], ip) + binary.BigEndian.PutUint16(dst[1+4:1+4+2], addr.Port) + return 1 + 4 + 2, nil + case socks5.AddressTypeIPv6: + ip := addr.IP.AsSlice() + copy(dst[1:1+16], ip) + binary.BigEndian.PutUint16(dst[1+16:1+16+2], addr.Port) + return 1 + 16 + 2, nil + case socks5.AddressTypeDomain: + domainLen := len(addr.Hostname) + dst[1] = byte(domainLen) + copy(dst[2:2+domainLen], addr.Hostname) + binary.BigEndian.PutUint16(dst[2+domainLen:2+domainLen+2], addr.Port) + return 1 + 1 + domainLen + 2, nil + default: + return 0, fmt.Errorf("unsupported address type: %v", addr.Type) + } +} + +func (c *TCPConn) ensureReadCipherBuf(size int) []byte { + if cap(c.readCipherBuf) < size { + c.readCipherBuf = make([]byte, size) + } + return c.readCipherBuf[:size] +} + +func (c *TCPConn) borrowWriteFrame(size int) []byte { + if size <= maxReusableWriteFrameSize { + if cap(c.writeFrame) < size { + c.writeFrame = make([]byte, size) + } + return c.writeFrame[:size] + } + return make([]byte, size) +} + +func (c *TCPConn) writeIdentityHeaderTo(dst []byte, offset int, salt []byte) (int, error) { + for i := 0; i < len(c.PSKList())-1; i++ { + if offset+aes.BlockSize > len(dst) { + return 0, io.ErrShortBuffer + } + identitySubkey := GenerateSubKey(c.PSKList()[i], salt, Shadowsocks2022IdentityHeaderInfo) + b, err := c.CipherConf().NewBlockCipher(identitySubkey) + if err != nil { + PutSubKey(identitySubkey) + return 0, err + } + plaintext := blake3.Sum512(c.PSKList()[i+1]) + b.Encrypt(dst[offset:offset+aes.BlockSize], plaintext[:aes.BlockSize]) + PutSubKey(identitySubkey) + offset += aes.BlockSize + } + return offset, nil +} + +func (c *TCPConn) sealPayload(dst []byte, payload []byte) int { + offset := 0 + var chunkLengthBuf [2]byte + for i := 0; i < len(payload); i += TCPChunkMaxLen { + chunkLength := common.Min(TCPChunkMaxLen, len(payload)-i) + binary.BigEndian.PutUint16(chunkLengthBuf[:], uint16(chunkLength)) + _ = c.cipherWrite.Seal(dst[offset:offset], c.nonceWrite, chunkLengthBuf[:], nil) + offset += 2 + c.CipherConf().TagLen + common.BytesIncLittleEndian(c.nonceWrite) + _ = c.cipherWrite.Seal(dst[offset:offset], c.nonceWrite, payload[i:i+chunkLength], nil) + offset += chunkLength + c.CipherConf().TagLen + common.BytesIncLittleEndian(c.nonceWrite) + } + return offset +} + +func (c *TCPConn) Read(b []byte) (n int, err error) { + c.readMutex.Lock() + defer c.readMutex.Unlock() + + if c.indexToRead < len(c.leftToRead) { + n = copy(b, c.leftToRead[c.indexToRead:]) + c.indexToRead += n + if c.indexToRead >= len(c.leftToRead) { + c.leftToRead = nil + c.indexToRead = 0 + } + return n, nil + } + + var payloadLength uint16 + + if !c.onceRead { + var saltBuf [32]byte + salt := saltBuf[:c.CipherConf().SaltLen] + n, err = io.ReadFull(c.Conn, salt) + if err != nil { + return 0, err + } + c.cipherRead, err = CreateCipher(c.UPSK(), salt, c.CipherConf()) + if err != nil { + return 0, oops.Wrapf(err, "fail to initiate cipher") + } + + var headerBuf [11 + 32 + 16]byte + header := headerBuf[:11+c.CipherConf().SaltLen+c.CipherConf().TagLen] + if _, err := io.ReadFull(c.Conn, header); err != nil { + return 0, err + } + header, err := c.cipherRead.Open(header[:0], c.nonceRead, header, nil) + if err != nil { + return 0, protocol.ErrFailAuth + } + common.BytesIncLittleEndian(c.nonceRead) + offset := 0 + typ := uint8(header[offset]) + offset += 1 + timestamp := time.Unix(int64(binary.BigEndian.Uint64(header[offset:offset+8])), 0) + offset += 8 + + if typ != HeaderTypeServerStream { + return 0, fmt.Errorf("received unexpected header type: %d", typ) + } + + if err := validateTimestamp(timestamp, time.Now()); err != nil { + return 0, err + } + + // Best-effort replay protection fallback for environments that provide bloom. + if c.bloom != nil { + if c.bloom.ExistOrAdd(salt) { + return 0, protocol.ErrReplayAttack + } + } + + // Skip request salt + offset += c.CipherConf().SaltLen + + payloadLength = binary.BigEndian.Uint16(header[offset : offset+2]) + + c.onceRead = true + } else { + var payloadLengthBuf [2 + 16]byte + payloadLengthRaw := payloadLengthBuf[:2+c.CipherConf().TagLen] + if _, err := io.ReadFull(c.Conn, payloadLengthRaw); err != nil { + return 0, err + } + payloadLengthPlain, err := c.cipherRead.Open(payloadLengthRaw[:0], c.nonceRead, payloadLengthRaw, nil) + if err != nil { + return 0, protocol.ErrFailAuth + } + common.BytesIncLittleEndian(c.nonceRead) + + payloadLength = binary.BigEndian.Uint16(payloadLengthPlain) + } + + if c.cipherRead == nil { + return 0, oops.Wrapf(err, "cipher is not initialized") + } + + payload := c.ensureReadCipherBuf(int(payloadLength) + c.CipherConf().TagLen) + if _, err = io.ReadFull(c.Conn, payload); err != nil { + return 0, err + } + payload, err = c.cipherRead.Open(payload[:0], c.nonceRead, payload, nil) + if err != nil { + return 0, protocol.ErrFailAuth + } + common.BytesIncLittleEndian(c.nonceRead) + + n = copy(b, payload) + if len(payload) > n { + c.leftToRead = payload + c.indexToRead = n + } else { + c.leftToRead = nil + c.indexToRead = 0 + } + return n, nil +} + +func (c *TCPConn) Write(b []byte) (n int, err error) { + n = len(b) + c.writeMutex.Lock() + defer c.writeMutex.Unlock() + if !c.onceWrite { + // Generate salt + salt := c.sg.Get() + defer pool.Put(salt) + + // Setup encryption + c.cipherWrite, err = CreateCipher(c.UPSK(), salt, c.CipherConf()) + if err != nil { + return 0, oops.Wrapf(err, "fail to initiate cipher") + } + + addrLen, err := addrInfoEncodedLen(c.addr) + if err != nil { + return 0, oops.Wrapf(err, "fail to calculate address length") + } + + initialPayloadMaxLength := TCPChunkMaxLen - (addrLen + 2) + initialPayloadLen := len(b) + if initialPayloadLen > initialPayloadMaxLength { + initialPayloadLen = initialPayloadMaxLength + } + firstVarHeaderLen := addrLen + 2 + initialPayloadLen + remainingPayload := b[initialPayloadLen:] + totalSize := len(salt) + + (len(c.PSKList())-1)*aes.BlockSize + + (11 + c.CipherConf().TagLen) + + (firstVarHeaderLen + c.CipherConf().TagLen) + + encryptedPayloadLen(len(remainingPayload), c.CipherConf().TagLen) + frame := c.borrowWriteFrame(totalSize) + offset := 0 + copy(frame[offset:], salt) + offset += len(salt) + + offset, err = c.writeIdentityHeaderTo(frame, offset, salt) + if err != nil { + return 0, oops.Wrapf(err, "fail to write identity header") + } + + fixedHeaderOffset := offset + fixedHeaderPlain := frame[offset : offset+11] + fixedHeaderPlain[0] = HeaderTypeClientStream + binary.BigEndian.PutUint64(fixedHeaderPlain[1:9], uint64(time.Now().Unix())) + binary.BigEndian.PutUint16(fixedHeaderPlain[9:11], uint16(firstVarHeaderLen)) + sealed := c.cipherWrite.Seal(frame[:fixedHeaderOffset], c.nonceWrite, fixedHeaderPlain, nil) + offset = len(sealed) + common.BytesIncLittleEndian(c.nonceWrite) + + varHeaderOffset := offset + varHeaderPlain := frame[offset : offset+firstVarHeaderLen] + addrWritten, err := writeAddrInfoTo(varHeaderPlain, c.addr) + if err != nil { + return 0, oops.Wrapf(err, "fail to encode request address") + } + binary.BigEndian.PutUint16(varHeaderPlain[addrWritten:addrWritten+2], 0) + copy(varHeaderPlain[addrWritten+2:], b[:initialPayloadLen]) + sealed = c.cipherWrite.Seal(frame[:varHeaderOffset], c.nonceWrite, varHeaderPlain, nil) + offset = len(sealed) + common.BytesIncLittleEndian(c.nonceWrite) + + offset += c.sealPayload(frame[offset:], remainingPayload) + c.onceWrite = true + _, err = c.Conn.Write(frame[:offset]) + return n, err + } + if c.cipherWrite == nil { + return 0, fmt.Errorf("cipher is not initialized") + } + frameSize := encryptedPayloadLen(len(b), c.CipherConf().TagLen) + frame := c.borrowWriteFrame(frameSize) + offset := c.sealPayload(frame, b) + _, err = c.Conn.Write(frame[:offset]) + return n, err +} diff --git a/protocol/shadowsocks_2022/udp_conn.go b/protocol/shadowsocks_2022/udp_conn.go new file mode 100644 index 00000000..0e0d1221 --- /dev/null +++ b/protocol/shadowsocks_2022/udp_conn.go @@ -0,0 +1,320 @@ +package shadowsocks_2022 + +import ( + "bytes" + "crypto/cipher" + "encoding/binary" + "fmt" + "io" + "net" + "net/netip" + "sync" + "sync/atomic" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/pkg/fastrand" + "github.com/daeuniverse/outbound/pool" + "github.com/daeuniverse/outbound/protocol" + "github.com/daeuniverse/outbound/protocol/socks5" + disk_bloom "github.com/mzz2017/disk-bloom" + "github.com/samber/oops" +) + +// UdpConn represents a Shadowsocks 2022 UDP connection. +// Design follows sing-box: cipher is created once at session initialization. +type UdpConn struct { + *SS2022Core + + net.Conn + + sessionID [8]byte + packetID atomic.Uint64 + + // cipher is derived from the local session ID and reused for outbound + // packets. Inbound packets must decrypt against the remote session ID + // carried in each packet, so they use decryptCiphers instead. + cipher cipher.AEAD + cipherOnce sync.Once + cipherErr error + + // decryptCiphers caches inbound AEAD instances by remote session ID. + // Keeping this per-UdpConn avoids the old process-wide cache while + // preserving the protocol requirement that receive-side decryption uses + // the sender's session ID, not the local one. + decryptCiphers sync.Map // map[[8]byte]cipher.AEAD + + bloom *disk_bloom.FilterGroup + + replayWindow sync.Map + replayCount atomic.Int64 + + cleanupCounter atomic.Int64 +} + +const ( + udpPacketReplayWindowSize = 1024 + maxTrackedUdpSessions = 128 +) + +type udpSessionReplayState struct { + filter *ciphers.SlidingWindowFilter + lastSeen atomic.Int64 +} + +// NewUdpConn creates a new UDP connection bound to a shared SS2022 profile. +func NewUdpConn(conn net.Conn, core *SS2022Core, bloom *disk_bloom.FilterGroup) (*UdpConn, error) { + u := &UdpConn{ + SS2022Core: core, + Conn: conn, + bloom: bloom, + } + + // Generate session ID + fastrand.Read(u.sessionID[:]) + return u, nil +} + +func (c *UdpConn) ensureCipher() error { + c.cipherOnce.Do(func() { + c.cipher, c.cipherErr = CreateCipher(c.UPSK(), c.sessionID[:], c.CipherConf()) + if c.cipherErr != nil { + c.cipherErr = fmt.Errorf("failed to create session cipher: %w", c.cipherErr) + } + }) + return c.cipherErr +} + +func (c *UdpConn) decryptCipherFor(sessionID [8]byte) (cipher.AEAD, error) { + if cached, ok := c.decryptCiphers.Load(sessionID); ok { + return cached.(cipher.AEAD), nil + } + + sessionCipher, err := CreateCipher(c.UPSK(), sessionID[:], c.CipherConf()) + if err != nil { + return nil, fmt.Errorf("failed to create decrypt cipher for remote session: %w", err) + } + actual, _ := c.decryptCiphers.LoadOrStore(sessionID, sessionCipher) + return actual.(cipher.AEAD), nil +} + +func (c *UdpConn) nextPacketID() uint64 { + return c.packetID.Add(1) +} + +func (c *UdpConn) checkAndUpdateReplay(sessionID [8]byte, packetID uint64, now time.Time) bool { + nowNano := now.UnixNano() + expireNano := ciphers.SaltStorageDuration.Nanoseconds() + + if v, ok := c.replayWindow.Load(sessionID); ok { + state := v.(*udpSessionReplayState) + lastSeen := state.lastSeen.Load() + if nowNano-lastSeen > expireNano { + if c.replayWindow.CompareAndDelete(sessionID, v) { + c.replayCount.Add(-1) + } + } else { + state.lastSeen.Store(nowNano) + return state.filter.CheckAndUpdate(packetID) + } + } + + if c.cleanupCounter.Add(1)%cleanupInterval == 0 { + go c.cleanupExpiredSessions(nowNano, expireNano) + } + + newState := &udpSessionReplayState{ + filter: ciphers.NewSlidingWindowFilter(udpPacketReplayWindowSize), + } + newState.lastSeen.Store(nowNano) + + actual, loaded := c.replayWindow.LoadOrStore(sessionID, newState) + state := actual.(*udpSessionReplayState) + + if loaded { + state.lastSeen.Store(nowNano) + } else { + c.replayCount.Add(1) + c.evictOldestIfNeeded() + } + + return state.filter.CheckAndUpdate(packetID) +} + +const cleanupInterval = 1000 + +func (c *UdpConn) cleanupExpiredSessions(nowNano, expireNano int64) { + c.replayWindow.Range(func(key, value interface{}) bool { + state := value.(*udpSessionReplayState) + if nowNano-state.lastSeen.Load() > expireNano { + if c.replayWindow.CompareAndDelete(key, value) { + c.replayCount.Add(-1) + } + } + return true + }) +} + +func (c *UdpConn) evictOldestIfNeeded() { + for c.replayCount.Load() > maxTrackedUdpSessions { + var ( + found bool + oldestKey [8]byte + oldestVal any + oldestNano int64 = ^int64(0) + ) + + c.replayWindow.Range(func(key, value interface{}) bool { + state := value.(*udpSessionReplayState) + seen := state.lastSeen.Load() + if !found || seen < oldestNano { + found = true + oldestKey = key.([8]byte) + oldestVal = value + oldestNano = seen + } + return true + }) + + if !found { + c.replayCount.Store(0) + return + } + if c.replayWindow.CompareAndDelete(oldestKey, oldestVal) { + c.replayCount.Add(-1) + continue + } + // Retry if the oldest entry changed concurrently. + } +} + +func (c *UdpConn) WriteTo(b []byte, addr string) (int, error) { + if err := c.ensureCipher(); err != nil { + return 0, err + } + + packetID := c.nextPacketID() + var separateHeader [16]byte + copy(separateHeader[:8], c.sessionID[:]) + binary.BigEndian.PutUint64(separateHeader[8:], packetID) + + var separateHeaderEncrypted [16]byte + c.BlockCipherEncrypt().Encrypt(separateHeaderEncrypted[:], separateHeader[:]) + + addrInfo, err := socks5.AddressFromString(addr) + if err != nil { + return 0, oops.Wrapf(err, "fail to parse target address") + } + addrLen, err := addrInfoEncodedLen(addrInfo) + if err != nil { + return 0, oops.Wrapf(err, "fail to calculate address length") + } + messageLen := 1 + 8 + 2 + addrLen + len(b) + totalPacketLen := len(separateHeaderEncrypted) + c.IdentityHeaderLen() + messageLen + c.CipherConf().TagLen + packet := pool.Get(totalPacketLen) + defer pool.Put(packet) + offset := 0 + copy(packet[offset:], separateHeaderEncrypted[:]) + offset += len(separateHeaderEncrypted) + + identityHeaderLen, err := c.WriteIdentityHeader(packet[offset:], separateHeader[:]) + if err != nil { + return 0, oops.Wrapf(err, "fail to write identity header") + } + offset += identityHeaderLen + + messageOffset := offset + message := packet[messageOffset : messageOffset+messageLen] + message[0] = HeaderTypeClientStream + binary.BigEndian.PutUint64(message[1:9], uint64(time.Now().Unix())) + binary.BigEndian.PutUint16(message[9:11], 0) + addrWritten, err := writeAddrInfoTo(message[11:], addrInfo) + if err != nil { + return 0, oops.Wrapf(err, "fail to encode request address") + } + copy(message[11+addrWritten:], b) + + // Use session-level cipher (no cache lookup needed) + packet = c.cipher.Seal(packet[:messageOffset], separateHeader[4:16], message, nil) + + _, err = c.Conn.Write(packet) + return len(b), err +} + +func (c *UdpConn) ReadFrom(b []byte) (n int, addr netip.AddrPort, err error) { + buf := pool.Get(len(b) + 16 + c.CipherConf().TagLen) + defer pool.Put(buf) + n, err = c.Conn.Read(buf) + if err != nil { + return 0, netip.AddrPort{}, err + } + if n < 16 { + return 0, netip.AddrPort{}, fmt.Errorf("short length to decrypt") + } + + c.BlockCipherDecrypt().Decrypt(buf[:16], buf[:16]) + var sessionID [8]byte + copy(sessionID[:], buf[:8]) + packetID := binary.BigEndian.Uint64(buf[8:16]) + now := time.Now() + if !c.checkAndUpdateReplay(sessionID, packetID, now) { + return 0, netip.AddrPort{}, protocol.ErrReplayAttack + } + + payload := buf[16:n] + sessionCipher, err := c.decryptCipherFor(sessionID) + if err != nil { + return 0, netip.AddrPort{}, err + } + payload, err = sessionCipher.Open(payload[:0], buf[4:16], payload, nil) + if err != nil { + return 0, netip.AddrPort{}, err + } + + reader := bytes.NewReader(payload) + + var typ uint8 + if err := binary.Read(reader, binary.BigEndian, &typ); err != nil { + return 0, netip.AddrPort{}, fmt.Errorf("failed to read header type: %w", err) + } + + var timestampRaw uint64 + if err := binary.Read(reader, binary.BigEndian, ×tampRaw); err != nil { + return 0, netip.AddrPort{}, fmt.Errorf("failed to read timestamp: %w", err) + } + timestamp := time.Unix(int64(timestampRaw), 0) + + if _, err := reader.Seek(8, io.SeekCurrent); err != nil { + return 0, netip.AddrPort{}, fmt.Errorf("failed to skip session ID: %w", err) + } + + var paddingLength uint16 + if err := binary.Read(reader, binary.BigEndian, &paddingLength); err != nil { + return 0, netip.AddrPort{}, fmt.Errorf("failed to read padding length: %w", err) + } + + if _, err := reader.Seek(int64(paddingLength), io.SeekCurrent); err != nil { + return 0, netip.AddrPort{}, fmt.Errorf("failed to skip padding: %w", err) + } + + if typ != HeaderTypeServerStream { + return 0, netip.AddrPort{}, fmt.Errorf("received unexpected header type: %d", typ) + } + + if err := validateTimestamp(timestamp, now); err != nil { + return 0, netip.AddrPort{}, err + } + + netAddr, err := socks5.ReadAddr(reader) + if err != nil { + return 0, netip.AddrPort{}, err + } + + if udpAddr, ok := netAddr.(*net.UDPAddr); ok { + ipAddr, _ := netip.AddrFromSlice(udpAddr.IP) + addr = netip.AddrPortFrom(ipAddr, uint16(udpAddr.Port)) + } + + n, err = reader.Read(b) + return +} diff --git a/protocol/shadowsocks_2022/udp_conn_race_test.go b/protocol/shadowsocks_2022/udp_conn_race_test.go new file mode 100644 index 00000000..e6a9fcf7 --- /dev/null +++ b/protocol/shadowsocks_2022/udp_conn_race_test.go @@ -0,0 +1,172 @@ +package shadowsocks_2022 + +import ( + "runtime" + "sync" + "testing" + "time" + + "github.com/daeuniverse/outbound/ciphers" +) + +func TestReplayWindowRace(t *testing.T) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + psk := make([]byte, 32) + + core, err := NewSS2022Core(conf, [][]byte{psk}, psk) + if err != nil { + t.Fatal(err) + } + + conn := &UdpConn{ + SS2022Core: core, + } + + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + sessionID := [8]byte{byte(id % 256), byte(id / 256)} + for j := 0; j < 100; j++ { + conn.checkAndUpdateReplay(sessionID, uint64(j), time.Now()) + } + }(i) + } + wg.Wait() +} + +func TestNewUdpConnCreatesCipher(t *testing.T) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + psk := make([]byte, 32) + + core, err := NewSS2022Core(conf, [][]byte{psk}, psk) + if err != nil { + t.Fatal(err) + } + + conn := &UdpConn{ + SS2022Core: core, + } + + // Verify cipher is nil before initialization + if conn.cipher != nil { + t.Error("cipher should be nil before NewUdpConn") + } + + // Create cipher (simulating NewUdpConn behavior) + sessionID := make([]byte, 8) + cipher, err := CreateCipher(psk, sessionID, conf) + if err != nil { + t.Fatal(err) + } + conn.cipher = cipher + + // Verify cipher is created + if conn.cipher == nil { + t.Error("cipher should be created") + } +} + +func TestNoGoroutineLeak(t *testing.T) { + before := runtime.NumGoroutine() + + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + psk := make([]byte, 32) + + for i := 0; i < 100; i++ { + _, err := NewSS2022Core(conf, [][]byte{psk}, psk) + if err != nil { + t.Fatal(err) + } + } + + time.Sleep(100 * time.Millisecond) + after := runtime.NumGoroutine() + + if after-before > 5 { + t.Errorf("Potential goroutine leak: before=%d, after=%d", before, after) + } +} + +func BenchmarkSessionCipherAccess(b *testing.B) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + psk := make([]byte, 32) + + core, err := NewSS2022Core(conf, [][]byte{psk}, psk) + if err != nil { + b.Fatal(err) + } + + sessionID := make([]byte, 8) + cipher, err := CreateCipher(psk, sessionID, conf) + if err != nil { + b.Fatal(err) + } + + conn := &UdpConn{ + SS2022Core: core, + cipher: cipher, + } + + plaintext := make([]byte, 1400) + nonce := make([]byte, 12) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + out := make([]byte, len(plaintext)+16) + _ = conn.cipher.Seal(out[:0], nonce, plaintext, nil) + } +} + +func BenchmarkSessionCipherAccessParallel(b *testing.B) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + psk := make([]byte, 32) + + core, err := NewSS2022Core(conf, [][]byte{psk}, psk) + if err != nil { + b.Fatal(err) + } + + sessionID := make([]byte, 8) + cipher, err := CreateCipher(psk, sessionID, conf) + if err != nil { + b.Fatal(err) + } + + conn := &UdpConn{ + SS2022Core: core, + cipher: cipher, + } + + plaintext := make([]byte, 1400) + nonce := make([]byte, 12) + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + out := make([]byte, len(plaintext)+16) + _ = conn.cipher.Seal(out[:0], nonce, plaintext, nil) + } + }) +} + +func BenchmarkReplayCheck(b *testing.B) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + psk := make([]byte, 32) + + core, err := NewSS2022Core(conf, [][]byte{psk}, psk) + if err != nil { + b.Fatal(err) + } + + conn := &UdpConn{ + SS2022Core: core, + } + sessionID := [8]byte{1, 2, 3, 4, 5, 6, 7, 8} + now := time.Now() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + conn.checkAndUpdateReplay(sessionID, uint64(i), now) + } +} diff --git a/protocol/shadowsocks_2022/udp_conn_test.go b/protocol/shadowsocks_2022/udp_conn_test.go new file mode 100644 index 00000000..8ec872ae --- /dev/null +++ b/protocol/shadowsocks_2022/udp_conn_test.go @@ -0,0 +1,216 @@ +package shadowsocks_2022 + +import ( + "encoding/binary" + "io" + "net" + "net/netip" + "sync" + "testing" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/protocol" + "github.com/daeuniverse/outbound/protocol/socks5" +) + +type udpReadBufferConn struct { + packet []byte + read bool +} + +func (c *udpReadBufferConn) Read(p []byte) (int, error) { + if c.read { + return 0, io.EOF + } + c.read = true + copy(p, c.packet) + return len(c.packet), nil +} + +func (c *udpReadBufferConn) Write(p []byte) (int, error) { return len(p), nil } +func (c *udpReadBufferConn) Close() error { return nil } +func (c *udpReadBufferConn) LocalAddr() net.Addr { return udpConnTestAddr("local") } +func (c *udpReadBufferConn) RemoteAddr() net.Addr { return udpConnTestAddr("remote") } +func (c *udpReadBufferConn) SetDeadline(_ time.Time) error { return nil } +func (c *udpReadBufferConn) SetReadDeadline(_ time.Time) error { return nil } +func (c *udpReadBufferConn) SetWriteDeadline(_ time.Time) error { return nil } + +type udpConnTestAddr string + +func (a udpConnTestAddr) Network() string { return "udp-test" } +func (a udpConnTestAddr) String() string { return string(a) } + +func buildServerPacket(t *testing.T, core *SS2022Core, serverSessionID, clientSessionID [8]byte, packetID uint64, addr string, payload []byte) []byte { + t.Helper() + + var separateHeader [16]byte + copy(separateHeader[:8], serverSessionID[:]) + binary.BigEndian.PutUint64(separateHeader[8:], packetID) + + var separateHeaderEncrypted [16]byte + core.BlockCipherDecrypt().Encrypt(separateHeaderEncrypted[:], separateHeader[:]) + + addrInfo, err := socks5.AddressFromString(addr) + if err != nil { + t.Fatalf("AddressFromString: %v", err) + } + addrLen, err := addrInfoEncodedLen(addrInfo) + if err != nil { + t.Fatalf("addrInfoEncodedLen: %v", err) + } + + messageLen := 1 + 8 + 8 + 2 + addrLen + len(payload) + message := make([]byte, messageLen) + message[0] = HeaderTypeServerStream + binary.BigEndian.PutUint64(message[1:9], uint64(time.Now().Unix())) + copy(message[9:17], clientSessionID[:]) + addrWritten, err := writeAddrInfoTo(message[19:], addrInfo) + if err != nil { + t.Fatalf("writeAddrInfoTo: %v", err) + } + copy(message[19+addrWritten:], payload) + + sessionCipher, err := CreateCipher(core.UPSK(), serverSessionID[:], core.CipherConf()) + if err != nil { + t.Fatalf("CreateCipher: %v", err) + } + return sessionCipher.Seal(separateHeaderEncrypted[:], separateHeader[4:16], message, nil) +} + +func TestValidateTimestamp(t *testing.T) { + now := time.Now() + if err := validateTimestamp(now, now); err != nil { + t.Fatalf("now should pass: %v", err) + } + if err := validateTimestamp(now.Add(ciphers.TimestampTolerance-time.Millisecond), now); err != nil { + t.Fatalf("near-future timestamp should pass: %v", err) + } + if err := validateTimestamp(now.Add(-ciphers.TimestampTolerance+time.Millisecond), now); err != nil { + t.Fatalf("near-past timestamp should pass: %v", err) + } + if err := validateTimestamp(now.Add(ciphers.TimestampTolerance+time.Millisecond), now); err != protocol.ErrReplayAttack { + t.Fatalf("too-far future timestamp should fail with replay, got: %v", err) + } + if err := validateTimestamp(now.Add(-ciphers.TimestampTolerance-time.Millisecond), now); err != protocol.ErrReplayAttack { + t.Fatalf("too-old timestamp should fail with replay, got: %v", err) + } +} + +func TestUdpConn_NextPacketID_ConcurrentUnique(t *testing.T) { + u := &UdpConn{} + const n = 2000 + + ids := make(chan uint64, n) + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + ids <- u.nextPacketID() + }() + } + wg.Wait() + close(ids) + + seen := make(map[uint64]struct{}, n) + var minID uint64 = ^uint64(0) + var maxID uint64 + for id := range ids { + if _, ok := seen[id]; ok { + t.Fatalf("duplicate packetID: %d", id) + } + seen[id] = struct{}{} + if id < minID { + minID = id + } + if id > maxID { + maxID = id + } + } + + if len(seen) != n { + t.Fatalf("unexpected unique count: got %d, want %d", len(seen), n) + } + if minID != 1 { + t.Fatalf("unexpected min packetID: got %d, want 1", minID) + } + if maxID != n { + t.Fatalf("unexpected max packetID: got %d, want %d", maxID, n) + } +} + +func TestUdpConn_ReplayWindow_PerSessionAndExpiry(t *testing.T) { + u := &UdpConn{} + now := time.Now() + + var sid1 [8]byte + copy(sid1[:], []byte{1, 1, 1, 1, 1, 1, 1, 1}) + var sid2 [8]byte + copy(sid2[:], []byte{2, 2, 2, 2, 2, 2, 2, 2}) + + if !u.checkAndUpdateReplay(sid1, 1, now) { + t.Fatalf("sid1 packet 1 should pass") + } + if u.checkAndUpdateReplay(sid1, 1, now) { + t.Fatalf("sid1 duplicate packet 1 should fail") + } + if !u.checkAndUpdateReplay(sid1, 2, now) { + t.Fatalf("sid1 packet 2 should pass") + } + if !u.checkAndUpdateReplay(sid2, 1, now) { + t.Fatalf("sid2 packet 1 should pass independently") + } + + if !u.checkAndUpdateReplay(sid1, 5000, now) { + t.Fatalf("sid1 packet 5000 should pass") + } + if u.checkAndUpdateReplay(sid1, 1, now) { + t.Fatalf("sid1 old packet should fail after large jump") + } + + future := now.Add(ciphers.SaltStorageDuration + time.Second) + if !u.checkAndUpdateReplay(sid1, 1, future) { + t.Fatalf("sid1 should reset after expiry and accept packet 1") + } +} + +func TestUdpConn_ReadFromUsesRemoteSessionCipher(t *testing.T) { + conf := ciphers.Aead2022CiphersConf["2022-blake3-aes-256-gcm"] + if conf == nil { + t.Fatal("missing ss2022 cipher config") + } + + psk := make([]byte, conf.KeyLen) + for i := range psk { + psk[i] = 0x23 + } + core, err := NewSS2022Core(conf, [][]byte{psk}, psk) + if err != nil { + t.Fatal(err) + } + + localSessionID := [8]byte{1, 2, 3, 4, 5, 6, 7, 8} + remoteSessionID := [8]byte{8, 7, 6, 5, 4, 3, 2, 1} + wantAddr := netip.MustParseAddrPort("203.0.113.9:853") + wantPayload := []byte("ss2022-udp-remote-session") + + packet := buildServerPacket(t, core, remoteSessionID, localSessionID, 1, wantAddr.String(), wantPayload) + conn, err := NewUdpConn(&udpReadBufferConn{packet: packet}, core, nil) + if err != nil { + t.Fatal(err) + } + conn.sessionID = localSessionID + + buf := make([]byte, 128) + n, addr, err := conn.ReadFrom(buf) + if err != nil { + t.Fatalf("ReadFrom: %v", err) + } + if addr != wantAddr { + t.Fatalf("unexpected addr: got %v want %v", addr, wantAddr) + } + if got := string(buf[:n]); got != string(wantPayload) { + t.Fatalf("unexpected payload: got %q want %q", got, string(wantPayload)) + } +} diff --git a/protocol/shadowsocks_2022/validation.go b/protocol/shadowsocks_2022/validation.go new file mode 100644 index 00000000..83d7659b --- /dev/null +++ b/protocol/shadowsocks_2022/validation.go @@ -0,0 +1,16 @@ +package shadowsocks_2022 + +import ( + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/protocol" +) + +func validateTimestamp(timestamp time.Time, now time.Time) error { + if timestamp.Before(now.Add(-ciphers.TimestampTolerance)) || + timestamp.After(now.Add(ciphers.TimestampTolerance)) { + return protocol.ErrReplayAttack + } + return nil +} diff --git a/protocol/shadowsocks_stream/udp_conn_test.go b/protocol/shadowsocks_stream/udp_conn_test.go new file mode 100644 index 00000000..f154448f --- /dev/null +++ b/protocol/shadowsocks_stream/udp_conn_test.go @@ -0,0 +1,107 @@ +package shadowsocks_stream + +import ( + "io" + "net/netip" + "sync" + "testing" + "time" + + "github.com/daeuniverse/outbound/ciphers" + "github.com/daeuniverse/outbound/protocol/infra/socks" +) + +type recordingPacketConn struct { + mu sync.Mutex + writes []recordedPacketWrite +} + +type recordedPacketWrite struct { + addr string + data []byte +} + +func (c *recordingPacketConn) Read([]byte) (int, error) { return 0, io.EOF } + +func (c *recordingPacketConn) Write(p []byte) (int, error) { + return c.WriteTo(p, "") +} + +func (c *recordingPacketConn) ReadFrom([]byte) (int, netip.AddrPort, error) { + return 0, netip.AddrPort{}, io.EOF +} + +func (c *recordingPacketConn) WriteTo(p []byte, addr string) (int, error) { + clone := append([]byte(nil), p...) + c.mu.Lock() + c.writes = append(c.writes, recordedPacketWrite{addr: addr, data: clone}) + c.mu.Unlock() + return len(p), nil +} + +func (c *recordingPacketConn) Close() error { return nil } +func (c *recordingPacketConn) SetDeadline(time.Time) error { return nil } +func (c *recordingPacketConn) SetReadDeadline(time.Time) error { return nil } +func (c *recordingPacketConn) SetWriteDeadline(time.Time) error { return nil } + +func TestUdpConnConcurrentWriteTo(t *testing.T) { + t.Helper() + + recorder := &recordingPacketConn{} + cipher, err := ciphers.NewStreamCipher("none", "password") + if err != nil { + t.Fatalf("NewStreamCipher failed: %v", err) + } + defaultAddr, err := socks.ParseAddr("1.1.1.1:53") + if err != nil { + t.Fatalf("ParseAddr default target failed: %v", err) + } + + conn := NewUdpConn(recorder, cipher, defaultAddr, "127.0.0.1:8388") + targets := []string{ + "1.1.1.1:53", + "8.8.8.8:53", + "9.9.9.9:53", + "208.67.222.222:53", + } + + var wg sync.WaitGroup + for _, target := range targets { + target := target + wg.Add(1) + go func() { + defer wg.Done() + if _, err := conn.WriteTo([]byte("payload-"+target), target); err != nil { + t.Errorf("WriteTo(%s) failed: %v", target, err) + } + }() + } + wg.Wait() + + if len(recorder.writes) != len(targets) { + t.Fatalf("unexpected write count: got %d want %d", len(recorder.writes), len(targets)) + } + + seen := make(map[string]bool, len(targets)) + for _, write := range recorder.writes { + if write.addr != "127.0.0.1:8388" { + t.Fatalf("unexpected proxy addr: got %q", write.addr) + } + addr := socks.SplitAddr(write.data) + if addr == nil { + t.Fatal("failed to parse encoded target addr") + } + target := addr.String() + payload := string(write.data[len(addr):]) + if payload != "payload-"+target { + t.Fatalf("payload mismatch for %s: got %q", target, payload) + } + seen[target] = true + } + + for _, target := range targets { + if !seen[target] { + t.Fatalf("missing target write for %s", target) + } + } +} diff --git a/protocol/socks5/addr.go b/protocol/socks5/addr.go new file mode 100644 index 00000000..8934eeb2 --- /dev/null +++ b/protocol/socks5/addr.go @@ -0,0 +1,174 @@ +package socks5 + +import ( + "bytes" + "encoding/binary" + "fmt" + "io" + "net" + "net/netip" + "strconv" + + "github.com/daeuniverse/outbound/pool" +) + +type AddressType uint8 + +// Address type constants for Shadowsocks protocol +const ( + AddressTypeIPv4 AddressType = 1 + AddressTypeDomain AddressType = 3 + AddressTypeIPv6 AddressType = 4 +) + +var ( + ErrInvalidAddress = fmt.Errorf("invalid address") +) + +// AddressInfo represents decoded address information +type AddressInfo struct { + Type AddressType + Hostname string + IP netip.Addr + Port uint16 +} + +func WriteAddr(addr string, buf *bytes.Buffer) error { + addressInfo, err := AddressFromString(addr) + if err != nil { + return err + } + return WriteAddrInfo(addressInfo, buf) +} + +// WriteAddrInfo writes address information to writer +func WriteAddrInfo(addr *AddressInfo, w io.Writer) error { + var typeBuf [1]byte + typeBuf[0] = byte(addr.Type) + if _, err := w.Write(typeBuf[:]); err != nil { + return err + } + + switch addr.Type { + case AddressTypeIPv4, AddressTypeIPv6: + if _, err := w.Write(addr.IP.AsSlice()); err != nil { + return err + } + var portBuf [2]byte + binary.BigEndian.PutUint16(portBuf[:], addr.Port) + _, err := w.Write(portBuf[:]) + return err + case AddressTypeDomain: + lenDN := len(addr.Hostname) + if lenDN > 255 { + return fmt.Errorf("domain name too long: %d bytes", lenDN) + } + var lenBuf [1]byte + lenBuf[0] = uint8(lenDN) + if _, err := w.Write(lenBuf[:]); err != nil { + return err + } + if _, err := io.WriteString(w, addr.Hostname); err != nil { + return err + } + var portBuf [2]byte + binary.BigEndian.PutUint16(portBuf[:], addr.Port) + _, err := w.Write(portBuf[:]) + return err + default: + return fmt.Errorf("unsupported address type: %v", addr.Type) + } + return nil +} + +func ReadAddr(data io.Reader) (net.Addr, error) { + addressInfo, err := ReadAddrInfo(data) + if err != nil { + return nil, err + } + + // Create address object (only support IP addresses for UDP) + switch addressInfo.Type { + case AddressTypeIPv4, AddressTypeIPv6: + return net.UDPAddrFromAddrPort(netip.AddrPortFrom(addressInfo.IP, addressInfo.Port)), nil + default: + return nil, fmt.Errorf("unsupported address type for UDP: %v", addressInfo.Type) + } +} + +// ReadAddr reads address from buffer +func ReadAddrInfo(data io.Reader) (*AddressInfo, error) { + var typ uint8 + if err := binary.Read(data, binary.BigEndian, &typ); err != nil { + return nil, fmt.Errorf("%w: too short", ErrInvalidAddress) + } + + info := &AddressInfo{Type: AddressType(typ)} + + switch info.Type { + case AddressTypeIPv4: + ip := pool.Get(4) + defer pool.Put(ip) + if _, err := data.Read(ip); err != nil { + return nil, fmt.Errorf("failed to read IP: %w", err) + } + info.IP = netip.AddrFrom4([4]byte(ip)) + if err := binary.Read(data, binary.BigEndian, &info.Port); err != nil { + return nil, fmt.Errorf("failed to read port: %w", err) + } + case AddressTypeIPv6: + ip := pool.Get(16) + defer pool.Put(ip) + if _, err := data.Read(ip); err != nil { + return nil, fmt.Errorf("failed to read IP: %w", err) + } + info.IP = netip.AddrFrom16([16]byte(ip)) + if err := binary.Read(data, binary.BigEndian, &info.Port); err != nil { + return nil, fmt.Errorf("failed to read port: %w", err) + } + case AddressTypeDomain: + var domainLen uint8 + if err := binary.Read(data, binary.BigEndian, &domainLen); err != nil { + return nil, fmt.Errorf("failed to read domain length: %w", err) + } + domain := pool.Get(int(domainLen)) + defer pool.Put(domain) + if _, err := data.Read(domain); err != nil { + return nil, fmt.Errorf("failed to read domain: %w", err) + } + info.Hostname = string(domain) + if err := binary.Read(data, binary.BigEndian, &info.Port); err != nil { + return nil, fmt.Errorf("failed to read port: %w", err) + } + default: + return nil, fmt.Errorf("%w: invalid type: %v", ErrInvalidAddress, info.Type) + } + return info, nil +} + +func AddressFromString(addr string) (*AddressInfo, error) { + hostname, port_, err := net.SplitHostPort(addr) + if err != nil { + return nil, err + } + port, err := strconv.ParseUint(port_, 10, 16) + if err != nil { + return nil, fmt.Errorf("invalid port: %v", port_) + } + + info := &AddressInfo{Port: uint16(port)} + + ip, err := netip.ParseAddr(hostname) + if err != nil { + info.Type = AddressTypeDomain + info.Hostname = hostname + } else { + info.IP = ip + if ip.Is4() { + info.Type = AddressTypeIPv4 + } else { + info.Type = AddressTypeIPv6 + } + } + return info, nil +} diff --git a/protocol/socks5/packet.go b/protocol/socks5/packet.go index 5e9701ca..bfd6f6b2 100644 --- a/protocol/socks5/packet.go +++ b/protocol/socks5/packet.go @@ -3,6 +3,7 @@ package socks5 import ( + "context" "errors" "fmt" "net" @@ -20,15 +21,18 @@ type PktConn struct { ctrlConn netproxy.Conn // tcp control conn target string proxyAddr string + cancel context.CancelFunc } // NewPktConn returns a PktConn, the writeAddr must be *net.UDPAddr or *net.UnixAddr. func NewPktConn(c netproxy.PacketConn, proxyAddr string, targetAddr string, ctrlConn netproxy.Conn) *PktConn { + ctx, cancel := context.WithCancel(context.Background()) pc := &PktConn{ PacketConn: c, target: targetAddr, proxyAddr: proxyAddr, ctrlConn: ctrlConn, + cancel: cancel, } if ctrlConn != nil { @@ -36,6 +40,11 @@ func NewPktConn(c netproxy.PacketConn, proxyAddr string, targetAddr string, ctrl buf := pool.Get(1) defer pool.Put(buf) for { + select { + case <-ctx.Done(): + return + default: + } _, err := ctrlConn.Read(buf) if err, ok := err.(net.Error); ok && err.Timeout() { continue @@ -114,6 +123,9 @@ func (pc *PktConn) WriteTo(b []byte, addr string) (int, error) { // Close . func (pc *PktConn) Close() error { + if pc.cancel != nil { + pc.cancel() + } if pc.ctrlConn != nil { pc.ctrlConn.Close() } diff --git a/protocol/socks5/packet_test.go b/protocol/socks5/packet_test.go new file mode 100644 index 00000000..cf1ee417 --- /dev/null +++ b/protocol/socks5/packet_test.go @@ -0,0 +1,104 @@ +package socks5 + +import ( + "io" + "net/netip" + "sync" + "testing" + "time" + + "github.com/daeuniverse/outbound/protocol/infra/socks" +) + +type recordingPacketConn struct { + mu sync.Mutex + writes []recordedPacketWrite +} + +type recordedPacketWrite struct { + addr string + data []byte +} + +func (c *recordingPacketConn) Read([]byte) (int, error) { return 0, io.EOF } + +func (c *recordingPacketConn) Write(p []byte) (int, error) { + return c.WriteTo(p, "") +} + +func (c *recordingPacketConn) ReadFrom([]byte) (int, netip.AddrPort, error) { + return 0, netip.AddrPort{}, io.EOF +} + +func (c *recordingPacketConn) WriteTo(p []byte, addr string) (int, error) { + clone := append([]byte(nil), p...) + c.mu.Lock() + c.writes = append(c.writes, recordedPacketWrite{addr: addr, data: clone}) + c.mu.Unlock() + return len(p), nil +} + +func (c *recordingPacketConn) Close() error { return nil } +func (c *recordingPacketConn) SetDeadline(time.Time) error { return nil } +func (c *recordingPacketConn) SetReadDeadline(time.Time) error { return nil } +func (c *recordingPacketConn) SetWriteDeadline(time.Time) error { return nil } + +func TestPktConnConcurrentWriteTo(t *testing.T) { + t.Helper() + + recorder := &recordingPacketConn{} + pc := NewPktConn(recorder, "127.0.0.1:1080", "1.1.1.1:53", nil) + + targets := []string{ + "1.1.1.1:53", + "8.8.8.8:53", + "9.9.9.9:53", + "208.67.222.222:53", + } + + var wg sync.WaitGroup + for _, target := range targets { + target := target + wg.Add(1) + go func() { + defer wg.Done() + if _, err := pc.WriteTo([]byte("payload-"+target), target); err != nil { + t.Errorf("WriteTo(%s) failed: %v", target, err) + } + }() + } + wg.Wait() + + if len(recorder.writes) != len(targets) { + t.Fatalf("unexpected write count: got %d want %d", len(recorder.writes), len(targets)) + } + + seen := make(map[string]bool, len(targets)) + for _, write := range recorder.writes { + if write.addr != "127.0.0.1:1080" { + t.Fatalf("unexpected proxy addr: got %q", write.addr) + } + if len(write.data) < 4 { + t.Fatalf("unexpected short socks5 packet: %d", len(write.data)) + } + if write.data[0] != 0 || write.data[1] != 0 || write.data[2] != 0 { + t.Fatalf("unexpected socks5 reserved header: %v", write.data[:3]) + } + addr := socks.SplitAddr(write.data[3:]) + if addr == nil { + t.Fatal("failed to parse encoded target address") + } + target := addr.String() + payload := string(write.data[3+len(addr):]) + if payload != "payload-"+target { + t.Fatalf("payload mismatch for %s: got %q", target, payload) + } + seen[target] = true + } + + for _, target := range targets { + if !seen[target] { + t.Fatalf("missing target write for %s", target) + } + } +} diff --git a/protocol/trojanc/conn.go b/protocol/trojanc/conn.go index 24cd8659..0590897c 100644 --- a/protocol/trojanc/conn.go +++ b/protocol/trojanc/conn.go @@ -9,6 +9,7 @@ import ( "encoding/hex" "fmt" "io" + "net" "sync" "time" @@ -19,6 +20,9 @@ import ( var ( CRLF = []byte{13, 10} FailAuthErr = fmt.Errorf("incorrect password") + + // passwordHashCache caches SHA224 hash results of passwords + passwordHashCache sync.Map ) type Conn struct { @@ -31,15 +35,34 @@ type Conn struct { onceRead sync.Once } -func NewConn(conn netproxy.Conn, metadata Metadata, password string) (c *Conn, err error) { +// getPasswordHash retrieves the SHA224 hash of a password (with caching) +// Optimization: Uses sync.Map to cache hash results, avoiding repeated computation +func getPasswordHash(password string) [56]byte { + // Try to get from cache + if cached, ok := passwordHashCache.Load(password); ok { + return cached.([56]byte) + } + + // Cache miss, calculate hash hash := sha256.New224() hash.Write([]byte(password)) + var result [56]byte + hex.Encode(result[:], hash.Sum(nil)) + + // Store in cache + passwordHashCache.Store(password, result) + return result +} + +func NewConn(conn netproxy.Conn, metadata Metadata, password string) (c *Conn, err error) { + // Use cached password hash for ~6x performance improvement + pass := getPasswordHash(password) + c = &Conn{ Conn: conn, metadata: metadata, - pass: [56]byte{}, + pass: pass, } - hex.Encode(c.pass[:], hash.Sum(nil)) if metadata.Network == "tcp" && metadata.IsClient { time.AfterFunc(100*time.Millisecond, func() { // avoid the situation where the server sends messages first @@ -51,31 +74,50 @@ func NewConn(conn netproxy.Conn, metadata Metadata, password string) (c *Conn, e return c, nil } -func (c *Conn) reqHeaderFromPool(payload []byte) (buf []byte) { +func (c *Conn) reqHeaderFromPool() (buf []byte) { reqLen := c.metadata.Len() - buf = pool.Get(56 + 2 + 1 + reqLen + 2 + len(payload)) + buf = pool.Get(56 + 2 + 1 + reqLen + 2) copy(buf, c.pass[:]) copy(buf[56:], CRLF) buf[58] = NetworkToByte(c.metadata.Network) c.metadata.PackTo(buf[59:]) copy(buf[59+reqLen:], CRLF) - copy(buf[61+reqLen:], payload) return buf } +func (c *Conn) writeRequestHeader(payload []byte) (n int, err error) { + header := c.reqHeaderFromPool() + defer pool.Put(header) + + buffers := net.Buffers{header} + if len(payload) > 0 { + buffers = append(buffers, payload) + } + written, err := buffers.WriteTo(c.Conn) + if err != nil { + if written <= int64(len(header)) { + return 0, fmt.Errorf("write header: %w", err) + } + return int(written) - len(header), fmt.Errorf("write header: %w", err) + } + if written < int64(len(header)) { + return 0, fmt.Errorf("write header: %w", io.ErrShortWrite) + } + return int(written) - len(header), nil +} + func (c *Conn) Write(b []byte) (n int, err error) { c.writeMutex.Lock() defer c.writeMutex.Unlock() if !c.onceWrite { if c.metadata.IsClient { - buf := c.reqHeaderFromPool(b) - defer pool.Put(buf) - if _, err = c.Conn.Write(buf); err != nil { - return 0, fmt.Errorf("write header: %w", err) + n, err = c.writeRequestHeader(b) + if err != nil { + return n, err } c.onceWrite = true - return len(b), nil + return n, nil } } return c.Conn.Write(b) diff --git a/protocol/trojanc/conn_bench_test.go b/protocol/trojanc/conn_bench_test.go new file mode 100644 index 00000000..075def4f --- /dev/null +++ b/protocol/trojanc/conn_bench_test.go @@ -0,0 +1,136 @@ +package trojanc + +import ( + "crypto/sha256" + "encoding/hex" + "sync" + "testing" +) + +// BenchmarkPasswordHashBaseline tests the current password hash implementation +func BenchmarkPasswordHashBaseline(b *testing.B) { + password := "test-password-12345" + b.ResetTimer() + + for i := 0; i < b.N; i++ { + hash := sha256.New224() + hash.Write([]byte(password)) + var result [56]byte + hex.Encode(result[:], hash.Sum(nil)) + _ = result + } +} + +// BenchmarkPasswordHashCached tests using cached password hash +func BenchmarkPasswordHashCached(b *testing.B) { + password := "test-password-12345" + cache := make(map[string][56]byte) + + // Pre-compute + hash := sha256.New224() + hash.Write([]byte(password)) + var result [56]byte + hex.Encode(result[:], hash.Sum(nil)) + cache[password] = result + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + if cached, ok := cache[password]; ok { + _ = cached + } + } +} + +// BenchmarkPasswordHashSyncMap tests using sync.Map for caching +func BenchmarkPasswordHashSyncMap(b *testing.B) { + password := "test-password-12345" + var cache sync.Map + + // 预计算 + hash := sha256.New224() + hash.Write([]byte(password)) + var result [56]byte + hex.Encode(result[:], hash.Sum(nil)) + cache.Store(password, result) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + if cached, ok := cache.Load(password); ok { + _ = cached.([56]byte) + } + } +} + +// BenchmarkNewConnComparison compares NewConn performance before and after optimization +func BenchmarkNewConnComparison(b *testing.B) { + password := "test-password-12345" + + b.Run("Original", func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + // 原始实现：每次都重新计算 + hash := sha256.New224() + hash.Write([]byte(password)) + var pass [56]byte + hex.Encode(pass[:], hash.Sum(nil)) + _ = pass + } + }) + + b.Run("OptimizedCached", func(b *testing.B) { + // 预热缓存 + _ = getPasswordHash(password) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // 优化后：使用缓存 + pass := getPasswordHash(password) + _ = pass + } + }) +} + +// BenchmarkMultiplePasswords tests scenarios with multiple different passwords +func BenchmarkMultiplePasswords(b *testing.B) { + passwords := []string{ + "password1", + "password2", + "password3", + "password4", + "password5", + } + + b.Run("Baseline", func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + password := passwords[i%len(passwords)] + hash := sha256.New224() + hash.Write([]byte(password)) + var result [56]byte + hex.Encode(result[:], hash.Sum(nil)) + _ = result + } + }) + + b.Run("SyncMap", func(b *testing.B) { + var cache sync.Map + + b.ResetTimer() + for i := 0; i < b.N; i++ { + password := passwords[i%len(passwords)] + + if cached, ok := cache.Load(password); ok { + _ = cached.([56]byte) + continue + } + + hash := sha256.New224() + hash.Write([]byte(password)) + var result [56]byte + hex.Encode(result[:], hash.Sum(nil)) + cache.Store(password, result) + } + }) +} diff --git a/protocol/trojanc/conn_optimized_test.go b/protocol/trojanc/conn_optimized_test.go new file mode 100644 index 00000000..9c79abbd --- /dev/null +++ b/protocol/trojanc/conn_optimized_test.go @@ -0,0 +1,97 @@ +package trojanc + +import ( + "crypto/sha256" + "encoding/hex" + "testing" + + "github.com/daeuniverse/outbound/netproxy" + "github.com/daeuniverse/outbound/protocol" +) + +// BenchmarkNewConnOptimized tests optimized NewConn performance +func BenchmarkNewConnOptimized(b *testing.B) { + // Mock network connection (nil is fine in benchmarks as we don't actually read/write) + var mockConn netproxy.Conn + + metadata := Metadata{ + Metadata: protocol.Metadata{ + Hostname: "example.com", + Port: 443, + }, + Network: "tcp", + } + password := "test-password-12345" + + // 预热缓存 + _, _ = NewConn(mockConn, metadata, password) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _, _ = NewConn(mockConn, metadata, password) + } +} + +// BenchmarkNewConnMultiplePasswords tests multiple password scenarios +func BenchmarkNewConnMultiplePasswords(b *testing.B) { + var mockConn netproxy.Conn + + passwords := []string{ + "password1", + "password2", + "password3", + "password4", + "password5", + } + + metadata := Metadata{ + Metadata: protocol.Metadata{ + Hostname: "example.com", + Port: 443, + }, + Network: "tcp", + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + password := passwords[i%len(passwords)] + _, _ = NewConn(mockConn, metadata, password) + } +} + +// TestPasswordHashConsistency tests password hash consistency +func TestPasswordHashConsistency(t *testing.T) { + password := "test-password" + + // First retrieval (computation) + hash1 := getPasswordHash(password) + + // 第二次获取（缓存） + hash2 := getPasswordHash(password) + + // 验证一致性 + if hash1 != hash2 { + t.Errorf("password hash inconsistency") + } +} + +// TestPasswordHashCorrectness tests password hash correctness +func TestPasswordHashCorrectness(t *testing.T) { + password := "test-password" + + // Compute using new function + hash := getPasswordHash(password) + + // 手动计算预期值 + expected := [56]byte{} + h := sha256.New224() + h.Write([]byte(password)) + hex.Encode(expected[:], h.Sum(nil)) + + // 验证正确性 + if hash != expected { + t.Errorf("password hash incorrect") + } +} diff --git a/protocol/trojanc/conn_write_bench_test.go b/protocol/trojanc/conn_write_bench_test.go new file mode 100644 index 00000000..178e9817 --- /dev/null +++ b/protocol/trojanc/conn_write_bench_test.go @@ -0,0 +1,43 @@ +package trojanc + +import ( + "io" + "testing" + "time" + + "github.com/daeuniverse/outbound/netproxy" + "github.com/daeuniverse/outbound/protocol" +) + +type discardConn struct{} + +func (c *discardConn) Read(_ []byte) (int, error) { return 0, io.EOF } +func (c *discardConn) Write(p []byte) (int, error) { return len(p), nil } +func (c *discardConn) Close() error { return nil } +func (c *discardConn) SetDeadline(_ time.Time) error { return nil } +func (c *discardConn) SetReadDeadline(_ time.Time) error { return nil } +func (c *discardConn) SetWriteDeadline(_ time.Time) error { return nil } + +var _ netproxy.Conn = (*discardConn)(nil) + +func BenchmarkConnFirstWrite(b *testing.B) { + baseMetadata, err := protocol.ParseMetadata("example.com:443") + if err != nil { + b.Fatalf("parse metadata failed: %v", err) + } + baseMetadata.IsClient = true + payload := make([]byte, 64<<10) + + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + conn, err := NewConn(&discardConn{}, Metadata{Metadata: baseMetadata, Network: "tcp"}, "test-password") + if err != nil { + b.Fatalf("new conn failed: %v", err) + } + if _, err := conn.Write(payload); err != nil { + b.Fatalf("write failed: %v", err) + } + } +} diff --git a/protocol/trojanc/conn_write_test.go b/protocol/trojanc/conn_write_test.go new file mode 100644 index 00000000..cabcb0d3 --- /dev/null +++ b/protocol/trojanc/conn_write_test.go @@ -0,0 +1,68 @@ +package trojanc + +import ( + "bytes" + "testing" + "time" + + "github.com/daeuniverse/outbound/netproxy" + "github.com/daeuniverse/outbound/protocol" +) + +type captureConn struct { + bytes.Buffer +} + +func (c *captureConn) Close() error { return nil } +func (c *captureConn) SetDeadline(_ time.Time) error { return nil } +func (c *captureConn) SetReadDeadline(_ time.Time) error { return nil } +func (c *captureConn) SetWriteDeadline(_ time.Time) error { return nil } + +var _ netproxy.Conn = (*captureConn)(nil) + +func TestConnWrite_WritesHeaderOnlyOnce(t *testing.T) { + raw := &captureConn{} + baseMetadata, err := protocol.ParseMetadata("example.com:443") + if err != nil { + t.Fatalf("parse metadata failed: %v", err) + } + baseMetadata.IsClient = true + conn, err := NewConn(raw, Metadata{Metadata: baseMetadata, Network: "tcp"}, "test-password") + if err != nil { + t.Fatalf("new conn failed: %v", err) + } + + firstPayload := []byte("hello") + n, err := conn.Write(firstPayload) + if err != nil { + t.Fatalf("first write failed: %v", err) + } + if n != len(firstPayload) { + t.Fatalf("unexpected first write length: got %d want %d", n, len(firstPayload)) + } + + headerMetadata := Metadata{Metadata: baseMetadata, Network: "tcp"} + headerLen := 56 + 2 + 1 + headerMetadata.Len() + 2 + if raw.Len() != headerLen+len(firstPayload) { + t.Fatalf("unexpected buffered length after first write: got %d want %d", raw.Len(), headerLen+len(firstPayload)) + } + if !bytes.Equal(raw.Bytes()[raw.Len()-len(firstPayload):], firstPayload) { + t.Fatal("first payload mismatch") + } + + secondPayload := []byte("world") + n, err = conn.Write(secondPayload) + if err != nil { + t.Fatalf("second write failed: %v", err) + } + if n != len(secondPayload) { + t.Fatalf("unexpected second write length: got %d want %d", n, len(secondPayload)) + } + + if raw.Len() != headerLen+len(firstPayload)+len(secondPayload) { + t.Fatalf("unexpected buffered length after second write: got %d want %d", raw.Len(), headerLen+len(firstPayload)+len(secondPayload)) + } + if !bytes.Equal(raw.Bytes()[raw.Len()-len(secondPayload):], secondPayload) { + t.Fatal("second payload mismatch") + } +} diff --git a/protocol/trojanc/dialer.go b/protocol/trojanc/dialer.go index bf98eb3b..2bf0067f 100644 --- a/protocol/trojanc/dialer.go +++ b/protocol/trojanc/dialer.go @@ -32,17 +32,6 @@ func NewDialer(nextDialer netproxy.Dialer, header protocol.Header) (netproxy.Dia }, nil } -func (d *Dialer) DialTcp(ctx context.Context, addr string) (c netproxy.Conn, err error) { - return d.DialContext(ctx, "tcp", addr) -} - -func (d *Dialer) DialUdp(ctx context.Context, addr string) (c netproxy.PacketConn, err error) { - pktConn, err := d.DialContext(ctx, "udp", addr) - if err != nil { - return nil, err - } - return pktConn.(netproxy.PacketConn), nil -} func (d *Dialer) DialContext(ctx context.Context, network string, addr string) (c netproxy.Conn, err error) { magicNetwork, err := netproxy.ParseMagicNetwork(network) diff --git a/protocol/trojanc/udp_bench_test.go b/protocol/trojanc/udp_bench_test.go new file mode 100644 index 00000000..153d7d0f --- /dev/null +++ b/protocol/trojanc/udp_bench_test.go @@ -0,0 +1,34 @@ +package trojanc + +import ( + "testing" +) + +// BenchmarkUDPPacketOverhead tests memory allocation for UDP packet processing +func BenchmarkUDPPacketOverhead(b *testing.B) { + b.Run("SmallPacket", func(b *testing.B) { + data := make([]byte, 100) + for i := range data { + data[i] = byte(i) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Simulate SealUDP allocation + _ = make([]byte, 100+4+100) + } + }) + + b.Run("LargePacket", func(b *testing.B) { + data := make([]byte, 1400) + for i := range data { + data[i] = byte(i) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Simulate SealUDP allocation + _ = make([]byte, 100+4+1400) + } + }) +} diff --git a/protocol/tuic/client.go b/protocol/tuic/client.go index 379d3ac1..04bcc584 100644 --- a/protocol/tuic/client.go +++ b/protocol/tuic/client.go @@ -7,16 +7,16 @@ import ( "crypto/tls" "fmt" "net" - "strings" "sync" "time" + outbounderrors "github.com/daeuniverse/outbound/common/errors" "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/pkg/fastrand" "github.com/daeuniverse/outbound/pool" "github.com/daeuniverse/outbound/protocol" "github.com/daeuniverse/outbound/protocol/tuic/common" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) const Ver5 = 0x5 @@ -171,11 +171,13 @@ func (t *clientImpl) handleMessage(quicConn quic.Connection) (err error) { t.deferQuicConn(quicConn, err) }() for { - // TODO: - ctx, cancel := context.WithTimeout(context.TODO(), 3*time.Minute) - message, err := quicConn.ReceiveDatagram(ctx) - cancel() + // Use context.Background() instead of fixed timeout + // QUIC's keepalive mechanism will handle connection health + message, err := quicConn.ReceiveDatagram(context.Background()) if err != nil { + if outbounderrors.IsTemporaryError(err) { + continue + } return err } go func(message []byte) (err error) { @@ -219,7 +221,10 @@ func (t *clientImpl) handleMessage(quicConn quic.Connection) (err error) { } func (t *clientImpl) deferQuicConn(quicConn quic.Connection, err error) { - if err != nil && !strings.Contains(err.Error(), common.ErrTooManyOpenStreams.Error()) { + // Only close connection on non-temporary errors + if err != nil && + !outbounderrors.IsTemporaryError(err) && + err != outbounderrors.ErrStreamExhausted { t.forceClose(quicConn, err) } } diff --git a/protocol/tuic/client_ring.go b/protocol/tuic/client_ring.go index d09f13d7..02695d63 100644 --- a/protocol/tuic/client_ring.go +++ b/protocol/tuic/client_ring.go @@ -3,8 +3,6 @@ package tuic import ( "container/list" "context" - "errors" - "strings" "sync" "github.com/daeuniverse/outbound/netproxy" @@ -90,10 +88,10 @@ func (r *clientRing) _tryNext(current **list.Element, f func(cli *clientRingNode *current = r.ring.Front() } } - if *current == r.current { - // Clients are exhausted. - if strings.Contains(err.Error(), common.ErrTooManyOpenStreams.Error()) || errors.Is(err, common.ErrClientClosed) || errors.Is(err, common.ErrHoldOn) { + if err == common.ErrTooManyOpenStreams || + err == common.ErrClientClosed || + err == common.ErrHoldOn { goto getNew } // Not the expected error. diff --git a/protocol/tuic/common/congestion.go b/protocol/tuic/common/congestion.go index f8676b9a..06cca9a3 100644 --- a/protocol/tuic/common/congestion.go +++ b/protocol/tuic/common/congestion.go @@ -2,7 +2,7 @@ package common import ( "github.com/daeuniverse/outbound/protocol/tuic/congestion" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) const ( diff --git a/protocol/tuic/common/stream.go b/protocol/tuic/common/stream.go index 3dead297..e3891a9b 100644 --- a/protocol/tuic/common/stream.go +++ b/protocol/tuic/common/stream.go @@ -6,7 +6,7 @@ import ( "time" "github.com/daeuniverse/outbound/netproxy" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) type safeStreamConn struct { diff --git a/protocol/tuic/common/type.go b/protocol/tuic/common/type.go index c1fcf427..083a5cf7 100644 --- a/protocol/tuic/common/type.go +++ b/protocol/tuic/common/type.go @@ -2,18 +2,18 @@ package common import ( "context" - "errors" "net" + outbounderrors "github.com/daeuniverse/outbound/common/errors" "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/protocol" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) var ( - ErrClientClosed = errors.New("client closed") - ErrTooManyOpenStreams = errors.New("too many open streams") - ErrHoldOn = errors.New("hold on") + ErrClientClosed = outbounderrors.ErrClientClosed + ErrTooManyOpenStreams = outbounderrors.ErrStreamExhausted + ErrHoldOn = outbounderrors.ErrOperationHold ) type DialFunc func(ctx context.Context, dialer netproxy.Dialer) (transport *quic.Transport, addr net.Addr, err error) @@ -31,3 +31,8 @@ const ( QUIC UdpRelayMode = iota NATIVE ) + +// IsTemporaryError checks if an error is temporary and should not close the connection +func IsTemporaryError(err error) bool { + return outbounderrors.IsTemporaryError(err) +} diff --git a/protocol/tuic/common/type_test.go b/protocol/tuic/common/type_test.go new file mode 100644 index 00000000..d34b9602 --- /dev/null +++ b/protocol/tuic/common/type_test.go @@ -0,0 +1,102 @@ +package common + +import ( + "context" + "errors" + "fmt" + "testing" +) + +// mockNetError implements net.Error for testing +type mockNetError struct { + msg string + timeout bool + temporary bool +} + +func (e *mockNetError) Error() string { return e.msg } +func (e *mockNetError) Timeout() bool { return e.timeout } +func (e *mockNetError) Temporary() bool { return e.temporary } + +func TestIsTemporaryError(t *testing.T) { + tests := []struct { + name string + err error + expected bool + }{ + { + name: "nil error", + err: nil, + expected: false, + }, + { + name: "context deadline exceeded", + err: context.DeadlineExceeded, + expected: true, + }, + { + name: "context canceled", + err: context.Canceled, + expected: true, + }, + { + name: "net temporary error", + err: &mockNetError{ + msg: "temporary network error", + timeout: false, + temporary: true, + }, + expected: true, + }, + { + name: "net timeout error (temporary)", + err: &mockNetError{ + msg: "timeout error", + timeout: true, + temporary: true, + }, + expected: true, + }, + { + name: "net permanent error", + err: &mockNetError{ + msg: "permanent error", + timeout: false, + temporary: false, + }, + expected: false, + }, + { + name: "generic error", + err: errors.New("some error"), + expected: false, + }, + { + name: "client closed error", + err: ErrClientClosed, + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := IsTemporaryError(tt.err) + if result != tt.expected { + t.Errorf("IsTemporaryError(%v) = %v, want %v", tt.err, result, tt.expected) + } + }) + } +} + +func TestIsTemporaryErrorWithWrappedErrors(t *testing.T) { + // Test wrapped context errors + wrappedDeadline := fmt.Errorf("wrapped: %w", context.DeadlineExceeded) + if !IsTemporaryError(wrappedDeadline) { + t.Error("IsTemporaryError should return true for wrapped context.DeadlineExceeded") + } + + wrappedCanceled := fmt.Errorf("wrapped: %w", context.Canceled) + if !IsTemporaryError(wrappedCanceled) { + t.Error("IsTemporaryError should return true for wrapped context.Canceled") + } +} diff --git a/protocol/tuic/congestion/bbr/bandwidth.go b/protocol/tuic/congestion/bbr/bandwidth.go index 23d870d9..572c9932 100644 --- a/protocol/tuic/congestion/bbr/bandwidth.go +++ b/protocol/tuic/congestion/bbr/bandwidth.go @@ -4,7 +4,7 @@ import ( "math" "time" - "github.com/daeuniverse/quic-go/congestion" + "github.com/olicesx/quic-go/congestion" ) const ( diff --git a/protocol/tuic/congestion/bbr/bandwidth_sampler.go b/protocol/tuic/congestion/bbr/bandwidth_sampler.go index 4b28d423..4915d123 100644 --- a/protocol/tuic/congestion/bbr/bandwidth_sampler.go +++ b/protocol/tuic/congestion/bbr/bandwidth_sampler.go @@ -4,7 +4,7 @@ import ( "math" "time" - "github.com/daeuniverse/quic-go/congestion" + "github.com/olicesx/quic-go/congestion" ) const ( diff --git a/protocol/tuic/congestion/bbr/bbr_sender.go b/protocol/tuic/congestion/bbr/bbr_sender.go index e63b11ce..19f1c343 100644 --- a/protocol/tuic/congestion/bbr/bbr_sender.go +++ b/protocol/tuic/congestion/bbr/bbr_sender.go @@ -7,7 +7,7 @@ import ( rand "github.com/daeuniverse/outbound/pkg/fastrand" "github.com/daeuniverse/outbound/protocol/tuic/congestion/common" - "github.com/daeuniverse/quic-go/congestion" + "github.com/olicesx/quic-go/congestion" ) // BbrSender implements BBR congestion control algorithm. BBR aims to estimate diff --git a/protocol/tuic/congestion/bbr/packet_number_indexed_queue.go b/protocol/tuic/congestion/bbr/packet_number_indexed_queue.go index e9fad5a8..e1a0c24a 100644 --- a/protocol/tuic/congestion/bbr/packet_number_indexed_queue.go +++ b/protocol/tuic/congestion/bbr/packet_number_indexed_queue.go @@ -1,7 +1,7 @@ package bbr import ( - "github.com/daeuniverse/quic-go/congestion" + "github.com/olicesx/quic-go/congestion" ) // packetNumberIndexedQueue is a queue of mostly continuous numbered entries diff --git a/protocol/tuic/congestion/brutal/brutal.go b/protocol/tuic/congestion/brutal/brutal.go index 15276b40..53e34088 100644 --- a/protocol/tuic/congestion/brutal/brutal.go +++ b/protocol/tuic/congestion/brutal/brutal.go @@ -8,7 +8,7 @@ import ( "github.com/daeuniverse/outbound/protocol/tuic/congestion/common" - "github.com/daeuniverse/quic-go/congestion" + "github.com/olicesx/quic-go/congestion" ) const ( diff --git a/protocol/tuic/congestion/common/pacer.go b/protocol/tuic/congestion/common/pacer.go index 092eb8fa..c48730f4 100644 --- a/protocol/tuic/congestion/common/pacer.go +++ b/protocol/tuic/congestion/common/pacer.go @@ -4,7 +4,7 @@ import ( "math" "time" - "github.com/daeuniverse/quic-go/congestion" + "github.com/olicesx/quic-go/congestion" ) const ( diff --git a/protocol/tuic/congestion/utils.go b/protocol/tuic/congestion/utils.go index ea83ae16..0849538b 100644 --- a/protocol/tuic/congestion/utils.go +++ b/protocol/tuic/congestion/utils.go @@ -3,7 +3,7 @@ package congestion import ( "github.com/daeuniverse/outbound/protocol/tuic/congestion/bbr" "github.com/daeuniverse/outbound/protocol/tuic/congestion/brutal" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) func UseBBR(conn quic.Connection) { diff --git a/protocol/tuic/dialer.go b/protocol/tuic/dialer.go index 2a42ef84..ce6ca087 100644 --- a/protocol/tuic/dialer.go +++ b/protocol/tuic/dialer.go @@ -9,8 +9,8 @@ import ( "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/protocol" "github.com/daeuniverse/outbound/protocol/tuic/common" - "github.com/daeuniverse/quic-go" "github.com/google/uuid" + "github.com/olicesx/quic-go" ) func init() { @@ -21,6 +21,7 @@ type Dialer struct { clientRing *clientRing proxyAddress string + proxyUDPAddr *net.UDPAddr nextDialer netproxy.Dialer metadata protocol.Metadata } @@ -41,6 +42,10 @@ func NewDialer(nextDialer netproxy.Dialer, header protocol.Header) (netproxy.Dia // FIXME: QUIC has severe performance problems. // udpRelayMode = common.QUIC } + proxyUDPAddr, err := net.ResolveUDPAddr("udp", header.ProxyAddress) + if err != nil { + return nil, err + } return &Dialer{ clientRing: newClientRing(func(capabilityCallback func(n int64)) *clientImpl { return &clientImpl{ @@ -69,23 +74,12 @@ func NewDialer(nextDialer netproxy.Dialer, header protocol.Header) (netproxy.Dia } }, 10), proxyAddress: header.ProxyAddress, + proxyUDPAddr: proxyUDPAddr, nextDialer: nextDialer, metadata: metadata, }, nil } -func (d *Dialer) DialTcp(ctx context.Context, addr string) (c netproxy.Conn, err error) { - return d.DialContext(ctx, "tcp", addr) -} - -func (d *Dialer) DialUdp(ctx context.Context, addr string) (c netproxy.PacketConn, err error) { - pktConn, err := d.DialContext(ctx, "udp", addr) - if err != nil { - return nil, err - } - return pktConn.(netproxy.PacketConn), nil -} - func (d *Dialer) dialFuncFactory(udpNetwork string, rAddr net.Addr) common.DialFunc { return func(ctx context.Context, dialer netproxy.Dialer) (transport *quic.Transport, addr net.Addr, err error) { conn, err := dialer.DialContext(ctx, udpNetwork, d.proxyAddress) @@ -114,10 +108,6 @@ func (d *Dialer) DialContext(ctx context.Context, network string, addr string) ( return nil, err } mdata.IsClient = d.metadata.IsClient - proxyAddr, err := net.ResolveUDPAddr("udp", d.proxyAddress) - if err != nil { - return nil, err - } udpNetwork := network if magicNetwork.Network == "tcp" { udpNetwork = netproxy.MagicNetwork{ @@ -125,7 +115,7 @@ func (d *Dialer) DialContext(ctx context.Context, network string, addr string) ( Mark: magicNetwork.Mark, }.Encode() tcpConn, err := d.clientRing.DialContextWithDialer(ctx, &mdata, d.nextDialer, - d.dialFuncFactory(udpNetwork, proxyAddr), + d.dialFuncFactory(udpNetwork, d.proxyUDPAddr), ) if err != nil { return nil, err @@ -133,7 +123,7 @@ func (d *Dialer) DialContext(ctx context.Context, network string, addr string) ( return tcpConn, nil } else { udpConn, err := d.clientRing.ListenPacketWithDialer(ctx, &mdata, d.nextDialer, - d.dialFuncFactory(udpNetwork, proxyAddr), + d.dialFuncFactory(udpNetwork, d.proxyUDPAddr), ) if err != nil { return nil, err diff --git a/protocol/tuic/frag.go b/protocol/tuic/frag.go index dfdb55d5..9e588b35 100644 --- a/protocol/tuic/frag.go +++ b/protocol/tuic/frag.go @@ -4,7 +4,7 @@ import ( "net/netip" "github.com/daeuniverse/outbound/pool/bytes" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) func fragWriteNative(quicConn quic.Connection, packet *Packet, buf *bytes.Buffer, fragSize int) (err error) { diff --git a/protocol/tuic/packet.go b/protocol/tuic/packet.go index 2188fb6e..1af476f5 100644 --- a/protocol/tuic/packet.go +++ b/protocol/tuic/packet.go @@ -7,6 +7,7 @@ import ( "net" "net/netip" "sync" + "sync/atomic" "time" "github.com/daeuniverse/outbound/netproxy" @@ -14,7 +15,7 @@ import ( "github.com/daeuniverse/outbound/pool" "github.com/daeuniverse/outbound/protocol" "github.com/daeuniverse/outbound/protocol/tuic/common" - "github.com/daeuniverse/quic-go" + "github.com/olicesx/quic-go" ) type Packets struct { @@ -22,7 +23,7 @@ type Packets struct { list *list.List isEmptyState context.Context cancelEmptyState func() - closed bool + closed atomic.Bool } func NewPackets() *Packets { @@ -48,7 +49,7 @@ func (p *Packets) PushBack(packet *Packet) { func (p *Packets) PopFrontBlock() (packet *Packet, closed bool) { <-p.isEmptyState.Done() - if p.closed { + if p.closed.Load() { return nil, true } p.mu.Lock() @@ -67,10 +68,10 @@ func (p *Packets) setEmpty() { func (p *Packets) Close() error { p.mu.Lock() defer p.mu.Unlock() - if p.closed { + if p.closed.Load() { return nil } - p.closed = true + p.closed.Store(true) select { case <-p.isEmptyState.Done(): default: @@ -96,7 +97,7 @@ type quicStreamPacketConn struct { closeOnce sync.Once closeErr error - closed bool + closed atomic.Bool // TODO: multiple defraggers for different PKT_ID deFraggers sync.Map @@ -107,7 +108,7 @@ type quicStreamPacketConn struct { func (q *quicStreamPacketConn) Close() error { q.closeOnce.Do(func() { - q.closed = true + q.closed.Store(true) q.closeErr = q.close() }) return q.closeErr @@ -179,37 +180,38 @@ func (q *quicStreamPacketConn) SetWriteDeadline(t time.Time) error { func (q *quicStreamPacketConn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) { q.mu.Lock() - defer q.mu.Unlock() - if q.incomingPackets != nil { - for { - packet, closed := q.incomingPackets.PopFrontBlock() - if closed { - err = net.ErrClosed - return - } - _d, _ := q.deFraggers.LoadOrStore(packet.PKT_ID, &deFragger{}) - d := _d.(*deFragger) - var assembled bool - // Feed packet into this deFragger. - // Return if this PKT_ID is ready and assembled. - if n, addr, assembled = d.Feed(packet, p); assembled { - q.deFraggers.Delete(packet.PKT_ID) - return - } else { - // FIXME: Timeout to clean deFraggers. - } + incomingPackets := q.incomingPackets + q.mu.Unlock() + + if incomingPackets == nil { + return 0, netip.AddrPort{}, net.ErrClosed + } + + for { + packet, closed := incomingPackets.PopFrontBlock() + if closed { + err = net.ErrClosed + return + } + _d, _ := q.deFraggers.LoadOrStore(packet.PKT_ID, &deFragger{}) + d := _d.(*deFragger) + var assembled bool + // Feed packet into this deFragger. + // Return if this PKT_ID is ready and assembled. + if n, addr, assembled = d.Feed(packet, p); assembled { + q.deFraggers.Delete(packet.PKT_ID) + return + } else { + // FIXME: Timeout to clean deFraggers. } - } else { - err = net.ErrClosed } - return } func (q *quicStreamPacketConn) WriteTo(p []byte, addr string) (n int, err error) { if len(p) > 0xffff { // uint16 max return 0, &quic.DatagramTooLargeError{MaxDataLen: 0xffff} } - if q.closed { + if q.closed.Load() { return 0, net.ErrClosed } if q.deferQuicConnFn != nil { diff --git a/protocol/tuic/packet_test.go b/protocol/tuic/packet_test.go new file mode 100644 index 00000000..889dff08 --- /dev/null +++ b/protocol/tuic/packet_test.go @@ -0,0 +1,137 @@ +package tuic + +import ( + "errors" + "net" + "sync" + "testing" + "time" +) + +func TestReadFromNoDeadlock(t *testing.T) { + packets := NewPackets() + q := &quicStreamPacketConn{ + incomingPackets: packets, + } + + var wg sync.WaitGroup + wg.Add(2) + + readDone := make(chan struct{}) + go func() { + defer wg.Done() + defer close(readDone) + _, _, _ = q.ReadFrom(make([]byte, 1024)) + }() + + time.Sleep(100 * time.Millisecond) + + go func() { + defer wg.Done() + packets.Close() + }() + + select { + case <-readDone: + t.Log("ReadFrom unblocked successfully - no deadlock") + case <-time.After(2 * time.Second): + t.Fatal("ReadFrom deadlocked - Close() could not unblock it") + } + + wg.Wait() +} + +func TestReadFromReturnsNilAfterClose(t *testing.T) { + packets := NewPackets() + q := &quicStreamPacketConn{ + incomingPackets: packets, + } + + packets.Close() + + _, _, err := q.ReadFrom(make([]byte, 1024)) + if err == nil { + t.Error("expected error after close, got nil") + } +} + +func TestPacketsCloseUnblocksPopFrontBlock(t *testing.T) { + p := NewPackets() + + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = p.PopFrontBlock() + }() + + time.Sleep(50 * time.Millisecond) + _ = p.Close() + + select { + case <-done: + t.Log("PopFrontBlock unblocked after Close") + case <-time.After(1 * time.Second): + t.Fatal("PopFrontBlock did not unblock after Close") + } +} + +func TestPacketsPushPop(t *testing.T) { + p := NewPackets() + + go func() { + time.Sleep(50 * time.Millisecond) + p.PushBack(&Packet{ + PKT_ID: 1, + FRAG_ID: 0, + FRAG_TOTAL: 1, + DATA: []byte("test data"), + ADDR: &Address{TYPE: AtypIPv4, ADDR: []byte{127, 0, 0, 1}, PORT: 8080}, + }) + }() + + packet, closed := p.PopFrontBlock() + if closed { + t.Fatal("expected packet, got closed") + } + if packet == nil { + t.Fatal("expected non-nil packet") + } + if string(packet.DATA) != "test data" { + t.Errorf("expected 'test data', got '%s'", packet.DATA) + } + _ = p.Close() +} + +func TestConcurrentPushClose(t *testing.T) { + p := NewPackets() + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + p.PushBack(&Packet{DATA: []byte("test")}) + }() + } + + time.Sleep(10 * time.Millisecond) + _ = p.Close() + + wg.Wait() + _, closed := p.PopFrontBlock() + if !closed { + t.Error("expected closed after Close") + } +} + +func TestQuicStreamPacketConnWriteToAfterClose(t *testing.T) { + q := &quicStreamPacketConn{} + + if err := q.Close(); err != nil { + t.Fatalf("Close returned error: %v", err) + } + + if _, err := q.WriteTo([]byte("test"), "127.0.0.1:53"); !errors.Is(err, net.ErrClosed) { + t.Fatalf("expected net.ErrClosed after Close, got %v", err) + } +} diff --git a/protocol/tuic/protocol.go b/protocol/tuic/protocol.go index 9eb63f8e..2e757d81 100644 --- a/protocol/tuic/protocol.go +++ b/protocol/tuic/protocol.go @@ -9,8 +9,8 @@ import ( "strconv" "github.com/daeuniverse/outbound/protocol" - "github.com/daeuniverse/quic-go" "github.com/google/uuid" + "github.com/olicesx/quic-go" ) type BufferedReader interface { diff --git a/protocol/vless/dialer.go b/protocol/vless/dialer.go index 9e704f49..a874a907 100644 --- a/protocol/vless/dialer.go +++ b/protocol/vless/dialer.go @@ -57,17 +57,6 @@ func NewDialer(nextDialer netproxy.Dialer, header protocol.Header) (netproxy.Dia }, nil } -func (d *Dialer) DialTcp(ctx context.Context, addr string) (c netproxy.Conn, err error) { - return d.DialContext(ctx, "tcp", addr) -} - -func (d *Dialer) DialUdp(ctx context.Context, addr string) (c netproxy.PacketConn, err error) { - pktConn, err := d.DialContext(ctx, "udp", addr) - if err != nil { - return nil, err - } - return pktConn.(netproxy.PacketConn), nil -} func (d *Dialer) DialContext(ctx context.Context, network string, addr string) (c netproxy.Conn, err error) { magicNetwork, err := netproxy.ParseMagicNetwork(network) diff --git a/protocol/vless/key.go b/protocol/vless/key.go index 3053225e..4e90b665 100644 --- a/protocol/vless/key.go +++ b/protocol/vless/key.go @@ -15,7 +15,7 @@ func Password2Key(password string) (id []byte, err error) { } password = strings.ReplaceAll(password, "-", "") if len(password) != 32 { - return nil, fmt.Errorf("invalid UUID: " + password) + return nil, fmt.Errorf("invalid UUID: %s", password) } id = make([]byte, 16) if _, err := hex.Decode(id[:], []byte(password)); err != nil { diff --git a/protocol/vless/vision/packet.go b/protocol/vless/vision/packet.go index 4498bca9..37c7932d 100644 --- a/protocol/vless/vision/packet.go +++ b/protocol/vless/vision/packet.go @@ -2,7 +2,6 @@ package vision import ( "encoding/binary" - "errors" "fmt" "io" "net/netip" @@ -20,26 +19,26 @@ type PacketConn struct { addr string } -func (c *PacketConn) Read(b []byte) (n int, err error) { - switch c.network { +func (pc *PacketConn) Read(b []byte) (n int, err error) { + switch pc.network { case "tcp": - return c.Conn.Read(b) + return pc.Conn.Read(b) case "udp": - n, _, err = c.ReadFrom(b) + n, _, err = pc.ReadFrom(b) return n, err default: - return 0, fmt.Errorf("unsupported network: %s", c.network) + return 0, fmt.Errorf("unsupported network: %s", pc.network) } } -func (c *PacketConn) Write(b []byte) (n int, err error) { - switch c.network { +func (pc *PacketConn) Write(b []byte) (n int, err error) { + switch pc.network { case "tcp": - return c.Conn.Write(b) + return pc.Conn.Write(b) case "udp": - return c.WriteTo(b, c.addr) + return pc.WriteTo(b, pc.addr) default: - return 0, fmt.Errorf("unsupported network: %s", c.network) + return 0, fmt.Errorf("unsupported network: %s", pc.network) } } @@ -50,30 +49,40 @@ func (c *PacketConn) Write(b []byte) (n int, err error) { // +-------------------+-------------------+ // | Length Data | Payload | // +-------------------+-------------------+ -func (c *PacketConn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) { +func (pc *PacketConn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) { // Read frame length (2 bytes) var frameLengthBytes [2]byte - if _, err = io.ReadFull(c.Conn, frameLengthBytes[:]); err != nil { + if _, err = io.ReadFull(pc.Conn, frameLengthBytes[:]); err != nil { return 0, netip.AddrPort{}, err } frameLength := binary.BigEndian.Uint16(frameLengthBytes[:]) + if frameLength < 4 { + return 0, netip.AddrPort{}, io.EOF + } + // Read frame header (4 bytes) var frameHeaderBytes [4]byte - if _, err = io.ReadFull(c.Conn, frameHeaderBytes[:]); err != nil { + if _, err = io.ReadFull(pc.Conn, frameHeaderBytes[:]); err != nil { return 0, netip.AddrPort{}, err } + discard := false switch frameHeaderBytes[2] { case 0x01: return 0, netip.AddrPort{}, fmt.Errorf("unexpected frame new") case 0x02: // Keep if frameLength > 4 { - addrData := make([]byte, frameLength-4) - if _, err = io.ReadFull(c.Conn, addrData); err != nil { + netInfo := make([]byte, frameLength-4) + if _, err = io.ReadFull(pc.Conn, netInfo); err != nil { return 0, netip.AddrPort{}, err } + netType := netInfo[0] + if netType != 0x02 { // net type udp + return 0, netip.AddrPort{}, fmt.Errorf("unsupported net type: %x", netType) + } + addrData := netInfo[1:] addr, err = ReadPacketAddr(addrData) if err != nil { return 0, netip.AddrPort{}, err @@ -83,27 +92,33 @@ func (c *PacketConn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) return 0, netip.AddrPort{}, io.EOF case 0x04: // KeepAlive + discard = true default: return 0, netip.AddrPort{}, fmt.Errorf("unsupported frame header: %x", frameHeaderBytes[2]) } if frameHeaderBytes[3]&1 != 1 { - return c.ReadFrom(p) + return pc.ReadFrom(p) } // Read length and payload var lengthBytes [2]byte - if _, err = io.ReadFull(c.Conn, lengthBytes[:]); err != nil { + if _, err = io.ReadFull(pc.Conn, lengthBytes[:]); err != nil { return 0, netip.AddrPort{}, err } length := binary.BigEndian.Uint16(lengthBytes[:]) - + if length == 0 { + return pc.ReadFrom(p) + } if length > uint16(len(p)) { - return 0, netip.AddrPort{}, fmt.Errorf("buffer too small") + return 0, netip.AddrPort{}, io.ErrShortBuffer } - n, err = io.ReadFull(c.Conn, p[:length]) - return n, addr, err + n, err = io.ReadFull(pc.Conn, p[:length]) + if !discard { + return n, addr, err + } + return pc.ReadFrom(p) } // +------------------------+------------------------+ @@ -169,52 +184,3 @@ func (pc *PacketConn) prefixPacket(addr string) (pool.PB, error) { return prefix, err } - -func IPAddrToPacketAddrLength(addr netip.AddrPort) int { - nip, ok := netip.AddrFromSlice(addr.Addr().AsSlice()) - if !ok { - return 0 - } - - if nip.Is4() { - return 1 + 4 + 2 - } else { - return 1 + 16 + 2 - } -} - -func PutPacketAddr(src []byte, addr netip.AddrPort) error { - nip, ok := netip.AddrFromSlice(addr.Addr().AsSlice()) - if !ok { - return errors.New("invalid IP") - } - - if nip.Is4() { - binary.BigEndian.PutUint16(src[0:2], addr.Port()) - src[2] = 1 - copy(src[3:7], nip.AsSlice()) - } else { - binary.BigEndian.PutUint16(src[0:2], addr.Port()) - src[2] = 3 - copy(src[3:19], nip.AsSlice()) - } - - return nil -} - -func ReadPacketAddr(p []byte) (addr netip.AddrPort, err error) { - p = p[1:] - port := binary.BigEndian.Uint16(p[0:2]) - ipType := p[2] - ip := p[3:] - if ipType == 1 { - ip = ip[:4] - } else { - ip = ip[:16] - } - ipAddr, ok := netip.AddrFromSlice(ip) - if !ok { - return netip.AddrPort{}, errors.New("invalid IP") - } - return netip.AddrPortFrom(ipAddr, port), nil -} diff --git a/protocol/vless/vision/packetaddr.go b/protocol/vless/vision/packetaddr.go new file mode 100644 index 00000000..615b28c6 --- /dev/null +++ b/protocol/vless/vision/packetaddr.go @@ -0,0 +1,73 @@ +package vision + +import ( + "encoding/binary" + "errors" + "net/netip" +) + +var ( + ErrInvalidPacketAddr = errors.New("invalid packet addr") + ErrInvalidAddrType = errors.New("invalid addr type") + ErrInvalidIP = errors.New("invalid IP") +) + +func IPAddrToPacketAddrLength(addr netip.AddrPort) int { + nip, ok := netip.AddrFromSlice(addr.Addr().AsSlice()) + if !ok { + return 0 + } + + if nip.Is4() { + return 1 + 4 + 2 + } else { + return 1 + 16 + 2 + } +} + +func PutPacketAddr(src []byte, addr netip.AddrPort) error { + nip, ok := netip.AddrFromSlice(addr.Addr().AsSlice()) + if !ok { + return ErrInvalidIP + } + + if nip.Is4() { + binary.BigEndian.PutUint16(src[0:2], addr.Port()) + src[2] = 1 + copy(src[3:7], nip.AsSlice()) + } else { + binary.BigEndian.PutUint16(src[0:2], addr.Port()) + src[2] = 3 + copy(src[3:19], nip.AsSlice()) + } + + return nil +} + +func ReadPacketAddr(p []byte) (addr netip.AddrPort, err error) { + if len(p) < 3 { + return netip.AddrPort{}, ErrInvalidPacketAddr + } + port := binary.BigEndian.Uint16(p[0:2]) + ipType := p[2] + ip := p[3:] + switch ipType { + case 1: + if len(ip) < 4 { + return netip.AddrPort{}, ErrInvalidPacketAddr + } + ip = ip[:4] + case 3: + if len(ip) < 16 { + return netip.AddrPort{}, ErrInvalidPacketAddr + } + ip = ip[:16] + default: + return netip.AddrPort{}, ErrInvalidAddrType + } + ipAddr, ok := netip.AddrFromSlice(ip) + if !ok { + return netip.AddrPort{}, ErrInvalidIP + } + return netip.AddrPortFrom(ipAddr, port), nil +} diff --git a/protocol/vmess/conn.go b/protocol/vmess/conn.go index fc550123..9ac0d8d5 100644 --- a/protocol/vmess/conn.go +++ b/protocol/vmess/conn.go @@ -1,7 +1,6 @@ package vmess import ( - "bytes" "crypto/cipher" "crypto/sha256" "encoding/binary" @@ -13,14 +12,19 @@ import ( "sync" "github.com/daeuniverse/outbound/common" + "github.com/daeuniverse/outbound/common/iout" "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/pkg/fastrand" "github.com/daeuniverse/outbound/pool" ) +var resolveUDPAddr = net.ResolveUDPAddr + const ( MaxChunkSize = 1 << 14 MaxUDPSize = 1 << 11 + + maxReusableSealFrameSize = 128 << 10 ) type Conn struct { @@ -31,6 +35,7 @@ type Conn struct { cmdKey []byte dialTgt string dialTgtAddrPort netip.AddrPort // lazy resolve + dialTgtMu sync.Mutex NewAEAD func(key []byte) (cipher.AEAD, error) @@ -56,6 +61,8 @@ type Conn struct { readMutex sync.Mutex leftToRead []byte indexToRead int + + writeSealFrame []byte } func NewConn(conn netproxy.Conn, metadata Metadata, dialTgt string, cmdKey []byte) (c *Conn, err error) { @@ -77,9 +84,36 @@ func NewConn(conn netproxy.Conn, metadata Metadata, dialTgt string, cmdKey []byt } func (c *Conn) Close() error { + c.readMutex.Lock() + if c.leftToRead != nil { + pool.Put(c.leftToRead) + c.leftToRead = nil + c.indexToRead = 0 + } + c.readMutex.Unlock() + + c.writeMutex.Lock() + c.writeSealFrame = nil + c.writeMutex.Unlock() return c.Conn.Close() } +func (c *Conn) dialTargetAddrPort() (netip.AddrPort, error) { + c.dialTgtMu.Lock() + defer c.dialTgtMu.Unlock() + + if c.dialTgtAddrPort.IsValid() { + return c.dialTgtAddrPort, nil + } + + tgt, err := resolveUDPAddr("udp", c.dialTgt) + if err != nil { + return netip.AddrPort{}, err + } + c.dialTgtAddrPort = tgt.AddrPort() + return c.dialTgtAddrPort, nil +} + func (c *Conn) chunks(size int) (payloadSize int, numChunks int) { payloadSize = MaxChunkSize - c.writeBodyCipher.Overhead() - int(c.writeChunkSizeParser.SizeBytes()) - int(c.writePaddingGenerator.MaxPaddingLen()) if size%payloadSize == 0 { @@ -104,8 +138,16 @@ func (c *Conn) sealFromPool(b []byte) (data []byte) { sizeSize := c.writeChunkSizeParser.SizeBytes() encryptedSize := int32(len(b) + c.writeBodyCipher.Overhead()) paddingSize := int32(c.writePaddingGenerator.NextPaddingLen()) + totalSize := int(sizeSize + encryptedSize + paddingSize) - data = pool.Get(int(sizeSize + encryptedSize + paddingSize)) + if totalSize <= maxReusableSealFrameSize { + if cap(c.writeSealFrame) < totalSize { + c.writeSealFrame = make([]byte, totalSize) + } + data = c.writeSealFrame[:totalSize] + } else { + data = make([]byte, totalSize) + } c.writeChunkSizeParser.Encode(uint16(encryptedSize+paddingSize), data) c.writeBodyCipher.Seal(data[sizeSize:sizeSize], c.writeNonceGenerator(), b, nil) @@ -114,17 +156,13 @@ func (c *Conn) sealFromPool(b []byte) (data []byte) { return data } -// writeStream splits mb into multiple FIXED size (payloadSize) chunks. -// Then seal the chunks and write separately. -// If the sum size of mb less than one payloadSize, seal and write it directly. func (c *Conn) writeStream(b []byte, preWrite []byte) (n int, err error) { payloadSize, numChunks := c.chunks(len(b)) var start = 0 if preWrite != nil { start++ data := c.sealFromPool(b[n:common.Min(n+payloadSize, len(b))]) - defer pool.Put(data) - if _, err = c.Conn.Write(bytes.Join([][]byte{preWrite, data}, nil)); err != nil { + if _, err = iout.MultiWrite(c.Conn, preWrite, data); err != nil { return 0, err } n += payloadSize @@ -134,7 +172,6 @@ func (c *Conn) writeStream(b []byte, preWrite []byte) (n int, err error) { if _, err = c.Conn.Write(data); err != nil { return n, err } - pool.Put(data) n += payloadSize } if n > len(b) { @@ -143,12 +180,10 @@ func (c *Conn) writeStream(b []byte, preWrite []byte) (n int, err error) { return n, nil } -// writePacket simply seal every buffer of mb and write. func (c *Conn) writePacket(b []byte, preWrite []byte) (n int, err error) { data := c.sealFromPool(b) - defer pool.Put(data) if preWrite != nil { - if _, err = c.Conn.Write(bytes.Join([][]byte{preWrite, data}, nil)); err != nil { + if _, err = iout.MultiWrite(c.Conn, preWrite, data); err != nil { return 0, err } } else { @@ -220,14 +255,11 @@ func (c *Conn) WriteReqHeader() (err error) { func (c *Conn) Write(b []byte) (n int, err error) { if c.metadata.IsPacketAddr() { - if !c.dialTgtAddrPort.IsValid() { - tgt, err := net.ResolveUDPAddr("udp", c.dialTgt) - if err != nil { - return 0, err - } - c.dialTgtAddrPort = tgt.AddrPort() + tgt, err := c.dialTargetAddrPort() + if err != nil { + return 0, err } - return c.WriteTo(b, c.dialTgtAddrPort.String()) + return c.WriteTo(b, tgt.String()) } else { return c.write(b) } @@ -272,7 +304,6 @@ func (c *Conn) write(b []byte) (n int, err error) { } if len(b) == 0 { data := c.sealFromPool(nil) - defer pool.Put(data) _, err = c.Conn.Write(data) return 0, err } @@ -432,6 +463,8 @@ func (c *Conn) read(b []byte) (n int, err error) { if c.indexToRead >= len(c.leftToRead) { // put the buf back pool.Put(c.leftToRead) + c.leftToRead = nil + c.indexToRead = 0 } return n, nil } @@ -449,6 +482,8 @@ func (c *Conn) read(b []byte) (n int, err error) { } else { // full reading. put the buf back pool.Put(chunk) + c.leftToRead = nil + c.indexToRead = 0 } return n, nil } diff --git a/protocol/vmess/conn_test.go b/protocol/vmess/conn_test.go new file mode 100644 index 00000000..63538543 --- /dev/null +++ b/protocol/vmess/conn_test.go @@ -0,0 +1,72 @@ +package vmess + +import ( + "net" + "net/netip" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestConnDialTargetAddrPortConcurrentSafe(t *testing.T) { + t.Helper() + + oldResolveUDPAddr := resolveUDPAddr + defer func() { + resolveUDPAddr = oldResolveUDPAddr + }() + + resolved := net.UDPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:8443")) + var calls atomic.Int32 + resolveUDPAddr = func(network, address string) (*net.UDPAddr, error) { + if network != "udp" { + t.Fatalf("unexpected network: %q", network) + } + if address != "example.com:8443" { + t.Fatalf("unexpected address: %q", address) + } + calls.Add(1) + time.Sleep(20 * time.Millisecond) + return resolved, nil + } + + c := &Conn{dialTgt: "example.com:8443"} + + const goroutines = 8 + results := make(chan netip.AddrPort, goroutines) + errs := make(chan error, goroutines) + + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + addr, err := c.dialTargetAddrPort() + if err != nil { + errs <- err + return + } + results <- addr + }() + } + wg.Wait() + close(results) + close(errs) + + for err := range errs { + if err != nil { + t.Fatalf("dialTargetAddrPort returned error: %v", err) + } + } + + for addr := range results { + if addr != resolved.AddrPort() { + t.Fatalf("unexpected resolved addr: got %v want %v", addr, resolved.AddrPort()) + } + } + + if got := calls.Load(); got != 1 { + t.Fatalf("unexpected resolver call count: got %d want 1", got) + } +} diff --git a/protocol/vmess/dialer.go b/protocol/vmess/dialer.go index e7ca14f0..7ab835e9 100644 --- a/protocol/vmess/dialer.go +++ b/protocol/vmess/dialer.go @@ -67,17 +67,6 @@ func NewDialerFactory(proto protocol.Protocol) func(nextDialer netproxy.Dialer, } } -func (d *Dialer) DialTcp(ctx context.Context, addr string) (c netproxy.Conn, err error) { - return d.DialContext(ctx, "tcp", addr) -} - -func (d *Dialer) DialUdp(ctx context.Context, addr string) (c netproxy.PacketConn, err error) { - pktConn, err := d.DialContext(ctx, "udp", addr) - if err != nil { - return nil, err - } - return pktConn.(netproxy.PacketConn), nil -} func (d *Dialer) DialContext(ctx context.Context, network string, addr string) (c netproxy.Conn, err error) { magicNetwork, err := netproxy.ParseMagicNetwork(network) diff --git a/protocol/vmess/packet.go b/protocol/vmess/packet.go index 58363b29..9d8e739c 100644 --- a/protocol/vmess/packet.go +++ b/protocol/vmess/packet.go @@ -25,15 +25,12 @@ func (c *Conn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) { copy(p, buf[addrLen:n]) return n - addrLen, address, err } else { - if !c.dialTgtAddrPort.IsValid() { - tgt, err := net.ResolveUDPAddr("udp", c.dialTgt) - if err != nil { - return 0, netip.AddrPort{}, err - } - c.dialTgtAddrPort = tgt.AddrPort() + tgt, err := c.dialTargetAddrPort() + if err != nil { + return 0, netip.AddrPort{}, err } copy(p, buf[:n]) - return n, c.dialTgtAddrPort, err + return n, tgt, err } } diff --git a/transport/shadowsocksr/obfs/obfs.go b/transport/shadowsocksr/obfs/obfs.go index 0d331834..fe0176e4 100644 --- a/transport/shadowsocksr/obfs/obfs.go +++ b/transport/shadowsocksr/obfs/obfs.go @@ -7,7 +7,7 @@ import ( type Creator func() IObfs type constructor struct { - New Creator + New Creator Overhead int } diff --git a/transport/shadowsocksr/proto/issue_validation_test.go b/transport/shadowsocksr/proto/issue_validation_test.go new file mode 100644 index 00000000..47178e41 --- /dev/null +++ b/transport/shadowsocksr/proto/issue_validation_test.go @@ -0,0 +1,594 @@ +// +build race + +// 验证 outbound 网络代码审查中发现的问题 +// 使用 -race 标志运行: go test -race -v -run TestIssueValidation +package proto + +import ( + "bytes" + "fmt" + "net" + "net/netip" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/daeuniverse/outbound/pool" + "github.com/daeuniverse/outbound/protocol/infra/socks" +) + +// ======================================== +// 问题 1: shadowsockr UDP 并发写入验证 +// ======================================== + +// MockProtocol 模拟 Protocol 接口 +type MockProtocol struct{} + +func (m *MockProtocol) EncodePkt(buf *bytes.Buffer) error { + // 模拟编码操作 + time.Sleep(1 * time.Microsecond) + return nil +} + +func (m *MockProtocol) DecodePkt(buf []byte) ([]byte, error) { + return buf, nil +} + +// MockPacketConn 模拟 netproxy.PacketConn +type MockPacketConn struct { + writeCount atomic.Int64 + dataCorruption atomic.Bool + lastData atomic.Value +} + +func (m *MockPacketConn) Read(b []byte) (n int, err error) { + return 0, nil +} + +func (m *MockPacketConn) Write(b []byte) (n int, err error) { + m.writeCount.Add(1) + return len(b), nil +} + +func (m *MockPacketConn) ReadFrom(p []byte) (n int, addr netip.AddrPort, err error) { + return 0, netip.AddrPort{}, nil +} + +func (m *MockPacketConn) WriteTo(p []byte, addr string) (n int, err error) { + m.writeCount.Add(1) + + // 模拟写入延迟 + time.Sleep(1 * time.Microsecond) + + // 检测数据竞争 + lastData := m.lastData.Load() + if lastData != nil { + oldData := lastData.([]byte) + if len(oldData) > 0 { + // 旧数据还在处理，可能有竞争 + m.dataCorruption.Store(true) + } + } + m.lastData.Store(p) + time.Sleep(1 * time.Microsecond) + m.lastData.Store([]byte{}) + + return len(p), nil +} + +func (m *MockPacketConn) Close() error { + return nil +} + +func (m *MockPacketConn) SetDeadline(t time.Time) error { + return nil +} + +func (m *MockPacketConn) SetReadDeadline(t time.Time) error { + return nil +} + +func (m *MockPacketConn) SetWriteDeadline(t time.Time) error { + return nil +} + +// SimulateShadowsockrPacketConn 模拟 shadowsockr 的 PacketConn（没有写锁） +type SimulateShadowsockrPacketConn struct { + inner *MockPacketConn + protocol *MockProtocol + tgt string + // 注意：这里没有 writeMu +} + +func (c *SimulateShadowsockrPacketConn) WriteTo(b []byte, to string) (int, error) { + // 模拟 shadowsockr 的 WriteTo 逻辑（没有写锁） + addr, err := socks.ParseAddr(to) + if err != nil { + return 0, err + } + + // 获取 buffer + pb := pool.GetMustBigger(len(addr) + len(b)) + defer pool.Put(pb) + + // 复制数据 + copy(pb, addr) + copy(pb[len(addr):], b) + + // 编码 + buf := bytes.NewBuffer(pb) + if err = c.protocol.EncodePkt(buf); err != nil { + return 0, err + } + + // 写入 - 这里没有锁保护 + _, err = c.inner.WriteTo(buf.Bytes(), c.tgt) + if err != nil { + return 0, err + } + + return len(b), nil +} + +// FixedShadowsockrPacketConn 修复后的 shadowsockr PacketConn（有写锁） +type FixedShadowsockrPacketConn struct { + inner *MockPacketConn + protocol *MockProtocol + tgt string + writeMu sync.Mutex // 添加写锁 +} + +func (c *FixedShadowsockrPacketConn) WriteTo(b []byte, to string) (int, error) { + c.writeMu.Lock() + defer c.writeMu.Unlock() + + addr, err := socks.ParseAddr(to) + if err != nil { + return 0, err + } + + pb := pool.GetMustBigger(len(addr) + len(b)) + defer pool.Put(pb) + + copy(pb, addr) + copy(pb[len(addr):], b) + + buf := bytes.NewBuffer(pb) + if err = c.protocol.EncodePkt(buf); err != nil { + return 0, err + } + + _, err = c.inner.WriteTo(buf.Bytes(), c.tgt) + if err != nil { + return 0, err + } + + return len(b), nil +} + +// TestIssue1_Outbound_ShadowsockrUDPRace 验证问题 1: shadowsockr UDP 并发写入 +func TestIssue1_Outbound_ShadowsockrUDPRace(t *testing.T) { + t.Log("🔍 验证问题 1: shadowsockr UDP 并发写入 (outbound)") + + // 测试没有锁的情况 + t.Run("WithoutLock", func(t *testing.T) { + inner := &MockPacketConn{} + protocol := &MockProtocol{} + + conn := &SimulateShadowsockrPacketConn{ + inner: inner, + protocol: protocol, + tgt: "127.0.0.1:8080", + } + + const goroutines = 10 + const writesPerGoroutine = 100 + + var wg sync.WaitGroup + wg.Add(goroutines) + + startTime := time.Now() + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + data := []byte(fmt.Sprintf("packet-%d-%d", id, j)) + _, err := conn.WriteTo(data, "127.0.0.1:8080") + if err != nil { + t.Errorf("WriteTo failed: %v", err) + } + } + }(i) + } + + wg.Wait() + elapsed := time.Since(startTime) + + writes := inner.writeCount.Load() + expected := int64(goroutines * writesPerGoroutine) + + t.Logf("✅ 无锁测试完成: %d 次写入，耗时 %v", writes, elapsed) + + if writes != expected { + t.Errorf("❌ 写入计数不匹配: got %d, expected %d", writes, expected) + } + + if inner.dataCorruption.Load() { + t.Log("⚠️ 检测到潜在的数据竞争迹象") + } + + t.Log("⚠️ 使用 'go test -race' 运行此测试以检测数据竞争") + }) + + // 测试有锁的情况 + t.Run("WithLock", func(t *testing.T) { + inner := &MockPacketConn{} + protocol := &MockProtocol{} + + conn := &FixedShadowsockrPacketConn{ + inner: inner, + protocol: protocol, + tgt: "127.0.0.1:8080", + } + + const goroutines = 10 + const writesPerGoroutine = 100 + + var wg sync.WaitGroup + wg.Add(goroutines) + + startTime := time.Now() + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + data := []byte(fmt.Sprintf("packet-%d-%d", id, j)) + _, err := conn.WriteTo(data, "127.0.0.1:8080") + if err != nil { + t.Errorf("WriteTo failed: %v", err) + } + } + }(i) + } + + wg.Wait() + elapsed := time.Since(startTime) + + writes := inner.writeCount.Load() + expected := int64(goroutines * writesPerGoroutine) + + t.Logf("✅ 有锁测试完成: %d 次写入，耗时 %v", writes, elapsed) + + if writes != expected { + t.Errorf("❌ 写入计数不匹配: got %d, expected %d", writes, expected) + } + }) +} + +// ======================================== +// 问题 2: directPacketConn 懒缓存竞争验证 +// ======================================== + +// SimulateDirectPacketConn 模拟 directPacketConn（没有写锁） +type SimulateDirectPacketConn struct { + conn *net.UDPConn + cachedDialTgt atomic.Pointer[netip.AddrPort] + cacheOnce atomic.Bool // 简化版，实际使用 sync.Once + dialTgt string + FullCone bool +} + +func (c *SimulateDirectPacketConn) resolveTarget() error { + // 模拟解析延迟 + time.Sleep(10 * time.Millisecond) + + target := netip.MustParseAddrPort(c.dialTgt) + c.cachedDialTgt.Store(&target) + return nil +} + +func (c *SimulateDirectPacketConn) Write(b []byte) (int, error) { + if !c.FullCone { + return c.conn.Write(b) + } + + // 没有锁保护的懒缓存 + cached := c.cachedDialTgt.Load() + if cached == nil { + if !c.cacheOnce.Swap(true) { + // 第一个 goroutine 解析 + c.resolveTarget() + } else { + // 其他 goroutine 等待解析完成 + for c.cachedDialTgt.Load() == nil { + time.Sleep(1 * time.Millisecond) + } + } + cached = c.cachedDialTgt.Load() + } + + // 写入 - 没有序列化 + return c.conn.WriteToUDPAddrPort(b, *cached) +} + +// FixedDirectPacketConn 修复后的 directPacketConn（有写锁） +type FixedDirectPacketConn struct { + conn *net.UDPConn + cachedDialTgt atomic.Pointer[netip.AddrPort] + resolveOnce sync.Once + resolveErr error + dialTgt string + FullCone bool + writeMu sync.Mutex +} + +func (c *FixedDirectPacketConn) resolveTarget() error { + c.resolveOnce.Do(func() { + time.Sleep(10 * time.Millisecond) + target := netip.MustParseAddrPort(c.dialTgt) + c.cachedDialTgt.Store(&target) + }) + return c.resolveErr +} + +func (c *FixedDirectPacketConn) Write(b []byte) (int, error) { + if !c.FullCone { + return c.conn.Write(b) + } + + // 确保目标已解析 + if c.cachedDialTgt.Load() == nil { + c.resolveTarget() + } + + // 有写锁保护 + c.writeMu.Lock() + defer c.writeMu.Unlock() + + cached := c.cachedDialTgt.Load() + return c.conn.WriteToUDPAddrPort(b, *cached) +} + +// TestIssue2_Outbound_DirectPacketConnLazyCache 验证问题 2: directPacketConn 懒缓存竞争 +func TestIssue2_Outbound_DirectPacketConnLazyCache(t *testing.T) { + t.Log("🔍 验证问题 2: directPacketConn 懒缓存竞争 (outbound)") + + // 创建真实的 UDP 连接 + conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + t.Fatalf("Failed to create UDP connection: %v", err) + } + defer conn.Close() + + // 测试没有锁的情况 + t.Run("WithoutLock", func(t *testing.T) { + directConn := &SimulateDirectPacketConn{ + conn: conn, + dialTgt: "127.0.0.1:8080", + FullCone: true, + } + + const goroutines = 10 + const writesPerGoroutine = 50 + + var wg sync.WaitGroup + wg.Add(goroutines) + + startTime := time.Now() + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + data := []byte(fmt.Sprintf("direct-%d-%d", id, j)) + _, err := directConn.Write(data) + if err != nil { + t.Logf("Write error: %v", err) + } + } + }(i) + } + + wg.Wait() + elapsed := time.Since(startTime) + + t.Logf("✅ 无锁测试完成，耗时 %v", elapsed) + t.Log("⚠️ 检查 UDP 连接是否有并发写入问题") + }) + + // 测试有锁的情况 + t.Run("WithLock", func(t *testing.T) { + directConn := &FixedDirectPacketConn{ + conn: conn, + dialTgt: "127.0.0.1:8080", + FullCone: true, + } + + const goroutines = 10 + const writesPerGoroutine = 50 + + var wg sync.WaitGroup + wg.Add(goroutines) + + startTime := time.Now() + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + data := []byte(fmt.Sprintf("direct-%d-%d", id, j)) + _, err := directConn.Write(data) + if err != nil { + t.Errorf("Write failed: %v", err) + } + } + }(i) + } + + wg.Wait() + elapsed := time.Since(startTime) + + t.Logf("✅ 有锁测试完成，耗时 %v", elapsed) + }) +} + +// ======================================== +// 问题 7: Pool.Put 边界检查验证 +// ======================================== + +// TestIssue7_Outbound_PoolPutBoundary 验证问题 7: Pool.Put 边界检查 +func TestIssue7_Outbound_PoolPutBoundary(t *testing.T) { + t.Log("🔍 验证问题 7: Pool.Put 边界检查 (outbound)") + + testCases := []struct { + name string + capacity int + shouldAccept bool + note string + }{ + {"64 bytes (too small)", 64, false, "should be rejected"}, + {"512 bytes (min)", 512, true, "bucket 9"}, + {"1024 bytes (2^10)", 1024, true, "bucket 10"}, + {"1536 bytes (not power of 2)", 1536, true, "⚠️ goes to bucket 10, not 11"}, + {"2048 bytes (2^11)", 2048, true, "bucket 11"}, + {"4096 bytes (2^12)", 4096, true, "bucket 12"}, + {"65536 bytes (max)", 65536, true, "bucket 16"}, + {"70000 bytes (too large)", 70000, false, "should be rejected"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + buf := make([]byte, tc.capacity) + + t.Logf("Testing: cap=%d, %s", tc.capacity, tc.note) + + // 调用 Put（不应该 panic） + pool.Put(buf) + + if tc.shouldAccept { + t.Logf("✅ Buffer accepted (cap=%d)", tc.capacity) + } else { + t.Logf("✅ Buffer rejected (cap=%d)", tc.capacity) + } + }) + } + + t.Log("⚠️ 问题确认: cap=1536 的 buffer 会被放入错误的 bucket") + t.Log(" 这会导致:") + t.Log(" 1. 内存浪费（大 buffer 放入小 bucket）") + t.Log(" 2. 性能下降（下次 Get 可能容量不足）") +} + +// ======================================== +// 综合对比测试 +// ======================================== + +// TestOutboundLockVsNoLock 对比有锁和无锁的性能 +func TestOutboundLockVsNoLock(t *testing.T) { + t.Log("🔍 对比测试: 有锁 vs 无锁") + + // 创建测试组件 + protocol := &MockProtocol{} + + const goroutines = 10 + const writesPerGoroutine = 100 + + t.Run("WithoutLock", func(t *testing.T) { + inner := &MockPacketConn{} + conn := &SimulateShadowsockrPacketConn{ + inner: inner, + protocol: protocol, + tgt: "127.0.0.1:8080", + } + + var wg sync.WaitGroup + wg.Add(goroutines) + + startTime := time.Now() + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + data := []byte(fmt.Sprintf("test-%d-%d", id, j)) + conn.WriteTo(data, "127.0.0.1:8080") + } + }(i) + } + + wg.Wait() + elapsed := time.Since(startTime) + + t.Logf("无锁: %v (%.2f ops/sec)", elapsed, float64(goroutines*writesPerGoroutine)/elapsed.Seconds()) + }) + + t.Run("WithLock", func(t *testing.T) { + inner := &MockPacketConn{} + conn := &FixedShadowsockrPacketConn{ + inner: inner, + protocol: protocol, + tgt: "127.0.0.1:8080", + } + + var wg sync.WaitGroup + wg.Add(goroutines) + + startTime := time.Now() + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + data := []byte(fmt.Sprintf("test-%d-%d", id, j)) + conn.WriteTo(data, "127.0.0.1:8080") + } + }(i) + } + + wg.Wait() + elapsed := time.Since(startTime) + + t.Logf("有锁: %v (%.2f ops/sec)", elapsed, float64(goroutines*writesPerGoroutine)/elapsed.Seconds()) + }) + + t.Log("⚠️ 注意: 锁的开销通常小于数据竞争修复的成本") +} + +// TestBufferPoolMemoryUsage 测试 buffer pool 的内存使用 +func TestBufferPoolMemoryUsage(t *testing.T) { + t.Log("🔍 Buffer Pool 内存使用测试") + + // 获取初始内存状态 + // var m1 runtime.MemStats + // runtime.ReadMemStats(&m1) + + const iterations = 10000 + + // 测试正常使用 + for i := 0; i < iterations; i++ { + buf := pool.Get(1500) + // 使用 buffer + _ = buf + pool.Put(buf) + } + + // var m2 runtime.MemStats + // runtime.ReadMemStats(&m2) + + // 测试问题场景：1536 字节的 buffer + for i := 0; i < iterations; i++ { + buf := make([]byte, 1536) + pool.Put(buf) // 会被放入错误的 bucket + } + + // var m3 runtime.MemStats + // runtime.ReadMemStats(&m3) + + t.Log("✅ 内存使用测试完成") + t.Log("⚠️ 使用 pprof 检查内存分配:") + t.Log(" go test -memprofile=mem.prof -bench=. -run=TestBufferPoolMemoryUsage") + t.Log(" go tool pprof mem.prof") +} diff --git a/transport/shadowsocksr/proto/udp_concurrent_test.go b/transport/shadowsocksr/proto/udp_concurrent_test.go new file mode 100644 index 00000000..772b81e5 --- /dev/null +++ b/transport/shadowsocksr/proto/udp_concurrent_test.go @@ -0,0 +1,230 @@ +package proto + +import ( + "net" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/daeuniverse/outbound/pool" +) + +// TestUdpConnConcurrentWriteWithRealUDP 使用真实 UDP 连接测试并发写入 +// 这个测试验证 shadowsocksr 的 PacketConn 在并发写入时是否存在竞争 +// 运行: go test -race -run TestUdpConnConcurrentWriteWithRealUDP +func TestUdpConnConcurrentWriteWithRealUDP(t *testing.T) { + // 创建服务器 + serverAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0") + if err != nil { + t.Fatalf("Failed to resolve server address: %v", err) + } + + serverConn, err := net.ListenUDP("udp", serverAddr) + if err != nil { + t.Fatalf("Failed to create server: %v", err) + } + defer serverConn.Close() + + // 接收计数器 + var receivedCount int64 + + // 启动服务器接收协程 + go func() { + buf := make([]byte, 1500) + for { + n, _, err := serverConn.ReadFromUDP(buf) + if err != nil { + return + } + if n > 0 { + atomic.AddInt64(&receivedCount, 1) + } + } + }() + + // 创建客户端 + clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + t.Fatalf("Failed to create client: %v", err) + } + defer clientConn.Close() + + target := serverConn.LocalAddr().(*net.UDPAddr).AddrPort() + + // 创建 PacketConn 包装器（模拟 shadowsocksr 的 PacketConn） + // 注意：这里直接使用 UDP 连接测试并发安全性 + // 问题：如果 PacketConn 的 WriteTo 方法没有锁保护， + // 多个 goroutine 并发调用会导致数据竞争 + + const goroutines = 10 + const writesPerGoroutine = 100 + + var wg sync.WaitGroup + wg.Add(goroutines) + + start := time.Now() + + for i := 0; i < goroutines; i++ { + go func(id int) { + defer wg.Done() + + for j := 0; j < writesPerGoroutine; j++ { + data := []byte("test data from goroutine") + // 直接写入 UDP 连接 + _, err := clientConn.WriteToUDPAddrPort(data, target) + if err != nil { + t.Errorf("Goroutine %d write %d failed: %v", id, j, err) + } + } + }(i) + } + + wg.Wait() + + // 等待数据被接收 + time.Sleep(100 * time.Millisecond) + + received := atomic.LoadInt64(&receivedCount) + expected := int64(goroutines * writesPerGoroutine) + + duration := time.Since(start) + + t.Logf("Sent %d packets, received %d packets in %v", expected, received, duration) + + if received < expected*9/10 { + t.Errorf("Packet loss detected: sent %d, received %d", expected, received) + } +} + +// TestUdpConnWriteToWithoutLock 演示没有锁保护的问题 +// 这个测试创建一个模拟的 PacketConn 来展示竞争条件 +func TestUdpConnWriteToWithoutLock(t *testing.T) { + // 模拟一个没有锁保护的写入器 + type unsafeWriter struct { + writes int64 + } + + writer := &unsafeWriter{} + + var wg sync.WaitGroup + const goroutines = 20 + const writesPerGoroutine = 1000 + + wg.Add(goroutines) + + // 模拟并发写入（没有锁保护） + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + // 模拟写入操作（非原子） + current := atomic.LoadInt64(&writer.writes) + time.Sleep(1 * time.Microsecond) // 增加竞争窗口 + atomic.StoreInt64(&writer.writes, current+1) + } + }() + } + + wg.Wait() + + writes := atomic.LoadInt64(&writer.writes) + expected := int64(goroutines * writesPerGoroutine) + + // 由于没有锁保护，写入次数可能不等于预期值 + t.Logf("Writes without lock: got %d, expected %d (loss: %d)", + writes, expected, expected-writes) + + if writes != expected { + t.Logf("⚠️ Race condition detected: %d writes lost", expected-writes) + } +} + +// TestUdpConnWriteToWithLock 演示有锁保护的情况 +func TestUdpConnWriteToWithLock(t *testing.T) { + type safeWriter struct { + writes int64 + mu sync.Mutex + } + + writer := &safeWriter{} + + var wg sync.WaitGroup + const goroutines = 20 + const writesPerGoroutine = 1000 + + wg.Add(goroutines) + + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + for j := 0; j < writesPerGoroutine; j++ { + writer.mu.Lock() + writer.writes++ + writer.mu.Unlock() + } + }() + } + + wg.Wait() + + writes := writer.writes + expected := int64(goroutines * writesPerGoroutine) + + t.Logf("Writes with lock: got %d, expected %d", writes, expected) + + if writes != expected { + t.Errorf("Lock protection failed: got %d, expected %d", writes, expected) + } +} + +// BenchmarkUdpConnWriteParallel 并发性能基准测试 +func BenchmarkUdpConnWriteParallel(b *testing.B) { + serverConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + b.Skip("Failed to create server") + } + defer serverConn.Close() + + clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + if err != nil { + b.Skip("Failed to create client") + } + defer clientConn.Close() + + target := serverConn.LocalAddr().(*net.UDPAddr).AddrPort() + data := []byte("benchmark test data") + + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + clientConn.WriteToUDPAddrPort(data, target) + } + }) +} + +// TestPoolConcurrentAccess 测试 buffer pool 的并发访问 +func TestPoolConcurrentAccess(t *testing.T) { + const goroutines = 20 + const opsPerGoroutine = 500 + + var wg sync.WaitGroup + wg.Add(goroutines) + + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + for j := 0; j < opsPerGoroutine; j++ { + buf := pool.Get(1500) + if len(buf) < 1500 { + t.Errorf("Buffer too small: %d", len(buf)) + } + pool.Put(buf) + } + }() + } + + wg.Wait() + t.Logf("Pool concurrent access test completed") +} diff --git a/transport/shadowsocksr/proto/udp_conn.go b/transport/shadowsocksr/proto/udp_conn.go index 0686706a..ec73a0ad 100644 --- a/transport/shadowsocksr/proto/udp_conn.go +++ b/transport/shadowsocksr/proto/udp_conn.go @@ -5,6 +5,7 @@ import ( "net/netip" "github.com/daeuniverse/outbound/ciphers" + "sync" "github.com/daeuniverse/outbound/netproxy" "github.com/daeuniverse/outbound/pool" "github.com/daeuniverse/outbound/pool/bytes" @@ -16,6 +17,7 @@ type PacketConn struct { netproxy.PacketConn Protocol IProtocol tgt string + writeMu sync.Mutex } func NewPacketConn(c netproxy.PacketConn, proto IProtocol, tgt string) (*PacketConn, error) { @@ -77,7 +79,15 @@ func (c *PacketConn) WriteTo(b []byte, to string) (n int, err error) { if err != nil { return 0, err } + + // Lock to protect concurrent writes. + // Critical section is minimized to only protect the shared buffer and write operation. + c.writeMu.Lock() + defer c.writeMu.Unlock() + pb := pool.GetMustBigger(len(addr) + len(b)) + defer pool.Put(pb) + copy(pb, addr) copy(pb[len(addr):], b) buf := bytes.NewBuffer(pb) @@ -88,5 +98,5 @@ func (c *PacketConn) WriteTo(b []byte, to string) (n int, err error) { return 0, err } - return len(b), err + return len(b), nil } diff --git a/transport/shadowsocksr/proto/udp_conn_test.go b/transport/shadowsocksr/proto/udp_conn_test.go new file mode 100644 index 00000000..d6ff3b3c --- /dev/null +++ b/transport/shadowsocksr/proto/udp_conn_test.go @@ -0,0 +1,117 @@ +package proto + +import ( + "io" + "net/netip" + "sync" + "testing" + "time" + + "github.com/daeuniverse/outbound/pool" + poolbytes "github.com/daeuniverse/outbound/pool/bytes" + "github.com/daeuniverse/outbound/protocol/infra/socks" +) + +type recordingPacketConn struct { + mu sync.Mutex + writes []recordedPacketWrite +} + +type recordedPacketWrite struct { + addr string + data []byte +} + +func (c *recordingPacketConn) Read([]byte) (int, error) { return 0, io.EOF } + +func (c *recordingPacketConn) Write(p []byte) (int, error) { + return c.WriteTo(p, "") +} + +func (c *recordingPacketConn) ReadFrom([]byte) (int, netip.AddrPort, error) { + return 0, netip.AddrPort{}, io.EOF +} + +func (c *recordingPacketConn) WriteTo(p []byte, addr string) (int, error) { + clone := append([]byte(nil), p...) + c.mu.Lock() + c.writes = append(c.writes, recordedPacketWrite{addr: addr, data: clone}) + c.mu.Unlock() + return len(p), nil +} + +func (c *recordingPacketConn) Close() error { return nil } +func (c *recordingPacketConn) SetDeadline(time.Time) error { return nil } +func (c *recordingPacketConn) SetReadDeadline(time.Time) error { return nil } +func (c *recordingPacketConn) SetWriteDeadline(time.Time) error { return nil } + +type noOpProtocol struct{} + +func (p *noOpProtocol) InitWithServerInfo(*ServerInfo) {} +func (p *noOpProtocol) Encode(data []byte) ([]byte, error) { return append([]byte(nil), data...), nil } +func (p *noOpProtocol) Decode(data []byte) ([]byte, int, error) { + return append([]byte(nil), data...), len(data), nil +} +func (p *noOpProtocol) EncodePkt(*poolbytes.Buffer) error { return nil } +func (p *noOpProtocol) SetData(data interface{}) {} +func (p *noOpProtocol) GetData() interface{} { return nil } +func (p *noOpProtocol) GetOverhead() int { return 0 } +func (p *noOpProtocol) DecodePkt(data []byte) (pool.Bytes, error) { + pb := pool.Get(len(data)) + copy(pb, data) + return pb, nil +} + +func TestPacketConnConcurrentWriteTo(t *testing.T) { + t.Helper() + + recorder := &recordingPacketConn{} + pc, err := NewPacketConn(recorder, &noOpProtocol{}, "1.1.1.1:53") + if err != nil { + t.Fatalf("NewPacketConn failed: %v", err) + } + + targets := []string{ + "1.1.1.1:53", + "8.8.8.8:53", + "9.9.9.9:53", + "208.67.222.222:53", + } + + var wg sync.WaitGroup + for _, target := range targets { + target := target + wg.Add(1) + go func() { + defer wg.Done() + if _, err := pc.WriteTo([]byte("payload-"+target), target); err != nil { + t.Errorf("WriteTo(%s) failed: %v", target, err) + } + }() + } + wg.Wait() + + if len(recorder.writes) != len(targets) { + t.Fatalf("unexpected write count: got %d want %d", len(recorder.writes), len(targets)) + } + + seen := make(map[string]bool, len(targets)) + for _, write := range recorder.writes { + addr := socks.SplitAddr(write.data) + if addr == nil { + t.Fatal("failed to parse encoded target addr") + } + target := addr.String() + payload := string(write.data[len(addr):]) + if payload != "payload-"+target { + t.Fatalf("payload mismatch for %s: got %q", target, payload) + } + seen[target] = true + } + + for _, target := range targets { + if !seen[target] { + t.Fatalf("missing target write for %s", target) + } + } +} diff --git a/transport/simpleobfs/http.go b/transport/simpleobfs/http.go index 76c965ea..26d81dd1 100644 --- a/transport/simpleobfs/http.go +++ b/transport/simpleobfs/http.go @@ -92,7 +92,6 @@ func (ho *HTTPObfs) Write(b []byte) (int, error) { return ho.Conn.Write(b) } -// NewHTTPObfs return a HTTPObfs func NewHTTPObfs(conn netproxy.Conn, host string, port string, path string) netproxy.Conn { if !strings.HasPrefix(path, "/") { path = "/" + path diff --git a/transport/simpleobfs/simpleobfs.go b/transport/simpleobfs/simpleobfs.go index 6fb35075..7e007c89 100644 --- a/transport/simpleobfs/simpleobfs.go +++ b/transport/simpleobfs/simpleobfs.go @@ -27,7 +27,6 @@ type SimpleObfs struct { host string } -// NewSimpleobfs returns a simpleobfs proxy. func NewSimpleObfs(option *dialer.ExtraOption, nextDialer netproxy.Dialer, link string) (netproxy.Dialer, *dialer.Property, error) { u, err := url.Parse(link) if err != nil { diff --git a/transport/simpleobfs/tls.go b/transport/simpleobfs/tls.go index df95e9cb..781c9b06 100644 --- a/transport/simpleobfs/tls.go +++ b/transport/simpleobfs/tls.go @@ -112,7 +112,6 @@ func (to *TLSObfs) write(b []byte) (int, error) { return len(b), err } -// NewTLSObfs return a SimpleObfs func NewTLSObfs(conn netproxy.Conn, server string) netproxy.Conn { return &TLSObfs{ Conn: conn, diff --git a/transport/tls/fragment.go b/transport/tls/fragment.go index 6c010d6a..83fca866 100644 --- a/transport/tls/fragment.go +++ b/transport/tls/fragment.go @@ -34,6 +34,8 @@ type FragmentConn struct { minInterval int64 } +const fragmentStackRecordScratch = 5 + 1024 + func NewFragmentConn(rawConn netproxy.Conn, minLength, maxLength, minInterval, maxInterval int64) *FragmentConn { return &FragmentConn{ rawConn: rawConn, @@ -56,33 +58,45 @@ func (f *FragmentConn) Write(b []byte) (n int, err error) { if len(b) < recordLen { return f.rawConn.Write(b) } + minChunkLen, maxChunkLen := normalizeFragmentBounds(f.minLength, f.maxLength) data := b[5:recordLen] - buf := make([]byte, 1024) - var hello []byte - for from := 0; ; { - to := common.Min(len(data), from+int(randBetween(f.minLength, f.maxLength))) - copy(buf[:3], b) - copy(buf[5:], data[from:to]) - l := to - from - from = to - buf[3] = byte(l >> 8) - buf[4] = byte(l) - if f.maxInterval == 0 { - hello = append(hello, buf[:5+l]...) - } else { - if _, err := f.rawConn.Write(buf[:5+l]); err != nil { - return 0, err - } - time.Sleep(time.Duration(randBetween(f.minInterval, f.maxInterval)) * time.Millisecond) - } - if from == len(data) { - break - } + var stackScratch [fragmentStackRecordScratch]byte + recordScratch := stackScratch[:] + if 5+maxChunkLen > len(recordScratch) { + recordScratch = make([]byte, 5+maxChunkLen) } - if len(hello) > 0 { + + if f.maxInterval == 0 { + hello := make([]byte, 0, fragmentAggregateCap(len(data), minChunkLen)) + for from := 0; from < len(data); { + to := common.Min(len(data), from+int(randBetween(int64(minChunkLen), int64(maxChunkLen)))) + chunkLen := to - from + start := len(hello) + hello = hello[:start+5+chunkLen] + copy(hello[start:start+3], b[:3]) + hello[start+3] = byte(chunkLen >> 8) + hello[start+4] = byte(chunkLen) + copy(hello[start+5:start+5+chunkLen], data[from:to]) + from = to + } if _, err := f.rawConn.Write(hello); err != nil { return 0, err } + } else { + frame := recordScratch + for from := 0; from < len(data); { + to := common.Min(len(data), from+int(randBetween(int64(minChunkLen), int64(maxChunkLen)))) + chunkLen := to - from + copy(frame[:3], b[:3]) + frame[3] = byte(chunkLen >> 8) + frame[4] = byte(chunkLen) + copy(frame[5:5+chunkLen], data[from:to]) + if _, err := f.rawConn.Write(frame[:5+chunkLen]); err != nil { + return 0, err + } + time.Sleep(time.Duration(randBetween(f.minInterval, f.maxInterval)) * time.Millisecond) + from = to + } } if len(b) > recordLen { if _, err := f.rawConn.Write(b[recordLen:]); err != nil { @@ -92,6 +106,29 @@ func (f *FragmentConn) Write(b []byte) (n int, err error) { return len(b), nil } +func normalizeFragmentBounds(minLength, maxLength int64) (minChunkLen, maxChunkLen int) { + minChunkLen = int(minLength) + maxChunkLen = int(maxLength) + if minChunkLen <= 0 { + minChunkLen = 1 + } + if maxChunkLen < minChunkLen { + maxChunkLen = minChunkLen + } + return minChunkLen, maxChunkLen +} + +func fragmentAggregateCap(dataLen, minChunkLen int) int { + chunks := dataLen / minChunkLen + if dataLen%minChunkLen != 0 { + chunks++ + } + if chunks == 0 { + chunks = 1 + } + return dataLen + chunks*5 +} + func (f *FragmentConn) Close() error { return f.rawConn.Close() } diff --git a/transport/tls/fragment_test.go b/transport/tls/fragment_test.go new file mode 100644 index 00000000..b578cdad --- /dev/null +++ b/transport/tls/fragment_test.go @@ -0,0 +1,113 @@ +package tls + +import ( + "bytes" + "io" + "testing" + "time" + + "github.com/daeuniverse/outbound/netproxy" +) + +type fragmentTestConn struct { + writes [][]byte +} + +func (c *fragmentTestConn) Read(_ []byte) (int, error) { return 0, io.EOF } +func (c *fragmentTestConn) Close() error { return nil } +func (c *fragmentTestConn) SetDeadline(_ time.Time) error { return nil } +func (c *fragmentTestConn) SetReadDeadline(_ time.Time) error { return nil } +func (c *fragmentTestConn) SetWriteDeadline(_ time.Time) error { return nil } + +func (c *fragmentTestConn) Write(p []byte) (int, error) { + c.writes = append(c.writes, append([]byte(nil), p...)) + return len(p), nil +} + +var _ netproxy.Conn = (*fragmentTestConn)(nil) + +func TestFragmentConnWrite_PassthroughNonTLSRecord(t *testing.T) { + raw := &fragmentTestConn{} + conn := NewFragmentConn(raw, 2, 2, 0, 0) + + payload := []byte("plain") + n, err := conn.Write(payload) + if err != nil { + t.Fatalf("write failed: %v", err) + } + if n != len(payload) { + t.Fatalf("unexpected write length: got %d want %d", n, len(payload)) + } + if len(raw.writes) != 1 { + t.Fatalf("unexpected write count: got %d want 1", len(raw.writes)) + } + if !bytes.Equal(raw.writes[0], payload) { + t.Fatalf("payload mismatch: got %x want %x", raw.writes[0], payload) + } +} + +func TestFragmentConnWrite_AggregatesFragmentsWithoutInterval(t *testing.T) { + raw := &fragmentTestConn{} + conn := NewFragmentConn(raw, 2, 2, 0, 0) + + record := []byte{ + 0x16, 0x03, 0x03, 0x00, 0x06, + 'a', 'b', 'c', 'd', 'e', 'f', + 'X', 'Y', + } + n, err := conn.Write(record) + if err != nil { + t.Fatalf("write failed: %v", err) + } + if n != len(record) { + t.Fatalf("unexpected write length: got %d want %d", n, len(record)) + } + if len(raw.writes) != 2 { + t.Fatalf("unexpected write count: got %d want 2", len(raw.writes)) + } + + wantHello := []byte{ + 0x16, 0x03, 0x03, 0x00, 0x02, 'a', 'b', + 0x16, 0x03, 0x03, 0x00, 0x02, 'c', 'd', + 0x16, 0x03, 0x03, 0x00, 0x02, 'e', 'f', + } + if !bytes.Equal(raw.writes[0], wantHello) { + t.Fatalf("fragmented hello mismatch:\n got %x\nwant %x", raw.writes[0], wantHello) + } + if !bytes.Equal(raw.writes[1], []byte{'X', 'Y'}) { + t.Fatalf("tail mismatch: got %x want %x", raw.writes[1], []byte{'X', 'Y'}) + } +} + +func TestFragmentConnWrite_WritesFragmentsIndividuallyWithInterval(t *testing.T) { + raw := &fragmentTestConn{} + conn := NewFragmentConn(raw, 2, 2, 1, 1) + + record := []byte{ + 0x16, 0x03, 0x03, 0x00, 0x06, + 'a', 'b', 'c', 'd', 'e', 'f', + 'X', 'Y', + } + n, err := conn.Write(record) + if err != nil { + t.Fatalf("write failed: %v", err) + } + if n != len(record) { + t.Fatalf("unexpected write length: got %d want %d", n, len(record)) + } + if len(raw.writes) != 4 { + t.Fatalf("unexpected write count: got %d want 4", len(raw.writes)) + } + + wantFragments := [][]byte{ + {0x16, 0x03, 0x03, 0x00, 0x02, 'a', 'b'}, + {0x16, 0x03, 0x03, 0x00, 0x02, 'c', 'd'}, + {0x16, 0x03, 0x03, 0x00, 0x02, 'e', 'f'}, + {'X', 'Y'}, + } + for i := range wantFragments { + if !bytes.Equal(raw.writes[i], wantFragments[i]) { + t.Fatalf("write %d mismatch:\n got %x\nwant %x", i, raw.writes[i], wantFragments[i]) + } + } +} diff --git a/transport/tls/reality.go b/transport/tls/reality.go index 5a2dfb30..f3b8f339 100644 --- a/transport/tls/reality.go +++ b/transport/tls/reality.go @@ -213,7 +213,12 @@ func (x *Reality) DialContext(ctx context.Context, network, addr string) (c netp // if config.Show { // logrus.Printf("REALITY hello.SessionId[:16]: %v\n", hello.SessionId[:16]) // } - if uConn.HandshakeState.State13.EcdheKey == nil { + // Use KeyShareKeys.Ecdhe (new API) with fallback to EcdheKey (deprecated) for compatibility + ecdheKey := uConn.HandshakeState.State13.KeyShareKeys.Ecdhe + if ecdheKey == nil { + ecdheKey = uConn.HandshakeState.State13.EcdheKey + } + if ecdheKey == nil { // logrus.Println("wtf", retry, addr) if retry > 2 { return nil, errors.New("nil ecdheKey") @@ -222,7 +227,7 @@ func (x *Reality) DialContext(ctx context.Context, network, addr string) (c netp goto retryHandshake // retry } // logrus.Println("OH YEAH", retry) - uConn.AuthKey, _ = uConn.HandshakeState.State13.EcdheKey.ECDH(x.publicKey) + uConn.AuthKey, _ = ecdheKey.ECDH(x.publicKey) if uConn.AuthKey == nil { return nil, errors.New("REALITY: SharedKey == nil") } diff --git a/transport/ws/conn.go b/transport/ws/conn.go index 71377272..ec939ed0 100644 --- a/transport/ws/conn.go +++ b/transport/ws/conn.go @@ -1,14 +1,20 @@ package ws import ( - "bytes" + "io" + "sync" + "github.com/gorilla/websocket" "time" ) type conn struct { *websocket.Conn - readBuffer bytes.Buffer + + readMu sync.Mutex + currentReader io.Reader + + writeMu sync.Mutex } func newConn(wsc *websocket.Conn) *conn { @@ -18,22 +24,53 @@ func newConn(wsc *websocket.Conn) *conn { } func (c *conn) Read(b []byte) (n int, err error) { - if c.readBuffer.Len() > 0 { - return c.readBuffer.Read(b) + c.readMu.Lock() + defer c.readMu.Unlock() + + for { + if c.currentReader == nil { + messageType, reader, err := c.Conn.NextReader() + if err != nil { + return 0, err + } + if messageType != websocket.BinaryMessage { + _, _ = io.Copy(io.Discard, reader) + continue + } + c.currentReader = reader + } + + n, err = c.currentReader.Read(b) + if err == nil { + return n, nil + } + if err == io.EOF { + c.currentReader = nil + if n > 0 { + return n, nil + } + continue + } + return n, err } - _, msg, err := c.Conn.ReadMessage() +} +func (c *conn) Write(b []byte) (n int, err error) { + c.writeMu.Lock() + defer c.writeMu.Unlock() + + writer, err := c.Conn.NextWriter(websocket.BinaryMessage) if err != nil { return 0, err } - n = copy(b, msg) - if n < len(msg) { - c.readBuffer.Write(msg[n:]) + n, err = writer.Write(b) + closeErr := writer.Close() + if err != nil { + return n, err + } + if closeErr != nil { + return n, closeErr } return n, nil - -} -func (c *conn) Write(b []byte) (n int, err error) { - return len(b), c.Conn.WriteMessage(websocket.BinaryMessage, b) } func (c *conn) SetDeadline(t time.Time) error { diff --git a/transport/ws/conn_test.go b/transport/ws/conn_test.go new file mode 100644 index 00000000..927bb9e9 --- /dev/null +++ b/transport/ws/conn_test.go @@ -0,0 +1,97 @@ +package ws + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gorilla/websocket" +) + +func newWSPair(t *testing.T) (*conn, *websocket.Conn, func()) { + t.Helper() + + upgrader := websocket.Upgrader{} + serverConnCh := make(chan *websocket.Conn, 1) + serverErrCh := make(chan error, 1) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + serverConnCh <- conn + })) + + clientConn, _, err := websocket.DefaultDialer.Dial("ws"+server.URL[len("http"):], nil) + if err != nil { + server.Close() + t.Fatalf("dial websocket failed: %v", err) + } + + var serverConn *websocket.Conn + select { + case err := <-serverErrCh: + clientConn.Close() + server.Close() + t.Fatalf("upgrade failed: %v", err) + case serverConn = <-serverConnCh: + } + + cleanup := func() { + _ = clientConn.Close() + _ = serverConn.Close() + server.Close() + } + return newConn(clientConn), serverConn, cleanup +} + +func TestConnReadStreamsLargeMessage(t *testing.T) { + client, server, cleanup := newWSPair(t) + defer cleanup() + + payload := bytes.Repeat([]byte("abcd"), 4096) + if err := server.WriteMessage(websocket.BinaryMessage, payload); err != nil { + t.Fatalf("server write failed: %v", err) + } + + got := make([]byte, len(payload)) + if _, err := io.ReadFull(client, got); err != nil { + t.Fatalf("client read failed: %v", err) + } + if !bytes.Equal(got, payload) { + t.Fatal("payload mismatch") + } +} + +func TestConnWriteStreamsBinaryMessage(t *testing.T) { + client, server, cleanup := newWSPair(t) + defer cleanup() + + payload := bytes.Repeat([]byte("hello"), 2048) + n, err := client.Write(payload) + if err != nil { + t.Fatalf("client write failed: %v", err) + } + if n != len(payload) { + t.Fatalf("unexpected write length: got %d want %d", n, len(payload)) + } + + messageType, reader, err := server.NextReader() + if err != nil { + t.Fatalf("server next reader failed: %v", err) + } + if messageType != websocket.BinaryMessage { + t.Fatalf("unexpected message type: got %d want %d", messageType, websocket.BinaryMessage) + } + got, err := io.ReadAll(reader) + if err != nil { + t.Fatalf("server read failed: %v", err) + } + if !bytes.Equal(got, payload) { + t.Fatal("payload mismatch") + } +}