chunk.go 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. package vmess
  2. import (
  3. "encoding/binary"
  4. "errors"
  5. "io"
  6. "github.com/metacubex/mihomo/common/pool"
  7. )
  8. const (
  9. lenSize = 2
  10. chunkSize = 1 << 14 // 2 ** 14 == 16 * 1024
  11. maxSize = 17 * 1024 // 2 + chunkSize + aead.Overhead()
  12. )
  13. type chunkReader struct {
  14. io.Reader
  15. buf []byte
  16. sizeBuf []byte
  17. offset int
  18. }
  19. func newChunkReader(reader io.Reader) *chunkReader {
  20. return &chunkReader{Reader: reader, sizeBuf: make([]byte, lenSize)}
  21. }
  22. func newChunkWriter(writer io.WriteCloser) *chunkWriter {
  23. return &chunkWriter{Writer: writer}
  24. }
  25. func (cr *chunkReader) Read(b []byte) (int, error) {
  26. if cr.buf != nil {
  27. n := copy(b, cr.buf[cr.offset:])
  28. cr.offset += n
  29. if cr.offset == len(cr.buf) {
  30. pool.Put(cr.buf)
  31. cr.buf = nil
  32. }
  33. return n, nil
  34. }
  35. _, err := io.ReadFull(cr.Reader, cr.sizeBuf)
  36. if err != nil {
  37. return 0, err
  38. }
  39. size := int(binary.BigEndian.Uint16(cr.sizeBuf))
  40. if size > maxSize {
  41. return 0, errors.New("buffer is larger than standard")
  42. }
  43. if len(b) >= size {
  44. _, err := io.ReadFull(cr.Reader, b[:size])
  45. if err != nil {
  46. return 0, err
  47. }
  48. return size, nil
  49. }
  50. buf := pool.Get(size)
  51. _, err = io.ReadFull(cr.Reader, buf)
  52. if err != nil {
  53. pool.Put(buf)
  54. return 0, err
  55. }
  56. n := copy(b, buf)
  57. cr.offset = n
  58. cr.buf = buf
  59. return n, nil
  60. }
  61. type chunkWriter struct {
  62. io.Writer
  63. }
  64. func (cw *chunkWriter) Write(b []byte) (n int, err error) {
  65. buf := pool.Get(pool.RelayBufferSize)
  66. defer pool.Put(buf)
  67. length := len(b)
  68. for {
  69. if length == 0 {
  70. break
  71. }
  72. readLen := chunkSize
  73. if length < chunkSize {
  74. readLen = length
  75. }
  76. payloadBuf := buf[lenSize : lenSize+chunkSize]
  77. copy(payloadBuf, b[n:n+readLen])
  78. binary.BigEndian.PutUint16(buf[:lenSize], uint16(readLen))
  79. _, err = cw.Writer.Write(buf[:lenSize+readLen])
  80. if err != nil {
  81. break
  82. }
  83. n += readLen
  84. length -= readLen
  85. }
  86. return
  87. }