Golang alloc.NewBuffer函数代码示例

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 TestTCPRequest(t *testing.T) {
	assert := assert.On(t)

	request := &protocol.RequestHeader{
		Version: Version,
		Command: protocol.RequestCommandTCP,
		Address: v2net.LocalHostIP,
		Option:  RequestOptionOneTimeAuth,
		Port:    1234,
		User: &protocol.User{
			Email: "[email protected]",
			Account: loader.NewTypedSettings(&Account{
				Password:   "tcp-password",
				CipherType: CipherType_CHACHA20,
			}),
		},
	}

	data := alloc.NewLocalBuffer(256).Clear().AppendString("test string")
	cache := alloc.NewBuffer().Clear()

	writer, err := WriteTCPRequest(request, cache)
	assert.Error(err).IsNil()

	writer.Write(data)

	decodedRequest, reader, err := ReadTCPSession(request.User, cache)
	assert.Error(err).IsNil()
	assert.Address(decodedRequest.Address).Equals(request.Address)
	assert.Port(decodedRequest.Port).Equals(request.Port)

	decodedData, err := reader.Read()
	assert.Error(err).IsNil()
	assert.Bytes(decodedData.Value).Equals([]byte("test string"))
}

func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
	return &BufferedWriter{
		writer: rawWriter,
		buffer: alloc.NewBuffer().Clear(),
		cached: true,
	}
}

func (this *ChainWriter) Write(payload []byte) (int, error) {
	this.Lock()
	defer this.Unlock()
	if this.writer == nil {
		return 0, io.ErrClosedPipe
	}

	bytesWritten := 0
	size := len(payload)
	for size > 0 {
		buffer := alloc.NewBuffer().Clear()
		if size > alloc.BufferSize {
			buffer.Append(payload[:alloc.BufferSize])
			size -= alloc.BufferSize
			payload = payload[alloc.BufferSize:]
			bytesWritten += alloc.BufferSize
		} else {
			buffer.Append(payload)
			bytesWritten += size
			size = 0
		}
		err := this.writer.Write(buffer)
		if err != nil {
			return bytesWritten, err
		}
	}

	return bytesWritten, nil
}

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

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

func (this *ChunkReader) Read() (*alloc.Buffer, error) {
	buffer := alloc.NewBuffer()
	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 length > alloc.BufferSize {
		// Theoretically the size of a chunk is 64K, but most Shadowsocks implementations used <4K buffer.
		buffer.Release()
		buffer = alloc.NewLocalBuffer(int(length) + 128)
	}
	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()
		return nil, errors.New("Shadowsocks|AuthenticationReader: Invalid auth.")
	}
	buffer.SliceFrom(AuthSize)

	return buffer, nil
}

func TestSwitchAccount(t *testing.T) {
	assert := assert.On(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()
	assert.Port(sa.Port).Equals(sa2.Port)
	assert.String(sa.ID.String()).Equals(sa2.ID.String())
	assert.Uint16(sa.AlterIds).Equals(sa2.AlterIds)
	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.ErrInvalidProtocolVersion)
}

func TestBufferedReader(t *testing.T) {
	assert := assert.On(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 NewBufferedReader(rawReader io.Reader) *BufferedReader {
	return &BufferedReader{
		reader: rawReader,
		buffer: alloc.NewBuffer().Clear(),
		cached: true,
	}
}

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))
	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})
	assert.Port(request.Port).Equals(8)
}

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

	buffer := alloc.NewBuffer().Clear()
	buffer.AppendBytes(1, 2, 3, 4)
	Authenticate(buffer)
	assert.Bytes(buffer.Value).Equals([]byte{0, 8, 87, 52, 168, 125, 1, 2, 3, 4})

	b2, err := NewAuthChunkReader(buffer).Read()
	assert.Error(err).IsNil()
	assert.Bytes(b2.Value).Equals([]byte{1, 2, 3, 4})
}

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

	user := &protocol.User{
		Account: loader.NewTypedSettings(&Account{
			Password:   "test-password",
			CipherType: CipherType_CHACHA20_IEFT,
		}),
	}
	cache := alloc.NewBuffer().Clear()
	writer := &UDPWriter{
		Writer: cache,
		Request: &protocol.RequestHeader{
			Version: Version,
			Address: v2net.DomainAddress("v2ray.com"),
			Port:    123,
			User:    user,
			Option:  RequestOptionOneTimeAuth,
		},
	}

	reader := &UDPReader{
		Reader: cache,
		User:   user,
	}

	err := writer.Write(alloc.NewBuffer().Clear().AppendString("test payload"))
	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"))
	assert.Error(err).IsNil()

	payload, err = reader.Read()
	assert.Error(err).IsNil()
	assert.String(payload.String()).Equals("test payload 2")
}

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

	tcpServer := &tcp.Server{
		Port: v2net.Port(50052),
		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()
	defer tcpServer.Close()

	assert.Error(InitializeServerSetOnce("test_6")).IsNil()

	for i := 0; i < 1; i++ {
		conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
			IP:   []byte{127, 0, 0, 1},
			Port: 50050,
		})
		assert.Error(err).IsNil()

		payload := "dokodemo request."
		nBytes, err := conn.Write([]byte(payload))
		assert.Error(err).IsNil()
		assert.Int(nBytes).Equals(len(payload))

		conn.CloseWrite()

		response := alloc.NewBuffer().Clear()
		finished := false
		for {
			_, err := response.FillFrom(conn)
			assert.Error(err).IsNil()
			if err != nil {
				break
			}
			if response.String() == "Processed: "+payload {
				finished = true
				break
			}
		}
		assert.Bool(finished).IsTrue()

		conn.Close()
	}

	CloseAllServers()
}

func NewBuffer() *Buffer {
	b := &Buffer{
		next:     0,
		released: 0,
		hold:     true,
		buffer:   alloc.NewBuffer(),
	}
	for idx := range b.distro {
		content := b.buffer.Value[idx*DistroSize : (idx+1)*DistroSize]
		b.distro[idx] = alloc.CreateBuffer(content, b)
	}
	return b
}