Golang message.NetMessage類代碼示例

// HandleMessage marshals and unmarshals net messages, forwarding them to the
// BitSwapMessage receiver
func (bsnet *impl) HandleMessage(
	ctx context.Context, incoming netmsg.NetMessage) netmsg.NetMessage {

	if bsnet.receiver == nil {
		return nil
	}

	received, err := bsmsg.FromNet(incoming)
	if err != nil {
		go bsnet.receiver.ReceiveError(err)
		return nil
	}

	p, bsmsg := bsnet.receiver.ReceiveMessage(ctx, incoming.Peer(), received)

	// TODO(brian): put this in a helper function
	if bsmsg == nil || p == nil {
		bsnet.receiver.ReceiveError(errors.New("ReceiveMessage returned nil peer or message"))
		return nil
	}

	outgoing, err := bsmsg.ToNet(p)
	if err != nil {
		go bsnet.receiver.ReceiveError(err)
		return nil
	}

	log.Debugf("Message size: %d", len(outgoing.Data()))
	return outgoing
}

func FromNet(nmsg netmsg.NetMessage) (BitSwapMessage, error) {
	pb := new(pb.Message)
	if err := proto.Unmarshal(nmsg.Data(), pb); err != nil {
		return nil, err
	}
	m := newMessageFromProto(*pb)
	return m, nil
}

// SendRequest sends a request message out and awaits a response.
func (s *service) SendRequest(ctx context.Context, m msg.NetMessage) (msg.NetMessage, error) {

	// check if we should bail given our contexts
	select {
	default:
	case <-s.Closing():
		return nil, fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// create a request
	r, err := NewRequest(m.Peer().ID())
	if err != nil {
		return nil, err
	}

	// register Request
	s.RequestsLock.Lock()
	s.Requests[r.Key()] = r
	s.RequestsLock.Unlock()

	// defer deleting this request
	defer func() {
		s.RequestsLock.Lock()
		delete(s.Requests, r.Key())
		s.RequestsLock.Unlock()
	}()

	// check if we should bail after waiting for mutex
	select {
	default:
	case <-s.Closing():
		return nil, fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// Send message
	s.sendMessage(ctx, m, r.ID)

	// wait for response
	m = nil
	err = nil
	select {
	case m = <-r.Response:
	case <-s.Closed():
		err = fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		err = ctx.Err()
	}

	if m == nil {
		return nil, ErrNoResponse
	}

	return m, err
}

func (t *ReverseHandler) HandleMessage(ctx context.Context, m msg.NetMessage) msg.NetMessage {

	d := m.Data()
	for i, j := 0, len(d)-1; i < j; i, j = i+1, j-1 {
		d[i], d[j] = d[j], d[i]
	}

	return msg.New(m.Peer(), d)
}

// HandleMessage implements the inet.Handler interface.
func (dht *IpfsDHT) HandleMessage(ctx context.Context, mes msg.NetMessage) msg.NetMessage {

	mData := mes.Data()
	if mData == nil {
		log.Error("Message contained nil data.")
		return nil
	}

	mPeer := mes.Peer()
	if mPeer == nil {
		log.Error("Message contained nil peer.")
		return nil
	}

	// deserialize msg
	pmes := new(pb.Message)
	err := proto.Unmarshal(mData, pmes)
	if err != nil {
		log.Error("Error unmarshaling data")
		return nil
	}

	// update the peer (on valid msgs only)
	dht.Update(ctx, mPeer)

	log.Event(ctx, "foo", dht.self, mPeer, pmes)

	// get handler for this msg type.
	handler := dht.handlerForMsgType(pmes.GetType())
	if handler == nil {
		log.Error("got back nil handler from handlerForMsgType")
		return nil
	}

	// dispatch handler.
	rpmes, err := handler(mPeer, pmes)
	if err != nil {
		log.Errorf("handle message error: %s", err)
		return nil
	}

	// if nil response, return it before serializing
	if rpmes == nil {
		log.Warning("Got back nil response from request.")
		return nil
	}

	// serialize response msg
	rmes, err := msg.FromObject(mPeer, rpmes)
	if err != nil {
		log.Errorf("serialze response error: %s", err)
		return nil
	}

	return rmes
}

// handleIncomingMessage routes message to the appropriate protocol.
func (m *Muxer) handleIncomingMessage(m1 msg.NetMessage) {
	defer m.Children().Done()

	m.bwiLock.Lock()
	// TODO: compensate for overhead
	m.bwIn += uint64(len(m1.Data()))
	m.bwiLock.Unlock()

	data, pid, err := unwrapData(m1.Data())
	if err != nil {
		log.Errorf("muxer de-serializing error: %v", err)
		return
	}
	conn.ReleaseBuffer(m1.Data())

	m2 := msg.New(m1.Peer(), data)
	proto, found := m.Protocols[pid]
	if !found {
		log.Errorf("muxer unknown protocol %v", pid)
		return
	}

	select {
	case proto.GetPipe().Incoming <- m2:
	case <-m.Closing():
		return
	}
}

func testWrappedMsg(t *testing.T, m msg.NetMessage, pid pb.ProtocolID, data []byte) {
	data2, pid2, err := unwrapData(m.Data())
	if err != nil {
		t.Error(err)
	}

	if pid != pid2 {
		t.Errorf("ProtocolIDs do not match: %v != %v", pid, pid2)
	}

	if !bytes.Equal(data, data2) {
		t.Errorf("Data does not match: %v != %v", data, data2)
	}
}

// sendMessage sends a message out (actual leg work. SendMessage is to export w/o rid)
func (s *service) sendMessage(ctx context.Context, m msg.NetMessage, rid RequestID) error {

	// serialize ServiceMessage wrapper
	data, err := wrapData(m.Data(), rid)
	if err != nil {
		return err
	}

	// log.Debug("Service send message [to = %s]", m.Peer())

	// send message
	m2 := msg.New(m.Peer(), data)
	select {
	case s.Outgoing <- m2:
	case <-ctx.Done():
		return ctx.Err()
	}

	return nil
}

// handleOutgoingMessage wraps out a message and sends it out the
func (m *Muxer) handleOutgoingMessage(pid pb.ProtocolID, m1 msg.NetMessage) {
	defer m.Children().Done()

	data, err := wrapData(m1.Data(), pid)
	if err != nil {
		log.Errorf("muxer serializing error: %v", err)
		return
	}

	m.bwoLock.Lock()
	// TODO: compensate for overhead
	// TODO(jbenet): switch this to a goroutine to prevent sync waiting.
	m.bwOut += uint64(len(data))
	m.bwoLock.Unlock()

	m2 := msg.New(m1.Peer(), data)
	select {
	case m.GetPipe().Outgoing <- m2:
	case <-m.Closing():
		return
	}
}

func (s *service) handleIncomingMessage(m msg.NetMessage) {
	defer s.Children().Done()

	// unwrap the incoming message
	data, rid, err := unwrapData(m.Data())
	if err != nil {
		log.Errorf("service de-serializing error: %v", err)
		return
	}

	m2 := msg.New(m.Peer(), data)

	// if it's a request (or has no RequestID), handle it
	if rid == nil || rid.IsRequest() {
		handler := s.GetHandler()
		if handler == nil {
			log.Errorf("service dropped msg: %v", m)
			return // no handler, drop it.
		}

		// should this be "go HandleMessage ... ?"
		r1 := handler.HandleMessage(s.Context(), m2)

		// if handler gave us a response, send it back out!
		if r1 != nil {
			err := s.sendMessage(s.Context(), r1, rid.Response())
			if err != nil {
				log.Errorf("error sending response message: %v", err)
			}
		}
		return
	}

	// Otherwise, it is a response. handle it.
	if !rid.IsResponse() {
		log.Errorf("RequestID should identify a response here.")
	}

	key := RequestKey(m.Peer().ID(), RequestID(rid))
	s.RequestsLock.RLock()
	r, found := s.Requests[key]
	s.RequestsLock.RUnlock()

	if !found {
		log.Errorf("no request key %v (timeout?)", []byte(key))
		return
	}

	select {
	case r.Response <- m2:
	case <-s.Closing():
	}
}

func testMsg(t *testing.T, m msg.NetMessage, data []byte) {
	if !bytes.Equal(data, m.Data()) {
		t.Errorf("Data does not match: %v != %v", data, m.Data())
	}
}