github.com/aquanetwork/aquachain@v1.7.8/p2p/message.go (about) 1 // Copyright 2014 The aquachain Authors 2 // This file is part of the aquachain library. 3 // 4 // The aquachain 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 aquachain 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 aquachain 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 "sync/atomic" 26 "time" 27 28 "gitlab.com/aquachain/aquachain/aqua/event" 29 "gitlab.com/aquachain/aquachain/p2p/discover" 30 "gitlab.com/aquachain/aquachain/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 // eofSignal wraps a reader with eof signaling. the eof channel is 114 // closed when the wrapped reader returns an error or when count bytes 115 // have been read. 116 type eofSignal struct { 117 wrapped io.Reader 118 count uint32 // number of bytes left 119 eof chan<- struct{} 120 } 121 122 // note: when using eofSignal to detect whether a message payload 123 // has been read, Read might not be called for zero sized messages. 124 func (r *eofSignal) Read(buf []byte) (int, error) { 125 if r.count == 0 { 126 if r.eof != nil { 127 r.eof <- struct{}{} 128 r.eof = nil 129 } 130 return 0, io.EOF 131 } 132 133 max := len(buf) 134 if int(r.count) < len(buf) { 135 max = int(r.count) 136 } 137 n, err := r.wrapped.Read(buf[:max]) 138 r.count -= uint32(n) 139 if (err != nil || r.count == 0) && r.eof != nil { 140 r.eof <- struct{}{} // tell Peer that msg has been consumed 141 r.eof = nil 142 } 143 return n, err 144 } 145 146 // MsgPipe creates a message pipe. Reads on one end are matched 147 // with writes on the other. The pipe is full-duplex, both ends 148 // implement MsgReadWriter. 149 func MsgPipe() (*MsgPipeRW, *MsgPipeRW) { 150 var ( 151 c1, c2 = make(chan Msg), make(chan Msg) 152 closing = make(chan struct{}) 153 closed = new(int32) 154 rw1 = &MsgPipeRW{c1, c2, closing, closed} 155 rw2 = &MsgPipeRW{c2, c1, closing, closed} 156 ) 157 return rw1, rw2 158 } 159 160 // ErrPipeClosed is returned from pipe operations after the 161 // pipe has been closed. 162 var ErrPipeClosed = errors.New("p2p: read or write on closed message pipe") 163 164 // MsgPipeRW is an endpoint of a MsgReadWriter pipe. 165 type MsgPipeRW struct { 166 w chan<- Msg 167 r <-chan Msg 168 closing chan struct{} 169 closed *int32 170 } 171 172 // WriteMsg sends a messsage on the pipe. 173 // It blocks until the receiver has consumed the message payload. 174 func (p *MsgPipeRW) WriteMsg(msg Msg) error { 175 if atomic.LoadInt32(p.closed) == 0 { 176 consumed := make(chan struct{}, 1) 177 msg.Payload = &eofSignal{msg.Payload, msg.Size, consumed} 178 select { 179 case p.w <- msg: 180 if msg.Size > 0 { 181 // wait for payload read or discard 182 select { 183 case <-consumed: 184 case <-p.closing: 185 } 186 } 187 return nil 188 case <-p.closing: 189 } 190 } 191 return ErrPipeClosed 192 } 193 194 // ReadMsg returns a message sent on the other end of the pipe. 195 func (p *MsgPipeRW) ReadMsg() (Msg, error) { 196 if atomic.LoadInt32(p.closed) == 0 { 197 select { 198 case msg := <-p.r: 199 return msg, nil 200 case <-p.closing: 201 } 202 } 203 return Msg{}, ErrPipeClosed 204 } 205 206 // Close unblocks any pending ReadMsg and WriteMsg calls on both ends 207 // of the pipe. They will return ErrPipeClosed. Close also 208 // interrupts any reads from a message payload. 209 func (p *MsgPipeRW) Close() error { 210 if atomic.AddInt32(p.closed, 1) != 1 { 211 // someone else is already closing 212 atomic.StoreInt32(p.closed, 1) // avoid overflow 213 return nil 214 } 215 close(p.closing) 216 return nil 217 } 218 219 // ExpectMsg reads a message from r and verifies that its 220 // code and encoded RLP content match the provided values. 221 // If content is nil, the payload is discarded and not verified. 222 func ExpectMsg(r MsgReader, code uint64, content interface{}) error { 223 msg, err := r.ReadMsg() 224 if err != nil { 225 return err 226 } 227 if msg.Code != code { 228 return fmt.Errorf("message code mismatch: got %d, expected %d", msg.Code, code) 229 } 230 if content == nil { 231 return msg.Discard() 232 } else { 233 contentEnc, err := rlp.EncodeToBytes(content) 234 if err != nil { 235 panic("content encode error: " + err.Error()) 236 } 237 if int(msg.Size) != len(contentEnc) { 238 return fmt.Errorf("message size mismatch: got %d, want %d", msg.Size, len(contentEnc)) 239 } 240 actualContent, err := ioutil.ReadAll(msg.Payload) 241 if err != nil { 242 return err 243 } 244 if !bytes.Equal(actualContent, contentEnc) { 245 return fmt.Errorf("message payload mismatch:\ngot: %x\nwant: %x", actualContent, contentEnc) 246 } 247 } 248 return nil 249 } 250 251 // msgEventer wraps a MsgReadWriter and sends events whenever a message is sent 252 // or received 253 type msgEventer struct { 254 MsgReadWriter 255 256 feed *event.Feed 257 peerID discover.NodeID 258 Protocol string 259 } 260 261 // newMsgEventer returns a msgEventer which sends message events to the given 262 // feed 263 func newMsgEventer(rw MsgReadWriter, feed *event.Feed, peerID discover.NodeID, proto string) *msgEventer { 264 return &msgEventer{ 265 MsgReadWriter: rw, 266 feed: feed, 267 peerID: peerID, 268 Protocol: proto, 269 } 270 } 271 272 // ReadMsg reads a message from the underlying MsgReadWriter and emits a 273 // "message received" event 274 func (self *msgEventer) ReadMsg() (Msg, error) { 275 msg, err := self.MsgReadWriter.ReadMsg() 276 if err != nil { 277 return msg, err 278 } 279 self.feed.Send(&PeerEvent{ 280 Type: PeerEventTypeMsgRecv, 281 Peer: self.peerID, 282 Protocol: self.Protocol, 283 MsgCode: &msg.Code, 284 MsgSize: &msg.Size, 285 }) 286 return msg, nil 287 } 288 289 // WriteMsg writes a message to the underlying MsgReadWriter and emits a 290 // "message sent" event 291 func (self *msgEventer) WriteMsg(msg Msg) error { 292 err := self.MsgReadWriter.WriteMsg(msg) 293 if err != nil { 294 return err 295 } 296 self.feed.Send(&PeerEvent{ 297 Type: PeerEventTypeMsgSend, 298 Peer: self.peerID, 299 Protocol: self.Protocol, 300 MsgCode: &msg.Code, 301 MsgSize: &msg.Size, 302 }) 303 return nil 304 } 305 306 // Close closes the underlying MsgReadWriter if it implements the io.Closer 307 // interface 308 func (self *msgEventer) Close() error { 309 if v, ok := self.MsgReadWriter.(io.Closer); ok { 310 return v.Close() 311 } 312 return nil 313 }