1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-23 02:26:49 -05:00
v2fly/common/buf/writer.go

271 lines
5.0 KiB
Go
Raw Normal View History

2016-12-09 07:17:34 -05:00
package buf
2017-11-03 20:33:35 -04:00
import (
"io"
2018-08-17 05:51:59 -04:00
"net"
2018-08-31 09:02:35 -04:00
"sync"
2017-11-03 20:33:35 -04:00
2018-04-11 10:15:29 -04:00
"v2ray.com/core/common"
2017-11-03 20:33:35 -04:00
"v2ray.com/core/common/errors"
)
2016-12-09 07:17:34 -05:00
// BufferToBytesWriter is a Writer that writes alloc.Buffer into underlying writer.
type BufferToBytesWriter struct {
2017-11-09 16:33:15 -05:00
io.Writer
2018-08-17 05:51:59 -04:00
cache [][]byte
2017-11-09 16:33:15 -05:00
}
// WriteMultiBuffer implements Writer. This method takes ownership of the given buffer.
func (w *BufferToBytesWriter) WriteMultiBuffer(mb MultiBuffer) error {
2018-11-17 16:45:07 -05:00
defer ReleaseMulti(mb)
2017-04-16 16:30:29 -04:00
2018-07-27 17:35:02 -04:00
size := mb.Len()
if size == 0 {
return nil
}
2018-08-16 06:40:55 -04:00
if len(mb) == 1 {
return WriteAllBytes(w.Writer, mb[0].Bytes())
}
2018-08-17 05:51:59 -04:00
bs := w.cache
for _, b := range mb {
bs = append(bs, b.Bytes())
}
w.cache = bs
defer func() {
for idx := range w.cache {
w.cache[idx] = nil
}
w.cache = w.cache[:0]
}()
2018-08-17 05:51:59 -04:00
nb := net.Buffers(bs)
2018-07-27 17:35:02 -04:00
for size > 0 {
2018-08-17 05:51:59 -04:00
n, err := nb.WriteTo(w.Writer)
2018-07-27 17:35:02 -04:00
if err != nil {
return err
}
size -= int32(n)
}
return nil
2016-12-09 07:17:34 -05:00
}
2017-11-21 10:55:22 -05:00
// ReadFrom implements io.ReaderFrom.
2017-11-09 16:33:15 -05:00
func (w *BufferToBytesWriter) ReadFrom(reader io.Reader) (int64, error) {
var sc SizeCounter
err := Copy(NewReader(reader), w, CountSize(&sc))
return sc.Size, err
2017-04-19 15:27:21 -04:00
}
2017-11-21 10:55:22 -05:00
// BufferedWriter is a Writer with internal buffer.
2017-11-09 16:33:15 -05:00
type BufferedWriter struct {
2018-08-31 09:02:35 -04:00
sync.Mutex
2017-11-25 15:58:47 -05:00
writer Writer
buffer *Buffer
buffered bool
2017-04-19 15:27:21 -04:00
}
2017-11-21 10:55:22 -05:00
// NewBufferedWriter creates a new BufferedWriter.
2017-11-09 16:33:15 -05:00
func NewBufferedWriter(writer Writer) *BufferedWriter {
2017-11-25 15:58:47 -05:00
return &BufferedWriter{
2017-11-09 16:33:15 -05:00
writer: writer,
buffer: New(),
buffered: true,
}
2017-04-20 05:00:15 -04:00
}
2018-02-13 05:15:04 -05:00
// WriteByte implements io.ByteWriter.
func (w *BufferedWriter) WriteByte(c byte) error {
2018-04-11 10:15:29 -04:00
return common.Error2(w.Write([]byte{c}))
}
2017-11-21 10:55:22 -05:00
// Write implements io.Writer.
2017-11-09 16:33:15 -05:00
func (w *BufferedWriter) Write(b []byte) (int, error) {
2018-09-01 15:19:33 -04:00
if len(b) == 0 {
return 0, nil
}
2018-08-31 09:02:35 -04:00
w.Lock()
defer w.Unlock()
2017-11-25 15:58:47 -05:00
if !w.buffered {
if writer, ok := w.writer.(io.Writer); ok {
return writer.Write(b)
}
2017-11-09 16:33:15 -05:00
}
2017-04-20 05:00:15 -04:00
2017-11-09 16:33:15 -05:00
totalBytes := 0
for len(b) > 0 {
2017-11-17 04:18:12 -05:00
if w.buffer == nil {
w.buffer = New()
}
2017-11-09 16:33:15 -05:00
nBytes, err := w.buffer.Write(b)
totalBytes += nBytes
if err != nil {
return totalBytes, err
}
2017-11-15 17:38:52 -05:00
if !w.buffered || w.buffer.IsFull() {
2018-08-31 09:02:35 -04:00
if err := w.flushInternal(); err != nil {
2017-11-09 16:33:15 -05:00
return totalBytes, err
}
2017-04-20 05:00:15 -04:00
}
2017-11-09 16:33:15 -05:00
b = b[nBytes:]
2017-04-20 05:00:15 -04:00
}
2017-11-15 17:38:52 -05:00
2017-11-09 16:33:15 -05:00
return totalBytes, nil
2017-04-20 05:00:15 -04:00
}
2017-11-21 10:55:22 -05:00
// WriteMultiBuffer implements Writer. It takes ownership of the given MultiBuffer.
2017-11-09 16:33:15 -05:00
func (w *BufferedWriter) WriteMultiBuffer(b MultiBuffer) error {
2018-09-01 15:19:33 -04:00
if b.IsEmpty() {
return nil
}
2018-08-31 09:02:35 -04:00
w.Lock()
defer w.Unlock()
2017-11-09 16:33:15 -05:00
if !w.buffered {
return w.writer.WriteMultiBuffer(b)
}
2017-04-21 08:51:09 -04:00
2018-11-18 13:36:36 -05:00
reader := MultiBufferContainer{
MultiBuffer: b,
}
defer reader.Close()
2017-04-21 08:51:09 -04:00
2018-11-18 13:36:36 -05:00
for !reader.MultiBuffer.IsEmpty() {
2017-12-15 20:02:48 -05:00
if w.buffer == nil {
w.buffer = New()
}
2018-11-18 13:36:36 -05:00
common.Must2(w.buffer.ReadFrom(&reader))
2017-11-09 16:33:15 -05:00
if w.buffer.IsFull() {
2018-08-31 09:02:35 -04:00
if err := w.flushInternal(); err != nil {
2017-11-09 16:33:15 -05:00
return err
}
}
2017-04-21 08:51:09 -04:00
}
return nil
}
2017-11-21 10:55:22 -05:00
// Flush flushes buffered content into underlying writer.
2017-11-09 16:33:15 -05:00
func (w *BufferedWriter) Flush() error {
2018-08-31 09:02:35 -04:00
w.Lock()
defer w.Unlock()
return w.flushInternal()
}
func (w *BufferedWriter) flushInternal() error {
2018-07-27 17:35:02 -04:00
if w.buffer.IsEmpty() {
return nil
}
2018-05-31 05:55:11 -04:00
2018-07-27 17:35:02 -04:00
b := w.buffer
w.buffer = nil
if writer, ok := w.writer.(io.Writer); ok {
2018-08-29 10:42:03 -04:00
err := WriteAllBytes(writer, b.Bytes())
b.Release()
return err
2017-11-09 16:33:15 -05:00
}
2018-07-27 17:35:02 -04:00
2018-11-16 05:08:12 -05:00
return w.writer.WriteMultiBuffer(MultiBuffer{b})
2016-12-09 07:17:34 -05:00
}
2018-02-13 05:15:04 -05:00
// SetBuffered sets whether the internal buffer is used. If set to false, Flush() will be called to clear the buffer.
2017-11-09 16:33:15 -05:00
func (w *BufferedWriter) SetBuffered(f bool) error {
2018-08-31 09:02:35 -04:00
w.Lock()
defer w.Unlock()
2017-11-09 16:33:15 -05:00
w.buffered = f
if !f {
2018-08-31 09:02:35 -04:00
return w.flushInternal()
2016-12-09 07:17:34 -05:00
}
2017-11-09 16:33:15 -05:00
return nil
2017-04-15 15:07:23 -04:00
}
2016-12-09 07:17:34 -05:00
2017-11-15 05:23:25 -05:00
// ReadFrom implements io.ReaderFrom.
2017-11-09 16:33:15 -05:00
func (w *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
if err := w.SetBuffered(false); err != nil {
return 0, err
2017-11-09 16:33:15 -05:00
}
var sc SizeCounter
2017-11-09 16:33:15 -05:00
err := Copy(NewReader(reader), w, CountSize(&sc))
return sc.Size, err
2016-12-09 07:17:34 -05:00
}
2017-04-17 16:35:20 -04:00
2018-04-16 08:57:13 -04:00
// Close implements io.Closable.
func (w *BufferedWriter) Close() error {
if err := w.Flush(); err != nil {
return err
}
return common.Close(w.writer)
}
2018-07-31 07:43:27 -04:00
// SequentialWriter is a Writer that writes MultiBuffer sequentially into the underlying io.Writer.
type SequentialWriter struct {
io.Writer
2017-11-09 16:33:15 -05:00
}
2018-07-31 07:43:27 -04:00
// WriteMultiBuffer implements Writer.
func (w *SequentialWriter) WriteMultiBuffer(mb MultiBuffer) error {
2018-11-17 16:45:07 -05:00
defer ReleaseMulti(mb)
2017-11-09 16:33:15 -05:00
for _, b := range mb {
2018-09-01 15:19:33 -04:00
if b.IsEmpty() {
continue
}
if err := WriteAllBytes(w.Writer, b.Bytes()); err != nil {
2017-11-09 16:33:15 -05:00
return err
2017-04-17 16:35:20 -04:00
}
}
2017-11-09 16:33:15 -05:00
return nil
2017-04-17 16:35:20 -04:00
}
2017-04-27 16:20:29 -04:00
2017-12-03 17:11:29 -05:00
type noOpWriter byte
2017-04-27 16:20:29 -04:00
2017-11-09 16:33:15 -05:00
func (noOpWriter) WriteMultiBuffer(b MultiBuffer) error {
2018-11-17 16:45:07 -05:00
ReleaseMulti(b)
2017-04-27 16:20:29 -04:00
return nil
}
2017-11-25 15:35:26 -05:00
func (noOpWriter) Write(b []byte) (int, error) {
2017-11-03 20:33:35 -04:00
return len(b), nil
}
2017-11-25 15:35:26 -05:00
func (noOpWriter) ReadFrom(reader io.Reader) (int64, error) {
2017-11-03 20:33:35 -04:00
b := New()
defer b.Release()
totalBytes := int64(0)
for {
2018-11-02 10:01:33 -04:00
b.Clear()
_, err := b.ReadFrom(reader)
2017-11-03 20:33:35 -04:00
totalBytes += int64(b.Len())
if err != nil {
if errors.Cause(err) == io.EOF {
return totalBytes, nil
}
return totalBytes, err
}
}
}
2017-04-27 16:20:29 -04:00
var (
2017-11-08 12:14:44 -05:00
// Discard is a Writer that swallows all contents written in.
2017-12-03 17:11:29 -05:00
Discard Writer = noOpWriter(0)
2017-11-08 12:14:44 -05:00
// DiscardBytes is an io.Writer that swallows all contents written in.
2017-12-03 17:11:29 -05:00
DiscardBytes io.Writer = noOpWriter(0)
2017-04-27 16:20:29 -04:00
)