Rewrote reading strategy

pull/150/head
obscuren 10 years ago
parent d02024929a
commit 74de0f1f2a
  1. 241
      ethwire/messaging.go

@ -4,7 +4,6 @@ package ethwire
import (
"bytes"
"errors"
"fmt"
"net"
"time"
@ -78,106 +77,10 @@ func NewMessage(msgType MsgType, data interface{}) *Msg {
type Messages []*Msg
// The connection object allows you to set up a connection to the Ethereum network.
// The Connection object takes care of all encoding and sending objects properly over
// the network.
type Connection struct {
conn net.Conn
nTimeout time.Duration
pendingMessages Messages
}
// Create a new connection to the Ethereum network
func New(conn net.Conn) *Connection {
return &Connection{conn: conn, nTimeout: 500}
}
// Read, reads from the network. It will block until the next message is received.
func (self *Connection) Read() *Msg {
if len(self.pendingMessages) == 0 {
self.readMessages()
}
ret := self.pendingMessages[0]
self.pendingMessages = self.pendingMessages[1:]
return ret
}
// Write to the Ethereum network specifying the type of the message and
// the data. Data can be of type RlpEncodable or []interface{}. Returns
// nil or if something went wrong an error.
func (self *Connection) Write(typ MsgType, v ...interface{}) error {
var pack []byte
slice := [][]interface{}{[]interface{}{byte(typ)}}
for _, value := range v {
if encodable, ok := value.(ethutil.RlpEncodeDecode); ok {
slice = append(slice, encodable.RlpValue())
} else if raw, ok := value.([]interface{}); ok {
slice = append(slice, raw)
} else {
panic(fmt.Sprintf("Unable to 'write' object of type %T", value))
}
}
// Encode the type and the (RLP encoded) data for sending over the wire
encoded := ethutil.NewValue(slice).Encode()
payloadLength := ethutil.NumberToBytes(uint32(len(encoded)), 32)
// Write magic token and payload length (first 8 bytes)
pack = append(MagicToken, payloadLength...)
pack = append(pack, encoded...)
// Write to the connection
_, err := self.conn.Write(pack)
if err != nil {
return err
}
return nil
}
func (self *Connection) readMessage(data []byte) (msg *Msg, remaining []byte, done bool, err error) {
if len(data) == 0 {
return nil, nil, true, nil
}
if len(data) <= 8 {
return nil, remaining, false, errors.New("Invalid message")
}
// Check if the received 4 first bytes are the magic token
if bytes.Compare(MagicToken, data[:4]) != 0 {
return nil, nil, false, fmt.Errorf("MagicToken mismatch. Received %v", data[:4])
}
messageLength := ethutil.BytesToNumber(data[4:8])
remaining = data[8+messageLength:]
if int(messageLength) > len(data[8:]) {
return nil, nil, false, fmt.Errorf("message length %d, expected %d", len(data[8:]), messageLength)
}
message := data[8 : 8+messageLength]
decoder := ethutil.NewValueFromBytes(message)
// Type of message
t := decoder.Get(0).Uint()
// Actual data
d := decoder.SliceFrom(1)
msg = &Msg{
Type: MsgType(t),
Data: d,
}
return
}
// The basic message reader waits for data on the given connection, decoding
// and doing a few sanity checks such as if there's a data type and
// unmarhals the given data
func (self *Connection) readMessages() (err error) {
func ReadMessages(conn net.Conn) (msgs []*Msg, err error) {
// The recovering function in case anything goes horribly wrong
defer func() {
if r := recover(); r != nil {
@ -185,137 +88,69 @@ func (self *Connection) readMessages() (err error) {
}
}()
// Buff for writing network message to
//buff := make([]byte, 1440)
var buff []byte
var totalBytes int
var (
buff []byte
messages [][]byte
msgLength int
)
for {
// Give buffering some time
self.conn.SetReadDeadline(time.Now().Add(self.nTimeout * time.Millisecond))
conn.SetReadDeadline(time.Now().Add(5 * time.Millisecond))
// Create a new temporarily buffer
b := make([]byte, 1440)
// Wait for a message from this peer
n, _ := self.conn.Read(b)
n, _ := conn.Read(b)
if err != nil && n == 0 {
if err.Error() != "EOF" {
fmt.Println("err now", err)
return err
return nil, err
} else {
break
}
// Messages can't be empty
} else if n == 0 {
break
}
buff = append(buff, b[:n]...)
totalBytes += n
}
// Reslice buffer
buff = buff[:totalBytes]
msg, remaining, done, err := self.readMessage(buff)
for ; done != true; msg, remaining, done, err = self.readMessage(remaining) {
//log.Println("rx", msg)
if msg != nil {
self.pendingMessages = append(self.pendingMessages, msg)
if n == 0 {
continue
}
}
return
}
func ReadMessage(data []byte) (msg *Msg, remaining []byte, done bool, err error) {
if len(data) == 0 {
return nil, nil, true, nil
}
if len(data) <= 8 {
return nil, remaining, false, errors.New("Invalid message")
}
// Check if the received 4 first bytes are the magic token
if bytes.Compare(MagicToken, data[:4]) != 0 {
return nil, nil, false, fmt.Errorf("MagicToken mismatch. Received %v", data[:4])
}
messageLength := ethutil.BytesToNumber(data[4:8])
remaining = data[8+messageLength:]
if int(messageLength) > len(data[8:]) {
return nil, nil, false, fmt.Errorf("message length %d, expected %d", len(data[8:]), messageLength)
}
if msgLength == 0 {
// Check if the received 4 first bytes are the magic token
if bytes.Compare(MagicToken, b[:4]) != 0 {
return nil, fmt.Errorf("MagicToken mismatch. Received %v", b[:4])
}
message := data[8 : 8+messageLength]
decoder := ethutil.NewValueFromBytes(message)
// Type of message
t := decoder.Get(0).Uint()
// Actual data
d := decoder.SliceFrom(1)
// Remove the token
b = b[4:]
// Read the length of the message
msgLength = int(ethutil.BytesToNumber(b[:4]))
msg = &Msg{
Type: MsgType(t),
Data: d,
}
// Remove the length
b = b[4:]
return
}
n -= 8
}
func bufferedRead(conn net.Conn) ([]byte, error) {
return nil, nil
}
buff = append(buff, b[:n]...)
// The basic message reader waits for data on the given connection, decoding
// and doing a few sanity checks such as if there's a data type and
// unmarhals the given data
func ReadMessages(conn net.Conn) (msgs []*Msg, err error) {
// The recovering function in case anything goes horribly wrong
defer func() {
if r := recover(); r != nil {
err = fmt.Errorf("ethwire.ReadMessage error: %v", r)
}
}()
if len(buff) >= msgLength {
messages = append(messages, buff[:msgLength])
buff = buff[msgLength:]
msgLength = 0
// Buff for writing network message to
//buff := make([]byte, 1440)
var buff []byte
var totalBytes int
for {
// This is a bit of a cheat actually to make buffering extremely fast.
defer recover()
// Give buffering some time
conn.SetReadDeadline(time.Now().Add(5 * time.Millisecond))
// Create a new temporarily buffer
b := make([]byte, 1440)
// Wait for a message from this peer
n, _ := conn.Read(b)
if err != nil && n == 0 {
if err.Error() != "EOF" {
fmt.Println("err now", err)
return nil, err
} else {
if len(buff) == 0 {
break
}
// Messages can't be empty
} else if n == 0 {
break
}
buff = append(buff, b[:n]...)
totalBytes += n
}
// Reslice buffer
buff = buff[:totalBytes]
msg, remaining, done, err := ReadMessage(buff)
for ; done != true; msg, remaining, done, err = ReadMessage(remaining) {
//log.Println("rx", msg)
for _, m := range messages {
decoder := ethutil.NewValueFromBytes(m)
// Type of message
t := decoder.Get(0).Uint()
// Actual data
d := decoder.SliceFrom(1)
if msg != nil {
msgs = append(msgs, msg)
}
msgs = append(msgs, &Msg{Type: MsgType(t), Data: d})
}
return

Loading…
Cancel
Save