func TestBufferedReader(t *testing.T) {
v2testing.Current(t)
content := alloc.NewLargeBuffer()
len := content.Len()
reader := NewBufferedReader(content)
assert.Bool(reader.Cached()).IsTrue()
payload := make([]byte, 16)
nBytes, err := reader.Read(payload)
assert.Int(nBytes).Equals(16)
assert.Error(err).IsNil()
len2 := content.Len()
assert.Int(len - len2).GreaterThan(16)
nBytes, err = reader.Read(payload)
assert.Int(nBytes).Equals(16)
assert.Error(err).IsNil()
assert.Int(content.Len()).Equals(len2)
reader.SetCached(false)
payload2 := alloc.NewBuffer()
reader.Read(payload2.Value)
assert.Int(content.Len()).Equals(len2)
reader.Read(payload2.Value)
assert.Int(content.Len()).LessThan(len2)
}
func (this *ChunkReader) Read() (*alloc.Buffer, error) {
buffer := alloc.NewLargeBuffer()
if _, err := io.ReadFull(this.reader, buffer.Value[:2]); err != nil {
buffer.Release()
return nil, err
}
// There is a potential buffer overflow here. Large buffer is 64K bytes,
// while uin16 + 10 will be more than that
length := serial.BytesToUint16(buffer.Value[:2]) + AuthSize
if _, err := io.ReadFull(this.reader, buffer.Value[:length]); err != nil {
buffer.Release()
return nil, err
}
buffer.Slice(0, int(length))
authBytes := buffer.Value[:AuthSize]
payload := buffer.Value[AuthSize:]
actualAuthBytes := this.auth.Authenticate(nil, payload)
if !bytes.Equal(authBytes, actualAuthBytes) {
buffer.Release()
log.Debug("AuthenticationReader: Unexpected auth: ", authBytes)
return nil, transport.ErrorCorruptedPacket
}
buffer.Value = payload
return buffer, nil
}
func (this *VMessInboundHandler) HandleConnection(connection *net.TCPConn) error {
defer connection.Close()
connReader := v2net.NewTimeOutReader(16, connection)
requestReader := protocol.NewVMessRequestReader(this.clients)
request, err := requestReader.Read(connReader)
if err != nil {
log.Access(connection.RemoteAddr().String(), "", log.AccessRejected, err.Error())
log.Warning("VMessIn: Invalid request from (%s): %v", connection.RemoteAddr().String(), err)
return err
}
log.Access(connection.RemoteAddr().String(), request.Address.String(), log.AccessAccepted, "")
log.Debug("VMessIn: Received request for %s", request.Address.String())
ray := this.space.PacketDispatcher().DispatchToOutbound(v2net.NewPacket(request.Destination(), nil, true))
input := ray.InboundInput()
output := ray.InboundOutput()
var readFinish, writeFinish sync.Mutex
readFinish.Lock()
writeFinish.Lock()
userSettings := vmess.GetUserSettings(request.User.Level)
connReader.SetTimeOut(userSettings.PayloadReadTimeout)
go handleInput(request, connReader, input, &readFinish)
responseKey := md5.Sum(request.RequestKey)
responseIV := md5.Sum(request.RequestIV)
aesStream, err := v2crypto.NewAesEncryptionStream(responseKey[:], responseIV[:])
if err != nil {
log.Error("VMessIn: Failed to create AES decryption stream: %v", err)
close(input)
return err
}
responseWriter := v2crypto.NewCryptionWriter(aesStream, connection)
// Optimize for small response packet
buffer := alloc.NewLargeBuffer().Clear()
defer buffer.Release()
buffer.AppendBytes(request.ResponseHeader[0] ^ request.ResponseHeader[1])
buffer.AppendBytes(request.ResponseHeader[2] ^ request.ResponseHeader[3])
buffer.AppendBytes(byte(0), byte(0))
if data, open := <-output; open {
buffer.Append(data.Value)
data.Release()
responseWriter.Write(buffer.Value)
go handleOutput(request, responseWriter, output, &writeFinish)
writeFinish.Lock()
}
connection.CloseWrite()
readFinish.Lock()
return nil
}
func (handler *VMessInboundHandler) HandleConnection(connection *net.TCPConn) error {
defer connection.Close()
connReader := v2net.NewTimeOutReader(16, connection)
requestReader := protocol.NewVMessRequestReader(handler.clients)
request, err := requestReader.Read(connReader)
if err != nil {
log.Access(connection.RemoteAddr().String(), "", log.AccessRejected, err.Error())
log.Warning("VMessIn: Invalid request from (%s): %v", connection.RemoteAddr().String(), err)
return err
}
log.Access(connection.RemoteAddr().String(), request.Address.String(), log.AccessAccepted, "")
log.Debug("VMessIn: Received request for %s", request.Address.String())
ray := handler.dispatcher.DispatchToOutbound(v2net.NewPacket(request.Destination(), nil, true))
input := ray.InboundInput()
output := ray.InboundOutput()
var readFinish, writeFinish sync.Mutex
readFinish.Lock()
writeFinish.Lock()
connReader.SetTimeOut(120)
go handleInput(request, connReader, input, &readFinish)
responseKey := md5.Sum(request.RequestKey)
responseIV := md5.Sum(request.RequestIV)
responseWriter, err := v2io.NewAesEncryptWriter(responseKey[:], responseIV[:], connection)
if err != nil {
log.Error("VMessIn: Failed to create encrypt writer: %v", err)
return err
}
// Optimize for small response packet
buffer := alloc.NewLargeBuffer().Clear()
buffer.Append(request.ResponseHeader)
if data, open := <-output; open {
buffer.Append(data.Value)
data.Release()
responseWriter.Write(buffer.Value)
buffer.Release()
go handleOutput(request, responseWriter, output, &writeFinish)
writeFinish.Lock()
}
connection.CloseWrite()
readFinish.Lock()
return nil
}
func TestAdaptiveWriter(t *testing.T) {
assert := assert.On(t)
lb := alloc.NewLargeBuffer()
rand.Read(lb.Value)
writeBuffer := make([]byte, 0, 1024*1024)
writer := NewAdaptiveWriter(NewBufferedWriter(bytes.NewBuffer(writeBuffer)))
err := writer.Write(lb)
assert.Error(err).IsNil()
assert.Bytes(lb.Bytes()).Equals(writeBuffer)
}
func TestBufferedWriter(t *testing.T) {
v2testing.Current(t)
content := alloc.NewLargeBuffer().Clear()
writer := NewBufferedWriter(content)
assert.Bool(writer.Cached()).IsTrue()
payload := make([]byte, 16)
nBytes, err := writer.Write(payload)
assert.Int(nBytes).Equals(16)
assert.Error(err).IsNil()
assert.Bool(content.IsEmpty()).IsTrue()
writer.SetCached(false)
assert.Int(content.Len()).Equals(16)
}
func NewValidationReader(reader io.Reader) *ValidationReader {
return &ValidationReader{
reader: reader,
buffer: alloc.NewLargeBuffer().Clear(),
}
}
func TestLargeIO(t *testing.T) {
assert := assert.On(t)
content := make([]byte, 1024*1024)
rand.Read(content)
chunckContent := bytes.NewBuffer(make([]byte, 0, len(content)*2))
writer := NewAuthChunkWriter(v2io.NewAdaptiveWriter(chunckContent))
writeSize := 0
for {
chunkSize := 7 * 1024
if chunkSize+writeSize > len(content) {
chunkSize = len(content) - writeSize
}
writer.Write(alloc.NewBuffer().Clear().Append(content[writeSize : writeSize+chunkSize]))
writeSize += chunkSize
if writeSize == len(content) {
break
}
chunkSize = 8 * 1024
if chunkSize+writeSize > len(content) {
chunkSize = len(content) - writeSize
}
writer.Write(alloc.NewLargeBuffer().Clear().Append(content[writeSize : writeSize+chunkSize]))
writeSize += chunkSize
if writeSize == len(content) {
break
}
chunkSize = 63 * 1024
if chunkSize+writeSize > len(content) {
chunkSize = len(content) - writeSize
}
writer.Write(alloc.NewLargeBuffer().Clear().Append(content[writeSize : writeSize+chunkSize]))
writeSize += chunkSize
if writeSize == len(content) {
break
}
chunkSize = 64*1024 - 16
if chunkSize+writeSize > len(content) {
chunkSize = len(content) - writeSize
}
writer.Write(alloc.NewLargeBuffer().Clear().Append(content[writeSize : writeSize+chunkSize]))
writeSize += chunkSize
if writeSize == len(content) {
break
}
}
writer.Write(alloc.NewBuffer().Clear())
writer.Release()
actualContent := make([]byte, 0, len(content))
reader := NewAuthChunkReader(chunckContent)
for {
buffer, err := reader.Read()
if err == io.EOF {
break
}
assert.Error(err).IsNil()
actualContent = append(actualContent, buffer.Value...)
}
assert.Int(len(actualContent)).Equals(len(content))
assert.Bytes(actualContent).Equals(content)
}
func (this *VMessInboundHandler) HandleConnection(connection *hub.TCPConn) {
defer connection.Close()
connReader := v2net.NewTimeOutReader(16, connection)
requestReader := protocol.NewVMessRequestReader(this.clients)
request, err := requestReader.Read(connReader)
if err != nil {
log.Access(connection.RemoteAddr(), serial.StringLiteral(""), log.AccessRejected, serial.StringLiteral(err.Error()))
log.Warning("VMessIn: Invalid request from ", connection.RemoteAddr(), ": ", err)
return
}
log.Access(connection.RemoteAddr(), request.Address, log.AccessAccepted, serial.StringLiteral(""))
log.Debug("VMessIn: Received request for ", request.Address)
ray := this.packetDispatcher.DispatchToOutbound(v2net.NewPacket(request.Destination(), nil, true))
input := ray.InboundInput()
output := ray.InboundOutput()
var readFinish, writeFinish sync.Mutex
readFinish.Lock()
writeFinish.Lock()
userSettings := vmess.GetUserSettings(request.User.Level)
connReader.SetTimeOut(userSettings.PayloadReadTimeout)
go handleInput(request, connReader, input, &readFinish)
responseKey := md5.Sum(request.RequestKey)
responseIV := md5.Sum(request.RequestIV)
aesStream, err := v2crypto.NewAesEncryptionStream(responseKey[:], responseIV[:])
if err != nil {
log.Error("VMessIn: Failed to create AES decryption stream: ", err)
close(input)
return
}
responseWriter := v2crypto.NewCryptionWriter(aesStream, connection)
// Optimize for small response packet
buffer := alloc.NewLargeBuffer().Clear()
defer buffer.Release()
buffer.AppendBytes(request.ResponseHeader, byte(0))
this.generateCommand(buffer)
if data, open := <-output; open {
if request.IsChunkStream() {
vmessio.Authenticate(data)
}
buffer.Append(data.Value)
data.Release()
responseWriter.Write(buffer.Value)
go func(finish *sync.Mutex) {
var writer v2io.Writer
writer = v2io.NewAdaptiveWriter(responseWriter)
if request.IsChunkStream() {
writer = vmessio.NewAuthChunkWriter(writer)
}
v2io.ChanToWriter(writer, output)
finish.Unlock()
}(&writeFinish)
writeFinish.Lock()
}
connection.CloseWrite()
readFinish.Lock()
}