diff --git a/message.go b/message.go index 08e0487..910360e 100644 --- a/message.go +++ b/message.go @@ -85,6 +85,32 @@ type Message struct { Raw []byte } +// MarshalBinary implements the encoding.BinaryMarshaler interface. +func (m Message) MarshalBinary() (data []byte, err error) { + // We can't return m.Raw, allocation is expected by implicit interface + // contract induced by other implementations. + b := make([]byte, len(m.Raw)) + copy(b, m.Raw) + return b, nil +} + +// UnmarshalBinary implements the encoding.BinaryUnmarshaler interface. +func (m *Message) UnmarshalBinary(data []byte) error { + // We can't retain data, copy is expected by interface contract. + m.Raw = append(m.Raw[:0], data...) + return m.Decode() +} + +// GobEncode implements the gob.GobEncoder interface. +func (m Message) GobEncode() ([]byte, error) { + return m.MarshalBinary() +} + +// GobDecode implements the gob.GobDecoder interface. +func (m *Message) GobDecode(data []byte) error { + return m.UnmarshalBinary(data) +} + // AddTo sets b.TransactionID to m.TransactionID. // // Implements Setter to aid in crafting responses. diff --git a/message_test.go b/message_test.go index 93b748c..2c913bc 100644 --- a/message_test.go +++ b/message_test.go @@ -1008,3 +1008,32 @@ func BenchmarkDecode(b *testing.B) { } } } + +func TestMessage_MarshalBinary(t *testing.T) { + m := MustBuild( + NewSoftware("software"), + &XORMappedAddress{ + IP: net.IPv4(213, 1, 223, 5), + }, + ) + data, err := m.MarshalBinary() + if err != nil { + t.Fatal(err) + } + + // Reset m.Raw to check retention. + for i := range m.Raw { + m.Raw[i] = 0 + } + if err := m.UnmarshalBinary(data); err != nil { + t.Fatal(err) + } + + // Reset data to check retention. + for i := range data { + data[i] = 0 + } + if err := m.Decode(); err != nil { + t.Fatal(err) + } +}