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
+18 -16
View File
@@ -122,28 +122,30 @@ func ReadUserPassRequest(reader io.Reader) (request Socks5UserPassRequest, err e
buffer := alloc.NewLocalBuffer(512)
defer buffer.Release()
_, err = reader.Read(buffer.Value[0:2])
_, err = buffer.FillFullFrom(reader, 2)
if err != nil {
return
}
request.version = buffer.Value[0]
nUsername := buffer.Value[1]
nBytes, err := reader.Read(buffer.Value[:nUsername])
if err != nil {
return
}
request.username = string(buffer.Value[:nBytes])
request.version = buffer.Byte(0)
nUsername := int(buffer.Byte(1))
_, err = reader.Read(buffer.Value[0:1])
buffer.Clear()
_, err = buffer.FillFullFrom(reader, nUsername)
if err != nil {
return
}
nPassword := buffer.Value[0]
nBytes, err = reader.Read(buffer.Value[:nPassword])
request.username = string(buffer.Bytes())
_, err = buffer.FillFullFrom(reader, 1)
if err != nil {
return
}
request.password = string(buffer.Value[:nBytes])
nPassword := int(buffer.Byte(0))
_, err = buffer.FillFullFrom(reader, nPassword)
if err != nil {
return
}
request.password = string(buffer.Bytes())
return
}
@@ -185,7 +187,7 @@ type Socks5Request struct {
}
func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
buffer := alloc.NewLocalBuffer(512).Clear()
buffer := alloc.NewLocalBuffer(512)
defer buffer.Release()
_, err = buffer.FillFullFrom(reader, 4)
@@ -194,10 +196,10 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
}
request = &Socks5Request{
Version: buffer.Value[0],
Command: buffer.Value[1],
Version: buffer.Byte(0),
Command: buffer.Byte(1),
// buffer[2] is a reserved field
AddrType: buffer.Value[3],
AddrType: buffer.Byte(3),
}
switch request.AddrType {
case AddrTypeIPv4:
+2 -2
View File
@@ -31,9 +31,9 @@ func TestSocks4AuthenticationResponseToBytes(t *testing.T) {
response := NewSocks4AuthenticationResponse(byte(0x10), 443, []byte{1, 2, 3, 4})
buffer := alloc.NewLocalBuffer(2048).Clear()
buffer := alloc.NewLocalBuffer(2048)
defer buffer.Release()
response.Write(buffer)
assert.Bytes(buffer.Value).Equals([]byte{0x00, 0x10, 0x01, 0xBB, 0x01, 0x02, 0x03, 0x04})
assert.Bytes(buffer.Bytes()).Equals([]byte{0x00, 0x10, 0x01, 0xBB, 0x01, 0x02, 0x03, 0x04})
}
+14 -9
View File
@@ -29,7 +29,8 @@ func TestHasAuthenticationMethod(t *testing.T) {
func TestAuthenticationRequestRead(t *testing.T) {
assert := assert.On(t)
buffer := alloc.NewBuffer().Clear().AppendBytes(
buffer := alloc.NewBuffer()
buffer.AppendBytes(
0x05, // version
0x01, // nMethods
0x02, // methods
@@ -83,7 +84,7 @@ func TestResponseWrite(t *testing.T) {
[16]byte{},
v2net.Port(53),
}
buffer := alloc.NewLocalBuffer(2048).Clear()
buffer := alloc.NewLocalBuffer(2048)
defer buffer.Release()
response.Write(buffer)
@@ -104,7 +105,7 @@ func TestSetIPv6(t *testing.T) {
response := NewSocks5Response()
response.SetIPv6([]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15})
buffer := alloc.NewLocalBuffer(2048).Clear()
buffer := alloc.NewLocalBuffer(2048)
defer buffer.Release()
response.Write(buffer)
assert.Bytes(buffer.Bytes()).Equals([]byte{
@@ -117,7 +118,7 @@ func TestSetDomain(t *testing.T) {
response := NewSocks5Response()
response.SetDomain("v2ray.com")
buffer := alloc.NewLocalBuffer(2048).Clear()
buffer := alloc.NewLocalBuffer(2048)
defer buffer.Release()
response.Write(buffer)
assert.Bytes(buffer.Bytes()).Equals([]byte{
@@ -127,7 +128,7 @@ func TestSetDomain(t *testing.T) {
func TestEmptyAuthRequest(t *testing.T) {
assert := assert.On(t)
_, _, err := ReadAuthentication(alloc.NewBuffer().Clear())
_, _, err := ReadAuthentication(alloc.NewBuffer())
assert.Error(err).Equals(io.EOF)
}
@@ -141,14 +142,16 @@ func TestSingleByteAuthRequest(t *testing.T) {
func TestZeroAuthenticationMethod(t *testing.T) {
assert := assert.On(t)
buffer := alloc.NewBuffer().Clear().AppendBytes(5, 0)
buffer := alloc.NewBuffer()
buffer.AppendBytes(5, 0)
_, _, err := ReadAuthentication(buffer)
assert.Error(err).Equals(proxy.ErrInvalidAuthentication)
}
func TestWrongProtocolVersion(t *testing.T) {
assert := assert.On(t)
buffer := alloc.NewBuffer().Clear().AppendBytes(6, 1, 0)
buffer := alloc.NewBuffer()
buffer.AppendBytes(6, 1, 0)
_, _, err := ReadAuthentication(buffer)
assert.Error(err).Equals(proxy.ErrInvalidProtocolVersion)
}
@@ -156,14 +159,16 @@ func TestWrongProtocolVersion(t *testing.T) {
func TestEmptyRequest(t *testing.T) {
assert := assert.On(t)
_, err := ReadRequest(alloc.NewBuffer().Clear())
_, err := ReadRequest(alloc.NewBuffer())
assert.Error(err).Equals(io.EOF)
}
func TestIPv6Request(t *testing.T) {
assert := assert.On(t)
request, err := ReadRequest(alloc.NewBuffer().Clear().AppendBytes(5, 1, 0, 4, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 0, 8))
b := alloc.NewBuffer()
b.AppendBytes(5, 1, 0, 4, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 0, 8)
request, err := ReadRequest(b)
assert.Error(err).IsNil()
assert.Byte(request.Command).Equals(1)
assert.Bytes(request.IPv6[:]).Equals([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6})
+12 -6
View File
@@ -5,6 +5,7 @@ import (
"v2ray.com/core/common/alloc"
"v2ray.com/core/common/errors"
v2net "v2ray.com/core/common/net"
"v2ray.com/core/common/serial"
)
var (
@@ -26,14 +27,17 @@ func (request *Socks5UDPRequest) Write(buffer *alloc.Buffer) {
buffer.AppendBytes(0, 0, request.Fragment)
switch request.Address.Family() {
case v2net.AddressFamilyIPv4:
buffer.AppendBytes(AddrTypeIPv4).Append(request.Address.IP())
buffer.AppendBytes(AddrTypeIPv4)
buffer.Append(request.Address.IP())
case v2net.AddressFamilyIPv6:
buffer.AppendBytes(AddrTypeIPv6).Append(request.Address.IP())
buffer.AppendBytes(AddrTypeIPv6)
buffer.Append(request.Address.IP())
case v2net.AddressFamilyDomain:
buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain()))).Append([]byte(request.Address.Domain()))
buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
buffer.Append([]byte(request.Address.Domain()))
}
buffer.AppendUint16(request.Port.Value())
buffer.Append(request.Data.Value)
buffer.AppendFunc(serial.WriteUint16(request.Port.Value()))
buffer.Append(request.Data.Bytes())
}
func ReadUDPRequest(packet []byte) (*Socks5UDPRequest, error) {
@@ -79,7 +83,9 @@ func ReadUDPRequest(packet []byte) (*Socks5UDPRequest, error) {
}
if len(packet) > dataBegin {
request.Data = alloc.NewBuffer().Clear().Append(packet[dataBegin:])
b := alloc.NewSmallBuffer()
b.Append(packet[dataBegin:])
request.Data = b
}
return request, nil
+1 -1
View File
@@ -32,5 +32,5 @@ func TestDomainAddressRequest(t *testing.T) {
assert.Byte(request.Fragment).Equals(1)
assert.Address(request.Address).EqualsString("v2ray.com")
assert.Port(request.Port).Equals(v2net.Port(80))
assert.Bytes(request.Data.Value).Equals([]byte("Actual payload"))
assert.String(request.Data.String()).Equals("Actual payload")
}
+3 -3
View File
@@ -26,7 +26,7 @@ func (v *Server) listenUDP() error {
func (v *Server) handleUDPPayload(payload *alloc.Buffer, session *proxy.SessionInfo) {
source := session.Source
log.Info("Socks: Client UDP connection from ", source)
request, err := protocol.ReadUDPRequest(payload.Value)
request, err := protocol.ReadUDPRequest(payload.Bytes())
payload.Release()
if err != nil {
@@ -55,7 +55,7 @@ func (v *Server) handleUDPPayload(payload *alloc.Buffer, session *proxy.SessionI
}
log.Info("Socks: Writing back UDP response with ", payload.Len(), " bytes to ", destination)
udpMessage := alloc.NewLocalBuffer(2048).Clear()
udpMessage := alloc.NewLocalBuffer(2048)
response.Write(udpMessage)
v.udpMutex.RLock()
@@ -63,7 +63,7 @@ func (v *Server) handleUDPPayload(payload *alloc.Buffer, session *proxy.SessionI
v.udpMutex.RUnlock()
return
}
nBytes, err := v.udpHub.WriteTo(udpMessage.Value, destination)
nBytes, err := v.udpHub.WriteTo(udpMessage.Bytes(), destination)
v.udpMutex.RUnlock()
udpMessage.Release()
response.Data.Release()