github.com/sunjiahui/go-ethereum@v1.10.31/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 "sync/atomic" 25 "time" 26 27 "github.com/sunjiahui/go-ethereum/event" 28 "github.com/sunjiahui/go-ethereum/p2p/enode" 29 "github.com/sunjiahui/go-ethereum/rlp" 30 ) 31 32 // Msg defines the structure of a p2p message. 33 // 34 // Note that a Msg can only be sent once since the Payload reader is 35 // consumed during sending. It is not possible to create a Msg and 36 // send it any number of times. If you want to reuse an encoded 37 // structure, encode the payload into a byte array and create a 38 // separate Msg with a bytes.Reader as Payload for each send. 39 type Msg struct { 40 Code uint64 41 Size uint32 // Size of the raw payload 42 Payload io.Reader 43 ReceivedAt time.Time 44 45 meterCap Cap // Protocol name and version for egress metering 46 meterCode uint64 // Message within protocol for egress metering 47 meterSize uint32 // Compressed message size for ingress metering 48 } 49 50 // Decode parses the RLP content of a message into 51 // the given value, which must be a pointer. 52 // 53 // For the decoding rules, please see package rlp. 54 func (msg Msg) Decode(val interface{}) error { 55 s := rlp.NewStream(msg.Payload, uint64(msg.Size)) 56 if err := s.Decode(val); err != nil { 57 return newPeerError(errInvalidMsg, "(code %x) (size %d) %v", msg.Code, msg.Size, err) 58 } 59 return nil 60 } 61 62 func (msg Msg) String() string { 63 return fmt.Sprintf("msg #%v (%v bytes)", msg.Code, msg.Size) 64 } 65 66 // Discard reads any remaining payload data into a black hole. 67 func (msg Msg) Discard() error { 68 _, err := io.Copy(io.Discard, msg.Payload) 69 return err 70 } 71 72 func (msg Msg) Time() time.Time { 73 return msg.ReceivedAt 74 } 75 76 type MsgReader interface { 77 ReadMsg() (Msg, error) 78 } 79 80 type MsgWriter interface { 81 // WriteMsg sends a message. It will block until the message's 82 // Payload has been consumed by the other end. 83 // 84 // Note that messages can be sent only once because their 85 // payload reader is drained. 86 WriteMsg(Msg) error 87 } 88 89 // MsgReadWriter provides reading and writing of encoded messages. 90 // Implementations should ensure that ReadMsg and WriteMsg can be 91 // called simultaneously from multiple goroutines. 92 type MsgReadWriter interface { 93 MsgReader 94 MsgWriter 95 } 96 97 // Send writes an RLP-encoded message with the given code. 98 // data should encode as an RLP list. 99 func Send(w MsgWriter, msgcode uint64, data interface{}) error { 100 size, r, err := rlp.EncodeToReader(data) 101 if err != nil { 102 return err 103 } 104 return w.WriteMsg(Msg{Code: msgcode, Size: uint32(size), Payload: r}) 105 } 106 107 // SendItems writes an RLP with the given code and data elements. 108 // For a call such as: 109 // 110 // SendItems(w, code, e1, e2, e3) 111 // 112 // the message payload will be an RLP list containing the items: 113 // 114 // [e1, e2, e3] 115 func SendItems(w MsgWriter, msgcode uint64, elems ...interface{}) error { 116 return Send(w, msgcode, elems) 117 } 118 119 // eofSignal wraps a reader with eof signaling. the eof channel is 120 // closed when the wrapped reader returns an error or when count bytes 121 // have been read. 122 type eofSignal struct { 123 wrapped io.Reader 124 count uint32 // number of bytes left 125 eof chan<- struct{} 126 } 127 128 // note: when using eofSignal to detect whether a message payload 129 // has been read, Read might not be called for zero sized messages. 130 func (r *eofSignal) Read(buf []byte) (int, error) { 131 if r.count == 0 { 132 if r.eof != nil { 133 r.eof <- struct{}{} 134 r.eof = nil 135 } 136 return 0, io.EOF 137 } 138 139 max := len(buf) 140 if int(r.count) < len(buf) { 141 max = int(r.count) 142 } 143 n, err := r.wrapped.Read(buf[:max]) 144 r.count -= uint32(n) 145 if (err != nil || r.count == 0) && r.eof != nil { 146 r.eof <- struct{}{} // tell Peer that msg has been consumed 147 r.eof = nil 148 } 149 return n, err 150 } 151 152 // MsgPipe creates a message pipe. Reads on one end are matched 153 // with writes on the other. The pipe is full-duplex, both ends 154 // implement MsgReadWriter. 155 func MsgPipe() (*MsgPipeRW, *MsgPipeRW) { 156 var ( 157 c1, c2 = make(chan Msg), make(chan Msg) 158 closing = make(chan struct{}) 159 closed = new(int32) 160 rw1 = &MsgPipeRW{c1, c2, closing, closed} 161 rw2 = &MsgPipeRW{c2, c1, closing, closed} 162 ) 163 return rw1, rw2 164 } 165 166 // ErrPipeClosed is returned from pipe operations after the 167 // pipe has been closed. 168 var ErrPipeClosed = errors.New("p2p: read or write on closed message pipe") 169 170 // MsgPipeRW is an endpoint of a MsgReadWriter pipe. 171 type MsgPipeRW struct { 172 w chan<- Msg 173 r <-chan Msg 174 closing chan struct{} 175 closed *int32 176 } 177 178 // WriteMsg sends a message on the pipe. 179 // It blocks until the receiver has consumed the message payload. 180 func (p *MsgPipeRW) WriteMsg(msg Msg) error { 181 if atomic.LoadInt32(p.closed) == 0 { 182 consumed := make(chan struct{}, 1) 183 msg.Payload = &eofSignal{msg.Payload, msg.Size, consumed} 184 select { 185 case p.w <- msg: 186 if msg.Size > 0 { 187 // wait for payload read or discard 188 select { 189 case <-consumed: 190 case <-p.closing: 191 } 192 } 193 return nil 194 case <-p.closing: 195 } 196 } 197 return ErrPipeClosed 198 } 199 200 // ReadMsg returns a message sent on the other end of the pipe. 201 func (p *MsgPipeRW) ReadMsg() (Msg, error) { 202 if atomic.LoadInt32(p.closed) == 0 { 203 select { 204 case msg := <-p.r: 205 return msg, nil 206 case <-p.closing: 207 } 208 } 209 return Msg{}, ErrPipeClosed 210 } 211 212 // Close unblocks any pending ReadMsg and WriteMsg calls on both ends 213 // of the pipe. They will return ErrPipeClosed. Close also 214 // interrupts any reads from a message payload. 215 func (p *MsgPipeRW) Close() error { 216 if atomic.AddInt32(p.closed, 1) != 1 { 217 // someone else is already closing 218 atomic.StoreInt32(p.closed, 1) // avoid overflow 219 return nil 220 } 221 close(p.closing) 222 return nil 223 } 224 225 // ExpectMsg reads a message from r and verifies that its 226 // code and encoded RLP content match the provided values. 227 // If content is nil, the payload is discarded and not verified. 228 func ExpectMsg(r MsgReader, code uint64, content interface{}) error { 229 msg, err := r.ReadMsg() 230 if err != nil { 231 return err 232 } 233 if msg.Code != code { 234 return fmt.Errorf("message code mismatch: got %d, expected %d", msg.Code, code) 235 } 236 if content == nil { 237 return msg.Discard() 238 } 239 contentEnc, err := rlp.EncodeToBytes(content) 240 if err != nil { 241 panic("content encode error: " + err.Error()) 242 } 243 if int(msg.Size) != len(contentEnc) { 244 return fmt.Errorf("message size mismatch: got %d, want %d", msg.Size, len(contentEnc)) 245 } 246 actualContent, err := io.ReadAll(msg.Payload) 247 if err != nil { 248 return err 249 } 250 if !bytes.Equal(actualContent, contentEnc) { 251 return fmt.Errorf("message payload mismatch:\ngot: %x\nwant: %x", actualContent, contentEnc) 252 } 253 return nil 254 } 255 256 // msgEventer wraps a MsgReadWriter and sends events whenever a message is sent 257 // or received 258 type msgEventer struct { 259 MsgReadWriter 260 261 feed *event.Feed 262 peerID enode.ID 263 Protocol string 264 localAddress string 265 remoteAddress string 266 } 267 268 // newMsgEventer returns a msgEventer which sends message events to the given 269 // feed 270 func newMsgEventer(rw MsgReadWriter, feed *event.Feed, peerID enode.ID, proto, remote, local string) *msgEventer { 271 return &msgEventer{ 272 MsgReadWriter: rw, 273 feed: feed, 274 peerID: peerID, 275 Protocol: proto, 276 remoteAddress: remote, 277 localAddress: local, 278 } 279 } 280 281 // ReadMsg reads a message from the underlying MsgReadWriter and emits a 282 // "message received" event 283 func (ev *msgEventer) ReadMsg() (Msg, error) { 284 msg, err := ev.MsgReadWriter.ReadMsg() 285 if err != nil { 286 return msg, err 287 } 288 ev.feed.Send(&PeerEvent{ 289 Type: PeerEventTypeMsgRecv, 290 Peer: ev.peerID, 291 Protocol: ev.Protocol, 292 MsgCode: &msg.Code, 293 MsgSize: &msg.Size, 294 LocalAddress: ev.localAddress, 295 RemoteAddress: ev.remoteAddress, 296 }) 297 return msg, nil 298 } 299 300 // WriteMsg writes a message to the underlying MsgReadWriter and emits a 301 // "message sent" event 302 func (ev *msgEventer) WriteMsg(msg Msg) error { 303 err := ev.MsgReadWriter.WriteMsg(msg) 304 if err != nil { 305 return err 306 } 307 ev.feed.Send(&PeerEvent{ 308 Type: PeerEventTypeMsgSend, 309 Peer: ev.peerID, 310 Protocol: ev.Protocol, 311 MsgCode: &msg.Code, 312 MsgSize: &msg.Size, 313 LocalAddress: ev.localAddress, 314 RemoteAddress: ev.remoteAddress, 315 }) 316 return nil 317 } 318 319 // Close closes the underlying MsgReadWriter if it implements the io.Closer 320 // interface 321 func (ev *msgEventer) Close() error { 322 if v, ok := ev.MsgReadWriter.(io.Closer); ok { 323 return v.Close() 324 } 325 return nil 326 }