1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 18:17:52 -05:00
v2fly/common/crypto/chunk.go

148 lines
2.9 KiB
Go
Raw Normal View History

2017-04-23 07:30:08 -04:00
package crypto
import (
"io"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/serial"
)
type ChunkSizeDecoder interface {
SizeBytes() int
Decode([]byte) (uint16, error)
}
type ChunkSizeEncoder interface {
SizeBytes() int
Encode(uint16, []byte) []byte
}
type PlainChunkSizeParser struct{}
func (PlainChunkSizeParser) SizeBytes() int {
return 2
}
func (PlainChunkSizeParser) Encode(size uint16, b []byte) []byte {
return serial.Uint16ToBytes(size, b)
}
func (PlainChunkSizeParser) Decode(b []byte) (uint16, error) {
return serial.BytesToUint16(b), nil
}
type ChunkStreamReader struct {
sizeDecoder ChunkSizeDecoder
reader buf.Reader
buffer []byte
leftOver buf.MultiBuffer
2017-05-02 16:50:27 -04:00
leftOverSize int
2017-04-23 07:30:08 -04:00
}
func NewChunkStreamReader(sizeDecoder ChunkSizeDecoder, reader io.Reader) *ChunkStreamReader {
return &ChunkStreamReader{
sizeDecoder: sizeDecoder,
reader: buf.NewReader(reader),
buffer: make([]byte, sizeDecoder.SizeBytes()),
}
}
2017-04-28 14:11:44 -04:00
func (r *ChunkStreamReader) readAtLeast(size int) error {
2017-04-23 07:30:08 -04:00
mb := r.leftOver
2017-05-02 16:50:27 -04:00
r.leftOver = nil
2017-04-23 07:30:08 -04:00
for mb.Len() < size {
extra, err := r.reader.Read()
if err != nil {
mb.Release()
2017-04-28 14:11:44 -04:00
return err
2017-04-23 07:30:08 -04:00
}
mb.AppendMulti(extra)
}
2017-04-28 14:11:44 -04:00
r.leftOver = mb
2017-04-23 07:30:08 -04:00
2017-04-28 14:11:44 -04:00
return nil
2017-04-23 07:30:08 -04:00
}
func (r *ChunkStreamReader) readSize() (uint16, error) {
if r.sizeDecoder.SizeBytes() > r.leftOver.Len() {
2017-04-28 14:11:44 -04:00
if err := r.readAtLeast(r.sizeDecoder.SizeBytes() - r.leftOver.Len()); err != nil {
2017-04-23 07:30:08 -04:00
return 0, err
}
}
r.leftOver.Read(r.buffer)
return r.sizeDecoder.Decode(r.buffer)
}
func (r *ChunkStreamReader) Read() (buf.MultiBuffer, error) {
2017-05-02 16:50:27 -04:00
size := r.leftOverSize
2017-04-23 07:30:08 -04:00
if size == 0 {
nextSize, err := r.readSize()
if err != nil {
return nil, err
}
if nextSize == 0 {
return nil, io.EOF
}
size = int(nextSize)
}
2017-04-28 14:11:44 -04:00
if r.leftOver.IsEmpty() {
if err := r.readAtLeast(1); err != nil {
2017-04-23 07:30:08 -04:00
return nil, err
}
}
2017-04-28 14:11:44 -04:00
if size >= r.leftOver.Len() {
mb := r.leftOver
2017-05-02 16:50:27 -04:00
r.leftOverSize = size - r.leftOver.Len()
2017-04-23 07:30:08 -04:00
r.leftOver = nil
2017-04-28 14:11:44 -04:00
return mb, nil
2017-04-23 07:30:08 -04:00
}
2017-04-28 14:11:44 -04:00
mb := r.leftOver.SliceBySize(size)
2017-04-23 07:30:08 -04:00
if mb.Len() != size {
b := buf.New()
2017-04-28 14:11:44 -04:00
b.AppendSupplier(buf.ReadFullFrom(&r.leftOver, size-mb.Len()))
2017-04-23 07:30:08 -04:00
mb.Append(b)
}
r.leftOverSize = 0
return mb, nil
}
type ChunkStreamWriter struct {
sizeEncoder ChunkSizeEncoder
writer buf.Writer
}
func NewChunkStreamWriter(sizeEncoder ChunkSizeEncoder, writer io.Writer) *ChunkStreamWriter {
return &ChunkStreamWriter{
sizeEncoder: sizeEncoder,
writer: buf.NewWriter(writer),
}
}
func (w *ChunkStreamWriter) Write(mb buf.MultiBuffer) error {
mb2Write := buf.NewMultiBuffer()
const sliceSize = 8192
for {
slice := mb.SliceBySize(sliceSize)
b := buf.New()
b.AppendSupplier(func(buffer []byte) (int, error) {
w.sizeEncoder.Encode(uint16(slice.Len()), buffer[:0])
return w.sizeEncoder.SizeBytes(), nil
})
mb2Write.Append(b)
mb2Write.AppendMulti(slice)
if mb.IsEmpty() {
break
}
}
return w.writer.Write(mb2Write)
}