Golang alloc.NewSmallBuffer函数代码示例

func (server *SocksServer) handleSocks4(reader io.Reader, writer io.Writer, auth protocol.Socks4AuthenticationRequest) error {
	result := protocol.Socks4RequestGranted
	if auth.Command == protocol.CmdBind {
		result = protocol.Socks4RequestRejected
	}
	socks4Response := protocol.NewSocks4AuthenticationResponse(result, auth.Port, auth.IP[:])

	responseBuffer := alloc.NewSmallBuffer().Clear()
	socks4Response.Write(responseBuffer)
	writer.Write(responseBuffer.Value)
	responseBuffer.Release()

	if result == protocol.Socks4RequestRejected {
		log.Warning("Unsupported socks 4 command %d", auth.Command)
		return UnsupportedSocksCommand
	}

	dest := v2net.NewTCPDestination(v2net.IPAddress(auth.IP[:], auth.Port))
	data, err := v2net.ReadFrom(reader, nil)
	if err != nil {
		return err
	}

	packet := v2net.NewPacket(dest, data, true)
	server.transport(reader, writer, packet)
	return nil
}

func TestInsufficientAddressRequest(t *testing.T) {
	v2testing.Current(t)

	buffer := alloc.NewSmallBuffer().Clear().AppendBytes(1, 1)
	_, err := ReadRequest(buffer, nil, false)
	assert.Error(err).Equals(transport.ErrorCorruptedPacket)

	buffer = alloc.NewSmallBuffer().Clear().AppendBytes(4, 1)
	_, err = ReadRequest(buffer, nil, false)
	assert.Error(err).Equals(transport.ErrorCorruptedPacket)

	buffer = alloc.NewSmallBuffer().Clear().AppendBytes(3, 255, 1)
	_, err = ReadRequest(buffer, nil, false)
	assert.Error(err).Equals(transport.ErrorCorruptedPacket)
}

func TestInsufficientPortRequest(t *testing.T) {
	assert := assert.On(t)

	buffer := alloc.NewSmallBuffer().Clear().AppendBytes(1, 1, 2, 3, 4, 5)
	_, err := ReadRequest(buffer, nil, false)
	assert.Error(err).Equals(transport.ErrCorruptedPacket)
}

func TestWrongAddressType(t *testing.T) {
	v2testing.Current(t)

	buffer := alloc.NewSmallBuffer().Clear().AppendBytes(5)
	_, err := ReadRequest(buffer, nil, false)
	assert.Error(err).Equals(transport.ErrorCorruptedPacket)
}

func TestSinglePacket(t *testing.T) {
	v2testing.Current(t)
	port := v2nettesting.PickPort()

	tcpServer := &tcp.Server{
		Port: port,
		MsgProcessor: func(data []byte) []byte {
			buffer := make([]byte, 0, 2048)
			buffer = append(buffer, []byte("Processed: ")...)
			buffer = append(buffer, data...)
			return buffer
		},
	}
	_, err := tcpServer.Start()
	assert.Error(err).IsNil()

	freedom := &FreedomConnection{}
	traffic := ray.NewRay()
	data2Send := "Data to be sent to remote"
	payload := alloc.NewSmallBuffer().Clear().Append([]byte(data2Send))
	packet := v2net.NewPacket(v2net.TCPDestination(v2net.IPAddress([]byte{127, 0, 0, 1}), port), payload, false)

	err = freedom.Dispatch(packet, traffic)
	assert.Error(err).IsNil()
	close(traffic.InboundInput())

	respPayload := <-traffic.InboundOutput()
	defer respPayload.Release()
	assert.Bytes(respPayload.Value).Equals([]byte("Processed: Data to be sent to remote"))

	_, open := <-traffic.InboundOutput()
	assert.Bool(open).IsFalse()

	tcpServer.Close()
}

func (handler *VMessOutboundHandler) Dispatch(firstPacket v2net.Packet, ray ray.OutboundRay) error {
	vNextList := handler.vNextList
	if firstPacket.Destination().IsUDP() {
		vNextList = handler.vNextListUDP
	}
	vNextAddress, vNextUser := pickVNext(vNextList)

	command := protocol.CmdTCP
	if firstPacket.Destination().IsUDP() {
		command = protocol.CmdUDP
	}
	request := &protocol.VMessRequest{
		Version: protocol.Version,
		User:    vNextUser,
		Command: command,
		Address: firstPacket.Destination().Address(),
	}

	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()
	v2net.ReadAllBytes(rand.Reader, buffer.Value[:36]) // 16 + 16 + 4
	request.RequestIV = buffer.Value[:16]
	request.RequestKey = buffer.Value[16:32]
	request.ResponseHeader = buffer.Value[32:36]

	return startCommunicate(request, vNextAddress, ray, firstPacket)
}

func TestSingleBytePayload(t *testing.T) {
	v2testing.Current(t)

	buffer := alloc.NewSmallBuffer().Clear().AppendBytes(1)
	_, err := ReadRequest(buffer, nil, false)
	assert.Error(err).Equals(transport.ErrorCorruptedPacket)
}

func (this *VMessOutboundHandler) Dispatch(firstPacket v2net.Packet, ray ray.OutboundRay) error {
	vNextAddress, vNextUser := this.receiverManager.PickReceiver()

	command := protocol.CmdTCP
	if firstPacket.Destination().IsUDP() {
		command = protocol.CmdUDP
	}
	request := &protocol.VMessRequest{
		Version: protocol.Version,
		User:    vNextUser,
		Command: command,
		Address: firstPacket.Destination().Address(),
		Port:    firstPacket.Destination().Port(),
	}
	if command == protocol.CmdUDP {
		request.Option |= protocol.OptionChunk
	}

	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()                      // Buffer is released after communication finishes.
	io.ReadFull(rand.Reader, buffer.Value[:33]) // 16 + 16 + 1
	request.RequestIV = buffer.Value[:16]
	request.RequestKey = buffer.Value[16:32]
	request.ResponseHeader = buffer.Value[32]

	return this.startCommunicate(request, vNextAddress, ray, firstPacket)
}

func TestEmptyPayload(t *testing.T) {
	assert := assert.On(t)

	buffer := alloc.NewSmallBuffer().Clear()
	_, err := ReadRequest(buffer, nil, false)
	assert.Error(err).Equals(io.EOF)
}

func (this *ClientSession) DecodeResponseHeader(reader io.Reader) (*protocol.ResponseHeader, error) {
	aesStream := crypto.NewAesDecryptionStream(this.responseBodyKey, this.responseBodyIV)
	this.responseReader = crypto.NewCryptionReader(aesStream, reader)

	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()

	_, err := io.ReadFull(this.responseReader, buffer.Value[:4])
	if err != nil {
		log.Error("Raw: Failed to read response header: ", err)
		return nil, err
	}

	if buffer.Value[0] != this.responseHeader {
		log.Warning("Raw: Unexpected response header. Expecting %d, but actually %d", this.responseHeader, buffer.Value[0])
		return nil, transport.ErrorCorruptedPacket
	}

	header := new(protocol.ResponseHeader)

	if buffer.Value[2] != 0 {
		cmdId := buffer.Value[2]
		dataLen := int(buffer.Value[3])
		_, err := io.ReadFull(this.responseReader, buffer.Value[:dataLen])
		if err != nil {
			log.Error("Raw: Failed to read response command: ", err)
			return nil, err
		}
		data := buffer.Value[:dataLen]
		command, err := UnmarshalCommand(cmdId, data)
		header.Command = command
	}

	return header, nil
}

func (server *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer) error {
	response := protocol.NewSocks5Response()
	response.Error = protocol.ErrorSuccess

	udpAddr := server.getUDPAddr()

	response.Port = udpAddr.Port()
	switch {
	case udpAddr.IsIPv4():
		response.SetIPv4(udpAddr.IP())
	case udpAddr.IsIPv6():
		response.SetIPv6(udpAddr.IP())
	case udpAddr.IsDomain():
		response.SetDomain(udpAddr.Domain())
	}

	responseBuffer := alloc.NewSmallBuffer().Clear()
	response.Write(responseBuffer)
	_, err := writer.Write(responseBuffer.Value)
	responseBuffer.Release()

	if err != nil {
		log.Error("Socks failed to write response: %v", err)
		return err
	}

	reader.SetTimeOut(300)      /* 5 minutes */
	v2net.ReadFrom(reader, nil) // Just in case of anything left in the socket
	// The TCP connection closes after this method returns. We need to wait until
	// the client closes it.
	// TODO: get notified from UDP part
	<-time.After(5 * time.Minute)

	return nil
}

func (this *VMessOutboundHandler) handleRequest(session *encoding.ClientSession, conn internet.Connection, request *protocol.RequestHeader, payload *alloc.Buffer, input v2io.Reader, finish *sync.Mutex) {
	defer finish.Unlock()

	writer := v2io.NewBufferedWriter(conn)
	defer writer.Release()
	session.EncodeRequestHeader(request, writer)

	bodyWriter := session.EncodeRequestBody(writer)
	var streamWriter v2io.Writer = v2io.NewAdaptiveWriter(bodyWriter)
	if request.Option.Has(protocol.RequestOptionChunkStream) {
		streamWriter = vmessio.NewAuthChunkWriter(streamWriter)
	}
	if err := streamWriter.Write(payload); err != nil {
		conn.SetReusable(false)
	}
	writer.SetCached(false)

	err := v2io.Pipe(input, streamWriter)
	if err != io.EOF {
		conn.SetReusable(false)
	}

	if request.Option.Has(protocol.RequestOptionChunkStream) {
		err := streamWriter.Write(alloc.NewSmallBuffer().Clear())
		if err != nil {
			conn.SetReusable(false)
		}
	}
	streamWriter.Release()
	return
}

func TestResponseWrite(t *testing.T) {
	assert := assert.On(t)

	response := Socks5Response{
		socksVersion,
		ErrorSuccess,
		AddrTypeIPv4,
		[4]byte{0x72, 0x72, 0x72, 0x72},
		"",
		[16]byte{},
		v2net.Port(53),
	}
	buffer := alloc.NewSmallBuffer().Clear()
	defer buffer.Release()

	response.Write(buffer)
	expectedBytes := []byte{
		socksVersion,
		ErrorSuccess,
		byte(0x00),
		AddrTypeIPv4,
		0x72, 0x72, 0x72, 0x72,
		byte(0x00), byte(0x035),
	}
	assert.Bytes(buffer.Value).Equals(expectedBytes)
}

func ReadUserPassRequest(reader io.Reader) (request Socks5UserPassRequest, err error) {
	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()

	_, err = reader.Read(buffer.Value[0: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])

	_, err = reader.Read(buffer.Value[0:1])
	if err != nil {
		return
	}
	nPassword := buffer.Value[0]
	nBytes, err = reader.Read(buffer.Value[:nPassword])
	if err != nil {
		return
	}
	request.password = string(buffer.Value[:nBytes])
	return
}

func (this *SocksServer) handlePacket(packet v2net.Packet, clientAddr *net.UDPAddr, targetAddr v2net.Address, port v2net.Port) {
	ray := this.space.PacketDispatcher().DispatchToOutbound(packet)
	close(ray.InboundInput())

	for data := range ray.InboundOutput() {
		response := &protocol.Socks5UDPRequest{
			Fragment: 0,
			Address:  targetAddr,
			Port:     port,
			Data:     data,
		}
		log.Info("Socks: Writing back UDP response with ", data.Len(), " bytes from ", targetAddr, " to ", clientAddr)

		udpMessage := alloc.NewSmallBuffer().Clear()
		response.Write(udpMessage)

		this.udpMutex.RLock()
		if !this.accepting {
			this.udpMutex.RUnlock()
			return
		}
		nBytes, err := this.udpConn.WriteToUDP(udpMessage.Value, clientAddr)
		this.udpMutex.RUnlock()
		udpMessage.Release()
		response.Data.Release()
		if err != nil {
			log.Error("Socks: failed to write UDP message (", nBytes, " bytes) to ", clientAddr, ": ", err)
		}
	}
}