addons.go 4.8 KB

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