Golang alloc.NewBuffer函数代码示例

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

	queue := NewSendingQueue(3)

	seg0 := alloc.NewBuffer()
	seg1 := alloc.NewBuffer()
	seg2 := alloc.NewBuffer()
	seg3 := alloc.NewBuffer()

	assert.Bool(queue.IsEmpty()).IsTrue()
	assert.Bool(queue.IsFull()).IsFalse()

	queue.Push(seg0)
	assert.Bool(queue.IsEmpty()).IsFalse()

	queue.Push(seg1)
	queue.Push(seg2)

	assert.Bool(queue.IsFull()).IsTrue()

	assert.Pointer(queue.Pop()).Equals(seg0)

	queue.Push(seg3)
	assert.Bool(queue.IsFull()).IsTrue()

	assert.Pointer(queue.Pop()).Equals(seg1)
	assert.Pointer(queue.Pop()).Equals(seg2)
	assert.Pointer(queue.Pop()).Equals(seg3)
	assert.Int(int(queue.Len())).Equals(0)
}

func (this *HttpProxyServer) handlePlainHTTP(request *http.Request, dest v2net.Destination, reader *bufio.Reader, writer io.Writer) {
	if len(request.URL.Host) <= 0 {
		hdr := http.Header(make(map[string][]string))
		hdr.Set("Connection", "close")
		response := &http.Response{
			Status:        "400 Bad Request",
			StatusCode:    400,
			Proto:         "HTTP/1.1",
			ProtoMajor:    1,
			ProtoMinor:    1,
			Header:        hdr,
			Body:          nil,
			ContentLength: 0,
			Close:         false,
		}

		buffer := alloc.NewSmallBuffer().Clear()
		response.Write(buffer)
		writer.Write(buffer.Value)
		buffer.Release()
		return
	}

	request.Host = request.URL.Host
	stripHopByHopHeaders(request)

	requestBuffer := alloc.NewBuffer().Clear()
	request.Write(requestBuffer)
	log.Info("Request to remote:\n%s", string(requestBuffer.Value))

	packet := v2net.NewPacket(dest, requestBuffer, true)
	ray := this.space.PacketDispatcher().DispatchToOutbound(packet)
	defer close(ray.InboundInput())

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		responseReader := bufio.NewReader(NewChanReader(ray.InboundOutput()))
		responseBuffer := alloc.NewBuffer()
		defer responseBuffer.Release()
		response, err := http.ReadResponse(responseReader, request)
		if err != nil {
			return
		}
		responseBuffer.Clear()
		response.Write(responseBuffer)
		writer.Write(responseBuffer.Value)
		response.Body.Close()
	}()
	wg.Wait()
}

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

	content := bytes.NewBuffer(make([]byte, 0, 1024*1024))

	writer := NewAuthChunkWriter(v2io.NewAdaptiveWriter(content))
	writer.Write(alloc.NewBuffer().Clear().AppendString("abcd"))
	writer.Write(alloc.NewBuffer().Clear())
	writer.Release()

	reader := NewAuthChunkReader(content)
	buffer, err := reader.Read()
	assert.Error(err).IsNil()
	assert.Bytes(buffer.Value).Equals([]byte("abcd"))
}

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

	user := protocol.NewUser(
		protocol.NewID(uuid.New()),
		protocol.UserLevelUntrusted,
		0,
		"[email protected]")

	expectedRequest := &protocol.RequestHeader{
		Version: 1,
		User:    user,
		Command: protocol.RequestCommandTCP,
		Option:  protocol.RequestOption(0),
		Address: v2net.DomainAddress("www.v2ray.com"),
		Port:    v2net.Port(443),
	}

	buffer := alloc.NewBuffer().Clear()
	client := NewClientSession(protocol.DefaultIDHash)
	client.EncodeRequestHeader(expectedRequest, buffer)

	userValidator := protocol.NewTimedUserValidator(protocol.DefaultIDHash)
	userValidator.Add(user)

	server := NewServerSession(userValidator)
	actualRequest, err := server.DecodeRequestHeader(buffer)
	assert.Error(err).IsNil()

	assert.Byte(expectedRequest.Version).Equals(actualRequest.Version)
	assert.Byte(byte(expectedRequest.Command)).Equals(byte(actualRequest.Command))
	assert.Byte(byte(expectedRequest.Option)).Equals(byte(actualRequest.Option))
	netassert.Address(expectedRequest.Address).Equals(actualRequest.Address)
	netassert.Port(expectedRequest.Port).Equals(actualRequest.Port)
}

func (handler *OutboundConnectionHandler) Dispatch(packet v2net.Packet, ray core.OutboundRay) error {
	input := ray.OutboundInput()
	output := ray.OutboundOutput()

	handler.Destination = packet.Destination()
	if packet.Chunk() != nil {
		handler.Data2Send.Write(packet.Chunk().Value)
	}

	go func() {
		for {
			data, open := <-input
			if !open {
				break
			}
			handler.Data2Send.Write(data.Value)
			data.Release()
		}
		response := alloc.NewBuffer()
		response.Clear()
		response.Append(handler.Data2Return)
		output <- response
		close(output)
	}()

	return nil
}

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 NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
	return &BufferedWriter{
		writer: rawWriter,
		buffer: alloc.NewBuffer().Clear(),
		cached: true,
	}
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := alloc.NewBuffer()
		nBytes, addr, err := conn.ReadFromUDP(buffer.Value)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			buffer.Release()
			continue
		}
		log.Info("Client UDP connection from %v", addr)
		request, err := protocol.ReadUDPRequest(buffer.Value[:nBytes])
		buffer.Release()
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			request.Data.Release()
			continue
		}
		if request.Fragment != 0 {
			log.Warning("Dropping fragmented UDP packets.")
			// TODO handle fragments
			request.Data.Release()
			continue
		}

		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
		log.Info("Send packet to %s with %d bytes", udpPacket.Destination().String(), request.Data.Len())
		go server.handlePacket(conn, udpPacket, addr, request.Address)
	}
}

func (handler *VMessInboundHandler) handlePacket(conn *net.UDPConn, request *protocol.VMessRequest, packet v2net.Packet, clientAddr *net.UDPAddr) {
	ray := handler.vPoint.DispatchToOutbound(packet)
	close(ray.InboundInput())

	responseKey := md5.Sum(request.RequestKey)
	responseIV := md5.Sum(request.RequestIV)

	buffer := alloc.NewBuffer().Clear()
	defer buffer.Release()

	responseWriter, err := v2io.NewAesEncryptWriter(responseKey[:], responseIV[:], buffer)
	if err != nil {
		log.Error("VMessIn: Failed to create encrypt writer: %v", err)
		return
	}
	responseWriter.Write(request.ResponseHeader)

	hasData := false

	if data, ok := <-ray.InboundOutput(); ok {
		hasData = true
		responseWriter.Write(data.Value)
		data.Release()
	}

	if hasData {
		conn.WriteToUDP(buffer.Value, clientAddr)
		log.Info("VMessIn sending %d bytes to %s", buffer.Len(), clientAddr.String())
	}
}

func NewBufferedReader(rawReader io.Reader) *BufferedReader {
	return &BufferedReader{
		reader: rawReader,
		buffer: alloc.NewBuffer().Clear(),
		cached: true,
	}
}

func ReadUDPRequest(packet []byte) (request Socks5UDPRequest, err error) {
	// packet[0] and packet[1] are reserved
	request.Fragment = packet[2]

	addrType := packet[3]
	var dataBegin int

	switch addrType {
	case AddrTypeIPv4:
		ip := packet[4:8]
		port := binary.BigEndian.Uint16(packet[8:10])
		request.Address = v2net.IPAddress(ip, port)
		dataBegin = 10
	case AddrTypeIPv6:
		ip := packet[4:20]
		port := binary.BigEndian.Uint16(packet[20:22])
		request.Address = v2net.IPAddress(ip, port)
		dataBegin = 22
	case AddrTypeDomain:
		domainLength := int(packet[4])
		domain := string(packet[5 : 5+domainLength])
		port := binary.BigEndian.Uint16(packet[5+domainLength : 5+domainLength+2])
		request.Address = v2net.DomainAddress(domain, port)
		dataBegin = 5 + domainLength + 2
	default:
		log.Warning("Unknown address type %d", addrType)
		err = ErrorUnknownAddressType
		return
	}

	request.Data = alloc.NewBuffer().Clear().Append(packet[dataBegin:])

	return
}

func (this *AuthChunkReader) Read() (*alloc.Buffer, error) {
	buffer := alloc.NewBuffer()
	if _, err := io.ReadFull(this.reader, buffer.Value[:2]); err != nil {
		buffer.Release()
		return nil, err
	}

	length := serial.BytesLiteral(buffer.Value[:2]).Uint16Value()
	if _, err := io.ReadFull(this.reader, buffer.Value[:length]); err != nil {
		buffer.Release()
		return nil, err
	}
	buffer.Slice(0, int(length))

	fnvHash := fnv.New32a()
	fnvHash.Write(buffer.Value[4:])
	expAuth := serial.BytesLiteral(fnvHash.Sum(nil))
	actualAuth := serial.BytesLiteral(buffer.Value[:4])
	if !actualAuth.Equals(expAuth) {
		buffer.Release()
		return nil, transport.ErrorCorruptedPacket
	}
	buffer.SliceFrom(4)
	return buffer, nil
}

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

	sa := &protocol.CommandSwitchAccount{
		Port:     1234,
		ID:       uuid.New(),
		AlterIds: 1024,
		Level:    128,
		ValidMin: 16,
	}

	buffer := alloc.NewBuffer().Clear()
	err := MarshalCommand(sa, buffer)
	assert.Error(err).IsNil()

	cmd, err := UnmarshalCommand(1, buffer.Value[2:])
	assert.Error(err).IsNil()

	sa2, ok := cmd.(*protocol.CommandSwitchAccount)
	assert.Bool(ok).IsTrue()
	assert.Pointer(sa.Host).IsNil()
	assert.Pointer(sa2.Host).IsNil()
	netassert.Port(sa.Port).Equals(sa2.Port)
	assert.String(sa.ID).Equals(sa2.ID.String())
	assert.Uint16(sa.AlterIds.Value()).Equals(sa2.AlterIds.Value())
	assert.Byte(byte(sa.Level)).Equals(byte(sa2.Level))
	assert.Byte(sa.ValidMin).Equals(sa2.ValidMin)
}

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

	buffer := alloc.NewBuffer().Clear().AppendBytes(6, 1, 0)
	_, _, err := ReadAuthentication(buffer)
	assert.Error(err).Equals(proxy.ErrorInvalidProtocolVersion)
}

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

	buffer := alloc.NewBuffer().Clear().AppendBytes(5, 0)
	_, _, err := ReadAuthentication(buffer)
	assert.Error(err).Equals(proxy.ErrorInvalidAuthentication)
}