FillFullFrom

pull/314/head
Darien Raymond 2016-12-05 17:05:47 +01:00
parent 10b75e8961
commit 186bf1f141
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
4 changed files with 40 additions and 37 deletions

View File

@ -220,7 +220,14 @@ func (b *Buffer) Read(data []byte) (int, error) {
func (b *Buffer) FillFrom(reader io.Reader) (int, error) { func (b *Buffer) FillFrom(reader io.Reader) (int, error) {
begin := b.Len() begin := b.Len()
nBytes, err := reader.Read(b.head[b.offset+begin:]) nBytes, err := reader.Read(b.head[b.offset+begin:])
b.Value = b.head[:b.offset+begin+nBytes] b.Value = b.head[b.offset : b.offset+begin+nBytes]
return nBytes, err
}
func (b *Buffer) FillFullFrom(reader io.Reader, amount int) (int, error) {
begin := b.Len()
nBytes, err := io.ReadFull(reader, b.head[b.offset+begin:b.offset+begin+amount])
b.Value = b.head[b.offset : b.offset+begin+nBytes]
return nBytes, err return nBytes, err
} }

View File

@ -71,8 +71,8 @@ func (v *ChunkReader) Release() {
} }
func (v *ChunkReader) Read() (*alloc.Buffer, error) { func (v *ChunkReader) Read() (*alloc.Buffer, error) {
buffer := alloc.NewBuffer() buffer := alloc.NewBuffer().Clear()
if _, err := io.ReadFull(v.reader, buffer.BytesTo(2)); err != nil { if _, err := buffer.FillFullFrom(v.reader, 2); err != nil {
buffer.Release() buffer.Release()
return nil, err return nil, err
} }
@ -84,11 +84,12 @@ func (v *ChunkReader) Read() (*alloc.Buffer, error) {
buffer.Release() buffer.Release()
buffer = alloc.NewLocalBuffer(int(length) + 128) buffer = alloc.NewLocalBuffer(int(length) + 128)
} }
if _, err := io.ReadFull(v.reader, buffer.BytesTo(int(length))); err != nil {
buffer.Clear()
if _, err := buffer.FillFullFrom(v.reader, int(length)); err != nil {
buffer.Release() buffer.Release()
return nil, err return nil, err
} }
buffer.Slice(0, int(length))
authBytes := buffer.BytesTo(AuthSize) authBytes := buffer.BytesTo(AuthSize)
payload := buffer.BytesFrom(AuthSize) payload := buffer.BytesFrom(AuthSize)

View File

@ -28,11 +28,11 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea
} }
account := rawAccount.(*ShadowsocksAccount) account := rawAccount.(*ShadowsocksAccount)
buffer := alloc.NewLocalBuffer(512) buffer := alloc.NewLocalBuffer(512).Clear()
defer buffer.Release() defer buffer.Release()
ivLen := account.Cipher.IVSize() ivLen := account.Cipher.IVSize()
_, err = io.ReadFull(reader, buffer.Bytes()[:ivLen]) _, err = buffer.FillFullFrom(reader, ivLen)
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read IV.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read IV.")
} }
@ -52,8 +52,8 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea
Command: protocol.RequestCommandTCP, Command: protocol.RequestCommandTCP,
} }
lenBuffer := 1 buffer.Clear()
_, err = io.ReadFull(reader, buffer.Bytes()[:1]) _, err = buffer.FillFullFrom(reader, 1)
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read address type.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read address type.")
} }
@ -73,53 +73,47 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea
switch addrType { switch addrType {
case AddrTypeIPv4: case AddrTypeIPv4:
_, err := io.ReadFull(reader, buffer.BytesRange(lenBuffer, lenBuffer+4)) _, err := buffer.FillFullFrom(reader, 4)
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read IPv4 address.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read IPv4 address.")
} }
request.Address = v2net.IPAddress(buffer.BytesRange(lenBuffer, lenBuffer+4)) request.Address = v2net.IPAddress(buffer.BytesFrom(-4))
lenBuffer += 4
case AddrTypeIPv6: case AddrTypeIPv6:
_, err := io.ReadFull(reader, buffer.BytesRange(lenBuffer, lenBuffer+16)) _, err := buffer.FillFullFrom(reader, 16)
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read IPv6 address.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read IPv6 address.")
} }
request.Address = v2net.IPAddress(buffer.BytesRange(lenBuffer, lenBuffer+16)) request.Address = v2net.IPAddress(buffer.BytesFrom(-16))
lenBuffer += 16
case AddrTypeDomain: case AddrTypeDomain:
_, err := io.ReadFull(reader, buffer.BytesRange(lenBuffer, lenBuffer+1)) _, err := buffer.FillFullFrom(reader, 1)
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read domain lenth.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read domain lenth.")
} }
domainLength := int(buffer.Bytes()[lenBuffer]) domainLength := int(buffer.BytesFrom(-1)[0])
lenBuffer++ _, err = buffer.FillFullFrom(reader, domainLength)
_, err = io.ReadFull(reader, buffer.BytesRange(lenBuffer, lenBuffer+domainLength))
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read domain.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read domain.")
} }
request.Address = v2net.DomainAddress(string(buffer.BytesRange(lenBuffer, lenBuffer+domainLength))) request.Address = v2net.DomainAddress(string(buffer.BytesFrom(-domainLength)))
lenBuffer += domainLength
default: default:
return nil, nil, errors.New("Shadowsocks|TCP: Unknown address type: ", addrType) return nil, nil, errors.New("Shadowsocks|TCP: Unknown address type: ", addrType)
} }
_, err = io.ReadFull(reader, buffer.BytesRange(lenBuffer, lenBuffer+2)) _, err = buffer.FillFullFrom(reader, 2)
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read port.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read port.")
} }
request.Port = v2net.PortFromBytes(buffer.BytesFrom(-2))
request.Port = v2net.PortFromBytes(buffer.BytesRange(lenBuffer, lenBuffer+2))
lenBuffer += 2
if request.Option.Has(RequestOptionOneTimeAuth) { if request.Option.Has(RequestOptionOneTimeAuth) {
authBytes := buffer.BytesRange(lenBuffer, lenBuffer+AuthSize) actualAuth := authenticator.Authenticate(nil, buffer.Bytes())
_, err = io.ReadFull(reader, authBytes)
_, err := buffer.FillFullFrom(reader, AuthSize)
if err != nil { if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read OTA.") return nil, nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to read OTA.")
} }
actualAuth := authenticator.Authenticate(nil, buffer.BytesTo(lenBuffer)) if !bytes.Equal(actualAuth, buffer.BytesFrom(-AuthSize)) {
if !bytes.Equal(actualAuth, authBytes) {
return nil, nil, errors.New("Shadowsocks|TCP: Invalid OTA") return nil, nil, errors.New("Shadowsocks|TCP: Invalid OTA")
} }
} }

View File

@ -185,10 +185,10 @@ type Socks5Request struct {
} }
func ReadRequest(reader io.Reader) (request *Socks5Request, err error) { func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
buffer := alloc.NewLocalBuffer(512) buffer := alloc.NewLocalBuffer(512).Clear()
defer buffer.Release() defer buffer.Release()
_, err = io.ReadFull(reader, buffer.Value[:4]) _, err = buffer.FillFullFrom(reader, 4)
if err != nil { if err != nil {
return return
} }
@ -206,17 +206,18 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
return return
} }
case AddrTypeDomain: case AddrTypeDomain:
_, err = io.ReadFull(reader, buffer.Value[0:1]) buffer.Clear()
_, err = buffer.FillFullFrom(reader, 1)
if err != nil { if err != nil {
return return
} }
domainLength := buffer.Value[0] domainLength := int(buffer.Byte(0))
_, err = io.ReadFull(reader, buffer.Value[:domainLength]) _, err = buffer.FillFullFrom(reader, domainLength)
if err != nil { if err != nil {
return return
} }
request.Domain = string(append([]byte(nil), buffer.Value[:domainLength]...)) request.Domain = string(buffer.BytesFrom(-domainLength))
case AddrTypeIPv6: case AddrTypeIPv6:
_, err = io.ReadFull(reader, request.IPv6[:]) _, err = io.ReadFull(reader, request.IPv6[:])
if err != nil { if err != nil {
@ -227,12 +228,12 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
return return
} }
_, err = io.ReadFull(reader, buffer.Value[:2]) _, err = buffer.FillFullFrom(reader, 2)
if err != nil { if err != nil {
return return
} }
request.Port = v2net.PortFromBytes(buffer.Value[:2]) request.Port = v2net.PortFromBytes(buffer.BytesFrom(-2))
return return
} }