diff --git a/app/dispatcher/impl/default.go b/app/dispatcher/impl/default.go index 1529c0aa7..ddf70e375 100644 --- a/app/dispatcher/impl/default.go +++ b/app/dispatcher/impl/default.go @@ -62,7 +62,7 @@ func (v *DefaultDispatcher) DispatchToOutbound(session *proxy.SessionInfo) ray.I } if session.Inbound != nil && session.Inbound.AllowPassiveConnection { - go dispatcher.Dispatch(destination, alloc.NewLocalBuffer(32).Clear(), direct) + go dispatcher.Dispatch(destination, alloc.NewLocalBuffer(32), direct) } else { go v.FilterPacketAndDispatch(destination, direct, dispatcher) } diff --git a/app/dispatcher/testing/dispatcher.go b/app/dispatcher/testing/dispatcher.go index 5aad956da..4f853c3d3 100644 --- a/app/dispatcher/testing/dispatcher.go +++ b/app/dispatcher/testing/dispatcher.go @@ -19,7 +19,8 @@ func NewTestPacketDispatcher(handler func(destination v2net.Destination, traffic if err != nil { break } - traffic.OutboundOutput().Write(payload.Prepend([]byte("Processed: "))) + payload.Prepend([]byte("Processed: ")) + traffic.OutboundOutput().Write(payload) } traffic.OutboundOutput().Close() } diff --git a/common/alloc/buffer.go b/common/alloc/buffer.go index 33c5a3d57..ad486b1a8 100644 --- a/common/alloc/buffer.go +++ b/common/alloc/buffer.go @@ -2,32 +2,32 @@ package alloc import ( - "hash" "io" - - "v2ray.com/core/common/serial" ) const ( defaultOffset = 16 ) +type BytesWriter func([]byte) int + // Buffer is a recyclable allocation of a byte array. Buffer.Release() recycles // the buffer into an internal buffer pool, in order to recreate a buffer more // quickly. type Buffer struct { - head []byte - pool Pool - Value []byte - offset int + head []byte + pool Pool + + start int + end int } func CreateBuffer(container []byte, parent Pool) *Buffer { b := new(Buffer) b.head = container b.pool = parent - b.Value = b.head[defaultOffset:] - b.offset = defaultOffset + b.start = defaultOffset + b.end = defaultOffset return b } @@ -40,140 +40,117 @@ func (b *Buffer) Release() { b.pool.Free(b) } b.head = nil - b.Value = nil b.pool = nil } // Clear clears the content of the buffer, results an empty buffer with // Len() = 0. -func (b *Buffer) Clear() *Buffer { - b.offset = defaultOffset - b.Value = b.head[b.offset:b.offset] - return b +func (b *Buffer) Clear() { + b.start = defaultOffset + b.end = defaultOffset } // Reset resets this Buffer into its original state. -func (b *Buffer) Reset() *Buffer { - b.offset = defaultOffset - b.Value = b.head[b.offset:] - return b +func (b *Buffer) Reset() { + b.start = defaultOffset + b.end = len(b.head) } // AppendBytes appends one or more bytes to the end of the buffer. -func (b *Buffer) AppendBytes(bytes ...byte) *Buffer { - b.Value = append(b.Value, bytes...) - return b +func (b *Buffer) AppendBytes(bytes ...byte) { + b.Append(bytes) } // Append appends a byte array to the end of the buffer. -func (b *Buffer) Append(data []byte) *Buffer { - b.Value = append(b.Value, data...) - return b +func (b *Buffer) Append(data []byte) { + nBytes := copy(b.head[b.end:], data) + b.end += nBytes } -// AppendString appends a given string to the end of the buffer. -func (b *Buffer) AppendString(s string) *Buffer { - b.Value = append(b.Value, s...) - return b -} - -func (b *Buffer) AppendUint16(val uint16) *Buffer { - b.Value = serial.Uint16ToBytes(val, b.Value) - return b -} - -func (b *Buffer) AppendUint32(val uint32) *Buffer { - b.Value = serial.Uint32ToBytes(val, b.Value) - return b +func (b *Buffer) AppendFunc(writer BytesWriter) { + nBytes := writer(b.head[b.end:]) + b.end += nBytes } // Prepend prepends bytes in front of the buffer. Caller must ensure total bytes prepended is // no more than 16 bytes. -func (b *Buffer) Prepend(data []byte) *Buffer { +func (b *Buffer) Prepend(data []byte) { b.SliceBack(len(data)) - copy(b.Value, data) - return b + copy(b.head[b.start:], data) } -func (b *Buffer) PrependBytes(data ...byte) *Buffer { - return b.Prepend(data) +func (b *Buffer) PrependBytes(data ...byte) { + b.Prepend(data) } -func (b *Buffer) PrependUint16(val uint16) *Buffer { - b.SliceBack(2) - serial.Uint16ToBytes(val, b.Value[:0]) - return b -} - -func (b *Buffer) PrependUint32(val uint32) *Buffer { - b.SliceBack(4) - serial.Uint32ToBytes(val, b.Value[:0]) - return b -} - -func (b *Buffer) PrependHash(h hash.Hash) *Buffer { - b.SliceBack(h.Size()) - h.Sum(b.Value[:0]) - return b +func (b *Buffer) PrependFunc(offset int, writer BytesWriter) { + b.SliceBack(offset) + writer(b.head[b.start:]) } func (b *Buffer) Byte(index int) byte { - return b.Value[index] + return b.head[b.start+index] } // Bytes returns the content bytes of this Buffer. func (b *Buffer) Bytes() []byte { - return b.Value + return b.head[b.start:b.end] } func (b *Buffer) BytesRange(from, to int) []byte { if from < 0 { - from += len(b.Value) + from += b.Len() } if to < 0 { - to += len(b.Value) + to += b.Len() } - return b.Value[from:to] + return b.head[b.start+from : b.start+to] } func (b *Buffer) BytesFrom(from int) []byte { if from < 0 { - from += len(b.Value) + from += b.Len() } - return b.Value[from:] + return b.head[b.start+from : b.end] } func (b *Buffer) BytesTo(to int) []byte { if to < 0 { - to += len(b.Value) + to += b.Len() } - return b.Value[:to] + return b.head[b.start : b.start+to] } // Slice cuts the buffer at the given position. -func (b *Buffer) Slice(from, to int) *Buffer { - b.offset += from - b.Value = b.Value[from:to] - return b +func (b *Buffer) Slice(from, to int) { + if from < 0 { + from += b.Len() + } + if to < 0 { + to += b.Len() + } + if to < from { + panic("Invalid slice") + } + b.end = b.start + to + b.start += from } // SliceFrom cuts the buffer at the given position. -func (b *Buffer) SliceFrom(from int) *Buffer { - b.offset += from - b.Value = b.Value[from:] - return b +func (b *Buffer) SliceFrom(from int) { + if from < 0 { + from += b.Len() + } + b.start += from } // SliceBack extends the Buffer to its front by offset bytes. // Caller must ensure cumulated offset is no more than 16. -func (b *Buffer) SliceBack(offset int) *Buffer { - newoffset := b.offset - offset - if newoffset < 0 { +func (b *Buffer) SliceBack(offset int) { + b.start -= offset + if b.start < 0 { panic("Negative buffer offset.") } - b.Value = b.head[newoffset : b.offset+len(b.Value)] - b.offset = newoffset - return b } // Len returns the length of the buffer content. @@ -181,7 +158,7 @@ func (b *Buffer) Len() int { if b == nil { return 0 } - return len(b.Value) + return b.end - b.start } func (b *Buffer) IsEmpty() bool { @@ -190,15 +167,13 @@ func (b *Buffer) IsEmpty() bool { // IsFull returns true if the buffer has no more room to grow. func (b *Buffer) IsFull() bool { - return len(b.Value) == cap(b.Value) + return b.end == len(b.head) } // Write implements Write method in io.Writer. func (b *Buffer) Write(data []byte) (int, error) { - begin := b.Len() - b.Value = b.Value[:cap(b.Value)] - nBytes := copy(b.Value[begin:], data) - b.Value = b.Value[:begin+nBytes] + nBytes := copy(b.head[b.end:], data) + b.end += nBytes return nBytes, nil } @@ -207,32 +182,29 @@ func (b *Buffer) Read(data []byte) (int, error) { if b.Len() == 0 { return 0, io.EOF } - nBytes := copy(data, b.Value) + nBytes := copy(data, b.head[b.start:b.end]) if nBytes == b.Len() { b.Clear() } else { - b.Value = b.Value[nBytes:] - b.offset += nBytes + b.start += nBytes } return nBytes, nil } func (b *Buffer) FillFrom(reader io.Reader) (int, error) { - begin := b.Len() - nBytes, err := reader.Read(b.head[b.offset+begin:]) - b.Value = b.head[b.offset : b.offset+begin+nBytes] + nBytes, err := reader.Read(b.head[b.end:]) + b.end += nBytes return nBytes, err } func (b *Buffer) FillFullFrom(reader io.Reader, amount int) (int, error) { - begin := b.Len() - nBytes, err := io.ReadFull(reader, b.head[b.offset+begin:b.offset+begin+amount]) - b.Value = b.head[b.offset : b.offset+begin+nBytes] + nBytes, err := io.ReadFull(reader, b.head[b.end:b.end+amount]) + b.end += nBytes return nBytes, err } func (b *Buffer) String() string { - return string(b.Value) + return string(b.head[b.start:b.end]) } // NewBuffer creates a Buffer with 8K bytes of arbitrary content. diff --git a/common/alloc/buffer_test.go b/common/alloc/buffer_test.go index 4c598ec73..8d00985de 100644 --- a/common/alloc/buffer_test.go +++ b/common/alloc/buffer_test.go @@ -4,13 +4,14 @@ import ( "testing" . "v2ray.com/core/common/alloc" + "v2ray.com/core/common/serial" "v2ray.com/core/testing/assert" ) func TestBufferClear(t *testing.T) { assert := assert.On(t) - buffer := NewBuffer().Clear() + buffer := NewBuffer() defer buffer.Release() payload := "Bytes" @@ -21,22 +22,19 @@ func TestBufferClear(t *testing.T) { assert.Int(buffer.Len()).Equals(0) } -func TestBufferIsFull(t *testing.T) { +func TestBufferIsEmpty(t *testing.T) { assert := assert.On(t) buffer := NewBuffer() defer buffer.Release() - assert.Bool(buffer.IsFull()).IsTrue() - - buffer.Clear() - assert.Bool(buffer.IsFull()).IsFalse() + assert.Bool(buffer.IsEmpty()).IsTrue() } func TestBufferPrepend(t *testing.T) { assert := assert.On(t) - buffer := NewBuffer().Clear() + buffer := NewBuffer() defer buffer.Release() buffer.Append([]byte{'a', 'b', 'c'}) @@ -52,17 +50,17 @@ func TestBufferPrepend(t *testing.T) { func TestBufferString(t *testing.T) { assert := assert.On(t) - buffer := NewBuffer().Clear() + buffer := NewBuffer() defer buffer.Release() - buffer.AppendString("Test String") + buffer.AppendFunc(serial.WriteString("Test String")) assert.String(buffer.String()).Equals("Test String") } func TestBufferWrite(t *testing.T) { assert := assert.On(t) - buffer := NewLocalBuffer(24).Clear() // 16 + 8 + buffer := NewLocalBuffer(24) // 16 + 8 nBytes, err := buffer.Write([]byte("abcd")) assert.Error(err).IsNil() assert.Int(nBytes).Equals(4) diff --git a/common/io/buffered_reader.go b/common/io/buffered_reader.go index 65d6244d2..1568283ac 100644 --- a/common/io/buffered_reader.go +++ b/common/io/buffered_reader.go @@ -17,7 +17,7 @@ type BufferedReader struct { func NewBufferedReader(rawReader io.Reader) *BufferedReader { return &BufferedReader{ reader: rawReader, - buffer: alloc.NewBuffer().Clear(), + buffer: alloc.NewBuffer(), cached: true, } } diff --git a/common/io/buffered_reader_test.go b/common/io/buffered_reader_test.go index d2ee2bffa..78e4e69f6 100644 --- a/common/io/buffered_reader_test.go +++ b/common/io/buffered_reader_test.go @@ -3,6 +3,7 @@ package io_test import ( "testing" + "crypto/rand" "v2ray.com/core/common/alloc" . "v2ray.com/core/common/io" "v2ray.com/core/testing/assert" @@ -12,6 +13,8 @@ func TestBufferedReader(t *testing.T) { assert := assert.On(t) content := alloc.NewBuffer() + content.FillFrom(rand.Reader) + len := content.Len() reader := NewBufferedReader(content) diff --git a/common/io/buffered_writer.go b/common/io/buffered_writer.go index 92bf84a16..3dc6e9744 100644 --- a/common/io/buffered_writer.go +++ b/common/io/buffered_writer.go @@ -17,7 +17,7 @@ type BufferedWriter struct { func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter { return &BufferedWriter{ writer: rawWriter, - buffer: alloc.NewSmallBuffer().Clear(), + buffer: alloc.NewSmallBuffer(), cached: true, } } diff --git a/common/io/buffered_writer_test.go b/common/io/buffered_writer_test.go index 9907578c3..0b8dac9c0 100644 --- a/common/io/buffered_writer_test.go +++ b/common/io/buffered_writer_test.go @@ -12,7 +12,7 @@ import ( func TestBufferedWriter(t *testing.T) { assert := assert.On(t) - content := alloc.NewBuffer().Clear() + content := alloc.NewBuffer() writer := NewBufferedWriter(content) assert.Bool(writer.Cached()).IsTrue() @@ -32,7 +32,7 @@ func TestBufferedWriter(t *testing.T) { func TestBufferedWriterLargePayload(t *testing.T) { assert := assert.On(t) - content := alloc.NewLocalBuffer(128 * 1024).Clear() + content := alloc.NewLocalBuffer(128 * 1024) writer := NewBufferedWriter(content) assert.Bool(writer.Cached()).IsTrue() diff --git a/common/io/chain_writer.go b/common/io/chain_writer.go index 325793c60..79d20acb3 100644 --- a/common/io/chain_writer.go +++ b/common/io/chain_writer.go @@ -28,7 +28,7 @@ func (v *ChainWriter) Write(payload []byte) (int, error) { bytesWritten := 0 size := len(payload) for size > 0 { - buffer := alloc.NewBuffer().Clear() + buffer := alloc.NewBuffer() if size > alloc.BufferSize { buffer.Append(payload[:alloc.BufferSize]) size -= alloc.BufferSize diff --git a/common/io/reader.go b/common/io/reader.go index d0c01ca3b..6554190f6 100644 --- a/common/io/reader.go +++ b/common/io/reader.go @@ -33,7 +33,7 @@ func NewAdaptiveReader(reader io.Reader) *AdaptiveReader { func (v *AdaptiveReader) Read() (*alloc.Buffer, error) { if v.highVolumn && v.largeBuffer.IsEmpty() { if v.largeBuffer == nil { - v.largeBuffer = alloc.NewLocalBuffer(256 * 1024).Clear() + v.largeBuffer = alloc.NewLocalBuffer(256 * 1024) } nBytes, err := v.largeBuffer.FillFrom(v.reader) if err != nil { @@ -44,7 +44,7 @@ func (v *AdaptiveReader) Read() (*alloc.Buffer, error) { } } - buffer := alloc.NewBuffer().Clear() + buffer := alloc.NewBuffer() if !v.largeBuffer.IsEmpty() { buffer.FillFrom(v.largeBuffer) return buffer, nil diff --git a/common/io/writer_test.go b/common/io/writer_test.go index d4c117608..ca07978ae 100644 --- a/common/io/writer_test.go +++ b/common/io/writer_test.go @@ -13,13 +13,15 @@ import ( func TestAdaptiveWriter(t *testing.T) { assert := assert.On(t) - lb := alloc.NewBuffer().Clear() + lb := alloc.NewBuffer() lb.FillFrom(rand.Reader) - writeBuffer := make([]byte, 0, 1024*1024) + expectedBytes := append([]byte(nil), lb.Bytes()...) - writer := NewAdaptiveWriter(NewBufferedWriter(bytes.NewBuffer(writeBuffer))) + writeBuffer := bytes.NewBuffer(make([]byte, 0, 1024*1024)) + + writer := NewAdaptiveWriter(NewBufferedWriter(writeBuffer)) err := writer.Write(lb) assert.Error(err).IsNil() - assert.Bytes(lb.Bytes()).Equals(writeBuffer) + assert.Bytes(expectedBytes).Equals(writeBuffer.Bytes()) } diff --git a/common/serial/hash.go b/common/serial/hash.go new file mode 100644 index 000000000..933cb9241 --- /dev/null +++ b/common/serial/hash.go @@ -0,0 +1,13 @@ +package serial + +import ( + "hash" + "v2ray.com/core/common/alloc" +) + +func WriteHash(h hash.Hash) alloc.BytesWriter { + return func(b []byte) int { + h.Sum(b[:0]) + return h.Size() + } +} diff --git a/common/serial/numbers.go b/common/serial/numbers.go index 883b2b642..2246076af 100644 --- a/common/serial/numbers.go +++ b/common/serial/numbers.go @@ -2,6 +2,7 @@ package serial import ( "strconv" + "v2ray.com/core/common/alloc" ) func Uint16ToBytes(value uint16, b []byte) []byte { @@ -12,6 +13,13 @@ func Uint16ToString(value uint16) string { return strconv.Itoa(int(value)) } +func WriteUint16(value uint16) alloc.BytesWriter { + return func(b []byte) int { + b = Uint16ToBytes(value, b[:0]) + return 2 + } +} + func Uint32ToBytes(value uint32, b []byte) []byte { return append(b, byte(value>>24), byte(value>>16), byte(value>>8), byte(value)) } @@ -20,6 +28,13 @@ func Uint32ToString(value uint32) string { return strconv.FormatUint(uint64(value), 10) } +func WriteUint32(value uint32) alloc.BytesWriter { + return func(b []byte) int { + b = Uint32ToBytes(value, b[:0]) + return 4 + } +} + func IntToBytes(value int, b []byte) []byte { return append(b, byte(value>>24), byte(value>>16), byte(value>>8), byte(value)) } diff --git a/common/serial/string.go b/common/serial/string.go index ef9ae61ba..4ca320685 100644 --- a/common/serial/string.go +++ b/common/serial/string.go @@ -3,6 +3,7 @@ package serial import ( "fmt" "strings" + "v2ray.com/core/common/alloc" ) func ToString(v interface{}) string { @@ -33,3 +34,10 @@ func Concat(v ...interface{}) string { } return strings.Join(values, "") } + +func WriteString(s string) alloc.BytesWriter { + return func(b []byte) int { + copy(b, []byte(s)) + return len(s) + } +} diff --git a/proxy/blackhole/config.go b/proxy/blackhole/config.go index b8804f0df..df1cf0fa3 100644 --- a/proxy/blackhole/config.go +++ b/proxy/blackhole/config.go @@ -6,6 +6,7 @@ import ( "github.com/golang/protobuf/ptypes" "github.com/golang/protobuf/ptypes/any" + "v2ray.com/core/common/serial" ) const ( @@ -31,7 +32,9 @@ func (v *NoneResponse) AsAny() *any.Any { } func (v *HTTPResponse) WriteTo(writer v2io.Writer) { - writer.Write(alloc.NewLocalBuffer(512).Clear().AppendString(http403response)) + b := alloc.NewLocalBuffer(512) + b.AppendFunc(serial.WriteString(http403response)) + writer.Write(b) } func (v *HTTPResponse) AsAny() *any.Any { diff --git a/proxy/blackhole/config_test.go b/proxy/blackhole/config_test.go index d765bd944..2884a5687 100644 --- a/proxy/blackhole/config_test.go +++ b/proxy/blackhole/config_test.go @@ -14,7 +14,7 @@ import ( func TestHTTPResponse(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewBuffer().Clear() + buffer := alloc.NewBuffer() httpResponse := new(HTTPResponse) httpResponse.WriteTo(v2io.NewAdaptiveWriter(buffer)) diff --git a/proxy/freedom/freedom_test.go b/proxy/freedom/freedom_test.go index c005fe5d5..fde289487 100644 --- a/proxy/freedom/freedom_test.go +++ b/proxy/freedom/freedom_test.go @@ -47,7 +47,8 @@ func TestSinglePacket(t *testing.T) { traffic := ray.NewRay() data2Send := "Data to be sent to remote" - payload := alloc.NewLocalBuffer(2048).Clear().Append([]byte(data2Send)) + payload := alloc.NewLocalBuffer(2048) + payload.Append([]byte(data2Send)) go freedom.Dispatch(v2net.TCPDestination(v2net.LocalHostIP, tcpServer.Port), payload, traffic) traffic.InboundInput().Close() diff --git a/proxy/shadowsocks/ota.go b/proxy/shadowsocks/ota.go index c6c3f367a..e3b2e6cbe 100644 --- a/proxy/shadowsocks/ota.go +++ b/proxy/shadowsocks/ota.go @@ -26,11 +26,14 @@ func NewAuthenticator(keygen KeyGenerator) *Authenticator { } } -func (v *Authenticator) Authenticate(auth []byte, data []byte) []byte { +func (v *Authenticator) Authenticate(data []byte) alloc.BytesWriter { hasher := hmac.New(sha1.New, v.key()) hasher.Write(data) res := hasher.Sum(nil) - return append(auth, res[:AuthSize]...) + return func(b []byte) int { + copy(b, res[:AuthSize]) + return AuthSize + } } func HeaderKeyGenerator(key []byte, iv []byte) func() []byte { @@ -71,7 +74,7 @@ func (v *ChunkReader) Release() { } func (v *ChunkReader) Read() (*alloc.Buffer, error) { - buffer := alloc.NewBuffer().Clear() + buffer := alloc.NewBuffer() if _, err := buffer.FillFullFrom(v.reader, 2); err != nil { buffer.Release() return nil, err @@ -94,7 +97,8 @@ func (v *ChunkReader) Read() (*alloc.Buffer, error) { authBytes := buffer.BytesTo(AuthSize) payload := buffer.BytesFrom(AuthSize) - actualAuthBytes := v.auth.Authenticate(nil, payload) + actualAuthBytes := make([]byte, AuthSize) + v.auth.Authenticate(payload)(actualAuthBytes) if !bytes.Equal(authBytes, actualAuthBytes) { buffer.Release() return nil, errors.New("Shadowsocks|AuthenticationReader: Invalid auth.") @@ -123,9 +127,8 @@ func (v *ChunkWriter) Release() { func (v *ChunkWriter) Write(payload *alloc.Buffer) error { totalLength := payload.Len() - payload.SliceBack(AuthSize) - v.auth.Authenticate(payload.BytesTo(0), payload.BytesFrom(AuthSize)) - payload.PrependUint16(uint16(totalLength)) + payload.PrependFunc(AuthSize, v.auth.Authenticate(payload.Bytes())) + payload.PrependFunc(2, serial.WriteUint16(uint16(totalLength))) _, err := v.writer.Write(payload.Bytes()) return err } diff --git a/proxy/shadowsocks/ota_test.go b/proxy/shadowsocks/ota_test.go index 957bcbac5..e5d001574 100644 --- a/proxy/shadowsocks/ota_test.go +++ b/proxy/shadowsocks/ota_test.go @@ -11,7 +11,8 @@ import ( func TestNormalChunkReading(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewBuffer().Clear().AppendBytes( + buffer := alloc.NewBuffer() + buffer.AppendBytes( 0, 8, 39, 228, 69, 96, 133, 39, 254, 26, 201, 70, 11, 12, 13, 14, 15, 16, 17, 18) reader := NewChunkReader(buffer, NewAuthenticator(ChunkKeyGenerator( []byte{21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36}))) @@ -26,11 +27,13 @@ func TestNormalChunkReading(t *testing.T) { func TestNormalChunkWriting(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewLocalBuffer(512).Clear() + buffer := alloc.NewLocalBuffer(512) writer := NewChunkWriter(buffer, NewAuthenticator(ChunkKeyGenerator( []byte{21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36}))) - err := writer.Write(alloc.NewLocalBuffer(256).Clear().Append([]byte{11, 12, 13, 14, 15, 16, 17, 18})) + b := alloc.NewLocalBuffer(256) + b.Append([]byte{11, 12, 13, 14, 15, 16, 17, 18}) + err := writer.Write(b) assert.Error(err).IsNil() assert.Bytes(buffer.Bytes()).Equals([]byte{0, 8, 39, 228, 69, 96, 133, 39, 254, 26, 201, 70, 11, 12, 13, 14, 15, 16, 17, 18}) } diff --git a/proxy/shadowsocks/protocol.go b/proxy/shadowsocks/protocol.go index fa40cd107..9f38942f5 100644 --- a/proxy/shadowsocks/protocol.go +++ b/proxy/shadowsocks/protocol.go @@ -10,6 +10,7 @@ import ( v2io "v2ray.com/core/common/io" v2net "v2ray.com/core/common/net" "v2ray.com/core/common/protocol" + "v2ray.com/core/common/serial" ) const ( @@ -28,7 +29,7 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea } account := rawAccount.(*ShadowsocksAccount) - buffer := alloc.NewLocalBuffer(512).Clear() + buffer := alloc.NewLocalBuffer(512) defer buffer.Release() ivLen := account.Cipher.IVSize() @@ -106,7 +107,8 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea request.Port = v2net.PortFromBytes(buffer.BytesFrom(-2)) if request.Option.Has(RequestOptionOneTimeAuth) { - actualAuth := authenticator.Authenticate(nil, buffer.Bytes()) + actualAuth := make([]byte, AuthSize) + authenticator.Authenticate(buffer.Bytes())(actualAuth) _, err := buffer.FillFullFrom(reader, AuthSize) if err != nil { @@ -150,7 +152,7 @@ func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (v2io.Wr writer = crypto.NewCryptionWriter(stream, writer) - header := alloc.NewLocalBuffer(512).Clear() + header := alloc.NewLocalBuffer(512) switch request.Address.Family() { case v2net.AddressFamilyIPv4: @@ -166,16 +168,16 @@ func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (v2io.Wr return nil, errors.New("Shadowsocks|TCP: Unsupported address type: ", request.Address.Family()) } - header.AppendUint16(uint16(request.Port)) + header.AppendFunc(serial.WriteUint16(uint16(request.Port))) if request.Option.Has(RequestOptionOneTimeAuth) { header.Bytes()[0] |= 0x10 authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv)) - header.Value = authenticator.Authenticate(header.Value, header.Value) + header.AppendFunc(authenticator.Authenticate(header.Bytes())) } - _, err = writer.Write(header.Value) + _, err = writer.Write(header.Bytes()) if err != nil { return nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to write header.") } @@ -243,9 +245,8 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a buffer := alloc.NewSmallBuffer() ivLen := account.Cipher.IVSize() - buffer.Slice(0, ivLen) - rand.Read(buffer.Value) - iv := buffer.Value + buffer.FillFullFrom(rand.Reader, ivLen) + iv := buffer.Bytes() switch request.Address.Family() { case v2net.AddressFamilyIPv4: @@ -261,14 +262,14 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a return nil, errors.New("Shadowsocks|UDP: Unsupported address type: ", request.Address.Family()) } - buffer.AppendUint16(uint16(request.Port)) - buffer.Append(payload.Value) + buffer.AppendFunc(serial.WriteUint16(uint16(request.Port))) + buffer.Append(payload.Bytes()) if request.Option.Has(RequestOptionOneTimeAuth) { authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv)) - buffer.Value[ivLen] |= 0x10 + buffer.Bytes()[ivLen] |= 0x10 - buffer.Value = authenticator.Authenticate(buffer.Value, buffer.Value[ivLen:]) + buffer.AppendFunc(authenticator.Authenticate(buffer.BytesFrom(ivLen))) } stream, err := account.Cipher.NewEncodingStream(account.Key, iv) @@ -276,7 +277,7 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a return nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to create encoding stream.") } - stream.XORKeyStream(buffer.Value[ivLen:], buffer.Value[ivLen:]) + stream.XORKeyStream(buffer.BytesFrom(ivLen), buffer.BytesFrom(ivLen)) return buffer, nil } @@ -288,14 +289,14 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ account := rawAccount.(*ShadowsocksAccount) ivLen := account.Cipher.IVSize() - iv := payload.Value[:ivLen] + iv := payload.BytesTo(ivLen) payload.SliceFrom(ivLen) stream, err := account.Cipher.NewDecodingStream(account.Key, iv) if err != nil { return nil, nil, errors.Base(err).Message("Shadowsocks|UDP: Failed to initialize decoding stream.") } - stream.XORKeyStream(payload.Value, payload.Value) + stream.XORKeyStream(payload.Bytes(), payload.Bytes()) authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv)) request := &protocol.RequestHeader{ @@ -304,8 +305,8 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ Command: protocol.RequestCommandUDP, } - addrType := (payload.Value[0] & 0x0F) - if (payload.Value[0] & 0x10) == 0x10 { + addrType := (payload.Byte(0) & 0x0F) + if (payload.Byte(0) & 0x10) == 0x10 { request.Option |= RequestOptionOneTimeAuth } @@ -319,9 +320,10 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ if request.Option.Has(RequestOptionOneTimeAuth) { payloadLen := payload.Len() - AuthSize - authBytes := payload.Value[payloadLen:] + authBytes := payload.BytesFrom(payloadLen) - actualAuth := authenticator.Authenticate(nil, payload.Value[0:payloadLen]) + actualAuth := make([]byte, AuthSize) + authenticator.Authenticate(payload.BytesTo(payloadLen))(actualAuth) if !bytes.Equal(actualAuth, authBytes) { return nil, nil, errors.New("Shadowsocks|UDP: Invalid OTA.") } @@ -333,20 +335,20 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ switch addrType { case AddrTypeIPv4: - request.Address = v2net.IPAddress(payload.Value[:4]) + request.Address = v2net.IPAddress(payload.BytesTo(4)) payload.SliceFrom(4) case AddrTypeIPv6: - request.Address = v2net.IPAddress(payload.Value[:16]) + request.Address = v2net.IPAddress(payload.BytesTo(16)) payload.SliceFrom(16) case AddrTypeDomain: - domainLength := int(payload.Value[0]) - request.Address = v2net.DomainAddress(string(payload.Value[1 : 1+domainLength])) + domainLength := int(payload.Byte(0)) + request.Address = v2net.DomainAddress(string(payload.BytesRange(1, 1+domainLength))) payload.SliceFrom(1 + domainLength) default: return nil, nil, errors.New("Shadowsocks|UDP: Unknown address type: ", addrType) } - request.Port = v2net.PortFromBytes(payload.Value[:2]) + request.Port = v2net.PortFromBytes(payload.BytesTo(2)) payload.SliceFrom(2) return request, payload, nil @@ -359,12 +361,11 @@ type UDPReader struct { func (v *UDPReader) Read() (*alloc.Buffer, error) { buffer := alloc.NewSmallBuffer() - nBytes, err := v.Reader.Read(buffer.Value) + _, err := buffer.FillFrom(v.Reader) if err != nil { buffer.Release() return nil, err } - buffer.Slice(0, nBytes) _, payload, err := DecodeUDPPacket(v.User, buffer) if err != nil { buffer.Release() @@ -386,7 +387,7 @@ func (v *UDPWriter) Write(buffer *alloc.Buffer) error { if err != nil { return err } - _, err = v.Writer.Write(payload.Value) + _, err = v.Writer.Write(payload.Bytes()) payload.Release() return err } diff --git a/proxy/shadowsocks/protocol_test.go b/proxy/shadowsocks/protocol_test.go index 3a3f7c3d1..ede36ceed 100644 --- a/proxy/shadowsocks/protocol_test.go +++ b/proxy/shadowsocks/protocol_test.go @@ -7,6 +7,7 @@ import ( "v2ray.com/core/common/loader" v2net "v2ray.com/core/common/net" "v2ray.com/core/common/protocol" + "v2ray.com/core/common/serial" . "v2ray.com/core/proxy/shadowsocks" "v2ray.com/core/testing/assert" ) @@ -29,7 +30,8 @@ func TestUDPEncoding(t *testing.T) { }, } - data := alloc.NewLocalBuffer(256).Clear().AppendString("test string") + data := alloc.NewLocalBuffer(256) + data.AppendFunc(serial.WriteString("test string")) encodedData, err := EncodeUDPPacket(request, data) assert.Error(err).IsNil() @@ -58,8 +60,9 @@ func TestTCPRequest(t *testing.T) { }, } - data := alloc.NewLocalBuffer(256).Clear().AppendString("test string") - cache := alloc.NewBuffer().Clear() + data := alloc.NewLocalBuffer(256) + data.AppendFunc(serial.WriteString("test string")) + cache := alloc.NewBuffer() writer, err := WriteTCPRequest(request, cache) assert.Error(err).IsNil() @@ -85,7 +88,7 @@ func TestUDPReaderWriter(t *testing.T) { CipherType: CipherType_CHACHA20_IEFT, }), } - cache := alloc.NewBuffer().Clear() + cache := alloc.NewBuffer() writer := &UDPWriter{ Writer: cache, Request: &protocol.RequestHeader{ @@ -102,14 +105,18 @@ func TestUDPReaderWriter(t *testing.T) { User: user, } - err := writer.Write(alloc.NewBuffer().Clear().AppendString("test payload")) + b := alloc.NewBuffer() + b.AppendFunc(serial.WriteString("test payload")) + err := writer.Write(b) assert.Error(err).IsNil() payload, err := reader.Read() assert.Error(err).IsNil() assert.String(payload.String()).Equals("test payload") - err = writer.Write(alloc.NewBuffer().Clear().AppendString("test payload 2")) + b = alloc.NewBuffer() + b.AppendFunc(serial.WriteString("test payload 2")) + err = writer.Write(b) assert.Error(err).IsNil() payload, err = reader.Read() diff --git a/proxy/socks/protocol/socks.go b/proxy/socks/protocol/socks.go index b5af88463..755dd79f7 100644 --- a/proxy/socks/protocol/socks.go +++ b/proxy/socks/protocol/socks.go @@ -122,28 +122,30 @@ func ReadUserPassRequest(reader io.Reader) (request Socks5UserPassRequest, err e buffer := alloc.NewLocalBuffer(512) defer buffer.Release() - _, err = reader.Read(buffer.Value[0:2]) + _, err = buffer.FillFullFrom(reader, 2) if err != nil { return } - request.version = buffer.Value[0] - nUsername := buffer.Value[1] - nBytes, err := reader.Read(buffer.Value[:nUsername]) - if err != nil { - return - } - request.username = string(buffer.Value[:nBytes]) + request.version = buffer.Byte(0) + nUsername := int(buffer.Byte(1)) - _, err = reader.Read(buffer.Value[0:1]) + buffer.Clear() + _, err = buffer.FillFullFrom(reader, nUsername) if err != nil { return } - nPassword := buffer.Value[0] - nBytes, err = reader.Read(buffer.Value[:nPassword]) + request.username = string(buffer.Bytes()) + + _, err = buffer.FillFullFrom(reader, 1) if err != nil { return } - request.password = string(buffer.Value[:nBytes]) + nPassword := int(buffer.Byte(0)) + _, err = buffer.FillFullFrom(reader, nPassword) + if err != nil { + return + } + request.password = string(buffer.Bytes()) return } @@ -185,7 +187,7 @@ type Socks5Request struct { } func ReadRequest(reader io.Reader) (request *Socks5Request, err error) { - buffer := alloc.NewLocalBuffer(512).Clear() + buffer := alloc.NewLocalBuffer(512) defer buffer.Release() _, err = buffer.FillFullFrom(reader, 4) @@ -194,10 +196,10 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) { } request = &Socks5Request{ - Version: buffer.Value[0], - Command: buffer.Value[1], + Version: buffer.Byte(0), + Command: buffer.Byte(1), // buffer[2] is a reserved field - AddrType: buffer.Value[3], + AddrType: buffer.Byte(3), } switch request.AddrType { case AddrTypeIPv4: diff --git a/proxy/socks/protocol/socks4_test.go b/proxy/socks/protocol/socks4_test.go index af43ada80..4a2852528 100644 --- a/proxy/socks/protocol/socks4_test.go +++ b/proxy/socks/protocol/socks4_test.go @@ -31,9 +31,9 @@ func TestSocks4AuthenticationResponseToBytes(t *testing.T) { response := NewSocks4AuthenticationResponse(byte(0x10), 443, []byte{1, 2, 3, 4}) - buffer := alloc.NewLocalBuffer(2048).Clear() + buffer := alloc.NewLocalBuffer(2048) defer buffer.Release() response.Write(buffer) - assert.Bytes(buffer.Value).Equals([]byte{0x00, 0x10, 0x01, 0xBB, 0x01, 0x02, 0x03, 0x04}) + assert.Bytes(buffer.Bytes()).Equals([]byte{0x00, 0x10, 0x01, 0xBB, 0x01, 0x02, 0x03, 0x04}) } diff --git a/proxy/socks/protocol/socks_test.go b/proxy/socks/protocol/socks_test.go index 5f800bd5d..f204c92be 100644 --- a/proxy/socks/protocol/socks_test.go +++ b/proxy/socks/protocol/socks_test.go @@ -29,7 +29,8 @@ func TestHasAuthenticationMethod(t *testing.T) { func TestAuthenticationRequestRead(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewBuffer().Clear().AppendBytes( + buffer := alloc.NewBuffer() + buffer.AppendBytes( 0x05, // version 0x01, // nMethods 0x02, // methods @@ -83,7 +84,7 @@ func TestResponseWrite(t *testing.T) { [16]byte{}, v2net.Port(53), } - buffer := alloc.NewLocalBuffer(2048).Clear() + buffer := alloc.NewLocalBuffer(2048) defer buffer.Release() response.Write(buffer) @@ -104,7 +105,7 @@ func TestSetIPv6(t *testing.T) { response := NewSocks5Response() response.SetIPv6([]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}) - buffer := alloc.NewLocalBuffer(2048).Clear() + buffer := alloc.NewLocalBuffer(2048) defer buffer.Release() response.Write(buffer) assert.Bytes(buffer.Bytes()).Equals([]byte{ @@ -117,7 +118,7 @@ func TestSetDomain(t *testing.T) { response := NewSocks5Response() response.SetDomain("v2ray.com") - buffer := alloc.NewLocalBuffer(2048).Clear() + buffer := alloc.NewLocalBuffer(2048) defer buffer.Release() response.Write(buffer) assert.Bytes(buffer.Bytes()).Equals([]byte{ @@ -127,7 +128,7 @@ func TestSetDomain(t *testing.T) { func TestEmptyAuthRequest(t *testing.T) { assert := assert.On(t) - _, _, err := ReadAuthentication(alloc.NewBuffer().Clear()) + _, _, err := ReadAuthentication(alloc.NewBuffer()) assert.Error(err).Equals(io.EOF) } @@ -141,14 +142,16 @@ func TestSingleByteAuthRequest(t *testing.T) { func TestZeroAuthenticationMethod(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewBuffer().Clear().AppendBytes(5, 0) + buffer := alloc.NewBuffer() + buffer.AppendBytes(5, 0) _, _, err := ReadAuthentication(buffer) assert.Error(err).Equals(proxy.ErrInvalidAuthentication) } func TestWrongProtocolVersion(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewBuffer().Clear().AppendBytes(6, 1, 0) + buffer := alloc.NewBuffer() + buffer.AppendBytes(6, 1, 0) _, _, err := ReadAuthentication(buffer) assert.Error(err).Equals(proxy.ErrInvalidProtocolVersion) } @@ -156,14 +159,16 @@ func TestWrongProtocolVersion(t *testing.T) { func TestEmptyRequest(t *testing.T) { assert := assert.On(t) - _, err := ReadRequest(alloc.NewBuffer().Clear()) + _, err := ReadRequest(alloc.NewBuffer()) assert.Error(err).Equals(io.EOF) } func TestIPv6Request(t *testing.T) { assert := assert.On(t) - request, err := ReadRequest(alloc.NewBuffer().Clear().AppendBytes(5, 1, 0, 4, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 0, 8)) + b := alloc.NewBuffer() + b.AppendBytes(5, 1, 0, 4, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 0, 8) + request, err := ReadRequest(b) assert.Error(err).IsNil() assert.Byte(request.Command).Equals(1) assert.Bytes(request.IPv6[:]).Equals([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}) diff --git a/proxy/socks/protocol/udp.go b/proxy/socks/protocol/udp.go index 51d9fdd61..84f6a2b7d 100644 --- a/proxy/socks/protocol/udp.go +++ b/proxy/socks/protocol/udp.go @@ -5,6 +5,7 @@ import ( "v2ray.com/core/common/alloc" "v2ray.com/core/common/errors" v2net "v2ray.com/core/common/net" + "v2ray.com/core/common/serial" ) var ( @@ -26,14 +27,17 @@ func (request *Socks5UDPRequest) Write(buffer *alloc.Buffer) { buffer.AppendBytes(0, 0, request.Fragment) switch request.Address.Family() { case v2net.AddressFamilyIPv4: - buffer.AppendBytes(AddrTypeIPv4).Append(request.Address.IP()) + buffer.AppendBytes(AddrTypeIPv4) + buffer.Append(request.Address.IP()) case v2net.AddressFamilyIPv6: - buffer.AppendBytes(AddrTypeIPv6).Append(request.Address.IP()) + buffer.AppendBytes(AddrTypeIPv6) + buffer.Append(request.Address.IP()) case v2net.AddressFamilyDomain: - buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain()))).Append([]byte(request.Address.Domain())) + buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain()))) + buffer.Append([]byte(request.Address.Domain())) } - buffer.AppendUint16(request.Port.Value()) - buffer.Append(request.Data.Value) + buffer.AppendFunc(serial.WriteUint16(request.Port.Value())) + buffer.Append(request.Data.Bytes()) } func ReadUDPRequest(packet []byte) (*Socks5UDPRequest, error) { @@ -79,7 +83,9 @@ func ReadUDPRequest(packet []byte) (*Socks5UDPRequest, error) { } if len(packet) > dataBegin { - request.Data = alloc.NewBuffer().Clear().Append(packet[dataBegin:]) + b := alloc.NewSmallBuffer() + b.Append(packet[dataBegin:]) + request.Data = b } return request, nil diff --git a/proxy/socks/protocol/udp_test.go b/proxy/socks/protocol/udp_test.go index 60f0d2c2f..576a29aaf 100644 --- a/proxy/socks/protocol/udp_test.go +++ b/proxy/socks/protocol/udp_test.go @@ -32,5 +32,5 @@ func TestDomainAddressRequest(t *testing.T) { assert.Byte(request.Fragment).Equals(1) assert.Address(request.Address).EqualsString("v2ray.com") assert.Port(request.Port).Equals(v2net.Port(80)) - assert.Bytes(request.Data.Value).Equals([]byte("Actual payload")) + assert.String(request.Data.String()).Equals("Actual payload") } diff --git a/proxy/socks/server_udp.go b/proxy/socks/server_udp.go index ef83e9db2..56efd7c15 100644 --- a/proxy/socks/server_udp.go +++ b/proxy/socks/server_udp.go @@ -26,7 +26,7 @@ func (v *Server) listenUDP() error { func (v *Server) handleUDPPayload(payload *alloc.Buffer, session *proxy.SessionInfo) { source := session.Source log.Info("Socks: Client UDP connection from ", source) - request, err := protocol.ReadUDPRequest(payload.Value) + request, err := protocol.ReadUDPRequest(payload.Bytes()) payload.Release() if err != nil { @@ -55,7 +55,7 @@ func (v *Server) handleUDPPayload(payload *alloc.Buffer, session *proxy.SessionI } log.Info("Socks: Writing back UDP response with ", payload.Len(), " bytes to ", destination) - udpMessage := alloc.NewLocalBuffer(2048).Clear() + udpMessage := alloc.NewLocalBuffer(2048) response.Write(udpMessage) v.udpMutex.RLock() @@ -63,7 +63,7 @@ func (v *Server) handleUDPPayload(payload *alloc.Buffer, session *proxy.SessionI v.udpMutex.RUnlock() return } - nBytes, err := v.udpHub.WriteTo(udpMessage.Value, destination) + nBytes, err := v.udpHub.WriteTo(udpMessage.Bytes(), destination) v.udpMutex.RUnlock() udpMessage.Release() response.Data.Release() diff --git a/proxy/vmess/encoding/commands.go b/proxy/vmess/encoding/commands.go index f26abc3c2..4898dd262 100644 --- a/proxy/vmess/encoding/commands.go +++ b/proxy/vmess/encoding/commands.go @@ -32,7 +32,7 @@ func MarshalCommand(command interface{}, writer io.Writer) error { return ErrUnknownCommand } - buffer := alloc.NewLocalBuffer(512).Clear() + buffer := alloc.NewLocalBuffer(512) defer buffer.Release() err := factory.Marshal(command, buffer) diff --git a/proxy/vmess/encoding/commands_test.go b/proxy/vmess/encoding/commands_test.go index 8382dc04b..94dc363cd 100644 --- a/proxy/vmess/encoding/commands_test.go +++ b/proxy/vmess/encoding/commands_test.go @@ -21,7 +21,7 @@ func TestSwitchAccount(t *testing.T) { ValidMin: 16, } - buffer := alloc.NewBuffer().Clear() + buffer := alloc.NewBuffer() err := MarshalCommand(sa, buffer) assert.Error(err).IsNil() diff --git a/proxy/vmess/encoding/encoding_test.go b/proxy/vmess/encoding/encoding_test.go index 9a79bf269..cd6ef93ab 100644 --- a/proxy/vmess/encoding/encoding_test.go +++ b/proxy/vmess/encoding/encoding_test.go @@ -35,7 +35,7 @@ func TestRequestSerialization(t *testing.T) { Port: v2net.Port(443), } - buffer := alloc.NewBuffer().Clear() + buffer := alloc.NewBuffer() client := NewClientSession(protocol.DefaultIDHash) client.EncodeRequestHeader(expectedRequest, buffer) diff --git a/proxy/vmess/inbound/inbound.go b/proxy/vmess/inbound/inbound.go index 88da6601b..5908a2696 100644 --- a/proxy/vmess/inbound/inbound.go +++ b/proxy/vmess/inbound/inbound.go @@ -236,7 +236,7 @@ func (v *VMessInboundHandler) HandleConnection(connection internet.Connection) { } output.Release() if request.Option.Has(protocol.RequestOptionChunkStream) { - if err := v2writer.Write(alloc.NewLocalBuffer(32).Clear()); err != nil { + if err := v2writer.Write(alloc.NewLocalBuffer(32)); err != nil { connection.SetReusable(false) } } diff --git a/proxy/vmess/io/io_test.go b/proxy/vmess/io/io_test.go index 2d1cd2a5c..16dffec0c 100644 --- a/proxy/vmess/io/io_test.go +++ b/proxy/vmess/io/io_test.go @@ -9,6 +9,7 @@ import ( "v2ray.com/core/common/alloc" "v2ray.com/core/common/errors" v2io "v2ray.com/core/common/io" + "v2ray.com/core/common/serial" . "v2ray.com/core/proxy/vmess/io" "v2ray.com/core/testing/assert" ) @@ -16,7 +17,7 @@ import ( func TestAuthenticate(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewBuffer().Clear() + buffer := alloc.NewBuffer() buffer.AppendBytes(1, 2, 3, 4) Authenticate(buffer) assert.Bytes(buffer.Bytes()).Equals([]byte{0, 8, 87, 52, 168, 125, 1, 2, 3, 4}) @@ -32,8 +33,10 @@ func TestSingleIO(t *testing.T) { content := bytes.NewBuffer(make([]byte, 0, 1024*1024)) writer := NewAuthChunkWriter(v2io.NewAdaptiveWriter(content)) - writer.Write(alloc.NewBuffer().Clear().AppendString("abcd")) - writer.Write(alloc.NewBuffer().Clear()) + b := alloc.NewBuffer() + b.AppendFunc(serial.WriteString("abcd")) + writer.Write(b) + writer.Write(alloc.NewBuffer()) writer.Release() reader := NewAuthChunkReader(content) @@ -56,13 +59,17 @@ func TestLargeIO(t *testing.T) { if chunkSize+writeSize > len(content) { chunkSize = len(content) - writeSize } - writer.Write(alloc.NewBuffer().Clear().Append(content[writeSize : writeSize+chunkSize])) + b := alloc.NewBuffer() + b.Append(content[writeSize : writeSize+chunkSize]) + writer.Write(b) + b.Release() + writeSize += chunkSize if writeSize == len(content) { break } } - writer.Write(alloc.NewBuffer().Clear()) + writer.Write(alloc.NewBuffer()) writer.Release() actualContent := make([]byte, 0, len(content)) diff --git a/proxy/vmess/io/reader.go b/proxy/vmess/io/reader.go index c71ddd721..f1b5e9c4c 100644 --- a/proxy/vmess/io/reader.go +++ b/proxy/vmess/io/reader.go @@ -50,7 +50,7 @@ func (v *AuthChunkReader) Read() (*alloc.Buffer, error) { buffer = v.last v.last = nil } else { - buffer = alloc.NewBuffer().Clear() + buffer = alloc.NewBuffer() } if v.chunkLength == -1 { @@ -96,7 +96,7 @@ func (v *AuthChunkReader) Read() (*alloc.Buffer, error) { } leftLength := buffer.Len() - v.chunkLength if leftLength > 0 { - v.last = alloc.NewBuffer().Clear() + v.last = alloc.NewBuffer() v.last.Append(buffer.BytesFrom(v.chunkLength)) buffer.Slice(0, v.chunkLength) } diff --git a/proxy/vmess/io/writer.go b/proxy/vmess/io/writer.go index 5798cd966..5c6ce82af 100644 --- a/proxy/vmess/io/writer.go +++ b/proxy/vmess/io/writer.go @@ -5,6 +5,7 @@ import ( "v2ray.com/core/common/alloc" v2io "v2ray.com/core/common/io" + "v2ray.com/core/common/serial" ) type AuthChunkWriter struct { @@ -30,7 +31,7 @@ func (v *AuthChunkWriter) Release() { func Authenticate(buffer *alloc.Buffer) { fnvHash := fnv.New32a() fnvHash.Write(buffer.Bytes()) - buffer.PrependHash(fnvHash) + buffer.PrependFunc(4, serial.WriteHash(fnvHash)) - buffer.PrependUint16(uint16(buffer.Len())) + buffer.PrependFunc(2, serial.WriteUint16(uint16(buffer.Len()))) } diff --git a/proxy/vmess/outbound/outbound.go b/proxy/vmess/outbound/outbound.go index 890ba4a25..923d31df8 100644 --- a/proxy/vmess/outbound/outbound.go +++ b/proxy/vmess/outbound/outbound.go @@ -111,7 +111,7 @@ func (v *VMessOutboundHandler) handleRequest(session *encoding.ClientSession, co } if request.Option.Has(protocol.RequestOptionChunkStream) { - err := streamWriter.Write(alloc.NewLocalBuffer(32).Clear()) + err := streamWriter.Write(alloc.NewLocalBuffer(32)) if err != nil { conn.SetReusable(false) } diff --git a/testing/scenarios/shadowsocks_test.go b/testing/scenarios/shadowsocks_test.go index 5c6c2d68d..490e522fd 100644 --- a/testing/scenarios/shadowsocks_test.go +++ b/testing/scenarios/shadowsocks_test.go @@ -42,7 +42,7 @@ func TestShadowsocksTCP(t *testing.T) { //conn.CloseWrite() - response := alloc.NewBuffer().Clear() + response := alloc.NewBuffer() finished := false expectedResponse := "Processed: " + payload for { @@ -56,7 +56,7 @@ func TestShadowsocksTCP(t *testing.T) { break } if response.Len() > len(expectedResponse) { - fmt.Printf("Unexpected response: %v\n", response.Value) + fmt.Printf("Unexpected response: %v\n", response.Bytes()) break } } diff --git a/transport/internet/authenticators/http/http.go b/transport/internet/authenticators/http/http.go index 306aeb0f8..c8c42e35b 100644 --- a/transport/internet/authenticators/http/http.go +++ b/transport/internet/authenticators/http/http.go @@ -5,10 +5,11 @@ import ( "io" "net" "net/http" + "strings" "time" - "v2ray.com/core/common/alloc" "v2ray.com/core/common/loader" + "v2ray.com/core/common/serial" "v2ray.com/core/transport/internet" ) @@ -17,6 +18,10 @@ const ( ENDING = CRLF + CRLF ) +var ( + writeCRLF = serial.WriteString(CRLF) +) + type Reader interface { Read(io.Reader) (*alloc.Buffer, error) } @@ -41,18 +46,18 @@ type HeaderReader struct { } func (*HeaderReader) Read(reader io.Reader) (*alloc.Buffer, error) { - buffer := alloc.NewSmallBuffer().Clear() + buffer := alloc.NewSmallBuffer() for { _, err := buffer.FillFrom(reader) if err != nil { return nil, err } - if n := bytes.Index(buffer.Value, []byte(ENDING)); n != -1 { + if n := bytes.Index(buffer.Bytes(), []byte(ENDING)); n != -1 { buffer.SliceFrom(n + len(ENDING)) break } if buffer.Len() >= len(ENDING) { - copy(buffer.Value, buffer.Value[buffer.Len()-len(ENDING):]) + copy(buffer.Bytes(), buffer.BytesFrom(buffer.Len()-len(ENDING))) buffer.Slice(0, len(ENDING)) } } @@ -77,7 +82,7 @@ func (v *HeaderWriter) Write(writer io.Writer) error { if v.header == nil { return nil } - _, err := writer.Write(v.header.Value) + _, err := writer.Write(v.header.Bytes()) v.header.Release() v.header = nil return err @@ -138,33 +143,39 @@ type HttpAuthenticator struct { } func (v HttpAuthenticator) GetClientWriter() *HeaderWriter { - header := alloc.NewSmallBuffer().Clear() + header := alloc.NewSmallBuffer() config := v.config.Request - header.AppendString(config.Method.GetValue()).AppendString(" ").AppendString(config.PickUri()).AppendString(" ").AppendString(config.GetFullVersion()).AppendString(CRLF) + header.AppendFunc(serial.WriteString(strings.Join([]string{config.Method.GetValue(), config.PickUri(), config.GetFullVersion()}, " "))) + header.AppendFunc(writeCRLF) headers := config.PickHeaders() for _, h := range headers { - header.AppendString(h).AppendString(CRLF) + header.AppendFunc(serial.WriteString(h)) + header.AppendFunc(writeCRLF) } - header.AppendString(CRLF) + header.AppendFunc(writeCRLF) return &HeaderWriter{ header: header, } } func (v HttpAuthenticator) GetServerWriter() *HeaderWriter { - header := alloc.NewSmallBuffer().Clear() + header := alloc.NewSmallBuffer() config := v.config.Response - header.AppendString(config.GetFullVersion()).AppendString(" ").AppendString(config.Status.GetCode()).AppendString(" ").AppendString(config.Status.GetReason()).AppendString(CRLF) + header.AppendFunc(serial.WriteString(strings.Join([]string{config.GetFullVersion(), config.Status.GetCode(), config.Status.GetReason()}, " "))) + header.AppendFunc(writeCRLF) headers := config.PickHeaders() for _, h := range headers { - header.AppendString(h).AppendString(CRLF) + header.AppendFunc(serial.WriteString(h)) + header.AppendFunc(writeCRLF) } if !config.HasHeader("Date") { - header.AppendString("Date: ").AppendString(time.Now().Format(http.TimeFormat)).AppendString(CRLF) + header.AppendFunc(serial.WriteString("Date: ")) + header.AppendFunc(serial.WriteString(time.Now().Format(http.TimeFormat))) + header.AppendFunc(writeCRLF) } - header.AppendString(CRLF) + header.AppendFunc(writeCRLF) return &HeaderWriter{ header: header, } diff --git a/transport/internet/authenticators/http/http_test.go b/transport/internet/authenticators/http/http_test.go index ea077b0f7..86e6ff667 100644 --- a/transport/internet/authenticators/http/http_test.go +++ b/transport/internet/authenticators/http/http_test.go @@ -4,6 +4,7 @@ import ( "testing" "v2ray.com/core/common/alloc" + "v2ray.com/core/common/serial" "v2ray.com/core/testing/assert" . "v2ray.com/core/transport/internet/authenticators/http" ) @@ -11,15 +12,17 @@ import ( func TestReaderWriter(t *testing.T) { assert := assert.On(t) - cache := alloc.NewBuffer().Clear() - writer := NewHeaderWriter(alloc.NewLocalBuffer(256).Clear().AppendString("abcd" + ENDING)) + cache := alloc.NewBuffer() + b := alloc.NewLocalBuffer(256) + b.AppendFunc(serial.WriteString("abcd" + ENDING)) + writer := NewHeaderWriter(b) writer.Write(cache) cache.Write([]byte{'e', 'f', 'g'}) reader := &HeaderReader{} buffer, err := reader.Read(cache) assert.Error(err).IsNil() - assert.Bytes(buffer.Value).Equals([]byte{'e', 'f', 'g'}) + assert.Bytes(buffer.Bytes()).Equals([]byte{'e', 'f', 'g'}) } func TestRequestHeader(t *testing.T) { @@ -38,7 +41,7 @@ func TestRequestHeader(t *testing.T) { }, }).(HttpAuthenticator) - cache := alloc.NewBuffer().Clear() + cache := alloc.NewBuffer() err := auth.GetClientWriter().Write(cache) assert.Error(err).IsNil() diff --git a/transport/internet/authenticators/srtp/srtp.go b/transport/internet/authenticators/srtp/srtp.go index 6b463498e..6caa8822c 100644 --- a/transport/internet/authenticators/srtp/srtp.go +++ b/transport/internet/authenticators/srtp/srtp.go @@ -5,6 +5,7 @@ import ( "v2ray.com/core/common/alloc" "v2ray.com/core/common/loader" + "v2ray.com/core/common/serial" "v2ray.com/core/transport/internet" ) @@ -24,8 +25,8 @@ func (v *SRTP) Open(payload *alloc.Buffer) bool { func (v *SRTP) Seal(payload *alloc.Buffer) { v.number++ - payload.PrependUint16(v.number) - payload.PrependUint16(v.header) + payload.PrependFunc(2, serial.WriteUint16(v.number)) + payload.PrependFunc(2, serial.WriteUint16(v.header)) } type SRTPFactory struct { diff --git a/transport/internet/authenticators/srtp/srtp_test.go b/transport/internet/authenticators/srtp/srtp_test.go index 77eeac28c..db4ef742e 100644 --- a/transport/internet/authenticators/srtp/srtp_test.go +++ b/transport/internet/authenticators/srtp/srtp_test.go @@ -12,7 +12,9 @@ func TestSRTPOpenSeal(t *testing.T) { assert := assert.On(t) content := []byte{'a', 'b', 'c', 'd', 'e', 'f', 'g'} - payload := alloc.NewLocalBuffer(2048).Clear().Append(content) + payload := alloc.NewLocalBuffer(2048) + payload.Append(content) + srtp := SRTP{} srtp.Seal(payload) assert.Int(payload.Len()).GreaterThan(len(content)) diff --git a/transport/internet/authenticators/utp/utp.go b/transport/internet/authenticators/utp/utp.go index 088749283..dcb1d4df5 100644 --- a/transport/internet/authenticators/utp/utp.go +++ b/transport/internet/authenticators/utp/utp.go @@ -5,6 +5,7 @@ import ( "v2ray.com/core/common/alloc" "v2ray.com/core/common/loader" + "v2ray.com/core/common/serial" "v2ray.com/core/transport/internet" ) @@ -24,7 +25,7 @@ func (v *UTP) Open(payload *alloc.Buffer) bool { } func (v *UTP) Seal(payload *alloc.Buffer) { - payload.PrependUint16(v.connectionId) + payload.PrependFunc(2, serial.WriteUint16(v.connectionId)) payload.PrependBytes(v.header, v.extension) } diff --git a/transport/internet/authenticators/utp/utp_test.go b/transport/internet/authenticators/utp/utp_test.go index c7bedde8a..a5d8db1c0 100644 --- a/transport/internet/authenticators/utp/utp_test.go +++ b/transport/internet/authenticators/utp/utp_test.go @@ -12,7 +12,9 @@ func TestUTPOpenSeal(t *testing.T) { assert := assert.On(t) content := []byte{'a', 'b', 'c', 'd', 'e', 'f', 'g'} - payload := alloc.NewLocalBuffer(2048).Clear().Append(content) + payload := alloc.NewLocalBuffer(2048) + payload.Append(content) + utp := UTP{} utp.Seal(payload) assert.Int(payload.Len()).GreaterThan(len(content)) diff --git a/transport/internet/kcp/crypt.go b/transport/internet/kcp/crypt.go index d0877bb9b..3c3450733 100644 --- a/transport/internet/kcp/crypt.go +++ b/transport/internet/kcp/crypt.go @@ -19,10 +19,10 @@ func (v *SimpleAuthenticator) Overhead() int { } func (v *SimpleAuthenticator) Seal(buffer *alloc.Buffer) { - buffer.PrependUint16(uint16(buffer.Len())) + buffer.PrependFunc(2, serial.WriteUint16(uint16(buffer.Len()))) fnvHash := fnv.New32a() fnvHash.Write(buffer.Bytes()) - buffer.PrependHash(fnvHash) + buffer.PrependFunc(4, serial.WriteHash(fnvHash)) len := buffer.Len() xtra := 4 - len%4 @@ -47,12 +47,12 @@ func (v *SimpleAuthenticator) Open(buffer *alloc.Buffer) bool { } fnvHash := fnv.New32a() - fnvHash.Write(buffer.Value[4:]) - if serial.BytesToUint32(buffer.Value[:4]) != fnvHash.Sum32() { + fnvHash.Write(buffer.BytesFrom(4)) + if serial.BytesToUint32(buffer.BytesTo(4)) != fnvHash.Sum32() { return false } - length := serial.BytesToUint16(buffer.Value[4:6]) + length := serial.BytesToUint16(buffer.BytesRange(4, 6)) if buffer.Len()-6 != int(length) { return false } diff --git a/transport/internet/kcp/crypt_test.go b/transport/internet/kcp/crypt_test.go index 28028bab6..4326455d3 100644 --- a/transport/internet/kcp/crypt_test.go +++ b/transport/internet/kcp/crypt_test.go @@ -12,7 +12,7 @@ import ( func TestSimpleAuthenticator(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewLocalBuffer(512).Clear() + buffer := alloc.NewLocalBuffer(512) buffer.AppendBytes('a', 'b', 'c', 'd', 'e', 'f', 'g') auth := NewSimpleAuthenticator() @@ -25,18 +25,18 @@ func TestSimpleAuthenticator(t *testing.T) { func TestSimpleAuthenticator2(t *testing.T) { assert := assert.On(t) - buffer := alloc.NewLocalBuffer(512).Clear() + buffer := alloc.NewLocalBuffer(512) buffer.AppendBytes('1', '2') auth := NewSimpleAuthenticator() auth.Seal(buffer) assert.Bool(auth.Open(buffer)).IsTrue() - assert.Bytes(buffer.Value).Equals([]byte{'1', '2'}) + assert.Bytes(buffer.Bytes()).Equals([]byte{'1', '2'}) } func BenchmarkSimpleAuthenticator(b *testing.B) { - buffer := alloc.NewLocalBuffer(2048).Clear() + buffer := alloc.NewLocalBuffer(2048) buffer.FillFullFrom(rand.Reader, 1024) auth := NewSimpleAuthenticator() diff --git a/transport/internet/kcp/dialer.go b/transport/internet/kcp/dialer.go index 3f6b0122a..c7f5d686a 100644 --- a/transport/internet/kcp/dialer.go +++ b/transport/internet/kcp/dialer.go @@ -51,15 +51,15 @@ func (o *ClientConnection) Run() { defer payload.Release() for { - nBytes, err := o.Conn.Read(payload.Value) + payload.Clear() + _, err := payload.FillFrom(o.Conn) if err != nil { payload.Release() return } - payload.Slice(0, nBytes) o.Lock() if o.input != nil && o.auth.Open(payload) { - o.input(payload.Value) + o.input(payload.Bytes()) } o.Unlock() payload.Reset() diff --git a/transport/internet/kcp/listener.go b/transport/internet/kcp/listener.go index b52ea43c8..b8c0ba662 100644 --- a/transport/internet/kcp/listener.go +++ b/transport/internet/kcp/listener.go @@ -53,7 +53,7 @@ func (o *ServerConnection) Input(b *alloc.Buffer) { defer b.Release() if o.auth.Open(b) { - o.input(b.Value) + o.input(b.Bytes()) } } @@ -153,8 +153,8 @@ func (v *Listener) OnReceive(payload *alloc.Buffer, session *proxy.SessionInfo) if payload.Len() < 4 { return } - conv := serial.BytesToUint16(payload.Value) - cmd := Command(payload.Value[2]) + conv := serial.BytesToUint16(payload.BytesTo(2)) + cmd := Command(payload.Byte(2)) id := ConnectionId{ Remote: src.Address, Port: src.Port, @@ -196,7 +196,7 @@ func (v *Listener) OnReceive(payload *alloc.Buffer, session *proxy.SessionInfo) } v.sessions[id] = conn } - conn.Input(payload.Value) + conn.Input(payload.Bytes()) } func (v *Listener) Remove(id ConnectionId) { diff --git a/transport/internet/kcp/output.go b/transport/internet/kcp/output.go index 1d52a12f7..d79dd40e7 100644 --- a/transport/internet/kcp/output.go +++ b/transport/internet/kcp/output.go @@ -37,10 +37,10 @@ func (v *BufferedSegmentWriter) Write(seg Segment) { } if v.buffer == nil { - v.buffer = alloc.NewSmallBuffer().Clear() + v.buffer = alloc.NewSmallBuffer() } - v.buffer.Value = seg.Bytes(v.buffer.Value) + v.buffer.AppendFunc(seg.Bytes()) } func (v *BufferedSegmentWriter) FlushWithoutLock() { @@ -69,7 +69,7 @@ func (v *AuthenticationWriter) Write(payload *alloc.Buffer) error { defer payload.Release() v.Authenticator.Seal(payload) - _, err := v.Writer.Write(payload.Value) + _, err := v.Writer.Write(payload.Bytes()) return err } diff --git a/transport/internet/kcp/receiving.go b/transport/internet/kcp/receiving.go index 66e589dd7..19032b002 100644 --- a/transport/internet/kcp/receiving.go +++ b/transport/internet/kcp/receiving.go @@ -196,7 +196,7 @@ func (v *ReceivingWorker) Read(b []byte) int { total := 0 if v.leftOver != nil { - nBytes := copy(b, v.leftOver.Value) + nBytes := copy(b, v.leftOver.Bytes()) if nBytes < v.leftOver.Len() { v.leftOver.SliceFrom(nBytes) return nBytes @@ -214,7 +214,7 @@ func (v *ReceivingWorker) Read(b []byte) int { v.window.Advance() v.nextNumber++ - nBytes := copy(b[total:], seg.Data.Value) + nBytes := copy(b[total:], seg.Data.Bytes()) total += nBytes if nBytes < seg.Data.Len() { seg.Data.SliceFrom(nBytes) diff --git a/transport/internet/kcp/segment.go b/transport/internet/kcp/segment.go index c0f74f146..62132d440 100644 --- a/transport/internet/kcp/segment.go +++ b/transport/internet/kcp/segment.go @@ -25,7 +25,7 @@ type Segment interface { common.Releasable Conversation() uint16 ByteSize() int - Bytes([]byte) []byte + Bytes() alloc.BytesWriter } const ( @@ -56,18 +56,21 @@ func (v *DataSegment) SetData(b []byte) { if v.Data == nil { v.Data = alloc.NewSmallBuffer() } - v.Data.Clear().Append(b) + v.Data.Clear() + v.Data.Append(b) } -func (v *DataSegment) Bytes(b []byte) []byte { - b = serial.Uint16ToBytes(v.Conv, b) - b = append(b, byte(CommandData), byte(v.Option)) - b = serial.Uint32ToBytes(v.Timestamp, b) - b = serial.Uint32ToBytes(v.Number, b) - b = serial.Uint32ToBytes(v.SendingNext, b) - b = serial.Uint16ToBytes(uint16(v.Data.Len()), b) - b = append(b, v.Data.Value...) - return b +func (v *DataSegment) Bytes() alloc.BytesWriter { + return func(b []byte) int { + b = serial.Uint16ToBytes(v.Conv, b[:0]) + b = append(b, byte(CommandData), byte(v.Option)) + b = serial.Uint32ToBytes(v.Timestamp, b) + b = serial.Uint32ToBytes(v.Number, b) + b = serial.Uint32ToBytes(v.SendingNext, b) + b = serial.Uint16ToBytes(uint16(v.Data.Len()), b) + b = append(b, v.Data.Bytes()...) + return v.ByteSize() + } } func (v *DataSegment) ByteSize() int { @@ -120,17 +123,19 @@ func (v *AckSegment) ByteSize() int { return 2 + 1 + 1 + 4 + 4 + 4 + 1 + int(v.Count)*4 } -func (v *AckSegment) Bytes(b []byte) []byte { - b = serial.Uint16ToBytes(v.Conv, b) - b = append(b, byte(CommandACK), byte(v.Option)) - b = serial.Uint32ToBytes(v.ReceivingWindow, b) - b = serial.Uint32ToBytes(v.ReceivingNext, b) - b = serial.Uint32ToBytes(v.Timestamp, b) - b = append(b, v.Count) - for i := byte(0); i < v.Count; i++ { - b = serial.Uint32ToBytes(v.NumberList[i], b) +func (v *AckSegment) Bytes() alloc.BytesWriter { + return func(b []byte) int { + b = serial.Uint16ToBytes(v.Conv, b[:0]) + b = append(b, byte(CommandACK), byte(v.Option)) + b = serial.Uint32ToBytes(v.ReceivingWindow, b) + b = serial.Uint32ToBytes(v.ReceivingNext, b) + b = serial.Uint32ToBytes(v.Timestamp, b) + b = append(b, v.Count) + for i := byte(0); i < v.Count; i++ { + b = serial.Uint32ToBytes(v.NumberList[i], b) + } + return v.ByteSize() } - return b } func (v *AckSegment) Release() { @@ -158,13 +163,15 @@ func (v *CmdOnlySegment) ByteSize() int { return 2 + 1 + 1 + 4 + 4 + 4 } -func (v *CmdOnlySegment) Bytes(b []byte) []byte { - b = serial.Uint16ToBytes(v.Conv, b) - b = append(b, byte(v.Command), byte(v.Option)) - b = serial.Uint32ToBytes(v.SendingNext, b) - b = serial.Uint32ToBytes(v.ReceivinNext, b) - b = serial.Uint32ToBytes(v.PeerRTO, b) - return b +func (v *CmdOnlySegment) Bytes() alloc.BytesWriter { + return func(b []byte) int { + b = serial.Uint16ToBytes(v.Conv, b[:0]) + b = append(b, byte(v.Command), byte(v.Option)) + b = serial.Uint32ToBytes(v.SendingNext, b) + b = serial.Uint32ToBytes(v.ReceivinNext, b) + b = serial.Uint32ToBytes(v.PeerRTO, b) + return v.ByteSize() + } } func (v *CmdOnlySegment) Release() { diff --git a/transport/internet/kcp/segment_test.go b/transport/internet/kcp/segment_test.go index 4fda45ff4..30627c76a 100644 --- a/transport/internet/kcp/segment_test.go +++ b/transport/internet/kcp/segment_test.go @@ -19,16 +19,19 @@ func TestBadSegment(t *testing.T) { func TestDataSegment(t *testing.T) { assert := assert.On(t) + b := alloc.NewLocalBuffer(512) + b.Append([]byte{'a', 'b', 'c', 'd'}) seg := &DataSegment{ Conv: 1, Timestamp: 3, Number: 4, SendingNext: 5, - Data: alloc.NewLocalBuffer(512).Clear().Append([]byte{'a', 'b', 'c', 'd'}), + Data: b, } nBytes := seg.ByteSize() - bytes := seg.Bytes(nil) + bytes := make([]byte, nBytes) + seg.Bytes()(bytes) assert.Int(len(bytes)).Equals(nBytes) @@ -54,7 +57,8 @@ func TestACKSegment(t *testing.T) { } nBytes := seg.ByteSize() - bytes := seg.Bytes(nil) + bytes := make([]byte, nBytes) + seg.Bytes()(bytes) assert.Int(len(bytes)).Equals(nBytes) @@ -83,7 +87,8 @@ func TestCmdSegment(t *testing.T) { } nBytes := seg.ByteSize() - bytes := seg.Bytes(nil) + bytes := make([]byte, nBytes) + seg.Bytes()(bytes) assert.Int(len(bytes)).Equals(nBytes) diff --git a/transport/internet/udp/hub.go b/transport/internet/udp/hub.go index 9b125e6a2..bc2babdc3 100644 --- a/transport/internet/udp/hub.go +++ b/transport/internet/udp/hub.go @@ -136,13 +136,22 @@ func (v *UDPHub) start() { oobBytes := make([]byte, 256) for v.Running() { buffer := alloc.NewSmallBuffer() - nBytes, noob, _, addr, err := ReadUDPMsg(v.conn, buffer.Bytes(), oobBytes) + var noob int + var addr *net.UDPAddr + var err error + buffer.AppendFunc(func(b []byte) int { + n, nb, _, a, e := ReadUDPMsg(v.conn, b, oobBytes) + noob = nb + addr = a + err = e + return n + }) + if err != nil { log.Info("UDP|Hub: Failed to read UDP msg: ", err) buffer.Release() continue } - buffer.Slice(0, nBytes) session := new(proxy.SessionInfo) session.Source = v2net.UDPDestination(v2net.IPAddress(addr.IP), v2net.Port(addr.Port))