1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2025-02-20 23:47:21 -05:00

simplify auth reader

This commit is contained in:
Darien Raymond 2017-12-03 13:23:24 +01:00
parent 6652edfa6f
commit bcfcba396b
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
2 changed files with 31 additions and 102 deletions

View File

@ -90,133 +90,58 @@ func (v *AEADAuthenticator) Seal(dst, plainText []byte) ([]byte, error) {
type AuthenticationReader struct { type AuthenticationReader struct {
auth Authenticator auth Authenticator
buffer *buf.Buffer reader *buf.BufferedReader
reader io.Reader
sizeParser ChunkSizeDecoder sizeParser ChunkSizeDecoder
size int
transferType protocol.TransferType transferType protocol.TransferType
} }
const (
readerBufferSize = 32 * 1024
)
func NewAuthenticationReader(auth Authenticator, sizeParser ChunkSizeDecoder, reader io.Reader, transferType protocol.TransferType) *AuthenticationReader { func NewAuthenticationReader(auth Authenticator, sizeParser ChunkSizeDecoder, reader io.Reader, transferType protocol.TransferType) *AuthenticationReader {
return &AuthenticationReader{ return &AuthenticationReader{
auth: auth, auth: auth,
buffer: buf.NewLocal(readerBufferSize), reader: buf.NewBufferedReader(buf.NewReader(reader)),
reader: reader,
sizeParser: sizeParser, sizeParser: sizeParser,
size: -1,
transferType: transferType, transferType: transferType,
} }
} }
func (r *AuthenticationReader) readSize() error { func (r *AuthenticationReader) readSize() (int, error) {
if r.size >= 0 { sizeBytes := make([]byte, r.sizeParser.SizeBytes())
return nil _, err := io.ReadFull(r.reader, sizeBytes)
}
sizeBytes := r.sizeParser.SizeBytes()
if r.buffer.Len() < sizeBytes {
if r.buffer.IsEmpty() {
r.buffer.Clear()
} else {
common.Must(r.buffer.Reset(buf.ReadFrom(r.buffer)))
}
delta := sizeBytes - r.buffer.Len()
if err := r.buffer.AppendSupplier(buf.ReadAtLeastFrom(r.reader, delta)); err != nil {
return err
}
}
size, err := r.sizeParser.Decode(r.buffer.BytesTo(sizeBytes))
if err != nil { if err != nil {
return err return 0, err
} }
r.size = int(size) size, err := r.sizeParser.Decode(sizeBytes)
r.buffer.SliceFrom(sizeBytes) return int(size), err
return nil
}
func (r *AuthenticationReader) readChunk(waitForData bool) ([]byte, error) {
if err := r.readSize(); err != nil {
return nil, err
}
if r.size > readerBufferSize-r.sizeParser.SizeBytes() {
return nil, newError("size too large ", r.size).AtWarning()
}
if r.size == r.auth.Overhead() {
return nil, io.EOF
}
if r.buffer.Len() < r.size {
if !waitForData {
return nil, io.ErrNoProgress
}
if r.buffer.IsEmpty() {
r.buffer.Clear()
} else {
common.Must(r.buffer.Reset(buf.ReadFrom(r.buffer)))
}
delta := r.size - r.buffer.Len()
if err := r.buffer.AppendSupplier(buf.ReadAtLeastFrom(r.reader, delta)); err != nil {
return nil, err
}
}
b, err := r.auth.Open(r.buffer.BytesTo(0), r.buffer.BytesTo(r.size))
if err != nil {
return nil, err
}
r.buffer.SliceFrom(r.size)
r.size = -1
return b, nil
} }
func (r *AuthenticationReader) ReadMultiBuffer() (buf.MultiBuffer, error) { func (r *AuthenticationReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
b, err := r.readChunk(true) size, err := r.readSize()
if err != nil { if err != nil {
return nil, err return nil, err
} }
var mb buf.MultiBuffer if size == r.auth.Overhead() {
if r.transferType == protocol.TransferTypeStream { return nil, io.EOF
mb.Write(b) }
var b *buf.Buffer
if size <= buf.Size {
b = buf.New()
} else { } else {
var bb *buf.Buffer b = buf.NewLocal(size)
if len(b) <= buf.Size { }
bb = buf.New() if err := b.Reset(buf.ReadFullFrom(r.reader, size)); err != nil {
} else { b.Release()
bb = buf.NewLocal(len(b)) return nil, err
}
bb.Append(b)
mb.Append(bb)
} }
for r.buffer.Len() >= r.sizeParser.SizeBytes() { rb, err := r.auth.Open(b.BytesTo(0), b.BytesTo(size))
b, err := r.readChunk(false) if err != nil {
if err != nil { b.Release()
break return nil, err
}
if r.transferType == protocol.TransferTypeStream {
mb.Write(b)
} else {
var bb *buf.Buffer
if len(b) <= buf.Size {
bb = buf.New()
} else {
bb = buf.NewLocal(len(b))
}
bb.Append(b)
mb.Append(bb)
}
} }
b.Slice(0, len(rb))
return mb, nil return buf.NewMultiBufferValue(b), nil
} }
type AuthenticationWriter struct { type AuthenticationWriter struct {

View File

@ -122,6 +122,10 @@ func TestAuthenticationReaderWriterPacket(t *testing.T) {
b1 := mb.SplitFirst() b1 := mb.SplitFirst()
assert(b1.String(), Equals, "abcd") assert(b1.String(), Equals, "abcd")
assert(mb.IsEmpty(), IsTrue)
mb, err = reader.ReadMultiBuffer()
assert(err, IsNil)
b2 := mb.SplitFirst() b2 := mb.SplitFirst()
assert(b2.String(), Equals, "efgh") assert(b2.String(), Equals, "efgh")
assert(mb.IsEmpty(), IsTrue) assert(mb.IsEmpty(), IsTrue)