rewrite alloc.buffer

This commit is contained in:
Darien Raymond
2016-12-06 11:03:42 +01:00
parent 4bf8288cc2
commit 8c8f0a53fd
51 changed files with 391 additions and 299 deletions
+10 -7
View File
@@ -26,11 +26,14 @@ func NewAuthenticator(keygen KeyGenerator) *Authenticator {
}
}
func (v *Authenticator) Authenticate(auth []byte, data []byte) []byte {
func (v *Authenticator) Authenticate(data []byte) alloc.BytesWriter {
hasher := hmac.New(sha1.New, v.key())
hasher.Write(data)
res := hasher.Sum(nil)
return append(auth, res[:AuthSize]...)
return func(b []byte) int {
copy(b, res[:AuthSize])
return AuthSize
}
}
func HeaderKeyGenerator(key []byte, iv []byte) func() []byte {
@@ -71,7 +74,7 @@ func (v *ChunkReader) Release() {
}
func (v *ChunkReader) Read() (*alloc.Buffer, error) {
buffer := alloc.NewBuffer().Clear()
buffer := alloc.NewBuffer()
if _, err := buffer.FillFullFrom(v.reader, 2); err != nil {
buffer.Release()
return nil, err
@@ -94,7 +97,8 @@ func (v *ChunkReader) Read() (*alloc.Buffer, error) {
authBytes := buffer.BytesTo(AuthSize)
payload := buffer.BytesFrom(AuthSize)
actualAuthBytes := v.auth.Authenticate(nil, payload)
actualAuthBytes := make([]byte, AuthSize)
v.auth.Authenticate(payload)(actualAuthBytes)
if !bytes.Equal(authBytes, actualAuthBytes) {
buffer.Release()
return nil, errors.New("Shadowsocks|AuthenticationReader: Invalid auth.")
@@ -123,9 +127,8 @@ func (v *ChunkWriter) Release() {
func (v *ChunkWriter) Write(payload *alloc.Buffer) error {
totalLength := payload.Len()
payload.SliceBack(AuthSize)
v.auth.Authenticate(payload.BytesTo(0), payload.BytesFrom(AuthSize))
payload.PrependUint16(uint16(totalLength))
payload.PrependFunc(AuthSize, v.auth.Authenticate(payload.Bytes()))
payload.PrependFunc(2, serial.WriteUint16(uint16(totalLength)))
_, err := v.writer.Write(payload.Bytes())
return err
}
+6 -3
View File
@@ -11,7 +11,8 @@ import (
func TestNormalChunkReading(t *testing.T) {
assert := assert.On(t)
buffer := alloc.NewBuffer().Clear().AppendBytes(
buffer := alloc.NewBuffer()
buffer.AppendBytes(
0, 8, 39, 228, 69, 96, 133, 39, 254, 26, 201, 70, 11, 12, 13, 14, 15, 16, 17, 18)
reader := NewChunkReader(buffer, NewAuthenticator(ChunkKeyGenerator(
[]byte{21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36})))
@@ -26,11 +27,13 @@ func TestNormalChunkReading(t *testing.T) {
func TestNormalChunkWriting(t *testing.T) {
assert := assert.On(t)
buffer := alloc.NewLocalBuffer(512).Clear()
buffer := alloc.NewLocalBuffer(512)
writer := NewChunkWriter(buffer, NewAuthenticator(ChunkKeyGenerator(
[]byte{21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36})))
err := writer.Write(alloc.NewLocalBuffer(256).Clear().Append([]byte{11, 12, 13, 14, 15, 16, 17, 18}))
b := alloc.NewLocalBuffer(256)
b.Append([]byte{11, 12, 13, 14, 15, 16, 17, 18})
err := writer.Write(b)
assert.Error(err).IsNil()
assert.Bytes(buffer.Bytes()).Equals([]byte{0, 8, 39, 228, 69, 96, 133, 39, 254, 26, 201, 70, 11, 12, 13, 14, 15, 16, 17, 18})
}
+29 -28
View File
@@ -10,6 +10,7 @@ import (
v2io "v2ray.com/core/common/io"
v2net "v2ray.com/core/common/net"
"v2ray.com/core/common/protocol"
"v2ray.com/core/common/serial"
)
const (
@@ -28,7 +29,7 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea
}
account := rawAccount.(*ShadowsocksAccount)
buffer := alloc.NewLocalBuffer(512).Clear()
buffer := alloc.NewLocalBuffer(512)
defer buffer.Release()
ivLen := account.Cipher.IVSize()
@@ -106,7 +107,8 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea
request.Port = v2net.PortFromBytes(buffer.BytesFrom(-2))
if request.Option.Has(RequestOptionOneTimeAuth) {
actualAuth := authenticator.Authenticate(nil, buffer.Bytes())
actualAuth := make([]byte, AuthSize)
authenticator.Authenticate(buffer.Bytes())(actualAuth)
_, err := buffer.FillFullFrom(reader, AuthSize)
if err != nil {
@@ -150,7 +152,7 @@ func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (v2io.Wr
writer = crypto.NewCryptionWriter(stream, writer)
header := alloc.NewLocalBuffer(512).Clear()
header := alloc.NewLocalBuffer(512)
switch request.Address.Family() {
case v2net.AddressFamilyIPv4:
@@ -166,16 +168,16 @@ func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (v2io.Wr
return nil, errors.New("Shadowsocks|TCP: Unsupported address type: ", request.Address.Family())
}
header.AppendUint16(uint16(request.Port))
header.AppendFunc(serial.WriteUint16(uint16(request.Port)))
if request.Option.Has(RequestOptionOneTimeAuth) {
header.Bytes()[0] |= 0x10
authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
header.Value = authenticator.Authenticate(header.Value, header.Value)
header.AppendFunc(authenticator.Authenticate(header.Bytes()))
}
_, err = writer.Write(header.Value)
_, err = writer.Write(header.Bytes())
if err != nil {
return nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to write header.")
}
@@ -243,9 +245,8 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a
buffer := alloc.NewSmallBuffer()
ivLen := account.Cipher.IVSize()
buffer.Slice(0, ivLen)
rand.Read(buffer.Value)
iv := buffer.Value
buffer.FillFullFrom(rand.Reader, ivLen)
iv := buffer.Bytes()
switch request.Address.Family() {
case v2net.AddressFamilyIPv4:
@@ -261,14 +262,14 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a
return nil, errors.New("Shadowsocks|UDP: Unsupported address type: ", request.Address.Family())
}
buffer.AppendUint16(uint16(request.Port))
buffer.Append(payload.Value)
buffer.AppendFunc(serial.WriteUint16(uint16(request.Port)))
buffer.Append(payload.Bytes())
if request.Option.Has(RequestOptionOneTimeAuth) {
authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
buffer.Value[ivLen] |= 0x10
buffer.Bytes()[ivLen] |= 0x10
buffer.Value = authenticator.Authenticate(buffer.Value, buffer.Value[ivLen:])
buffer.AppendFunc(authenticator.Authenticate(buffer.BytesFrom(ivLen)))
}
stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
@@ -276,7 +277,7 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a
return nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to create encoding stream.")
}
stream.XORKeyStream(buffer.Value[ivLen:], buffer.Value[ivLen:])
stream.XORKeyStream(buffer.BytesFrom(ivLen), buffer.BytesFrom(ivLen))
return buffer, nil
}
@@ -288,14 +289,14 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ
account := rawAccount.(*ShadowsocksAccount)
ivLen := account.Cipher.IVSize()
iv := payload.Value[:ivLen]
iv := payload.BytesTo(ivLen)
payload.SliceFrom(ivLen)
stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
if err != nil {
return nil, nil, errors.Base(err).Message("Shadowsocks|UDP: Failed to initialize decoding stream.")
}
stream.XORKeyStream(payload.Value, payload.Value)
stream.XORKeyStream(payload.Bytes(), payload.Bytes())
authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
request := &protocol.RequestHeader{
@@ -304,8 +305,8 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ
Command: protocol.RequestCommandUDP,
}
addrType := (payload.Value[0] & 0x0F)
if (payload.Value[0] & 0x10) == 0x10 {
addrType := (payload.Byte(0) & 0x0F)
if (payload.Byte(0) & 0x10) == 0x10 {
request.Option |= RequestOptionOneTimeAuth
}
@@ -319,9 +320,10 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ
if request.Option.Has(RequestOptionOneTimeAuth) {
payloadLen := payload.Len() - AuthSize
authBytes := payload.Value[payloadLen:]
authBytes := payload.BytesFrom(payloadLen)
actualAuth := authenticator.Authenticate(nil, payload.Value[0:payloadLen])
actualAuth := make([]byte, AuthSize)
authenticator.Authenticate(payload.BytesTo(payloadLen))(actualAuth)
if !bytes.Equal(actualAuth, authBytes) {
return nil, nil, errors.New("Shadowsocks|UDP: Invalid OTA.")
}
@@ -333,20 +335,20 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ
switch addrType {
case AddrTypeIPv4:
request.Address = v2net.IPAddress(payload.Value[:4])
request.Address = v2net.IPAddress(payload.BytesTo(4))
payload.SliceFrom(4)
case AddrTypeIPv6:
request.Address = v2net.IPAddress(payload.Value[:16])
request.Address = v2net.IPAddress(payload.BytesTo(16))
payload.SliceFrom(16)
case AddrTypeDomain:
domainLength := int(payload.Value[0])
request.Address = v2net.DomainAddress(string(payload.Value[1 : 1+domainLength]))
domainLength := int(payload.Byte(0))
request.Address = v2net.DomainAddress(string(payload.BytesRange(1, 1+domainLength)))
payload.SliceFrom(1 + domainLength)
default:
return nil, nil, errors.New("Shadowsocks|UDP: Unknown address type: ", addrType)
}
request.Port = v2net.PortFromBytes(payload.Value[:2])
request.Port = v2net.PortFromBytes(payload.BytesTo(2))
payload.SliceFrom(2)
return request, payload, nil
@@ -359,12 +361,11 @@ type UDPReader struct {
func (v *UDPReader) Read() (*alloc.Buffer, error) {
buffer := alloc.NewSmallBuffer()
nBytes, err := v.Reader.Read(buffer.Value)
_, err := buffer.FillFrom(v.Reader)
if err != nil {
buffer.Release()
return nil, err
}
buffer.Slice(0, nBytes)
_, payload, err := DecodeUDPPacket(v.User, buffer)
if err != nil {
buffer.Release()
@@ -386,7 +387,7 @@ func (v *UDPWriter) Write(buffer *alloc.Buffer) error {
if err != nil {
return err
}
_, err = v.Writer.Write(payload.Value)
_, err = v.Writer.Write(payload.Bytes())
payload.Release()
return err
}
+13 -6
View File
@@ -7,6 +7,7 @@ import (
"v2ray.com/core/common/loader"
v2net "v2ray.com/core/common/net"
"v2ray.com/core/common/protocol"
"v2ray.com/core/common/serial"
. "v2ray.com/core/proxy/shadowsocks"
"v2ray.com/core/testing/assert"
)
@@ -29,7 +30,8 @@ func TestUDPEncoding(t *testing.T) {
},
}
data := alloc.NewLocalBuffer(256).Clear().AppendString("test string")
data := alloc.NewLocalBuffer(256)
data.AppendFunc(serial.WriteString("test string"))
encodedData, err := EncodeUDPPacket(request, data)
assert.Error(err).IsNil()
@@ -58,8 +60,9 @@ func TestTCPRequest(t *testing.T) {
},
}
data := alloc.NewLocalBuffer(256).Clear().AppendString("test string")
cache := alloc.NewBuffer().Clear()
data := alloc.NewLocalBuffer(256)
data.AppendFunc(serial.WriteString("test string"))
cache := alloc.NewBuffer()
writer, err := WriteTCPRequest(request, cache)
assert.Error(err).IsNil()
@@ -85,7 +88,7 @@ func TestUDPReaderWriter(t *testing.T) {
CipherType: CipherType_CHACHA20_IEFT,
}),
}
cache := alloc.NewBuffer().Clear()
cache := alloc.NewBuffer()
writer := &UDPWriter{
Writer: cache,
Request: &protocol.RequestHeader{
@@ -102,14 +105,18 @@ func TestUDPReaderWriter(t *testing.T) {
User: user,
}
err := writer.Write(alloc.NewBuffer().Clear().AppendString("test payload"))
b := alloc.NewBuffer()
b.AppendFunc(serial.WriteString("test payload"))
err := writer.Write(b)
assert.Error(err).IsNil()
payload, err := reader.Read()
assert.Error(err).IsNil()
assert.String(payload.String()).Equals("test payload")
err = writer.Write(alloc.NewBuffer().Clear().AppendString("test payload 2"))
b = alloc.NewBuffer()
b.AppendFunc(serial.WriteString("test payload 2"))
err = writer.Write(b)
assert.Error(err).IsNil()
payload, err = reader.Read()