Skip to content

Commit

Permalink
🥪 ss2022/tcp: change padding behavior
Browse files Browse the repository at this point in the history 
Previously, for Shadowsocks 2022 TCP, padding is only added when there's no initial payload. This commit changes the behavior so that padding is added as long as the size of the initial payload does not exceed MaxPaddingLength.

Along with the above change is a bit of refactoring and a fix for large (>128k) initial payload handling.
  • Loading branch information
@database64128
database64128 committed Oct 12, 2022
1 parent 7151714 commit d98e94d
Show file tree
Hide file tree
Showing 4 changed files with 141 additions and 139 deletions.
43 changes: 18 additions & 25 deletions ss2022/header.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,29 +193,27 @@ func ParseTCPRequestVariableLengthHeader(b []byte) (targetAddr conn.Addr, payloa
return
}

// WriteTCPRequestVariableLengthHeader writes a TCP request variable-length header into the buffer
// and returns the number of bytes written.
// WriteTCPRequestVariableLengthHeader writes a TCP request variable-length header into the buffer.
//
// The buffer must be at least
// socks5.LengthOfAddrFromConnAddr(targetAddr) + 2 + MaxPaddingLength bytes long if there's no payload, or
// socks5.LengthOfAddrFromConnAddr(targetAddr) + 2 + len(payload) bytes long if there's initial payload.
// The total header length must not exceed MaxPayloadSize.
func WriteTCPRequestVariableLengthHeader(b []byte, targetAddr conn.Addr, payload []byte) (n int) {
// The header fills the whole buffer. Excess bytes are used as padding.
//
// The buffer size can be calculated with:
//
// socks5.LengthOfAddrFromConnAddr(targetAddr) + 2 + len(payload) + paddingLen
//
// The buffer size must not exceed [MaxPayloadSize].
// The excess space in the buffer must not be larger than [MaxPaddingLength] bytes.
func WriteTCPRequestVariableLengthHeader(b []byte, targetAddr conn.Addr, payload []byte) {
// SOCKS address
n = socks5.WriteAddrFromConnAddr(b, targetAddr)
n := socks5.WriteAddrFromConnAddr(b, targetAddr)

// Padding length
var paddingLen int
if len(payload) == 0 {
paddingLen = rand.Intn(MaxPaddingLength) + 1
}
paddingLen := len(b) - n - 2 - len(payload)
binary.BigEndian.PutUint16(b[n:], uint16(paddingLen))
n += 2 + paddingLen

// Initial payload
n += copy(b[n:], payload)

return
copy(b[n:], payload)
}

// ParseTCPResponseHeader parses a TCP response fixed-length header and returns the length
Expand Down Expand Up @@ -256,26 +254,21 @@ func ParseTCPResponseHeader(b []byte, requestSalt []byte) (n int, err error) {
return
}

// WriteTCPResponseHeader writes a TCP response fixed-length header into the buffer
// and returns the number of bytes written.
// WriteTCPResponseHeader writes a TCP response fixed-length header into the buffer.
//
// This function does not check buffer length.
// The buffer must be at least 1 + 8 + salt length + 2 bytes long.
func WriteTCPResponseHeader(b []byte, requestSalt []byte, length uint16) (n int) {
// The buffer size must be exactly 1 + 8 + len(requestSalt) + 2 bytes.
func WriteTCPResponseHeader(b []byte, requestSalt []byte, length uint16) {
// Type
b[0] = HeaderTypeServerStream

// Timestamp
binary.BigEndian.PutUint64(b[1:], uint64(time.Now().Unix()))

// Request salt
n = 1 + 8
n += copy(b[n:], requestSalt)
copy(b[1+8:], requestSalt)

// Length
binary.BigEndian.PutUint16(b[n:], length)
n += 2
return
binary.BigEndian.PutUint16(b[1+8+len(requestSalt):], length)
}

// ParseSessionIDAndPacketID parses the session ID and packet ID segment of a decrypted UDP packet.
Expand Down
80 changes: 36 additions & 44 deletions ss2022/header_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,26 +66,23 @@ func TestWriteAndParseTCPRequestFixedLengthHeader(t *testing.T) {
}

func TestWriteAndParseTCPRequestVariableLengthHeader(t *testing.T) {
payloadLen := mrand.Intn(1024)
payloadLen := 1 + mrand.Intn(1024)
payload := make([]byte, payloadLen)
_, err := rand.Read(payload)
if err != nil {
t.Fatal(err)
}
targetAddr := conn.AddrFromIPPort(netip.AddrPortFrom(netip.IPv6Unspecified(), 443))
targetAddrLen := socks5.LengthOfAddrFromConnAddr(targetAddr)
expectedHeaderWithPayloadLength := targetAddrLen + 2 + payloadLen
bufLen := TCPRequestVariableLengthHeaderNoPayloadMaxLength + payloadLen
noPayloadLen := targetAddrLen + 2 + 1 + mrand.Intn(MaxPaddingLength)
noPaddingLen := targetAddrLen + 2 + payloadLen
bufLen := noPaddingLen + MaxPaddingLength
b := make([]byte, bufLen)

// 1. Good header (with initial payload)
n := WriteTCPRequestVariableLengthHeader(b, targetAddr, payload)
if n != expectedHeaderWithPayloadLength {
t.Fatalf("Expected n %d, got %d", expectedHeaderWithPayloadLength, n)
}
header := b[:n]
// 1. Good header (padding + initial payload)
WriteTCPRequestVariableLengthHeader(b, targetAddr, payload)

ta, p, err := ParseTCPRequestVariableLengthHeader(header)
ta, p, err := ParseTCPRequestVariableLengthHeader(b)
if err != nil {
t.Fatal(err)
}
Expand All @@ -96,92 +93,87 @@ func TestWriteAndParseTCPRequestVariableLengthHeader(t *testing.T) {
t.Fatalf("Expected target address %s, got %s", targetAddr, ta)
}

// 2. Good header (no payload)
n = WriteTCPRequestVariableLengthHeader(b, targetAddr, nil)
if n <= targetAddrLen+2 {
t.Fatalf("Header should have been padded!\nActual length: %d", n)
}
header = b[:n]
// 2. Good header (initial payload)
b = b[:noPaddingLen]
WriteTCPRequestVariableLengthHeader(b, targetAddr, payload)

ta, p, err = ParseTCPRequestVariableLengthHeader(header)
ta, p, err = ParseTCPRequestVariableLengthHeader(b)
if err != nil {
t.Fatal(err)
}
if len(p) > 0 {
t.Fatalf("Expected empty initial payload, got length %d", len(p))
if !bytes.Equal(p, payload) {
t.Fatalf("Expected payload %v\nGot: %v", payload, p)
}
if ta != targetAddr {
t.Fatalf("Expected target address %s, got %s", targetAddr, ta)
}

// 3. Good header (padding + payload)
n += copy(b[n:], payload)
header = b[:n]
// 3. Good header (padding)
b = b[:noPayloadLen]
WriteTCPRequestVariableLengthHeader(b, targetAddr, nil)

ta, p, err = ParseTCPRequestVariableLengthHeader(header)
ta, p, err = ParseTCPRequestVariableLengthHeader(b)
if err != nil {
t.Fatal(err)
}
if !bytes.Equal(p, payload) {
t.Fatalf("Expected payload %v\nGot: %v", payload, p)
if len(p) > 0 {
t.Fatalf("Expected empty initial payload, got length %d", len(p))
}
if ta != targetAddr {
t.Fatalf("Expected target address %s, got %s", targetAddr, ta)
}

// 4. Bad header (incomplete padding)
n -= payloadLen
n -= 1
header = b[:n]
b = b[:noPayloadLen-1]

_, _, err = ParseTCPRequestVariableLengthHeader(header)
_, _, err = ParseTCPRequestVariableLengthHeader(b)
if !errors.Is(err, ErrPaddingExceedChunkBorder) {
t.Fatalf("Expected: %s\nGot: %s", ErrPaddingExceedChunkBorder, err)
}

// 5. Bad header (padding length out of range)
binary.BigEndian.PutUint16(b[targetAddrLen:], MaxPaddingLength+1)

_, _, err = ParseTCPRequestVariableLengthHeader(header)
_, _, err = ParseTCPRequestVariableLengthHeader(b)
if !errors.Is(err, ErrPaddingLengthOutOfRange) {
t.Fatalf("Expected: %s\nGot: %s", ErrPaddingLengthOutOfRange, err)
}

// 6. Bad header (incomplete padding length)
n = targetAddrLen + 1
header = b[:n]
b = b[:targetAddrLen+1]

_, _, err = ParseTCPRequestVariableLengthHeader(header)
_, _, err = ParseTCPRequestVariableLengthHeader(b)
if !errors.Is(err, ErrIncompleteHeaderInFirstChunk) {
t.Fatalf("Expected: %s\nGot: %s", ErrIncompleteHeaderInFirstChunk, err)
}

// 7. Bad header (incomplete SOCKS address)
n = targetAddrLen - 1
header = b[:n]
b = b[:targetAddrLen-1]

_, _, err = ParseTCPRequestVariableLengthHeader(header)
_, _, err = ParseTCPRequestVariableLengthHeader(b)
if err == nil {
t.Fatal("Expected error, got nil")
}
}

func TestWriteAndParseTCPResponseHeader(t *testing.T) {
b := make([]byte, TCPResponseHeaderMaxLength)
const (
saltLen = 32
bufLen = 1 + 8 + saltLen + 2
)

b := make([]byte, bufLen)
length := mrand.Intn(math.MaxUint16)
requestSalt := make([]byte, 32)
requestSalt := make([]byte, saltLen)
_, err := rand.Read(requestSalt)
if err != nil {
t.Fatal(err)
}

// 1. Good header
n := WriteTCPResponseHeader(b, requestSalt, uint16(length))
if n != TCPResponseHeaderMaxLength {
t.Fatalf("Expected: %d\nGot: %d", TCPResponseHeaderMaxLength, n)
}
WriteTCPResponseHeader(b, requestSalt, uint16(length))

n, err = ParseTCPResponseHeader(b, requestSalt)
n, err := ParseTCPResponseHeader(b, requestSalt)
if err != nil {
t.Fatal(err)
}
Expand All @@ -190,7 +182,7 @@ func TestWriteAndParseTCPResponseHeader(t *testing.T) {
}

// 2. Bad request salt
_, err = rand.Read(b[1+8 : 1+8+32])
_, err = rand.Read(b[1+8 : 1+8+saltLen])
if err != nil {
t.Fatal(err)
}
Expand Down
55 changes: 30 additions & 25 deletions ss2022/stream.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"encoding/binary"
"errors"
"io"
mrand "math/rand"

"github.com/database64128/shadowsocks-go/conn"
"github.com/database64128/shadowsocks-go/socks5"
Expand Down Expand Up @@ -166,14 +167,14 @@ func (rw *ShadowStreamServerReadWriter) WriteZeroCopy(b []byte, payloadStart, pa
if rw.w == nil { // first write
urspLen := len(rw.unsafeResponseStreamPrefix)
saltLen := len(rw.cipherConfig.PSK)
fixedLengthHeaderStart := urspLen + saltLen
responseHeaderPlaintextEnd := fixedLengthHeaderStart + TCPRequestFixedLengthHeaderLength + saltLen
payloadBufStart := responseHeaderPlaintextEnd + 16
responseHeaderStart := urspLen + saltLen
responseHeaderEnd := responseHeaderStart + TCPRequestFixedLengthHeaderLength + saltLen
payloadBufStart := responseHeaderEnd + 16
bufferLen := payloadBufStart + payloadLen + 16
hb := make([]byte, bufferLen)
ursp := hb[:urspLen]
salt := hb[urspLen:fixedLengthHeaderStart]
responseHeader := hb[fixedLengthHeaderStart:responseHeaderPlaintextEnd]
salt := hb[urspLen:responseHeaderStart]
responseHeader := hb[responseHeaderStart:responseHeaderEnd]

// Write unsafe response stream prefix.
copy(ursp, rw.unsafeResponseStreamPrefix)
Expand All @@ -185,7 +186,7 @@ func (rw *ShadowStreamServerReadWriter) WriteZeroCopy(b []byte, payloadStart, pa
}

// Write response header.
_ = WriteTCPResponseHeader(responseHeader, rw.requestSalt, uint16(payloadLen))
WriteTCPResponseHeader(responseHeader, rw.requestSalt, uint16(payloadLen))

// Create AEAD cipher.
shadowStreamCipher := rw.cipherConfig.NewShadowStreamCipher(salt)
Expand Down Expand Up @@ -254,21 +255,25 @@ type ShadowStreamClientReadWriter struct {
// NewShadowStreamClientReadWriter writes request headers to rw and returns a Shadowsocks stream client ready for reads and writes.
func NewShadowStreamClientReadWriter(rwo zerocopy.DirectReadWriteCloserOpener, cipherConfig *CipherConfig, eihPSKHashes [][IdentityHeaderLength]byte, targetAddr conn.Addr, payload, unsafeRequestStreamPrefix, unsafeResponseStreamPrefix []byte) (sscRW *ShadowStreamClientReadWriter, rawRW zerocopy.DirectReadWriteCloser, err error) {
var (
payloadOrPaddingMaxLen int
excessivePayload []byte
paddingPayloadLen int
excessPayload []byte
)

targetAddrLen := socks5.LengthOfAddrFromConnAddr(targetAddr)
payloadLen := len(payload)
roomForPayload := MaxPayloadSize - targetAddrLen - 2

switch {
case len(payload) > roomForPayload:
payloadOrPaddingMaxLen = roomForPayload
excessivePayload = payload[roomForPayload:]
case len(payload) > 0:
payloadOrPaddingMaxLen = len(payload)
case payloadLen > roomForPayload:
paddingPayloadLen = roomForPayload
excessPayload = payload[roomForPayload:]
payload = payload[:roomForPayload]
case payloadLen >= MaxPaddingLength:
paddingPayloadLen = payloadLen
case payloadLen > 0:
paddingPayloadLen = payloadLen + mrand.Intn(MaxPaddingLength-payloadLen+1)
default:
payloadOrPaddingMaxLen = MaxPaddingLength
paddingPayloadLen = 1 + mrand.Intn(MaxPaddingLength)
}

urspLen := len(unsafeRequestStreamPrefix)
Expand All @@ -278,7 +283,8 @@ func NewShadowStreamClientReadWriter(rwo zerocopy.DirectReadWriteCloserOpener, c
fixedLengthHeaderStart := identityHeadersStart + identityHeadersLen
fixedLengthHeaderEnd := fixedLengthHeaderStart + TCPRequestFixedLengthHeaderLength
variableLengthHeaderStart := fixedLengthHeaderEnd + 16
variableLengthHeaderEnd := variableLengthHeaderStart + targetAddrLen + 2 + payloadOrPaddingMaxLen
variableLengthHeaderLen := targetAddrLen + 2 + paddingPayloadLen
variableLengthHeaderEnd := variableLengthHeaderStart + variableLengthHeaderLen
bufferLen := variableLengthHeaderEnd + 16
b := make([]byte, bufferLen)
ursp := b[:urspLen]
Expand All @@ -305,10 +311,10 @@ func NewShadowStreamClientReadWriter(rwo zerocopy.DirectReadWriteCloserOpener, c
}

// Write variable-length header.
n := WriteTCPRequestVariableLengthHeader(variableLengthHeaderPlaintext, targetAddr, payload)
WriteTCPRequestVariableLengthHeader(variableLengthHeaderPlaintext, targetAddr, payload)

// Write fixed-length header.
WriteTCPRequestFixedLengthHeader(fixedLengthHeaderPlaintext, uint16(n))
WriteTCPRequestFixedLengthHeader(fixedLengthHeaderPlaintext, uint16(variableLengthHeaderLen))

// Create AEAD cipher.
shadowStreamCipher := cipherConfig.NewShadowStreamCipher(salt)
Expand All @@ -317,11 +323,10 @@ func NewShadowStreamClientReadWriter(rwo zerocopy.DirectReadWriteCloserOpener, c
shadowStreamCipher.EncryptInPlace(fixedLengthHeaderPlaintext)

// Seal variable-length header.
shadowStreamCipher.EncryptInPlace(variableLengthHeaderPlaintext[:n])
shadowStreamCipher.EncryptInPlace(variableLengthHeaderPlaintext)

// Write out.
n += variableLengthHeaderStart + 16
rawRW, err = rwo.Open(b[:n])
rawRW, err = rwo.Open(b)
if err != nil {
return
}
Expand All @@ -331,11 +336,11 @@ func NewShadowStreamClientReadWriter(rwo zerocopy.DirectReadWriteCloserOpener, c
ssc: shadowStreamCipher,
}

// Write excessive payload.
if len(excessivePayload) > 0 {
copy(payload[2+16:], excessivePayload)
_, err = w.WriteZeroCopy(payload, 2+16, len(excessivePayload))
if err != nil {
// Write excess payload, reusing the variable-length header buffer.
for len(excessPayload) > 0 {
n := copy(variableLengthHeaderPlaintext, excessPayload)
excessPayload = excessPayload[n:]
if _, err = w.WriteZeroCopy(b, variableLengthHeaderStart, n); err != nil {
rawRW.Close()
return
}
Expand Down
Loading

0 comments on commit d98e94d

Please sign in to comment.