diff --git a/fakemessage_test.go b/fakemessage_test.go
index 36eac01a..c1676ee9 100644
--- a/fakemessage_test.go
+++ b/fakemessage_test.go
@@ -15,6 +15,7 @@ type fakeMessage struct {
command string
payload []byte
forceEncodeErr bool
+ forceLenErr bool
}
// BtcDecode doesn't do anything. It just satisfies the btcwire.Message
@@ -45,8 +46,14 @@ func (msg *fakeMessage) Command() string {
return msg.command
}
-// MaxPayloadLength simply returns 0. It is only here to satisfy the
-// btcwire.Message interface.
+// MaxPayloadLength returns the length of the payload field of fake message
+// or a smaller value if the forceLenErr flag of the fake message is set. It
+// satisfies the btcwire.Message interface.
func (msg *fakeMessage) MaxPayloadLength(pver uint32) uint32 {
- return 0
+ lenp := uint32(len(msg.payload))
+ if msg.forceLenErr {
+ return lenp - 1
+ }
+
+ return lenp
}
diff --git a/message.go b/message.go
index 80b748a9..06c65f1a 100644
--- a/message.go
+++ b/message.go
@@ -15,7 +15,7 @@ import (
// header. Shorter commands must be zero padded.
const commandSize = 12
-// maxMessagePayload is the maximum byes a message can be regardless of other
+// maxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const maxMessagePayload = (1024 * 1024 * 32) // 32MB
@@ -158,6 +158,7 @@ func discardInput(r io.Reader, n uint32) {
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
var command [commandSize]byte
+ // Enforce max command size.
cmd := msg.Command()
if len(cmd) > commandSize {
str := fmt.Sprintf("command [%s] is too long [max %v]",
@@ -166,6 +167,7 @@ func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) erro
}
copy(command[:], []byte(cmd))
+ // Encode the message payload.
var bw bytes.Buffer
err := msg.BtcEncode(&bw, pver)
if err != nil {
@@ -174,7 +176,7 @@ func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) erro
payload := bw.Bytes()
lenp := len(payload)
- // Enforce maximum message payload.
+ // Enforce maximum overall message payload.
if lenp > maxMessagePayload {
str := fmt.Sprintf("message payload is too large - encoded "+
"%d bytes, but maximum message payload is %d bytes",
@@ -182,6 +184,15 @@ func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) erro
return messageError("WriteMessage", str)
}
+ // Enforce maximum message payload based on the message type.
+ mpl := msg.MaxPayloadLength(pver)
+ if uint32(lenp) > mpl {
+ str := fmt.Sprintf("message payload is too large - encoded "+
+ "%d bytes, but maximum message payload size for "+
+ "messages of type [%s] is %d.", lenp, cmd, mpl)
+ return messageError("WriteMessage", str)
+ }
+
// Create header for the message.
hdr := messageHeader{}
hdr.magic = btcnet
diff --git a/message_test.go b/message_test.go
index e3483fc2..9186bdbd 100644
--- a/message_test.go
+++ b/message_test.go
@@ -317,10 +317,16 @@ func TestWriteMessageWireErrors(t *testing.T) {
// Fake message with a problem during encoding
encodeErrMsg := &fakeMessage{forceEncodeErr: true}
- // Fake message that has a payload which exceed max.
- exceedPayload := make([]byte, btcwire.MaxMessagePayload+1)
- exceedPayloadErrMsg := &fakeMessage{payload: exceedPayload}
+ // Fake message that has payload which exceeds max overall message size.
+ exceedOverallPayload := make([]byte, btcwire.MaxMessagePayload+1)
+ exceedOverallPayloadErrMsg := &fakeMessage{payload: exceedOverallPayload}
+ // Fake message that has payload which exceeds max allowed per message.
+ exceedPayload := make([]byte, 1)
+ exceedPayloadErrMsg := &fakeMessage{payload: exceedPayload, forceLenErr: true}
+
+ // Fake message that is used to force errors in the header and payload
+ // writes.
bogusPayload := []byte{0x01, 0x02, 0x03, 0x04}
bogusMsg := &fakeMessage{command: "bogus", payload: bogusPayload}
@@ -335,7 +341,9 @@ func TestWriteMessageWireErrors(t *testing.T) {
{badCommandMsg, pver, btcnet, 0, btcwireErr},
// Force error in payload encode.
{encodeErrMsg, pver, btcnet, 0, btcwireErr},
- // Force error due to exceeding max payload size.
+ // Force error due to exceeding max overall message payload size.
+ {exceedOverallPayloadErrMsg, pver, btcnet, 0, btcwireErr},
+ // Force error due to exceeding max payload for message type.
{exceedPayloadErrMsg, pver, btcnet, 0, btcwireErr},
// Force error in header write.
{bogusMsg, pver, btcnet, 0, io.ErrShortWrite},
diff --git a/msgblock.go b/msgblock.go
index b87496c2..e3ecc488 100644
--- a/msgblock.go
+++ b/msgblock.go
@@ -12,6 +12,9 @@ import (
// MaxBlocksPerMsg is the maximum number of blocks allowed per message.
const MaxBlocksPerMsg = 500
+// maxBlockPayload is the maximum bytes a block message can be.
+const maxBlockPayload = (1024 * 1024) // 1MB
+
// TxLoc holds locator data for the offset and length of where a transaction is
// located within a MsgBlock data buffer.
type TxLoc struct {
@@ -131,8 +134,8 @@ func (msg *MsgBlock) Command() string {
// receiver. This is part of the Message interface implementation.
func (msg *MsgBlock) MaxPayloadLength(pver uint32) uint32 {
// Block header at 81 bytes + max transactions which can vary up to the
- // max message payload.
- return maxMessagePayload
+ // maxBlockPayload (including the block header).
+ return maxBlockPayload
}
// BlockSha computes the block identifier hash for this block.
diff --git a/msgblock_test.go b/msgblock_test.go
index bd5082a5..780a657e 100644
--- a/msgblock_test.go
+++ b/msgblock_test.go
@@ -35,7 +35,7 @@ func TestBlock(t *testing.T) {
// Ensure max payload is expected value for latest protocol version.
// Num addresses (varInt) + max allowed addresses.
- wantPayload := uint32(1024 * 1024 * 32)
+ wantPayload := uint32(1024 * 1024)
maxPayload := msg.MaxPayloadLength(pver)
if maxPayload != wantPayload {
t.Errorf("MaxPayloadLength: wrong max payload length for "+
diff --git a/test_coverage.txt b/test_coverage.txt
index 69c5b942..fce68621 100644
--- a/test_coverage.txt
+++ b/test_coverage.txt
@@ -1,6 +1,6 @@
github.com/conformal/btcwire/message.go ReadMessage 100.00% (37/37)
-github.com/conformal/btcwire/message.go WriteMessage 100.00% (27/27)
+github.com/conformal/btcwire/message.go WriteMessage 100.00% (31/31)
github.com/conformal/btcwire/msgtx.go MsgTx.BtcDecode 100.00% (25/25)
github.com/conformal/btcwire/msgversion.go MsgVersion.BtcDecode 100.00% (25/25)
github.com/conformal/btcwire/msgtx.go MsgTx.Copy 100.00% (24/24)
@@ -149,5 +149,5 @@ github.com/conformal/btcwire/msgversion.go NewMsgVersion 100.00% (1/1)
github.com/conformal/btcwire/netaddress.go NetAddress.AddService 100.00% (1/1)
github.com/conformal/btcwire/shahash.go ShaHash.IsEqual 100.00% (1/1)
github.com/conformal/btcwire/invvect.go NewInvVect 100.00% (1/1)
-github.com/conformal/btcwire --------------------------------- 100.00% (923/923)
+github.com/conformal/btcwire --------------------------------- 100.00% (927/927)