github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/common/crypto/chunk.go (about)

     1  package crypto
     2  
     3  import (
     4  	"encoding/binary"
     5  	"io"
     6  
     7  	"github.com/v2fly/v2ray-core/v5/common"
     8  	"github.com/v2fly/v2ray-core/v5/common/buf"
     9  )
    10  
    11  // ChunkSizeDecoder is a utility class to decode size value from bytes.
    12  type ChunkSizeDecoder interface {
    13  	// SizeBytes must be stable, return the same value across all calls
    14  	SizeBytes() int32
    15  	Decode([]byte) (uint16, error)
    16  }
    17  
    18  type ChunkSizeDecoderWithOffset interface {
    19  	ChunkSizeDecoder
    20  	// HasConstantOffset set the constant offset of Decode
    21  	// The effective size should be HasConstantOffset() + Decode(_).[0](uint64)
    22  	HasConstantOffset() uint16
    23  }
    24  
    25  // ChunkSizeEncoder is a utility class to encode size value into bytes.
    26  type ChunkSizeEncoder interface {
    27  	SizeBytes() int32
    28  	Encode(uint16, []byte) []byte
    29  }
    30  
    31  type PaddingLengthGenerator interface {
    32  	MaxPaddingLen() uint16
    33  	NextPaddingLen() uint16
    34  }
    35  
    36  type PlainChunkSizeParser struct{}
    37  
    38  func (PlainChunkSizeParser) SizeBytes() int32 {
    39  	return 2
    40  }
    41  
    42  func (PlainChunkSizeParser) Encode(size uint16, b []byte) []byte {
    43  	binary.BigEndian.PutUint16(b, size)
    44  	return b[:2]
    45  }
    46  
    47  func (PlainChunkSizeParser) Decode(b []byte) (uint16, error) {
    48  	return binary.BigEndian.Uint16(b), nil
    49  }
    50  
    51  type AEADChunkSizeParser struct {
    52  	Auth *AEADAuthenticator
    53  }
    54  
    55  func (p *AEADChunkSizeParser) SizeBytes() int32 {
    56  	return 2 + int32(p.Auth.Overhead())
    57  }
    58  
    59  func (p *AEADChunkSizeParser) Encode(size uint16, b []byte) []byte {
    60  	binary.BigEndian.PutUint16(b, size-uint16(p.Auth.Overhead()))
    61  	b, err := p.Auth.Seal(b[:0], b[:2])
    62  	common.Must(err)
    63  	return b
    64  }
    65  
    66  func (p *AEADChunkSizeParser) Decode(b []byte) (uint16, error) {
    67  	b, err := p.Auth.Open(b[:0], b)
    68  	if err != nil {
    69  		return 0, err
    70  	}
    71  	return binary.BigEndian.Uint16(b) + uint16(p.Auth.Overhead()), nil
    72  }
    73  
    74  type ChunkStreamReader struct {
    75  	sizeDecoder ChunkSizeDecoder
    76  	reader      *buf.BufferedReader
    77  
    78  	buffer       []byte
    79  	leftOverSize int32
    80  	maxNumChunk  uint32
    81  	numChunk     uint32
    82  }
    83  
    84  func NewChunkStreamReader(sizeDecoder ChunkSizeDecoder, reader io.Reader) *ChunkStreamReader {
    85  	return NewChunkStreamReaderWithChunkCount(sizeDecoder, reader, 0)
    86  }
    87  
    88  func NewChunkStreamReaderWithChunkCount(sizeDecoder ChunkSizeDecoder, reader io.Reader, maxNumChunk uint32) *ChunkStreamReader {
    89  	r := &ChunkStreamReader{
    90  		sizeDecoder: sizeDecoder,
    91  		buffer:      make([]byte, sizeDecoder.SizeBytes()),
    92  		maxNumChunk: maxNumChunk,
    93  	}
    94  	if breader, ok := reader.(*buf.BufferedReader); ok {
    95  		r.reader = breader
    96  	} else {
    97  		r.reader = &buf.BufferedReader{Reader: buf.NewReader(reader)}
    98  	}
    99  
   100  	return r
   101  }
   102  
   103  func (r *ChunkStreamReader) readSize() (uint16, error) {
   104  	if _, err := io.ReadFull(r.reader, r.buffer); err != nil {
   105  		return 0, err
   106  	}
   107  	return r.sizeDecoder.Decode(r.buffer)
   108  }
   109  
   110  func (r *ChunkStreamReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
   111  	size := r.leftOverSize
   112  	if size == 0 {
   113  		r.numChunk++
   114  		if r.maxNumChunk > 0 && r.numChunk > r.maxNumChunk {
   115  			return nil, io.EOF
   116  		}
   117  		nextSize, err := r.readSize()
   118  		if err != nil {
   119  			return nil, err
   120  		}
   121  		if nextSize == 0 {
   122  			return nil, io.EOF
   123  		}
   124  		size = int32(nextSize)
   125  	}
   126  	r.leftOverSize = size
   127  
   128  	mb, err := r.reader.ReadAtMost(size)
   129  	if !mb.IsEmpty() {
   130  		r.leftOverSize -= mb.Len()
   131  		return mb, nil
   132  	}
   133  	return nil, err
   134  }
   135  
   136  type ChunkStreamWriter struct {
   137  	sizeEncoder ChunkSizeEncoder
   138  	writer      buf.Writer
   139  }
   140  
   141  func NewChunkStreamWriter(sizeEncoder ChunkSizeEncoder, writer io.Writer) *ChunkStreamWriter {
   142  	return &ChunkStreamWriter{
   143  		sizeEncoder: sizeEncoder,
   144  		writer:      buf.NewWriter(writer),
   145  	}
   146  }
   147  
   148  func (w *ChunkStreamWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
   149  	const sliceSize = 8192
   150  	mbLen := mb.Len()
   151  	mb2Write := make(buf.MultiBuffer, 0, mbLen/buf.Size+mbLen/sliceSize+2)
   152  
   153  	for {
   154  		mb2, slice := buf.SplitSize(mb, sliceSize)
   155  		mb = mb2
   156  
   157  		b := buf.New()
   158  		w.sizeEncoder.Encode(uint16(slice.Len()), b.Extend(w.sizeEncoder.SizeBytes()))
   159  		mb2Write = append(mb2Write, b)
   160  		mb2Write = append(mb2Write, slice...)
   161  
   162  		if mb.IsEmpty() {
   163  			break
   164  		}
   165  	}
   166  
   167  	return w.writer.WriteMultiBuffer(mb2Write)
   168  }