github.com/ylsGit/go-ethereum@v1.6.5/p2p/message.go (about)

     1  // Copyright 2014 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum 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-ethereum 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-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package p2p
    18  
    19  import (
    20  	"bytes"
    21  	"errors"
    22  	"fmt"
    23  	"io"
    24  	"io/ioutil"
    25  	"net"
    26  	"sync"
    27  	"sync/atomic"
    28  	"time"
    29  
    30  	"github.com/ethereum/go-ethereum/rlp"
    31  )
    32  
    33  // Msg defines the structure of a p2p message.
    34  //
    35  // Note that a Msg can only be sent once since the Payload reader is
    36  // consumed during sending. It is not possible to create a Msg and
    37  // send it any number of times. If you want to reuse an encoded
    38  // structure, encode the payload into a byte array and create a
    39  // separate Msg with a bytes.Reader as Payload for each send.
    40  type Msg struct {
    41  	Code       uint64
    42  	Size       uint32 // size of the paylod
    43  	Payload    io.Reader
    44  	ReceivedAt time.Time
    45  }
    46  
    47  // Decode parses the RLP content of a message into
    48  // the given value, which must be a pointer.
    49  //
    50  // For the decoding rules, please see package rlp.
    51  func (msg Msg) Decode(val interface{}) error {
    52  	s := rlp.NewStream(msg.Payload, uint64(msg.Size))
    53  	if err := s.Decode(val); err != nil {
    54  		return newPeerError(errInvalidMsg, "(code %x) (size %d) %v", msg.Code, msg.Size, err)
    55  	}
    56  	return nil
    57  }
    58  
    59  func (msg Msg) String() string {
    60  	return fmt.Sprintf("msg #%v (%v bytes)", msg.Code, msg.Size)
    61  }
    62  
    63  // Discard reads any remaining payload data into a black hole.
    64  func (msg Msg) Discard() error {
    65  	_, err := io.Copy(ioutil.Discard, msg.Payload)
    66  	return err
    67  }
    68  
    69  type MsgReader interface {
    70  	ReadMsg() (Msg, error)
    71  }
    72  
    73  type MsgWriter interface {
    74  	// WriteMsg sends a message. It will block until the message's
    75  	// Payload has been consumed by the other end.
    76  	//
    77  	// Note that messages can be sent only once because their
    78  	// payload reader is drained.
    79  	WriteMsg(Msg) error
    80  }
    81  
    82  // MsgReadWriter provides reading and writing of encoded messages.
    83  // Implementations should ensure that ReadMsg and WriteMsg can be
    84  // called simultaneously from multiple goroutines.
    85  type MsgReadWriter interface {
    86  	MsgReader
    87  	MsgWriter
    88  }
    89  
    90  // Send writes an RLP-encoded message with the given code.
    91  // data should encode as an RLP list.
    92  func Send(w MsgWriter, msgcode uint64, data interface{}) error {
    93  	size, r, err := rlp.EncodeToReader(data)
    94  	if err != nil {
    95  		return err
    96  	}
    97  	return w.WriteMsg(Msg{Code: msgcode, Size: uint32(size), Payload: r})
    98  }
    99  
   100  // SendItems writes an RLP with the given code and data elements.
   101  // For a call such as:
   102  //
   103  //    SendItems(w, code, e1, e2, e3)
   104  //
   105  // the message payload will be an RLP list containing the items:
   106  //
   107  //    [e1, e2, e3]
   108  //
   109  func SendItems(w MsgWriter, msgcode uint64, elems ...interface{}) error {
   110  	return Send(w, msgcode, elems)
   111  }
   112  
   113  // netWrapper wraps a MsgReadWriter with locks around
   114  // ReadMsg/WriteMsg and applies read/write deadlines.
   115  type netWrapper struct {
   116  	rmu, wmu sync.Mutex
   117  
   118  	rtimeout, wtimeout time.Duration
   119  	conn               net.Conn
   120  	wrapped            MsgReadWriter
   121  }
   122  
   123  func (rw *netWrapper) ReadMsg() (Msg, error) {
   124  	rw.rmu.Lock()
   125  	defer rw.rmu.Unlock()
   126  	rw.conn.SetReadDeadline(time.Now().Add(rw.rtimeout))
   127  	return rw.wrapped.ReadMsg()
   128  }
   129  
   130  func (rw *netWrapper) WriteMsg(msg Msg) error {
   131  	rw.wmu.Lock()
   132  	defer rw.wmu.Unlock()
   133  	rw.conn.SetWriteDeadline(time.Now().Add(rw.wtimeout))
   134  	return rw.wrapped.WriteMsg(msg)
   135  }
   136  
   137  // eofSignal wraps a reader with eof signaling. the eof channel is
   138  // closed when the wrapped reader returns an error or when count bytes
   139  // have been read.
   140  type eofSignal struct {
   141  	wrapped io.Reader
   142  	count   uint32 // number of bytes left
   143  	eof     chan<- struct{}
   144  }
   145  
   146  // note: when using eofSignal to detect whether a message payload
   147  // has been read, Read might not be called for zero sized messages.
   148  func (r *eofSignal) Read(buf []byte) (int, error) {
   149  	if r.count == 0 {
   150  		if r.eof != nil {
   151  			r.eof <- struct{}{}
   152  			r.eof = nil
   153  		}
   154  		return 0, io.EOF
   155  	}
   156  
   157  	max := len(buf)
   158  	if int(r.count) < len(buf) {
   159  		max = int(r.count)
   160  	}
   161  	n, err := r.wrapped.Read(buf[:max])
   162  	r.count -= uint32(n)
   163  	if (err != nil || r.count == 0) && r.eof != nil {
   164  		r.eof <- struct{}{} // tell Peer that msg has been consumed
   165  		r.eof = nil
   166  	}
   167  	return n, err
   168  }
   169  
   170  // MsgPipe creates a message pipe. Reads on one end are matched
   171  // with writes on the other. The pipe is full-duplex, both ends
   172  // implement MsgReadWriter.
   173  func MsgPipe() (*MsgPipeRW, *MsgPipeRW) {
   174  	var (
   175  		c1, c2  = make(chan Msg), make(chan Msg)
   176  		closing = make(chan struct{})
   177  		closed  = new(int32)
   178  		rw1     = &MsgPipeRW{c1, c2, closing, closed}
   179  		rw2     = &MsgPipeRW{c2, c1, closing, closed}
   180  	)
   181  	return rw1, rw2
   182  }
   183  
   184  // ErrPipeClosed is returned from pipe operations after the
   185  // pipe has been closed.
   186  var ErrPipeClosed = errors.New("p2p: read or write on closed message pipe")
   187  
   188  // MsgPipeRW is an endpoint of a MsgReadWriter pipe.
   189  type MsgPipeRW struct {
   190  	w       chan<- Msg
   191  	r       <-chan Msg
   192  	closing chan struct{}
   193  	closed  *int32
   194  }
   195  
   196  // WriteMsg sends a messsage on the pipe.
   197  // It blocks until the receiver has consumed the message payload.
   198  func (p *MsgPipeRW) WriteMsg(msg Msg) error {
   199  	if atomic.LoadInt32(p.closed) == 0 {
   200  		consumed := make(chan struct{}, 1)
   201  		msg.Payload = &eofSignal{msg.Payload, msg.Size, consumed}
   202  		select {
   203  		case p.w <- msg:
   204  			if msg.Size > 0 {
   205  				// wait for payload read or discard
   206  				select {
   207  				case <-consumed:
   208  				case <-p.closing:
   209  				}
   210  			}
   211  			return nil
   212  		case <-p.closing:
   213  		}
   214  	}
   215  	return ErrPipeClosed
   216  }
   217  
   218  // ReadMsg returns a message sent on the other end of the pipe.
   219  func (p *MsgPipeRW) ReadMsg() (Msg, error) {
   220  	if atomic.LoadInt32(p.closed) == 0 {
   221  		select {
   222  		case msg := <-p.r:
   223  			return msg, nil
   224  		case <-p.closing:
   225  		}
   226  	}
   227  	return Msg{}, ErrPipeClosed
   228  }
   229  
   230  // Close unblocks any pending ReadMsg and WriteMsg calls on both ends
   231  // of the pipe. They will return ErrPipeClosed. Close also
   232  // interrupts any reads from a message payload.
   233  func (p *MsgPipeRW) Close() error {
   234  	if atomic.AddInt32(p.closed, 1) != 1 {
   235  		// someone else is already closing
   236  		atomic.StoreInt32(p.closed, 1) // avoid overflow
   237  		return nil
   238  	}
   239  	close(p.closing)
   240  	return nil
   241  }
   242  
   243  // ExpectMsg reads a message from r and verifies that its
   244  // code and encoded RLP content match the provided values.
   245  // If content is nil, the payload is discarded and not verified.
   246  func ExpectMsg(r MsgReader, code uint64, content interface{}) error {
   247  	msg, err := r.ReadMsg()
   248  	if err != nil {
   249  		return err
   250  	}
   251  	if msg.Code != code {
   252  		return fmt.Errorf("message code mismatch: got %d, expected %d", msg.Code, code)
   253  	}
   254  	if content == nil {
   255  		return msg.Discard()
   256  	} else {
   257  		contentEnc, err := rlp.EncodeToBytes(content)
   258  		if err != nil {
   259  			panic("content encode error: " + err.Error())
   260  		}
   261  		if int(msg.Size) != len(contentEnc) {
   262  			return fmt.Errorf("message size mismatch: got %d, want %d", msg.Size, len(contentEnc))
   263  		}
   264  		actualContent, err := ioutil.ReadAll(msg.Payload)
   265  		if err != nil {
   266  			return err
   267  		}
   268  		if !bytes.Equal(actualContent, contentEnc) {
   269  			return fmt.Errorf("message payload mismatch:\ngot:  %x\nwant: %x", actualContent, contentEnc)
   270  		}
   271  	}
   272  	return nil
   273  }