github.com/xxf098/lite-proxy@v0.15.1-0.20230422081941-12c69f323218/transport/vmess/chunk.go (about)

     1  package vmess
     2  
     3  import (
     4  	"encoding/binary"
     5  	"errors"
     6  	"io"
     7  
     8  	"github.com/xxf098/lite-proxy/common/pool"
     9  )
    10  
    11  const (
    12  	lenSize   = 2
    13  	chunkSize = 1 << 14   // 2 ** 14 == 16 * 1024
    14  	maxSize   = 17 * 1024 // 2 + chunkSize + aead.Overhead()
    15  )
    16  
    17  type chunkReader struct {
    18  	io.Reader
    19  	buf     []byte
    20  	sizeBuf []byte
    21  	offset  int
    22  }
    23  
    24  func newChunkReader(reader io.Reader) *chunkReader {
    25  	return &chunkReader{Reader: reader, sizeBuf: make([]byte, lenSize)}
    26  }
    27  
    28  func newChunkWriter(writer io.WriteCloser) *chunkWriter {
    29  	return &chunkWriter{Writer: writer}
    30  }
    31  
    32  func (cr *chunkReader) Read(b []byte) (int, error) {
    33  	if cr.buf != nil {
    34  		n := copy(b, cr.buf[cr.offset:])
    35  		cr.offset += n
    36  		if cr.offset == len(cr.buf) {
    37  			pool.Put(cr.buf)
    38  			cr.buf = nil
    39  		}
    40  		return n, nil
    41  	}
    42  
    43  	_, err := io.ReadFull(cr.Reader, cr.sizeBuf)
    44  	if err != nil {
    45  		return 0, err
    46  	}
    47  
    48  	size := int(binary.BigEndian.Uint16(cr.sizeBuf))
    49  	if size > maxSize {
    50  		return 0, errors.New("buffer is larger than standard")
    51  	}
    52  
    53  	if len(b) >= size {
    54  		_, err := io.ReadFull(cr.Reader, b[:size])
    55  		if err != nil {
    56  			return 0, err
    57  		}
    58  
    59  		return size, nil
    60  	}
    61  
    62  	buf := pool.Get(size)
    63  	_, err = io.ReadFull(cr.Reader, buf)
    64  	if err != nil {
    65  		pool.Put(buf)
    66  		return 0, err
    67  	}
    68  	n := copy(b, buf)
    69  	cr.offset = n
    70  	cr.buf = buf
    71  	return n, nil
    72  }
    73  
    74  type chunkWriter struct {
    75  	io.Writer
    76  }
    77  
    78  func (cw *chunkWriter) Write(b []byte) (n int, err error) {
    79  	buf := pool.Get(pool.RelayBufferSize)
    80  	defer pool.Put(buf)
    81  	length := len(b)
    82  	for {
    83  		if length == 0 {
    84  			break
    85  		}
    86  		readLen := chunkSize
    87  		if length < chunkSize {
    88  			readLen = length
    89  		}
    90  		payloadBuf := buf[lenSize : lenSize+chunkSize]
    91  		copy(payloadBuf, b[n:n+readLen])
    92  
    93  		binary.BigEndian.PutUint16(buf[:lenSize], uint16(readLen))
    94  		_, err = cw.Writer.Write(buf[:lenSize+readLen])
    95  		if err != nil {
    96  			break
    97  		}
    98  		n += readLen
    99  		length -= readLen
   100  	}
   101  	return
   102  }