mihomo/component/vmess/chunk.go

package vmess

import (
	"encoding/binary"
	"errors"
	"io"
	"sync"
)

const (
	lenSize   = 2
	chunkSize = 1 << 14   // 2 ** 14 == 16 * 1024
	maxSize   = 17 * 1024 // 2 + chunkSize + aead.Overhead()
)

var bufPool = sync.Pool{New: func() interface{} { return make([]byte, maxSize) }}

type chunkReader struct {
	io.Reader
	buf     []byte
	sizeBuf []byte
	offset  int
}

func newChunkReader(reader io.Reader) *chunkReader {
	return &chunkReader{Reader: reader, sizeBuf: make([]byte, lenSize)}
}

func newChunkWriter(writer io.WriteCloser) *chunkWriter {
	return &chunkWriter{Writer: writer}
}

func (cr *chunkReader) Read(b []byte) (int, error) {
	if cr.buf != nil {
		n := copy(b, cr.buf[cr.offset:])
		cr.offset += n
		if cr.offset == len(cr.buf) {
			bufPool.Put(cr.buf[:cap(cr.buf)])
			cr.buf = nil
		}
		return n, nil
	}

	_, err := io.ReadFull(cr.Reader, cr.sizeBuf)
	if err != nil {
		return 0, err
	}

	size := int(binary.BigEndian.Uint16(cr.sizeBuf))
	if size > maxSize {
		return 0, errors.New("Buffer is larger than standard")
	}

	if len(b) >= size {
		_, err := io.ReadFull(cr.Reader, b[:size])
		if err != nil {
			return 0, err
		}

		return size, nil
	}

	buf := bufPool.Get().([]byte)
	_, err = io.ReadFull(cr.Reader, buf[:size])
	if err != nil {
		bufPool.Put(buf[:cap(buf)])
		return 0, err
	}
	n := copy(b, cr.buf[:])
	cr.offset = n
	cr.buf = buf[:size]
	return n, nil
}

type chunkWriter struct {
	io.Writer
}

func (cw *chunkWriter) Write(b []byte) (n int, err error) {
	buf := bufPool.Get().([]byte)
	defer bufPool.Put(buf[:cap(buf)])
	length := len(b)
	for {
		if length == 0 {
			break
		}
		readLen := chunkSize
		if length < chunkSize {
			readLen = length
		}
		payloadBuf := buf[lenSize : lenSize+chunkSize]
		copy(payloadBuf, b[n:n+readLen])

		binary.BigEndian.PutUint16(buf[:lenSize], uint16(readLen))
		_, err = cw.Writer.Write(buf[:lenSize+readLen])
		if err != nil {
			break
		}
		n += readLen
		length -= readLen
	}
	return
}