simplify buffer extension

proxy/shadowsocks/config.go

@@ -198,13 +198,10 @@ 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
-	})
+
+	b.Extend(int32(auth.Overhead()))
+	_, err := auth.Seal(b.BytesTo(ivLen), b.BytesRange(ivLen, payloadLen))
+	return err
 }
 
 func (c *AEADCipher) DecodePacket(key []byte, b *buf.Buffer) error {
@@ -214,16 +211,12 @@ 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 {
+
+	bbb, err := auth.Open(b.BytesTo(ivLen), b.BytesRange(ivLen, payloadLen))
+	if err != nil {
 		return err
 	}
-	b.Advance(ivLen)
+	b.Resize(ivLen, ivLen+int32(len(bbb)))
 	return nil
 }
 

proxy/shadowsocks/ota.go

@@ -30,13 +30,11 @@ func NewAuthenticator(keygen KeyGenerator) *Authenticator {
 	}
 }
 
-func (v *Authenticator) Authenticate(data []byte) buf.Supplier {
+func (v *Authenticator) Authenticate(data []byte, dest []byte) {
 	hasher := hmac.New(sha1.New, v.key())
 	common.Must2(hasher.Write(data))
 	res := hasher.Sum(nil)
-	return func(b []byte) (int, error) {
-		return copy(b, res[:AuthSize]), nil
-	}
+	copy(dest, res[:AuthSize])
 }
 
 func HeaderKeyGenerator(key []byte, iv []byte) func() []byte {
@@ -89,7 +87,7 @@ func (v *ChunkReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
 	payload := buffer[AuthSize:size]
 
 	actualAuthBytes := make([]byte, AuthSize)
-	v.auth.Authenticate(payload)(actualAuthBytes)
+	v.auth.Authenticate(payload, actualAuthBytes)
 	if !bytes.Equal(authBytes, actualAuthBytes) {
 		return nil, newError("invalid auth")
 	}
@@ -121,7 +119,7 @@ func (w *ChunkWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 	for {
 		payloadLen, _ := mb.Read(w.buffer[2+AuthSize:])
 		binary.BigEndian.PutUint16(w.buffer, uint16(payloadLen))
-		w.auth.Authenticate(w.buffer[2+AuthSize : 2+AuthSize+payloadLen])(w.buffer[2:])
+		w.auth.Authenticate(w.buffer[2+AuthSize:2+AuthSize+payloadLen], w.buffer[2:])
 		if err := buf.WriteAllBytes(w.writer, w.buffer[:2+AuthSize+payloadLen]); err != nil {
 			return err
 		}

proxy/shadowsocks/protocol.go

@@ -142,7 +142,8 @@ func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (buf.Wri
 		header.SetByte(0, header.Byte(0)|0x10)
 
 		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
-		common.Must(header.AppendSupplier(authenticator.Authenticate(header.Bytes())))
+		authBuffer := header.Extend(AuthSize)
+		authenticator.Authenticate(header.Bytes(), authBuffer)
 	}
 
 	if err := w.WriteMultiBuffer(buf.NewMultiBufferValue(header)); err != nil {
@@ -210,7 +211,9 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload []byte) (*buf.Buff
 		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
 		buffer.SetByte(ivLen, buffer.Byte(ivLen)|0x10)
 
-		common.Must(buffer.AppendSupplier(authenticator.Authenticate(buffer.BytesFrom(ivLen))))
+		authPayload := buffer.BytesFrom(ivLen)
+		authBuffer := buffer.Extend(AuthSize)
+		authenticator.Authenticate(authPayload, authBuffer)
 	}
 	if err := account.Cipher.EncodePacket(account.Key, buffer); err != nil {
 		return nil, newError("failed to encrypt UDP payload").Base(err)

proxy/shadowsocks/protocol_test.go

@@ -7,7 +7,6 @@ import (
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/protocol"
-	"v2ray.com/core/common/serial"
 	. "v2ray.com/core/proxy/shadowsocks"
 	. "v2ray.com/ext/assert"
 )
@@ -37,7 +36,7 @@ func TestUDPEncoding(t *testing.T) {
 	}
 
 	data := buf.New()
-	data.AppendSupplier(serial.WriteString("test string"))
+	common.Must2(data.WriteString("test string"))
 	encodedData, err := EncodeUDPPacket(request, data.Bytes())
 	assert(err, IsNil)
 
@@ -168,7 +167,7 @@ func TestUDPReaderWriter(t *testing.T) {
 
 	{
 		b := buf.New()
-		b.AppendSupplier(serial.WriteString("test payload"))
+		common.Must2(b.WriteString("test payload"))
 		err := writer.WriteMultiBuffer(buf.NewMultiBufferValue(b))
 		assert(err, IsNil)
 
@@ -179,7 +178,7 @@ func TestUDPReaderWriter(t *testing.T) {
 
 	{
 		b := buf.New()
-		b.AppendSupplier(serial.WriteString("test payload 2"))
+		common.Must2(b.WriteString("test payload 2"))
 		err := writer.WriteMultiBuffer(buf.NewMultiBufferValue(b))
 		assert(err, IsNil)