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