diff --git a/app/proxyman/mux/frame.go b/app/proxyman/mux/frame.go index ba43d19b5..f6c64bc9e 100644 --- a/app/proxyman/mux/frame.go +++ b/app/proxyman/mux/frame.go @@ -59,21 +59,21 @@ type FrameMetadata struct { func (f FrameMetadata) WriteTo(b *buf.Buffer) error { lenBytes := b.Bytes() - common.Must2(b.AppendBytes(0x00, 0x00)) + common.Must2(b.WriteBytes(0x00, 0x00)) len0 := b.Len() if err := b.AppendSupplier(serial.WriteUint16(f.SessionID)); err != nil { return err } - common.Must2(b.AppendBytes(byte(f.SessionStatus), byte(f.Option))) + common.Must2(b.WriteBytes(byte(f.SessionStatus), byte(f.Option))) if f.SessionStatus == SessionStatusNew { switch f.Target.Network { case net.Network_TCP: - common.Must2(b.AppendBytes(byte(TargetNetworkTCP))) + common.Must2(b.WriteBytes(byte(TargetNetworkTCP))) case net.Network_UDP: - common.Must2(b.AppendBytes(byte(TargetNetworkUDP))) + common.Must2(b.WriteBytes(byte(TargetNetworkUDP))) } if err := addrParser.WriteAddressPort(b, f.Target.Address, f.Target.Port); err != nil { diff --git a/app/proxyman/mux/reader.go b/app/proxyman/mux/reader.go index 5cc843ad7..111bbee2b 100644 --- a/app/proxyman/mux/reader.go +++ b/app/proxyman/mux/reader.go @@ -17,7 +17,7 @@ func ReadMetadata(reader io.Reader) (*FrameMetadata, error) { return nil, newError("invalid metalen ", metaLen).AtError() } - b := buf.NewSize(int32(metaLen)) + b := buf.New() defer b.Release() if err := b.Reset(buf.ReadFullFrom(reader, int32(metaLen))); err != nil { diff --git a/common/buf/buffer.go b/common/buf/buffer.go index 3f7287033..bc1743c68 100644 --- a/common/buf/buffer.go +++ b/common/buf/buffer.go @@ -35,11 +35,6 @@ func (b *Buffer) Clear() { b.end = 0 } -// AppendBytes appends one or more bytes to the end of the buffer. -func (b *Buffer) AppendBytes(bytes ...byte) (int, error) { - return b.Write(bytes) -} - // AppendSupplier appends the content of a BytesWriter to the buffer. func (b *Buffer) AppendSupplier(writer Supplier) error { nBytes, err := writer(b.v[b.end:]) @@ -145,6 +140,11 @@ func (b *Buffer) Write(data []byte) (int, error) { return nBytes, nil } +// WriteBytes appends one or more bytes to the end of the buffer. +func (b *Buffer) WriteBytes(bytes ...byte) (int, error) { + return b.Write(bytes) +} + // Read implements io.Reader.Read(). func (b *Buffer) Read(data []byte) (int, error) { if b.Len() == 0 { diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index a93aa1196..3816d6264 100644 --- a/common/buf/multi_buffer.go +++ b/common/buf/multi_buffer.go @@ -21,26 +21,6 @@ func ReadAllToMultiBuffer(reader io.Reader) (MultiBuffer, error) { return mb, nil } -// ReadSizeToMultiBuffer reads specific number of bytes from reader into a MultiBuffer. -func ReadSizeToMultiBuffer(reader io.Reader, size int32) (MultiBuffer, error) { - mb := NewMultiBufferCap(32) - - for size > 0 { - bSize := size - if bSize > Size { - bSize = Size - } - b := NewSize(bSize) - if err := b.Reset(ReadFullFrom(reader, bSize)); err != nil { - mb.Release() - return nil, err - } - size -= bSize - mb.Append(b) - } - return mb, nil -} - // ReadAllToBytes reads all content from the reader into a byte array, until EOF. func ReadAllToBytes(reader io.Reader) ([]byte, error) { mb, err := ReadAllToMultiBuffer(reader) @@ -100,7 +80,7 @@ func (mb *MultiBuffer) ReadFrom(reader io.Reader) (int64, error) { for { b := New() - err := b.Reset(ReadFrom(reader)) + err := b.Reset(ReadFullFrom(reader, Size)) if b.IsEmpty() { b.Release() } else { @@ -108,7 +88,7 @@ func (mb *MultiBuffer) ReadFrom(reader io.Reader) (int64, error) { } totalBytes += int64(b.Len()) if err != nil { - if errors.Cause(err) == io.EOF { + if errors.Cause(err) == io.EOF || errors.Cause(err) == io.ErrUnexpectedEOF { return totalBytes, nil } return totalBytes, err @@ -178,6 +158,9 @@ func (mb *MultiBuffer) Write(b []byte) (int, error) { // WriteMultiBuffer implements Writer. func (mb *MultiBuffer) WriteMultiBuffer(b MultiBuffer) error { *mb = append(*mb, b...) + for i := range b { + b[i] = nil + } return nil } diff --git a/common/buf/multi_buffer_test.go b/common/buf/multi_buffer_test.go index 8a55e5562..a3982704f 100644 --- a/common/buf/multi_buffer_test.go +++ b/common/buf/multi_buffer_test.go @@ -14,10 +14,10 @@ func TestMultiBufferRead(t *testing.T) { assert := With(t) b1 := New() - b1.AppendBytes('a', 'b') + b1.WriteBytes('a', 'b') b2 := New() - b2.AppendBytes('c', 'd') + b2.WriteBytes('c', 'd') mb := NewMultiBufferValue(b1, b2) bs := make([]byte, 32) @@ -32,7 +32,7 @@ func TestMultiBufferAppend(t *testing.T) { var mb MultiBuffer b := New() - b.AppendBytes('a', 'b') + b.WriteBytes('a', 'b') mb.Append(b) assert(mb.Len(), Equals, int32(2)) } diff --git a/common/buf/reader_test.go b/common/buf/reader_test.go index 59029e96f..3a4d8d5d5 100644 --- a/common/buf/reader_test.go +++ b/common/buf/reader_test.go @@ -37,9 +37,9 @@ func TestBytesReaderWriteTo(t *testing.T) { pReader, pWriter := pipe.New(pipe.WithSizeLimit(1024)) reader := &BufferedReader{Reader: pReader} b1 := New() - b1.AppendBytes('a', 'b', 'c') + b1.WriteBytes('a', 'b', 'c') b2 := New() - b2.AppendBytes('e', 'f', 'g') + b2.WriteBytes('e', 'f', 'g') assert(pWriter.WriteMultiBuffer(NewMultiBufferValue(b1, b2)), IsNil) pWriter.Close() @@ -64,9 +64,9 @@ func TestBytesReaderMultiBuffer(t *testing.T) { pReader, pWriter := pipe.New(pipe.WithSizeLimit(1024)) reader := &BufferedReader{Reader: pReader} b1 := New() - b1.AppendBytes('a', 'b', 'c') + b1.WriteBytes('a', 'b', 'c') b2 := New() - b2.AppendBytes('e', 'f', 'g') + b2.WriteBytes('e', 'f', 'g') assert(pWriter.WriteMultiBuffer(NewMultiBufferValue(b1, b2)), IsNil) pWriter.Close() diff --git a/common/crypto/chunk_test.go b/common/crypto/chunk_test.go index 38a3ec438..c8df17cab 100644 --- a/common/crypto/chunk_test.go +++ b/common/crypto/chunk_test.go @@ -19,11 +19,11 @@ func TestChunkStreamIO(t *testing.T) { reader := NewChunkStreamReader(PlainChunkSizeParser{}, cache) b := buf.New() - b.AppendBytes('a', 'b', 'c', 'd') + b.WriteBytes('a', 'b', 'c', 'd') common.Must(writer.WriteMultiBuffer(buf.NewMultiBufferValue(b))) b = buf.New() - b.AppendBytes('e', 'f', 'g') + b.WriteBytes('e', 'f', 'g') common.Must(writer.WriteMultiBuffer(buf.NewMultiBufferValue(b))) common.Must(writer.WriteMultiBuffer(buf.MultiBuffer{})) diff --git a/proxy/http/server.go b/proxy/http/server.go index de96e9298..5f6a05f76 100755 --- a/proxy/http/server.go +++ b/proxy/http/server.go @@ -182,7 +182,8 @@ func (s *Server) handleConnect(ctx context.Context, request *http.Request, reade } if reader.Buffered() > 0 { - payload, err := buf.ReadSizeToMultiBuffer(reader, int32(reader.Buffered())) + var payload buf.MultiBuffer + _, err := payload.ReadFrom(&io.LimitedReader{R: reader, N: int64(reader.Buffered())}) if err != nil { return err } diff --git a/proxy/shadowsocks/ota_test.go b/proxy/shadowsocks/ota_test.go index 3281e09c0..fc8c44228 100644 --- a/proxy/shadowsocks/ota_test.go +++ b/proxy/shadowsocks/ota_test.go @@ -12,7 +12,7 @@ func TestNormalChunkReading(t *testing.T) { assert := With(t) buffer := buf.New() - buffer.AppendBytes( + buffer.WriteBytes( 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}))) diff --git a/proxy/socks/protocol.go b/proxy/socks/protocol.go index c7d1fe5b2..a3c378503 100644 --- a/proxy/socks/protocol.go +++ b/proxy/socks/protocol.go @@ -241,7 +241,7 @@ func writeSocks5Response(writer io.Writer, errCode byte, address net.Address, po buffer := buf.New() defer buffer.Release() - common.Must2(buffer.AppendBytes(socks5Version, errCode, 0x00 /* reserved */)) + common.Must2(buffer.WriteBytes(socks5Version, errCode, 0x00 /* reserved */)) if err := addrParser.WriteAddressPort(buffer, address, port); err != nil { return err } @@ -253,7 +253,7 @@ func writeSocks4Response(writer io.Writer, errCode byte, address net.Address, po buffer := buf.New() defer buffer.Release() - common.Must2(buffer.AppendBytes(0x00, errCode)) + common.Must2(buffer.WriteBytes(0x00, errCode)) common.Must(buffer.AppendSupplier(serial.WriteUint16(port.Value()))) common.Must2(buffer.Write(address.IP())) return buf.WriteAllBytes(writer, buffer.Bytes()) @@ -286,7 +286,7 @@ func DecodeUDPPacket(packet *buf.Buffer) (*protocol.RequestHeader, error) { func EncodeUDPPacket(request *protocol.RequestHeader, data []byte) (*buf.Buffer, error) { b := buf.New() - common.Must2(b.AppendBytes(0, 0, 0 /* Fragment */)) + common.Must2(b.WriteBytes(0, 0, 0 /* Fragment */)) if err := addrParser.WriteAddressPort(b, request.Address, request.Port); err != nil { b.Release() return nil, err @@ -348,7 +348,7 @@ func ClientHandshake(request *protocol.RequestHeader, reader io.Reader, writer i b := buf.New() defer b.Release() - common.Must2(b.AppendBytes(socks5Version, 0x01, authByte)) + common.Must2(b.WriteBytes(socks5Version, 0x01, authByte)) if authByte == authPassword { rawAccount, err := request.User.GetTypedAccount() if err != nil { @@ -356,9 +356,9 @@ func ClientHandshake(request *protocol.RequestHeader, reader io.Reader, writer i } account := rawAccount.(*Account) - common.Must2(b.AppendBytes(0x01, byte(len(account.Username)))) + common.Must2(b.WriteBytes(0x01, byte(len(account.Username)))) common.Must2(b.Write([]byte(account.Username))) - common.Must2(b.AppendBytes(byte(len(account.Password)))) + common.Must2(b.WriteBytes(byte(len(account.Password)))) common.Must2(b.Write([]byte(account.Password))) } @@ -392,7 +392,7 @@ func ClientHandshake(request *protocol.RequestHeader, reader io.Reader, writer i if request.Command == protocol.RequestCommandUDP { command = byte(cmdUDPPort) } - common.Must2(b.AppendBytes(socks5Version, command, 0x00 /* reserved */)) + common.Must2(b.WriteBytes(socks5Version, command, 0x00 /* reserved */)) if err := addrParser.WriteAddressPort(b, request.Address, request.Port); err != nil { return nil, err } diff --git a/proxy/vmess/encoding/client.go b/proxy/vmess/encoding/client.go index 5b0fc5fc5..cf4475000 100644 --- a/proxy/vmess/encoding/client.go +++ b/proxy/vmess/encoding/client.go @@ -69,14 +69,14 @@ func (c *ClientSession) EncodeRequestHeader(header *protocol.RequestHeader, writ buffer := buf.New() defer buffer.Release() - buffer.AppendBytes(Version) - buffer.Write(c.requestBodyIV[:]) - buffer.Write(c.requestBodyKey[:]) - buffer.AppendBytes(c.responseHeader, byte(header.Option)) + common.Must2(buffer.WriteBytes(Version)) + common.Must2(buffer.Write(c.requestBodyIV[:])) + common.Must2(buffer.Write(c.requestBodyKey[:])) + common.Must2(buffer.WriteBytes(c.responseHeader, byte(header.Option))) padingLen := dice.Roll(16) security := byte(padingLen<<4) | byte(header.Security) - buffer.AppendBytes(security, byte(0), byte(header.Command)) + common.Must2(buffer.WriteBytes(security, byte(0), byte(header.Command))) if header.Command != protocol.RequestCommandMux { if err := addrParser.WriteAddressPort(buffer, header.Address, header.Port); err != nil { diff --git a/transport/internet/kcp/segment_test.go b/transport/internet/kcp/segment_test.go index 7dfe60552..9d2a60e61 100644 --- a/transport/internet/kcp/segment_test.go +++ b/transport/internet/kcp/segment_test.go @@ -50,7 +50,7 @@ func Test1ByteDataSegment(t *testing.T) { Number: 4, SendingNext: 5, } - seg.Data().AppendBytes('a') + seg.Data().WriteBytes('a') nBytes := seg.ByteSize() bytes := make([]byte, nBytes) diff --git a/transport/internet/udp/dispatcher_test.go b/transport/internet/udp/dispatcher_test.go index d642e36ea..60e7d6900 100644 --- a/transport/internet/udp/dispatcher_test.go +++ b/transport/internet/udp/dispatcher_test.go @@ -58,7 +58,7 @@ func TestSameDestinationDispatching(t *testing.T) { dest := net.UDPDestination(net.LocalHostIP, 53) b := buf.New() - b.AppendBytes('a', 'b', 'c', 'd') + b.WriteBytes('a', 'b', 'c', 'd') var msgCount uint32 dispatcher := NewDispatcher(td, func(ctx context.Context, payload *buf.Buffer) { diff --git a/transport/pipe/pipe_test.go b/transport/pipe/pipe_test.go index 3a27c27de..964043342 100644 --- a/transport/pipe/pipe_test.go +++ b/transport/pipe/pipe_test.go @@ -103,7 +103,7 @@ func TestPipeWriteMultiThread(t *testing.T) { wg.Add(1) go func() { b := buf.New() - b.AppendBytes('a', 'b', 'c', 'd') + b.WriteBytes('a', 'b', 'c', 'd') pWriter.WriteMultiBuffer(buf.NewMultiBufferValue(b)) wg.Done() }()