github.com/yaling888/clash@v1.53.0/transport/vmess/chunk.go (about)

     1  package vmess
     2  
     3  import (
     4  	"encoding/binary"
     5  	"errors"
     6  	"io"
     7  
     8  	"github.com/yaling888/clash/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  	bufP    *[]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.bufP != nil {
    34  		n := copy(b, (*cr.bufP)[cr.offset:])
    35  		cr.offset += n
    36  		if cr.offset == len(*cr.bufP) {
    37  			pool.PutNetBuf(cr.bufP)
    38  			cr.bufP = 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  	bufP := pool.GetNetBuf()
    63  	_, err = io.ReadFull(cr.Reader, (*bufP)[:size])
    64  	if err != nil {
    65  		pool.PutNetBuf(bufP)
    66  		return 0, err
    67  	}
    68  	n := copy(b, (*bufP)[:size])
    69  	*bufP = (*bufP)[:size]
    70  	cr.offset = n
    71  	cr.bufP = bufP
    72  	return n, nil
    73  }
    74  
    75  type chunkWriter struct {
    76  	io.Writer
    77  }
    78  
    79  func (cw *chunkWriter) Write(b []byte) (n int, err error) {
    80  	bufP := pool.GetNetBuf()
    81  	defer pool.PutNetBuf(bufP)
    82  	length := len(b)
    83  	for {
    84  		if length == 0 {
    85  			break
    86  		}
    87  		readLen := chunkSize
    88  		if length < chunkSize {
    89  			readLen = length
    90  		}
    91  		payloadBuf := (*bufP)[lenSize : lenSize+chunkSize]
    92  		copy(payloadBuf, b[n:n+readLen])
    93  
    94  		binary.BigEndian.PutUint16((*bufP)[:lenSize], uint16(readLen))
    95  		_, err = cw.Writer.Write((*bufP)[:lenSize+readLen])
    96  		if err != nil {
    97  			break
    98  		}
    99  		n += readLen
   100  		length -= readLen
   101  	}
   102  	return
   103  }