github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/proxy/vless/encoding/addons.go (about)

     1  //go:build !confonly
     2  // +build !confonly
     3  
     4  package encoding
     5  
     6  import (
     7  	"io"
     8  
     9  	"google.golang.org/protobuf/proto"
    10  
    11  	"github.com/v2fly/v2ray-core/v5/common/buf"
    12  	"github.com/v2fly/v2ray-core/v5/common/errors"
    13  	"github.com/v2fly/v2ray-core/v5/common/protocol"
    14  )
    15  
    16  // EncodeHeaderAddons Add addons byte to the header
    17  func EncodeHeaderAddons(buffer *buf.Buffer, addons *Addons) error {
    18  	if err := buffer.WriteByte(0); err != nil {
    19  		return newError("failed to write addons protobuf length").Base(err)
    20  	}
    21  	return nil
    22  }
    23  
    24  func DecodeHeaderAddons(buffer *buf.Buffer, reader io.Reader) (*Addons, error) {
    25  	addons := new(Addons)
    26  	buffer.Clear()
    27  	if _, err := buffer.ReadFullFrom(reader, 1); err != nil {
    28  		return nil, newError("failed to read addons protobuf length").Base(err)
    29  	}
    30  
    31  	if length := int32(buffer.Byte(0)); length != 0 {
    32  		buffer.Clear()
    33  		if _, err := buffer.ReadFullFrom(reader, length); err != nil {
    34  			return nil, newError("failed to read addons protobuf value").Base(err)
    35  		}
    36  
    37  		if err := proto.Unmarshal(buffer.Bytes(), addons); err != nil {
    38  			return nil, newError("failed to unmarshal addons protobuf value").Base(err)
    39  		}
    40  	}
    41  
    42  	return addons, nil
    43  }
    44  
    45  // EncodeBodyAddons returns a Writer that auto-encrypt content written by caller.
    46  func EncodeBodyAddons(writer io.Writer, request *protocol.RequestHeader, addons *Addons) buf.Writer {
    47  	if request.Command == protocol.RequestCommandUDP {
    48  		return NewMultiLengthPacketWriter(writer.(buf.Writer))
    49  	}
    50  	return buf.NewWriter(writer)
    51  }
    52  
    53  // DecodeBodyAddons returns a Reader from which caller can fetch decrypted body.
    54  func DecodeBodyAddons(reader io.Reader, request *protocol.RequestHeader, addons *Addons) buf.Reader {
    55  	if request.Command == protocol.RequestCommandUDP {
    56  		return NewLengthPacketReader(reader)
    57  	}
    58  	return buf.NewReader(reader)
    59  }
    60  
    61  func NewMultiLengthPacketWriter(writer buf.Writer) *MultiLengthPacketWriter {
    62  	return &MultiLengthPacketWriter{
    63  		Writer: writer,
    64  	}
    65  }
    66  
    67  type MultiLengthPacketWriter struct {
    68  	buf.Writer
    69  }
    70  
    71  func (w *MultiLengthPacketWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
    72  	defer buf.ReleaseMulti(mb)
    73  
    74  	if len(mb)+1 > 64*1024*1024 {
    75  		return errors.New("value too large")
    76  	}
    77  	sliceSize := len(mb) + 1
    78  	mb2Write := make(buf.MultiBuffer, 0, sliceSize)
    79  	for _, b := range mb {
    80  		length := b.Len()
    81  		if length == 0 || length+2 > buf.Size {
    82  			continue
    83  		}
    84  		eb := buf.New()
    85  		if err := eb.WriteByte(byte(length >> 8)); err != nil {
    86  			eb.Release()
    87  			continue
    88  		}
    89  		if err := eb.WriteByte(byte(length)); err != nil {
    90  			eb.Release()
    91  			continue
    92  		}
    93  		if _, err := eb.Write(b.Bytes()); err != nil {
    94  			eb.Release()
    95  			continue
    96  		}
    97  		mb2Write = append(mb2Write, eb)
    98  	}
    99  	if mb2Write.IsEmpty() {
   100  		return nil
   101  	}
   102  	return w.Writer.WriteMultiBuffer(mb2Write)
   103  }
   104  
   105  type LengthPacketWriter struct {
   106  	io.Writer
   107  	cache []byte
   108  }
   109  
   110  func (w *LengthPacketWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
   111  	length := mb.Len() // none of mb is nil
   112  	if length == 0 {
   113  		return nil
   114  	}
   115  	defer func() {
   116  		w.cache = w.cache[:0]
   117  	}()
   118  	w.cache = append(w.cache, byte(length>>8), byte(length))
   119  	for i, b := range mb {
   120  		w.cache = append(w.cache, b.Bytes()...)
   121  		b.Release()
   122  		mb[i] = nil
   123  	}
   124  	if _, err := w.Write(w.cache); err != nil {
   125  		return newError("failed to write a packet").Base(err)
   126  	}
   127  	return nil
   128  }
   129  
   130  func NewLengthPacketReader(reader io.Reader) *LengthPacketReader {
   131  	return &LengthPacketReader{
   132  		Reader: reader,
   133  		cache:  make([]byte, 2),
   134  	}
   135  }
   136  
   137  type LengthPacketReader struct {
   138  	io.Reader
   139  	cache []byte
   140  }
   141  
   142  func (r *LengthPacketReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
   143  	if _, err := io.ReadFull(r.Reader, r.cache); err != nil { // maybe EOF
   144  		return nil, newError("failed to read packet length").Base(err)
   145  	}
   146  	length := int32(r.cache[0])<<8 | int32(r.cache[1])
   147  	mb := make(buf.MultiBuffer, 0, length/buf.Size+1)
   148  	for length > 0 {
   149  		size := length
   150  		if size > buf.Size {
   151  			size = buf.Size
   152  		}
   153  		length -= size
   154  		b := buf.New()
   155  		if _, err := b.ReadFullFrom(r.Reader, size); err != nil {
   156  			return nil, newError("failed to read packet payload").Base(err)
   157  		}
   158  		mb = append(mb, b)
   159  	}
   160  	return mb, nil
   161  }