github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/bufio/chunk.go (about)

     1  package bufio
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/sagernet/sing/common"
     7  	"github.com/sagernet/sing/common/buf"
     8  	N "github.com/sagernet/sing/common/network"
     9  )
    10  
    11  type ChunkReader struct {
    12  	upstream     N.ExtendedReader
    13  	maxChunkSize int
    14  	cache        *buf.Buffer
    15  }
    16  
    17  func NewChunkReader(upstream io.Reader, maxChunkSize int) *ChunkReader {
    18  	return &ChunkReader{
    19  		upstream:     NewExtendedReader(upstream),
    20  		maxChunkSize: maxChunkSize,
    21  	}
    22  }
    23  
    24  func (c *ChunkReader) ReadBuffer(buffer *buf.Buffer) error {
    25  	if buffer.FreeLen() >= c.maxChunkSize {
    26  		return c.upstream.ReadBuffer(buffer)
    27  	}
    28  	if c.cache == nil {
    29  		c.cache = buf.NewSize(c.maxChunkSize)
    30  	} else if !c.cache.IsEmpty() {
    31  		return common.Error(buffer.ReadFrom(c.cache))
    32  	}
    33  	c.cache.Reset()
    34  	err := c.upstream.ReadBuffer(c.cache)
    35  	if err != nil {
    36  		c.cache.Release()
    37  		c.cache = nil
    38  		return err
    39  	}
    40  	return common.Error(buffer.ReadFrom(c.cache))
    41  }
    42  
    43  func (c *ChunkReader) Read(p []byte) (n int, err error) {
    44  	if c.cache == nil {
    45  		c.cache = buf.NewSize(c.maxChunkSize)
    46  	} else if !c.cache.IsEmpty() {
    47  		return c.cache.Read(p)
    48  	}
    49  	c.cache.Reset()
    50  	err = c.upstream.ReadBuffer(c.cache)
    51  	if err != nil {
    52  		c.cache.Release()
    53  		c.cache = nil
    54  		return
    55  	}
    56  	return c.cache.Read(p)
    57  }
    58  
    59  func (c *ChunkReader) ReadByte() (byte, error) {
    60  	buffer, err := c.ReadChunk()
    61  	if err != nil {
    62  		return 0, err
    63  	}
    64  	return buffer.ReadByte()
    65  }
    66  
    67  func (c *ChunkReader) ReadChunk() (*buf.Buffer, error) {
    68  	if c.cache == nil {
    69  		c.cache = buf.NewSize(c.maxChunkSize)
    70  	} else if !c.cache.IsEmpty() {
    71  		return c.cache, nil
    72  	}
    73  	c.cache.Reset()
    74  	err := c.upstream.ReadBuffer(c.cache)
    75  	if err != nil {
    76  		c.cache.Release()
    77  		c.cache = nil
    78  		return nil, err
    79  	}
    80  	return c.cache, nil
    81  }
    82  
    83  func (c *ChunkReader) MTU() int {
    84  	return c.maxChunkSize
    85  }
    86  
    87  type ChunkWriter struct {
    88  	upstream     N.ExtendedWriter
    89  	maxChunkSize int
    90  }
    91  
    92  func NewChunkWriter(writer io.Writer, maxChunkSize int) *ChunkWriter {
    93  	return &ChunkWriter{
    94  		upstream:     NewExtendedWriter(writer),
    95  		maxChunkSize: maxChunkSize,
    96  	}
    97  }
    98  
    99  func (w *ChunkWriter) Write(p []byte) (n int, err error) {
   100  	for pLen := len(p); pLen > 0; {
   101  		var data []byte
   102  		if pLen > w.maxChunkSize {
   103  			data = p[:w.maxChunkSize]
   104  			p = p[w.maxChunkSize:]
   105  			pLen -= w.maxChunkSize
   106  		} else {
   107  			data = p
   108  			pLen = 0
   109  		}
   110  		var writeN int
   111  		writeN, err = w.upstream.Write(data)
   112  		n += writeN
   113  		if err != nil {
   114  			return
   115  		}
   116  	}
   117  	return
   118  }
   119  
   120  func (w *ChunkWriter) WriteBuffer(buffer *buf.Buffer) error {
   121  	if buffer.Len() > w.maxChunkSize {
   122  		defer buffer.Release()
   123  		return common.Error(w.Write(buffer.Bytes()))
   124  	}
   125  	return w.upstream.WriteBuffer(buffer)
   126  }
   127  
   128  func (w *ChunkWriter) Upstream() any {
   129  	return w.upstream
   130  }
   131  
   132  func (w *ChunkWriter) MTU() int {
   133  	return w.maxChunkSize
   134  }