@@ -11,8 +11,6 @@ import (
11
11
"github.com/ethereum/go-ethereum/ethutil"
12
12
)
13
13
14
- type MsgCode uint64
15
-
16
14
// Msg defines the structure of a p2p message.
17
15
//
18
16
// Note that a Msg can only be sent once since the Payload reader is
@@ -21,13 +19,13 @@ type MsgCode uint64
21
19
// structure, encode the payload into a byte array and create a
22
20
// separate Msg with a bytes.Reader as Payload for each send.
23
21
type Msg struct {
24
- Code MsgCode
22
+ Code uint64
25
23
Size uint32 // size of the paylod
26
24
Payload io.Reader
27
25
}
28
26
29
27
// NewMsg creates an RLP-encoded message with the given code.
30
- func NewMsg (code MsgCode , params ... interface {}) Msg {
28
+ func NewMsg (code uint64 , params ... interface {}) Msg {
31
29
buf := new (bytes.Buffer )
32
30
for _ , p := range params {
33
31
buf .Write (ethutil .Encode (p ))
@@ -63,6 +61,52 @@ func (msg Msg) Discard() error {
63
61
return err
64
62
}
65
63
64
+ type MsgReader interface {
65
+ ReadMsg () (Msg , error )
66
+ }
67
+
68
+ type MsgWriter interface {
69
+ // WriteMsg sends an existing message.
70
+ // The Payload reader of the message is consumed.
71
+ // Note that messages can be sent only once.
72
+ WriteMsg (Msg ) error
73
+
74
+ // EncodeMsg writes an RLP-encoded message with the given
75
+ // code and data elements.
76
+ EncodeMsg (code uint64 , data ... interface {}) error
77
+ }
78
+
79
+ // MsgReadWriter provides reading and writing of encoded messages.
80
+ type MsgReadWriter interface {
81
+ MsgReader
82
+ MsgWriter
83
+ }
84
+
85
+ // MsgLoop reads messages off the given reader and
86
+ // calls the handler function for each decoded message until
87
+ // it returns an error or the peer connection is closed.
88
+ //
89
+ // If a message is larger than the given maximum size,
90
+ // MsgLoop returns an appropriate error.
91
+ func MsgLoop (r MsgReader , maxsize uint32 , f func (code uint64 , data * ethutil.Value ) error ) error {
92
+ for {
93
+ msg , err := r .ReadMsg ()
94
+ if err != nil {
95
+ return err
96
+ }
97
+ if msg .Size > maxsize {
98
+ return newPeerError (errInvalidMsg , "size %d exceeds maximum size of %d" , msg .Size , maxsize )
99
+ }
100
+ value , err := msg .Data ()
101
+ if err != nil {
102
+ return err
103
+ }
104
+ if err := f (msg .Code , value ); err != nil {
105
+ return err
106
+ }
107
+ }
108
+ }
109
+
66
110
var magicToken = []byte {34 , 64 , 8 , 145 }
67
111
68
112
func writeMsg (w io.Writer , msg Msg ) error {
@@ -103,10 +147,10 @@ func readMsg(r byteReader) (msg Msg, err error) {
103
147
// read magic and payload size
104
148
start := make ([]byte , 8 )
105
149
if _ , err = io .ReadFull (r , start ); err != nil {
106
- return msg , NewPeerError ( ReadError , "%v" , err )
150
+ return msg , newPeerError ( errRead , "%v" , err )
107
151
}
108
152
if ! bytes .HasPrefix (start , magicToken ) {
109
- return msg , NewPeerError ( MagicTokenMismatch , "got %x, want %x" , start [:4 ], magicToken )
153
+ return msg , newPeerError ( errMagicTokenMismatch , "got %x, want %x" , start [:4 ], magicToken )
110
154
}
111
155
size := binary .BigEndian .Uint32 (start [4 :])
112
156
@@ -152,13 +196,13 @@ func readListHeader(r byteReader) (len uint64, hdrlen uint32, err error) {
152
196
}
153
197
154
198
// readUint reads an RLP-encoded unsigned integer from r.
155
- func readMsgCode (r byteReader ) (code MsgCode , codelen uint32 , err error ) {
199
+ func readMsgCode (r byteReader ) (code uint64 , codelen uint32 , err error ) {
156
200
b , err := r .ReadByte ()
157
201
if err != nil {
158
202
return 0 , 0 , err
159
203
}
160
204
if b < 0x80 {
161
- return MsgCode (b ), 1 , nil
205
+ return uint64 (b ), 1 , nil
162
206
} else if b < 0x89 { // max length for uint64 is 8 bytes
163
207
codelen = uint32 (b - 0x80 )
164
208
if codelen == 0 {
@@ -168,7 +212,7 @@ func readMsgCode(r byteReader) (code MsgCode, codelen uint32, err error) {
168
212
if _ , err := io .ReadFull (r , buf [8 - codelen :]); err != nil {
169
213
return 0 , 0 , err
170
214
}
171
- return MsgCode ( binary .BigEndian .Uint64 (buf ) ), codelen , nil
215
+ return binary .BigEndian .Uint64 (buf ), codelen , nil
172
216
}
173
217
return 0 , 0 , fmt .Errorf ("bad RLP type for message code: %x" , b )
174
218
}
0 commit comments