diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index 9cd31dc3..43ea2b85 100644 --- a/common/buf/multi_buffer.go +++ b/common/buf/multi_buffer.go @@ -35,6 +35,7 @@ func MergeMulti(dest MultiBuffer, src MultiBuffer) (MultiBuffer, MultiBuffer) { return dest, src[:0] } +// MergeBytes merges the given bytes into MultiBuffer and return the new address of the merged MultiBuffer. func MergeBytes(dest MultiBuffer, src []byte) MultiBuffer { n := len(dest) if n > 0 && !(dest)[n-1].IsFull() { @@ -94,6 +95,8 @@ func ReadFrom(reader io.Reader) (MultiBuffer, error) { } } +// SplitBytes splits the given amount of bytes from the beginning of the MultiBuffer. +// It returns the new address of MultiBuffer leftover, and number of bytes written into the input byte slice. func SplitBytes(mb MultiBuffer, b []byte) (MultiBuffer, int) { totalBytes := 0 @@ -112,6 +115,18 @@ func SplitBytes(mb MultiBuffer, b []byte) (MultiBuffer, int) { return mb, totalBytes } +// SplitFirst splits the first Buffer from the beginning of the MultiBuffer. +func SplitFirst(mb MultiBuffer) (MultiBuffer, *Buffer) { + if len(mb) == 0 { + return mb, nil + } + + b := mb[0] + mb[0] = nil + mb = mb[1:] + return mb, b +} + // Len returns the total number of bytes in the MultiBuffer. func (mb MultiBuffer) Len() int32 { if mb == nil { @@ -167,21 +182,12 @@ func (mb *MultiBuffer) SliceBySize(size int32) MultiBuffer { return slice } -// SplitFirst splits out the first Buffer in this MultiBuffer. -func (mb *MultiBuffer) SplitFirst() *Buffer { - if len(*mb) == 0 { - return nil - } - b := (*mb)[0] - (*mb)[0] = nil - *mb = (*mb)[1:] - return b -} - +// MultiBufferContainer is a ReadWriteCloser wrapper over MultiBuffer. type MultiBufferContainer struct { MultiBuffer } +// Read implements io.Reader. func (c *MultiBufferContainer) Read(b []byte) (int, error) { if c.MultiBuffer.IsEmpty() { return 0, io.EOF @@ -192,23 +198,27 @@ func (c *MultiBufferContainer) Read(b []byte) (int, error) { return nBytes, nil } +// ReadMultiBuffer implements Reader. func (c *MultiBufferContainer) ReadMultiBuffer() (MultiBuffer, error) { mb := c.MultiBuffer c.MultiBuffer = nil return mb, nil } +// Write implements io.Writer. func (c *MultiBufferContainer) Write(b []byte) (int, error) { c.MultiBuffer = MergeBytes(c.MultiBuffer, b) return len(b), nil } +// WriteMultiBuffer implement Writer. func (c *MultiBufferContainer) WriteMultiBuffer(b MultiBuffer) error { mb, _ := MergeMulti(c.MultiBuffer, b) c.MultiBuffer = mb return nil } +// Close implement io.Closer. func (c *MultiBufferContainer) Close() error { c.MultiBuffer = ReleaseMulti(c.MultiBuffer) return nil diff --git a/common/crypto/auth_test.go b/common/crypto/auth_test.go index b79a15be..d06f5b01 100644 --- a/common/crypto/auth_test.go +++ b/common/crypto/auth_test.go @@ -116,10 +116,10 @@ func TestAuthenticationReaderWriterPacket(t *testing.T) { mb, err := reader.ReadMultiBuffer() assert(err, IsNil) - b1 := mb.SplitFirst() + mb, b1 := buf.SplitFirst(mb) assert(b1.String(), Equals, "abcd") - b2 := mb.SplitFirst() + mb, b2 := buf.SplitFirst(mb) assert(b2.String(), Equals, "efgh") assert(mb.IsEmpty(), IsTrue) diff --git a/common/mux/writer.go b/common/mux/writer.go index 2dc3904c..5259bedb 100644 --- a/common/mux/writer.go +++ b/common/mux/writer.go @@ -96,7 +96,9 @@ func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error { if w.transferType == protocol.TransferTypeStream { chunk = mb.SliceBySize(8 * 1024) } else { - chunk = buf.MultiBuffer{mb.SplitFirst()} + mb2, b := buf.SplitFirst(mb) + mb = mb2 + chunk = buf.MultiBuffer{b} } if err := w.writeData(chunk); err != nil { return err