github.com/EagleQL/Xray-core@v1.4.3/common/buf/buffer.go (about)

     1  package buf
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/xtls/xray-core/common/bytespool"
     7  	"github.com/xtls/xray-core/common/net"
     8  )
     9  
    10  const (
    11  	// Size of a regular buffer.
    12  	Size = 8192
    13  )
    14  
    15  var pool = bytespool.GetPool(Size)
    16  
    17  // Buffer is a recyclable allocation of a byte array. Buffer.Release() recycles
    18  // the buffer into an internal buffer pool, in order to recreate a buffer more
    19  // quickly.
    20  type Buffer struct {
    21  	v     []byte
    22  	start int32
    23  	end   int32
    24  	UDP   *net.Destination
    25  }
    26  
    27  // New creates a Buffer with 0 length and 2K capacity.
    28  func New() *Buffer {
    29  	return &Buffer{
    30  		v: pool.Get().([]byte),
    31  	}
    32  }
    33  
    34  func NewExisted(b []byte) *Buffer {
    35  	if cap(b) < Size {
    36  		panic("Invalid buffer")
    37  	}
    38  
    39  	oLen := len(b)
    40  	if oLen < Size {
    41  		b = b[:Size]
    42  	}
    43  
    44  	return &Buffer{
    45  		v:   b,
    46  		end: int32(oLen),
    47  	}
    48  }
    49  
    50  // StackNew creates a new Buffer object on stack.
    51  // This method is for buffers that is released in the same function.
    52  func StackNew() Buffer {
    53  	return Buffer{
    54  		v: pool.Get().([]byte),
    55  	}
    56  }
    57  
    58  // Release recycles the buffer into an internal buffer pool.
    59  func (b *Buffer) Release() {
    60  	if b == nil || b.v == nil {
    61  		return
    62  	}
    63  
    64  	p := b.v
    65  	b.v = nil
    66  	b.Clear()
    67  	pool.Put(p)
    68  	b.UDP = nil
    69  }
    70  
    71  // Clear clears the content of the buffer, results an empty buffer with
    72  // Len() = 0.
    73  func (b *Buffer) Clear() {
    74  	b.start = 0
    75  	b.end = 0
    76  }
    77  
    78  // Byte returns the bytes at index.
    79  func (b *Buffer) Byte(index int32) byte {
    80  	return b.v[b.start+index]
    81  }
    82  
    83  // SetByte sets the byte value at index.
    84  func (b *Buffer) SetByte(index int32, value byte) {
    85  	b.v[b.start+index] = value
    86  }
    87  
    88  // Bytes returns the content bytes of this Buffer.
    89  func (b *Buffer) Bytes() []byte {
    90  	return b.v[b.start:b.end]
    91  }
    92  
    93  // Extend increases the buffer size by n bytes, and returns the extended part.
    94  // It panics if result size is larger than buf.Size.
    95  func (b *Buffer) Extend(n int32) []byte {
    96  	end := b.end + n
    97  	if end > int32(len(b.v)) {
    98  		panic("extending out of bound")
    99  	}
   100  	ext := b.v[b.end:end]
   101  	b.end = end
   102  	return ext
   103  }
   104  
   105  // BytesRange returns a slice of this buffer with given from and to boundary.
   106  func (b *Buffer) BytesRange(from, to int32) []byte {
   107  	if from < 0 {
   108  		from += b.Len()
   109  	}
   110  	if to < 0 {
   111  		to += b.Len()
   112  	}
   113  	return b.v[b.start+from : b.start+to]
   114  }
   115  
   116  // BytesFrom returns a slice of this Buffer starting from the given position.
   117  func (b *Buffer) BytesFrom(from int32) []byte {
   118  	if from < 0 {
   119  		from += b.Len()
   120  	}
   121  	return b.v[b.start+from : b.end]
   122  }
   123  
   124  // BytesTo returns a slice of this Buffer from start to the given position.
   125  func (b *Buffer) BytesTo(to int32) []byte {
   126  	if to < 0 {
   127  		to += b.Len()
   128  	}
   129  	if to < 0 {
   130  		to = 0
   131  	}
   132  	return b.v[b.start : b.start+to]
   133  }
   134  
   135  // Resize cuts the buffer at the given position.
   136  func (b *Buffer) Resize(from, to int32) {
   137  	if from < 0 {
   138  		from += b.Len()
   139  	}
   140  	if to < 0 {
   141  		to += b.Len()
   142  	}
   143  	if to < from {
   144  		panic("Invalid slice")
   145  	}
   146  	b.end = b.start + to
   147  	b.start += from
   148  }
   149  
   150  // Advance cuts the buffer at the given position.
   151  func (b *Buffer) Advance(from int32) {
   152  	if from < 0 {
   153  		from += b.Len()
   154  	}
   155  	b.start += from
   156  }
   157  
   158  // Len returns the length of the buffer content.
   159  func (b *Buffer) Len() int32 {
   160  	if b == nil {
   161  		return 0
   162  	}
   163  	return b.end - b.start
   164  }
   165  
   166  // IsEmpty returns true if the buffer is empty.
   167  func (b *Buffer) IsEmpty() bool {
   168  	return b.Len() == 0
   169  }
   170  
   171  // IsFull returns true if the buffer has no more room to grow.
   172  func (b *Buffer) IsFull() bool {
   173  	return b != nil && b.end == int32(len(b.v))
   174  }
   175  
   176  // Write implements Write method in io.Writer.
   177  func (b *Buffer) Write(data []byte) (int, error) {
   178  	nBytes := copy(b.v[b.end:], data)
   179  	b.end += int32(nBytes)
   180  	return nBytes, nil
   181  }
   182  
   183  // WriteByte writes a single byte into the buffer.
   184  func (b *Buffer) WriteByte(v byte) error {
   185  	if b.IsFull() {
   186  		return newError("buffer full")
   187  	}
   188  	b.v[b.end] = v
   189  	b.end++
   190  	return nil
   191  }
   192  
   193  // WriteString implements io.StringWriter.
   194  func (b *Buffer) WriteString(s string) (int, error) {
   195  	return b.Write([]byte(s))
   196  }
   197  
   198  // Read implements io.Reader.Read().
   199  func (b *Buffer) Read(data []byte) (int, error) {
   200  	if b.Len() == 0 {
   201  		return 0, io.EOF
   202  	}
   203  	nBytes := copy(data, b.v[b.start:b.end])
   204  	if int32(nBytes) == b.Len() {
   205  		b.Clear()
   206  	} else {
   207  		b.start += int32(nBytes)
   208  	}
   209  	return nBytes, nil
   210  }
   211  
   212  // ReadFrom implements io.ReaderFrom.
   213  func (b *Buffer) ReadFrom(reader io.Reader) (int64, error) {
   214  	n, err := reader.Read(b.v[b.end:])
   215  	b.end += int32(n)
   216  	return int64(n), err
   217  }
   218  
   219  // ReadFullFrom reads exact size of bytes from given reader, or until error occurs.
   220  func (b *Buffer) ReadFullFrom(reader io.Reader, size int32) (int64, error) {
   221  	end := b.end + size
   222  	if end > int32(len(b.v)) {
   223  		v := end
   224  		return 0, newError("out of bound: ", v)
   225  	}
   226  	n, err := io.ReadFull(reader, b.v[b.end:end])
   227  	b.end += int32(n)
   228  	return int64(n), err
   229  }
   230  
   231  // String returns the string form of this Buffer.
   232  func (b *Buffer) String() string {
   233  	return string(b.Bytes())
   234  }