fix aead udp

proxy/shadowsocks/config.go

@@ -103,6 +103,8 @@ type Cipher interface {
 	NewEncryptionWriter(key []byte, iv []byte, writer io.Writer) (buf.Writer, error)
 	NewDecryptionReader(key []byte, iv []byte, reader io.Reader) (buf.Reader, error)
 	IsAEAD() bool
+	EncodePacket(key []byte, b *buf.Buffer) error
+	DecodePacket(key []byte, b *buf.Buffer) error
 }
 
 type AesCfb struct {
@@ -131,6 +133,21 @@ func (v *AesCfb) NewDecryptionReader(key []byte, iv []byte, reader io.Reader) (b
 	return buf.NewReader(crypto.NewCryptionReader(stream, reader)), nil
 }
 
+func (v *AesCfb) EncodePacket(key []byte, b *buf.Buffer) error {
+	iv := b.BytesTo(v.IVSize())
+	stream := crypto.NewAesEncryptionStream(key, iv)
+	stream.XORKeyStream(b.BytesFrom(v.IVSize()), b.BytesFrom(v.IVSize()))
+	return nil
+}
+
+func (v *AesCfb) DecodePacket(key []byte, b *buf.Buffer) error {
+	iv := b.BytesTo(v.IVSize())
+	stream := crypto.NewAesEncryptionStream(key, iv)
+	stream.XORKeyStream(b.BytesFrom(v.IVSize()), b.BytesFrom(v.IVSize()))
+	b.SliceFrom(v.IVSize())
+	return nil
+}
+
 type AEADCipher struct {
 	KeyBytes        int
 	IVBytes         int
@@ -149,32 +166,60 @@ func (c *AEADCipher) IVSize() int {
 	return c.IVBytes
 }
 
-func (c *AEADCipher) NewEncryptionWriter(key []byte, iv []byte, writer io.Writer) (buf.Writer, error) {
+func (c *AEADCipher) createAuthenticator(key []byte, iv []byte) *crypto.AEADAuthenticator {
 	nonce := crypto.NewIncreasingAEADNonceGenerator()
 	subkey := make([]byte, c.KeyBytes)
 	hkdfSHA1(key, iv, subkey)
-	auth := &crypto.AEADAuthenticator{
+	return &crypto.AEADAuthenticator{
 		AEAD:           c.AEADAuthCreator(subkey),
 		NonceGenerator: nonce,
 	}
+}
+
+func (c *AEADCipher) NewEncryptionWriter(key []byte, iv []byte, writer io.Writer) (buf.Writer, error) {
+	auth := c.createAuthenticator(key, iv)
 	return crypto.NewAuthenticationWriter(auth, &crypto.AEADChunkSizeParser{
 		Auth: auth,
 	}, writer, protocol.TransferTypeStream), nil
 }
 
 func (c *AEADCipher) NewDecryptionReader(key []byte, iv []byte, reader io.Reader) (buf.Reader, error) {
-	nonce := crypto.NewIncreasingAEADNonceGenerator()
-	subkey := make([]byte, c.KeyBytes)
-	hkdfSHA1(key, iv, subkey)
-	auth := &crypto.AEADAuthenticator{
-		AEAD:           c.AEADAuthCreator(subkey),
-		NonceGenerator: nonce,
-	}
+	auth := c.createAuthenticator(key, iv)
 	return crypto.NewAuthenticationReader(auth, &crypto.AEADChunkSizeParser{
 		Auth: auth,
 	}, reader, protocol.TransferTypeStream), nil
 }
 
+func (c *AEADCipher) EncodePacket(key []byte, b *buf.Buffer) error {
+	ivLen := c.IVSize()
+	payloadLen := b.Len()
+	auth := c.createAuthenticator(key, b.BytesTo(ivLen))
+	return b.Reset(func(bb []byte) (int, error) {
+		bbb, err := auth.Seal(bb[:ivLen], bb[ivLen:payloadLen])
+		if err != nil {
+			return 0, err
+		}
+		return len(bbb), nil
+	})
+}
+
+func (c *AEADCipher) DecodePacket(key []byte, b *buf.Buffer) error {
+	ivLen := c.IVSize()
+	payloadLen := b.Len()
+	auth := c.createAuthenticator(key, b.BytesTo(ivLen))
+	if err := b.Reset(func(bb []byte) (int, error) {
+		bbb, err := auth.Open(bb[:ivLen], bb[ivLen:payloadLen])
+		if err != nil {
+			return 0, err
+		}
+		return len(bbb), nil
+	}); err != nil {
+		return err
+	}
+	b.SliceFrom(ivLen)
+	return nil
+}
+
 type ChaCha20 struct {
 	IVBytes int
 }
@@ -201,6 +246,21 @@ func (v *ChaCha20) NewDecryptionReader(key []byte, iv []byte, reader io.Reader)
 	return buf.NewReader(crypto.NewCryptionReader(stream, reader)), nil
 }
 
+func (v *ChaCha20) EncodePacket(key []byte, b *buf.Buffer) error {
+	iv := b.BytesTo(v.IVSize())
+	stream := crypto.NewChaCha20Stream(key, iv)
+	stream.XORKeyStream(b.BytesFrom(v.IVSize()), b.BytesFrom(v.IVSize()))
+	return nil
+}
+
+func (v *ChaCha20) DecodePacket(key []byte, b *buf.Buffer) error {
+	iv := b.BytesTo(v.IVSize())
+	stream := crypto.NewChaCha20Stream(key, iv)
+	stream.XORKeyStream(b.BytesFrom(v.IVSize()), b.BytesFrom(v.IVSize()))
+	b.SliceFrom(v.IVSize())
+	return nil
+}
+
 func PasswordToCipherKey(password string, keySize int) []byte {
 	pwdBytes := []byte(password)
 	key := make([]byte, 0, keySize)

proxy/shadowsocks/protocol.go

@@ -247,39 +247,31 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload []byte) (*buf.Buff
 	common.Must(buffer.Reset(buf.ReadFullFrom(rand.Reader, ivLen)))
 	iv := buffer.Bytes()
 
-	payloadBuffer := buf.New()
-	defer payloadBuffer.Release()
-
 	switch request.Address.Family() {
 	case net.AddressFamilyIPv4:
-		payloadBuffer.AppendBytes(AddrTypeIPv4)
-		payloadBuffer.Append([]byte(request.Address.IP()))
+		buffer.AppendBytes(AddrTypeIPv4)
+		buffer.Append([]byte(request.Address.IP()))
 	case net.AddressFamilyIPv6:
-		payloadBuffer.AppendBytes(AddrTypeIPv6)
-		payloadBuffer.Append([]byte(request.Address.IP()))
+		buffer.AppendBytes(AddrTypeIPv6)
+		buffer.Append([]byte(request.Address.IP()))
 	case net.AddressFamilyDomain:
-		payloadBuffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
-		payloadBuffer.Append([]byte(request.Address.Domain()))
+		buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
+		buffer.Append([]byte(request.Address.Domain()))
 	default:
 		return nil, newError("unsupported address type: ", request.Address.Family()).AtError()
 	}
 
-	common.Must(payloadBuffer.AppendSupplier(serial.WriteUint16(uint16(request.Port))))
-	payloadBuffer.Append(payload)
+	common.Must(buffer.AppendSupplier(serial.WriteUint16(uint16(request.Port))))
+	buffer.Append(payload)
 
 	if !account.Cipher.IsAEAD() && request.Option.Has(RequestOptionOneTimeAuth) {
 		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
-		payloadBuffer.SetByte(0, payloadBuffer.Byte(0)|0x10)
+		buffer.SetByte(ivLen, buffer.Byte(ivLen)|0x10)
 
-		common.Must(payloadBuffer.AppendSupplier(authenticator.Authenticate(payloadBuffer.Bytes())))
+		common.Must(buffer.AppendSupplier(authenticator.Authenticate(buffer.BytesFrom(ivLen))))
 	}
-
-	w, err := account.Cipher.NewEncryptionWriter(account.Key, iv, buffer)
-	if err != nil {
-		return nil, newError("failed to create encoding stream").Base(err).AtError()
-	}
-	if err := w.WriteMultiBuffer(buf.NewMultiBufferValue(payloadBuffer)); err != nil {
-		return nil, newError("failed to encrypt UDP payload").Base(err).AtWarning()
+	if err := account.Cipher.EncodePacket(account.Key, buffer); err != nil {
+		return nil, newError("failed to encrypt UDP payload").Base(err)
 	}
 
 	return buffer, nil
@@ -292,24 +284,15 @@ func DecodeUDPPacket(user *protocol.User, payload *buf.Buffer) (*protocol.Reques
 	}
 	account := rawAccount.(*ShadowsocksAccount)
 
-	ivLen := account.Cipher.IVSize()
-	iv := make([]byte, ivLen)
-	copy(iv, payload.BytesTo(ivLen))
-	payload.SliceFrom(ivLen)
-
-	r, err := account.Cipher.NewDecryptionReader(account.Key, iv, payload)
-	if err != nil {
-		return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError()
+	var authenticator *Authenticator
+	if !account.Cipher.IsAEAD() {
+		authenticator = NewAuthenticator(HeaderKeyGenerator(account.Key, payload.BytesTo(account.Cipher.IVSize())))
 	}
-	mb, err := r.ReadMultiBuffer()
-	if err != nil {
-		return nil, nil, newError("failed to decrypt UDP payload").Base(err).AtWarning()
+
+	if err := account.Cipher.DecodePacket(account.Key, payload); err != nil {
+		return nil, nil, newError("failed to decrypt UDP payload").Base(err)
 	}
-	payload.Release()
-	payload = mb.SplitFirst()
-	mb.Release()
 
-	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
 	request := &protocol.RequestHeader{
 		Version: Version,
 		User:    user,

testing/scenarios/shadowsocks_test.go

@@ -10,6 +10,7 @@ import (
 	"v2ray.com/core"
 	"v2ray.com/core/app/log"
 	"v2ray.com/core/app/proxyman"
+	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/protocol"
 	"v2ray.com/core/common/serial"
@@ -690,3 +691,85 @@ func TestShadowsocksAES256GCMConformance(t *testing.T) {
 
 	CloseAllServers(servers)
 }
+
+func TestShadowsocksChacha20Poly1305UDPConformance(t *testing.T) {
+	assert := With(t)
+
+	udpServer := udp.Server{
+		MsgProcessor: xor,
+	}
+	dest, err := udpServer.Start()
+	assert(err, IsNil)
+	defer udpServer.Close()
+
+	account := serial.ToTypedMessage(&shadowsocks.Account{
+		Password:   "ss-password",
+		CipherType: shadowsocks.CipherType_CHACHA20_POLY1305,
+	})
+
+	serverPort := pickPort()
+	serverConfig := &core.Config{
+		App: []*serial.TypedMessage{
+			serial.ToTypedMessage(&log.Config{
+				ErrorLogLevel: log.LogLevel_Debug,
+				ErrorLogType:  log.LogType_Console,
+			}),
+		},
+		Inbound: []*proxyman.InboundHandlerConfig{
+			{
+				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
+					PortRange: net.SinglePortRange(serverPort),
+					Listen:    net.NewIPOrDomain(net.LocalHostIP),
+				}),
+				ProxySettings: serial.ToTypedMessage(&shadowsocks.ServerConfig{
+					UdpEnabled: true,
+					User: &protocol.User{
+						Account: account,
+						Level:   1,
+					},
+				}),
+			},
+		},
+		Outbound: []*proxyman.OutboundHandlerConfig{
+			{
+				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
+			},
+		},
+	}
+
+	servers, err := InitializeServerConfigs(serverConfig)
+	assert(err, IsNil)
+
+	cipher, err := ss.PickCipher("CHACHA20-IETF-POLY1305", nil, "ss-password")
+	assert(err, IsNil)
+	conn, err := ss.ListenPacket("udp", ":0", cipher)
+	assert(err, IsNil)
+
+	for i := 0; i < 100; i++ {
+
+		payload := buf.New()
+		payload.AppendBytes(1, 127, 0, 0, 1)
+		payload.AppendSupplier(serial.WriteUint16(dest.Port.Value()))
+
+		payload.AppendSupplier(buf.ReadFullFrom(rand.Reader, 10))
+
+		nBytes, err := conn.WriteTo(payload.Bytes(), &net.UDPAddr{
+			IP:   []byte{127, 0, 0, 1},
+			Port: int(serverPort),
+		})
+		assert(err, IsNil)
+		assert(nBytes, Equals, payload.Len())
+
+		conn.SetReadDeadline(time.Now().Add(time.Second * 10))
+		response := make([]byte, 10240)
+		nBytes, _, err = conn.ReadFrom(response)
+		assert(err, IsNil)
+		assert(response[:7], Equals, payload.BytesTo(7))
+		assert(response[7:nBytes], Equals, xor(payload.BytesFrom(7)))
+
+	}
+
+	assert(conn.Close(), IsNil)
+
+	CloseAllServers(servers)
+}