Merge weaveworks/519_remove_outermost_gob (PR #1189)

router/consts.go

@@ -25,6 +25,7 @@ const (
 	MaxDuration         = time.Duration(math.MaxInt64)
 	MaxMissedHeartbeats = 6
 	HeartbeatTimeout    = MaxMissedHeartbeats * SlowHeartbeat
+	MaxTCPMsgSize       = 10 * 1024 * 1024
 )
 
 var (

router/crypto.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"encoding/gob"
 	"fmt"
+	"io"
 	"sync"
 
 	"github.com/andybalholm/go-bit"
@@ -332,28 +333,48 @@ type TCPSender interface {
 	Send([]byte) error
 }
 
-type SimpleTCPSender struct {
+type GobTCPSender struct {
 	encoder *gob.Encoder
 }
 
+type LengthPrefixTCPSender struct {
+	writer io.Writer
+}
+
 type EncryptedTCPSender struct {
-	SimpleTCPSender
 	sync.RWMutex
-	state *TCPCryptoState
+	sender TCPSender
+	state  *TCPCryptoState
 }
 
-func NewSimpleTCPSender(encoder *gob.Encoder) *SimpleTCPSender {
-	return &SimpleTCPSender{encoder: encoder}
+func NewGobTCPSender(encoder *gob.Encoder) *GobTCPSender {
+	return &GobTCPSender{encoder: encoder}
 }
 
-func (sender *SimpleTCPSender) Send(msg []byte) error {
+func (sender *GobTCPSender) Send(msg []byte) error {
 	return sender.encoder.Encode(msg)
 }
 
-func NewEncryptedTCPSender(encoder *gob.Encoder, sessionKey *[32]byte, outbound bool) *EncryptedTCPSender {
-	return &EncryptedTCPSender{
-		SimpleTCPSender: *NewSimpleTCPSender(encoder),
-		state:           NewTCPCryptoState(sessionKey, outbound)}
+func NewLengthPrefixTCPSender(writer io.Writer) *LengthPrefixTCPSender {
+	return &LengthPrefixTCPSender{writer: writer}
+}
+
+func (sender *LengthPrefixTCPSender) Send(msg []byte) error {
+	l := len(msg)
+	if l > MaxTCPMsgSize {
+		return fmt.Errorf("outgoing message exceeds maximum size: %d > %d", l, MaxTCPMsgSize)
+	}
+	// We copy the message so we can send it in a single Write
+	// operation, thus making this thread-safe without locking.
+	prefixedMsg := make([]byte, 4+l)
+	binary.BigEndian.PutUint32(prefixedMsg, uint32(l))
+	copy(prefixedMsg[4:], msg)
+	_, err := sender.writer.Write(prefixedMsg)
+	return err
+}
+
+func NewEncryptedTCPSender(sender TCPSender, sessionKey *[32]byte, outbound bool) *EncryptedTCPSender {
+	return &EncryptedTCPSender{sender: sender, state: NewTCPCryptoState(sessionKey, outbound)}
 }
 
 func (sender *EncryptedTCPSender) Send(msg []byte) error {
@@ -361,40 +382,60 @@ func (sender *EncryptedTCPSender) Send(msg []byte) error {
 	defer sender.Unlock()
 	encodedMsg := secretbox.Seal(nil, msg, &sender.state.nonce, sender.state.sessionKey)
 	sender.state.advance()
-	return sender.SimpleTCPSender.Send(encodedMsg)
+	return sender.sender.Send(encodedMsg)
 }
 
 type TCPReceiver interface {
 	Receive() ([]byte, error)
 }
 
-type SimpleTCPReceiver struct {
+type GobTCPReceiver struct {
 	decoder *gob.Decoder
 }
 
+type LengthPrefixTCPReceiver struct {
+	reader io.Reader
+}
+
 type EncryptedTCPReceiver struct {
-	SimpleTCPReceiver
-	state *TCPCryptoState
+	receiver TCPReceiver
+	state    *TCPCryptoState
 }
 
-func NewSimpleTCPReceiver(decoder *gob.Decoder) *SimpleTCPReceiver {
-	return &SimpleTCPReceiver{decoder: decoder}
+func NewGobTCPReceiver(decoder *gob.Decoder) *GobTCPReceiver {
+	return &GobTCPReceiver{decoder: decoder}
 }
 
-func (receiver *SimpleTCPReceiver) Receive() ([]byte, error) {
+func (receiver *GobTCPReceiver) Receive() ([]byte, error) {
 	var msg []byte
 	err := receiver.decoder.Decode(&msg)
 	return msg, err
 }
 
-func NewEncryptedTCPReceiver(decoder *gob.Decoder, sessionKey *[32]byte, outbound bool) *EncryptedTCPReceiver {
-	return &EncryptedTCPReceiver{
-		SimpleTCPReceiver: *NewSimpleTCPReceiver(decoder),
-		state:             NewTCPCryptoState(sessionKey, !outbound)}
+func NewLengthPrefixTCPReceiver(reader io.Reader) *LengthPrefixTCPReceiver {
+	return &LengthPrefixTCPReceiver{reader: reader}
+}
+
+func (receiver *LengthPrefixTCPReceiver) Receive() ([]byte, error) {
+	lenPrefix := make([]byte, 4)
+	if _, err := io.ReadFull(receiver.reader, lenPrefix); err != nil {
+		return nil, err
+	}
+	l := binary.BigEndian.Uint32(lenPrefix)
+	if l > MaxTCPMsgSize {
+		return nil, fmt.Errorf("incoming message exceeds maximum size: %d > %d", l, MaxTCPMsgSize)
+	}
+	msg := make([]byte, l)
+	_, err := io.ReadFull(receiver.reader, msg)
+	return msg, err
+}
+
+func NewEncryptedTCPReceiver(receiver TCPReceiver, sessionKey *[32]byte, outbound bool) *EncryptedTCPReceiver {
+	return &EncryptedTCPReceiver{receiver: receiver, state: NewTCPCryptoState(sessionKey, !outbound)}
 }
 
 func (receiver *EncryptedTCPReceiver) Receive() ([]byte, error) {
-	msg, err := receiver.SimpleTCPReceiver.Receive()
+	msg, err := receiver.receiver.Receive()
 	if err != nil {
 		return nil, err
 	}

router/protocol.go

@@ -180,12 +180,13 @@ func (res *ProtocolIntroResults) doIntroV1(params ProtocolIntroParams, pubKey, p
 		return err
 	}
 
+	res.Sender = NewGobTCPSender(enc)
+	res.Receiver = NewGobTCPReceiver(dec)
+
 	if pubKey == nil {
 		if _, present := res.Features["PublicKey"]; present {
 			return ErrExpectedNoCrypto
 		}
-
-		res.setupNoCrypto(enc, dec)
 	} else {
 		remotePubKeyStr, ok := res.Features["PublicKey"]
 		if !ok {
@@ -197,7 +198,7 @@ func (res *ProtocolIntroResults) doIntroV1(params ProtocolIntroParams, pubKey, p
 			return err
 		}
 
-		res.setupCrypto(params, enc, dec, remotePubKey, privKey)
+		res.setupCrypto(params, remotePubKey, privKey)
 	}
 
 	res.Features = filterV1Features(res.Features)
@@ -227,12 +228,11 @@ func filterV1Features(intro map[string]string) map[string]string {
 // - When the connection is encrypted, 32 bytes follow containing the
 // public key.
 //
-// - Then a stream of gobified values.
+// - Then a stream of length-prefixed messages, which are encrypted
+// for an encrypted connection.
 //
-// The gobified values are the messages on the connection (encrypted
-// for an encrypted connection).  The first message contains the
-// encoded features map (so in contrast to V1, it will be encrypted on
-// an encrypted connection).
+// The first message contains the encoded features map (so in contrast
+// to V1, it will be encrypted on an encrypted connection).
 func (res *ProtocolIntroResults) doIntroV2(params ProtocolIntroParams, pubKey, privKey *[32]byte) error {
 	// Public key exchange
 	var wbuf []byte
@@ -265,7 +265,8 @@ func (res *ProtocolIntroResults) doIntroV2(params ProtocolIntroParams, pubKey, p
 			return ErrExpectedCrypto
 		}
 
-		res.setupNoCrypto(gob.NewEncoder(params.Conn), gob.NewDecoder(params.Conn))
+		res.Sender = NewLengthPrefixTCPSender(params.Conn)
+		res.Receiver = NewLengthPrefixTCPReceiver(params.Conn)
 
 	case 1:
 		if pubKey == nil {
@@ -277,7 +278,9 @@ func (res *ProtocolIntroResults) doIntroV2(params ProtocolIntroParams, pubKey, p
 			return err
 		}
 
-		res.setupCrypto(params, gob.NewEncoder(params.Conn), gob.NewDecoder(params.Conn), rbuf, privKey)
+		res.Sender = NewLengthPrefixTCPSender(params.Conn)
+		res.Receiver = NewLengthPrefixTCPReceiver(params.Conn)
+		res.setupCrypto(params, rbuf, privKey)
 
 	default:
 		return fmt.Errorf("Bad encryption flag %d", rbuf[0])
@@ -314,18 +317,12 @@ func (res *ProtocolIntroResults) doIntroV2(params ProtocolIntroParams, pubKey, p
 	return nil
 }
 
-func (res *ProtocolIntroResults) setupNoCrypto(enc *gob.Encoder, dec *gob.Decoder) {
-	res.Sender = NewSimpleTCPSender(enc)
-	res.Receiver = NewSimpleTCPReceiver(dec)
-}
-
-func (res *ProtocolIntroResults) setupCrypto(params ProtocolIntroParams,
-	enc *gob.Encoder, dec *gob.Decoder, remotePubKey []byte, privKey *[32]byte) {
+func (res *ProtocolIntroResults) setupCrypto(params ProtocolIntroParams, remotePubKey []byte, privKey *[32]byte) {
 	var remotePubKeyArr [32]byte
 	copy(remotePubKeyArr[:], remotePubKey)
 	res.SessionKey = FormSessionKey(&remotePubKeyArr, privKey, params.Password)
-	res.Sender = NewEncryptedTCPSender(enc, res.SessionKey, params.Outbound)
-	res.Receiver = NewEncryptedTCPReceiver(dec, res.SessionKey, params.Outbound)
+	res.Sender = NewEncryptedTCPSender(res.Sender, res.SessionKey, params.Outbound)
+	res.Receiver = NewEncryptedTCPReceiver(res.Receiver, res.SessionKey, params.Outbound)
 }
 
 type ProtocolTag byte