package p2p import ( "bufio" "bytes" "encoding/binary" "errors" "fmt" "io" "io/ioutil" "math/big" "net" "sync" "sync/atomic" "time" "github.com/ethereum/go-ethereum/ethutil" "github.com/ethereum/go-ethereum/rlp" ) // Msg defines the structure of a p2p message. // // Note that a Msg can only be sent once since the Payload reader is // consumed during sending. It is not possible to create a Msg and // send it any number of times. If you want to reuse an encoded // structure, encode the payload into a byte array and create a // separate Msg with a bytes.Reader as Payload for each send. type Msg struct { Code uint64 Size uint32 // size of the paylod Payload io.Reader } // NewMsg creates an RLP-encoded message with the given code. func NewMsg(code uint64, params ...interface{}) Msg { buf := new(bytes.Buffer) for _, p := range params { buf.Write(ethutil.Encode(p)) } return Msg{Code: code, Size: uint32(buf.Len()), Payload: buf} } func encodePayload(params ...interface{}) []byte { buf := new(bytes.Buffer) for _, p := range params { buf.Write(ethutil.Encode(p)) } return buf.Bytes() } // Decode parse the RLP content of a message into // the given value, which must be a pointer. // // For the decoding rules, please see package rlp. func (msg Msg) Decode(val interface{}) error { s := rlp.NewListStream(msg.Payload, uint64(msg.Size)) if err := s.Decode(val); err != nil { return newPeerError(errInvalidMsg, "(code %#x) (size %d) %v", msg.Code, msg.Size, err) } return nil } func (msg Msg) String() string { return fmt.Sprintf("msg #%v (%v bytes)", msg.Code, msg.Size) } // Discard reads any remaining payload data into a black hole. func (msg Msg) Discard() error { _, err := io.Copy(ioutil.Discard, msg.Payload) return err } type MsgReader interface { ReadMsg() (Msg, error) } type MsgWriter interface { // WriteMsg sends a message. It will block until the message's // Payload has been consumed by the other end. // // Note that messages can be sent only once because their // payload reader is drained. WriteMsg(Msg) error } // MsgReadWriter provides reading and writing of encoded messages. // Implementations should ensure that ReadMsg and WriteMsg can be // called simultaneously from multiple goroutines. type MsgReadWriter interface { MsgReader MsgWriter } // EncodeMsg writes an RLP-encoded message with the given code and // data elements. func EncodeMsg(w MsgWriter, code uint64, data ...interface{}) error { return w.WriteMsg(NewMsg(code, data...)) } // frameRW is a MsgReadWriter that reads and writes devp2p message frames. // As required by the interface, ReadMsg and WriteMsg can be called from // multiple goroutines. type frameRW struct { net.Conn // make Conn methods available. be careful. bufconn *bufio.ReadWriter // this channel is used to 'lend' bufconn to a caller of ReadMsg // until the message payload has been consumed. the channel // receives a value when EOF is reached on the payload, unblocking // a pending call to ReadMsg. rsync chan struct{} // this mutex guards writes to bufconn. writeMu sync.Mutex } func newFrameRW(conn net.Conn, timeout time.Duration) *frameRW { rsync := make(chan struct{}, 1) rsync <- struct{}{} return &frameRW{ Conn: conn, bufconn: bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)), rsync: rsync, } } var magicToken = []byte{34, 64, 8, 145} func (rw *frameRW) WriteMsg(msg Msg) error { rw.writeMu.Lock() defer rw.writeMu.Unlock() rw.SetWriteDeadline(time.Now().Add(msgWriteTimeout)) if err := writeMsg(rw.bufconn, msg); err != nil { return err } return rw.bufconn.Flush() } func writeMsg(w io.Writer, msg Msg) error { // TODO: handle case when Size + len(code) + len(listhdr) overflows uint32 code := ethutil.Encode(uint32(msg.Code)) listhdr := makeListHeader(msg.Size + uint32(len(code))) payloadLen := uint32(len(listhdr)) + uint32(len(code)) + msg.Size start := make([]byte, 8) copy(start, magicToken) binary.BigEndian.PutUint32(start[4:], payloadLen) for _, b := range [][]byte{start, listhdr, code} { if _, err := w.Write(b); err != nil { return err } } _, err := io.CopyN(w, msg.Payload, int64(msg.Size)) return err } func makeListHeader(length uint32) []byte { if length < 56 { return []byte{byte(length + 0xc0)} } enc := big.NewInt(int64(length)).Bytes() lenb := byte(len(enc)) + 0xf7 return append([]byte{lenb}, enc...) } func (rw *frameRW) ReadMsg() (msg Msg, err error) { <-rw.rsync // wait until bufconn is ours // this read timeout applies also to the payload. // TODO: proper read timeout rw.SetReadDeadline(time.Now().Add(msgReadTimeout)) // read magic and payload size start := make([]byte, 8) if _, err = io.ReadFull(rw.bufconn, start); err != nil { return msg, err } if !bytes.HasPrefix(start, magicToken) { return msg, fmt.Errorf("bad magic token %x", start[:4], magicToken) } size := binary.BigEndian.Uint32(start[4:]) // decode start of RLP message to get the message code posr := &postrack{rw.bufconn, 0} s := rlp.NewStream(posr) if _, err := s.List(); err != nil { return msg, err } msg.Code, err = s.Uint() if err != nil { return msg, err } msg.Size = size - posr.p if msg.Size <= wholePayloadSize { // msg is small, read all of it and move on to the next message. pbuf := make([]byte, msg.Size) if _, err := io.ReadFull(rw.bufconn, pbuf); err != nil { return msg, err } rw.rsync <- struct{}{} // bufconn is available again msg.Payload = bytes.NewReader(pbuf) } else { // lend bufconn to the caller until it has // consumed the payload. eofSignal will send a value // on rw.rsync when EOF is reached. pr := &eofSignal{rw.bufconn, msg.Size, rw.rsync} msg.Payload = pr } return msg, nil } // postrack wraps an rlp.ByteReader with a position counter. type postrack struct { r rlp.ByteReader p uint32 } func (r *postrack) Read(buf []byte) (int, error) { n, err := r.r.Read(buf) r.p += uint32(n) return n, err } func (r *postrack) ReadByte() (byte, error) { b, err := r.r.ReadByte() if err == nil { r.p++ } return b, err } // eofSignal wraps a reader with eof signaling. the eof channel is // closed when the wrapped reader returns an error or when count bytes // have been read. type eofSignal struct { wrapped io.Reader count uint32 // number of bytes left eof chan<- struct{} } // note: when using eofSignal to detect whether a message payload // has been read, Read might not be called for zero sized messages. func (r *eofSignal) Read(buf []byte) (int, error) { if r.count == 0 { if r.eof != nil { r.eof <- struct{}{} r.eof = nil } return 0, io.EOF } max := len(buf) if int(r.count) < len(buf) { max = int(r.count) } n, err := r.wrapped.Read(buf[:max]) r.count -= uint32(n) if (err != nil || r.count == 0) && r.eof != nil { r.eof <- struct{}{} // tell Peer that msg has been consumed r.eof = nil } return n, err } // MsgPipe creates a message pipe. Reads on one end are matched // with writes on the other. The pipe is full-duplex, both ends // implement MsgReadWriter. func MsgPipe() (*MsgPipeRW, *MsgPipeRW) { var ( c1, c2 = make(chan Msg), make(chan Msg) closing = make(chan struct{}) closed = new(int32) rw1 = &MsgPipeRW{c1, c2, closing, closed} rw2 = &MsgPipeRW{c2, c1, closing, closed} ) return rw1, rw2 } // ErrPipeClosed is returned from pipe operations after the // pipe has been closed. var ErrPipeClosed = errors.New("p2p: read or write on closed message pipe") // MsgPipeRW is an endpoint of a MsgReadWriter pipe. type MsgPipeRW struct { w chan<- Msg r <-chan Msg closing chan struct{} closed *int32 } // WriteMsg sends a messsage on the pipe. // It blocks until the receiver has consumed the message payload. func (p *MsgPipeRW) WriteMsg(msg Msg) error { if atomic.LoadInt32(p.closed) == 0 { consumed := make(chan struct{}, 1) msg.Payload = &eofSignal{msg.Payload, msg.Size, consumed} select { case p.w <- msg: if msg.Size > 0 { // wait for payload read or discard <-consumed } return nil case <-p.closing: } } return ErrPipeClosed } // ReadMsg returns a message sent on the other end of the pipe. func (p *MsgPipeRW) ReadMsg() (Msg, error) { if atomic.LoadInt32(p.closed) == 0 { select { case msg := <-p.r: return msg, nil case <-p.closing: } } return Msg{}, ErrPipeClosed } // Close unblocks any pending ReadMsg and WriteMsg calls on both ends // of the pipe. They will return ErrPipeClosed. Note that Close does // not interrupt any reads from a message payload. func (p *MsgPipeRW) Close() error { if atomic.AddInt32(p.closed, 1) != 1 { // someone else is already closing atomic.StoreInt32(p.closed, 1) // avoid overflow return nil } close(p.closing) return nil }