github.com/n1ghtfa1l/go-vnt@v0.6.4-alpha.6/vntp2p/protocol.go (about)

     1  // Copyright 2019 The go-vnt Authors
     2  // This file is part of the go-vnt library.
     3  //
     4  // The go-vnt library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-vnt library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-vnt library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package vntp2p
    18  
    19  import (
    20  	"encoding/binary"
    21  	"encoding/json"
    22  	"io"
    23  	"time"
    24  
    25  	inet "github.com/libp2p/go-libp2p-net"
    26  	libp2p "github.com/libp2p/go-libp2p-peer"
    27  	"github.com/vntchain/go-vnt/log"
    28  	"github.com/vntchain/go-vnt/rlp"
    29  )
    30  
    31  // 目前依然沿用原有的子协议结构，减少上层的改动
    32  type Protocol struct {
    33  	Name     string
    34  	Version  uint
    35  	Length   uint64
    36  	Run      func(peer *Peer, rw MsgReadWriter) error
    37  	NodeInfo func() interface{}
    38  	PeerInfo func(id libp2p.ID) interface{}
    39  }
    40  
    41  // HandleStream handle all message which is from anywhere
    42  // 主、被动连接都走的流程
    43  func (server *Server) HandleStream(s inet.Stream) {
    44  	// if peer is blacklisted, ignore it
    45  	if blacklist.exists(s.Conn().RemotePeer()) {
    46  		log.Trace("HandleStream: related peer is blacklisted", "pid", s.Conn().RemotePeer())
    47  		s.Conn().Close()
    48  		return
    49  	}
    50  
    51  	// peer信息只获取1次即可
    52  	log.Debug("Stream data coming...")
    53  	peer := server.GetPeerByRemoteID(s)
    54  	if peer == nil {
    55  		log.Debug("HandleStream", "remotePeerID", s.Conn().RemotePeer(), "this remote peer is nil, don't handle it")
    56  		_ = s.Reset()
    57  		return
    58  	}
    59  
    60  	// 发生错误时才会退出
    61  	defer func() {
    62  		peer.log.Debug("HandleStream reset stream before exit")
    63  		peer.Reset()
    64  	}()
    65  
    66  	// stream未关闭则连接正常可持续读取消息
    67  	for {
    68  		// 读取消息
    69  		msgHeaderByte := make([]byte, MessageHeaderLength)
    70  		_, err := io.ReadFull(s, msgHeaderByte)
    71  		if err != nil {
    72  			peer.log.Warn("HandleStream", "read msg header error", err)
    73  			notifyError(peer.msgers, err)
    74  			return
    75  		}
    76  		bodySize := binary.LittleEndian.Uint32(msgHeaderByte)
    77  
    78  		msgBodyByte := make([]byte, bodySize)
    79  		_, err = io.ReadFull(s, msgBodyByte)
    80  		if err != nil {
    81  			peer.log.Warn("HandleStream", "read msg Body error", err)
    82  			notifyError(peer.msgers, err)
    83  			return
    84  		}
    85  		msgBody := &MsgBody{Payload: &rlp.EncReader{}}
    86  		err = json.Unmarshal(msgBodyByte, msgBody)
    87  		if err != nil {
    88  			peer.log.Warn("HandleStream", "unmarshal msg Body error", err)
    89  			notifyError(peer.msgers, err)
    90  			return
    91  		}
    92  		msgBody.ReceivedAt = time.Now()
    93  
    94  		// 传递给msger
    95  		var msgHeader MsgHeader
    96  		copy(msgHeader[:], msgHeaderByte)
    97  
    98  		msg := Msg{
    99  			Header: msgHeader,
   100  			Body:   *msgBody,
   101  		}
   102  		if msger, ok := peer.msgers[msgBody.ProtocolID]; ok { // this node support protocolID
   103  			// 非阻塞向上层协议传递消息，如果2s还未被读取，认为上层协议有故障
   104  			select {
   105  			case msger.in <- msg:
   106  				peer.log.Trace("HandleStream send message to messager success")
   107  			case <-time.NewTimer(time.Second * 2).C:
   108  				peer.log.Trace("HandleStream send message to messager timeout")
   109  			}
   110  		} else {
   111  			peer.log.Warn("HandleStream", "receive unknown message", msg)
   112  		}
   113  	}
   114  }
   115  
   116  func notifyError(msgers map[string]*VNTMsger, err error) {
   117  	for _, m := range msgers {
   118  		m.err <- err
   119  	}
   120  }