Golang typed.NewReadBuffer函数代码示例

// Handles an incoming InitReq.  If we are waiting for the peer to send us an
// InitReq, and the InitReq is valid, send a corresponding InitRes and mark
// ourselves as active
func (c *Connection) handleInitReq(frame *Frame) {
	if err := c.withStateRLock(func() error {
		return nil
	}); err != nil {
		c.connectionError(err)
		return
	}

	var req initReq
	rbuf := typed.NewReadBuffer(frame.SizedPayload())
	if err := req.read(rbuf); err != nil {
		// TODO(mmihic): Technically probably a protocol error
		c.connectionError(err)
		return
	}

	if req.Version != CurrentProtocolVersion {
		// TODO(mmihic): Send protocol error
		c.protocolError(fmt.Errorf("Unsupported protocol version %d from peer", req.Version))
		return
	}

	var ok bool
	if c.remotePeerInfo.HostPort, ok = req.initParams[InitParamHostPort]; !ok {
		c.protocolError(fmt.Errorf("Header %v is required", InitParamHostPort))
		return
	}
	if c.remotePeerInfo.ProcessName, ok = req.initParams[InitParamProcessName]; !ok {
		c.protocolError(fmt.Errorf("Header %v is required", InitParamProcessName))
		return
	}
	if c.remotePeerInfo.IsEphemeral() {
		// TODO(prashant): Add an IsEphemeral bool to the peer info.
		c.remotePeerInfo.HostPort = c.conn.RemoteAddr().String()
	}

	res := initRes{initMessage{id: frame.Header.ID}}
	res.initParams = initParams{
		InitParamHostPort:    c.localPeerInfo.HostPort,
		InitParamProcessName: c.localPeerInfo.ProcessName,
	}
	res.Version = CurrentProtocolVersion
	if err := c.sendMessage(&res); err != nil {
		c.connectionError(err)
		return
	}

	c.withStateLock(func() error {
		switch c.state {
		case connectionWaitingToRecvInitReq:
			c.state = connectionActive
		}

		return nil
	})

	c.callOnActive()
}

func (ch fragmentChannel) recvNextFragment(initial bool) (*readableFragment, error) {
	rbuf := typed.NewReadBuffer(<-ch)
	fragment := new(readableFragment)
	fragment.done = func() {}
	fragment.flags = rbuf.ReadSingleByte()
	fragment.checksumType = ChecksumType(rbuf.ReadSingleByte())
	fragment.checksum = rbuf.ReadBytes(fragment.checksumType.ChecksumSize())
	fragment.contents = rbuf
	return fragment, rbuf.Err()
}

// parseInboundFragment parses an incoming fragment based on the given message
func parseInboundFragment(framePool FramePool, frame *Frame, message message) (*readableFragment, error) {
	rbuf := typed.NewReadBuffer(frame.SizedPayload())
	fragment := new(readableFragment)
	fragment.flags = rbuf.ReadByte()
	if err := message.read(rbuf); err != nil {
		return nil, err
	}

	fragment.checksumType = ChecksumType(rbuf.ReadByte())
	fragment.checksum = rbuf.ReadBytes(fragment.checksumType.ChecksumSize())
	fragment.contents = rbuf
	fragment.done = func() {
		framePool.Release(frame)
	}
	return fragment, rbuf.Err()
}

// handleError andles an error coming back from the peer. If the error is a
// protocol level error, the entire connection will be closed.  If the error is
// a request specific error, it will be written to the request's response
// channel and converted into a SystemError returned from the next reader or
// access call.
func (c *Connection) handleError(frame *Frame) {
	var errorMessage errorMessage
	rbuf := typed.NewReadBuffer(frame.SizedPayload())
	if err := errorMessage.read(rbuf); err != nil {
		c.log.Warnf("Unable to read Error frame from %s: %v", c.remotePeerInfo, err)
		c.connectionError(err)
		return
	}

	if errorMessage.errCode == ErrCodeProtocol {
		c.log.Warnf("Peer %s reported protocol error: %s", c.remotePeerInfo, errorMessage.message)
		c.connectionError(errorMessage.AsSystemError())
		return
	}

	if err := c.outbound.forwardPeerFrame(frame); err != nil {
		c.outbound.removeExchange(frame.Header.ID)
	}
}

// TODO(prashant): Allow typed.ReadBuffer to read directly from the reader.
func readHeaders(r io.Reader) (map[string]string, error) {
	bs, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	buffer := typed.NewReadBuffer(bs)
	numHeaders := buffer.ReadUint16()
	if numHeaders == 0 {
		return nil, nil
	}

	headers := make(map[string]string)
	for i := 0; i < int(numHeaders); i++ {
		k := buffer.ReadLen16String()
		v := buffer.ReadLen16String()
		headers[k] = v
	}
	return headers, buffer.Err()
}

func TestFraming(t *testing.T) {
	fh := FrameHeader{
		size:        uint16(0xFF34),
		messageType: messageTypeCallReq,
		ID:          0xDEADBEEF,
	}

	wbuf := typed.NewWriteBufferWithSize(1024)
	require.Nil(t, fh.write(wbuf))

	var b bytes.Buffer
	if _, err := wbuf.FlushTo(&b); err != nil {
		require.Nil(t, err)
	}

	rbuf := typed.NewReadBuffer(b.Bytes())

	var fh2 FrameHeader
	require.Nil(t, fh2.read(rbuf))

	assert.Equal(t, fh, fh2)
}