packet: mrt: commiting support and test for a split function for use in MRT messages under the bufio interface

2 files changed, 47 insertions, 0 deletions

diff --git a/packet/mrt.go b/packet/mrt.go
index 5e4f2597..369e8063 100644
--- a/packet/mrt.go
+++ b/packet/mrt.go
@@ -697,6 +697,27 @@ func (m *BGP4MPMessage) String() string {
 	return fmt.Sprintf("%s: PeerAS [%d] LocalAS [%d] InterfaceIndex [%d] PeerIP [%s] LocalIP [%s] BGPMessage [%s]", title, m.PeerAS, m.LocalAS, m.InterfaceIndex, m.PeerIpAddress, m.LocalIpAddress, m.BGPMessage)
 }
 
+//This function can be passed into a bufio.Scanner.Split() to read buffered mrt msgs
+func SplitMrt(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	if atEOF && len(data) == 0 {
+		return 0, nil, nil
+	}
+	if cap(data) < MRT_COMMON_HEADER_LEN { // read more
+		return 0, nil, nil
+	}
+	//this reads the data
+	hdr := &MRTHeader{}
+	errh := hdr.DecodeFromBytes(data[:MRT_COMMON_HEADER_LEN])
+	if errh != nil {
+		return 0, nil, errh
+	}
+	totlen := int(hdr.Len + MRT_COMMON_HEADER_LEN)
+	if len(data) < totlen { //need to read more
+		return 0, nil, nil
+	}
+	return totlen, data[0:totlen], nil
+}
+
 func ParseMRTBody(h *MRTHeader, data []byte) (*MRTMessage, error) {
 	if len(data) < int(h.Len) {
 		return nil, fmt.Errorf("Not all MRT message bytes available. expected: %d, actual: %d", int(h.Len), len(data))
diff --git a/packet/mrt_test.go b/packet/mrt_test.go
index 7aa073c4..c43d45e0 100644
--- a/packet/mrt_test.go
+++ b/packet/mrt_test.go
@@ -16,6 +16,8 @@
 package bgp
 
 import (
+	"bufio"
+	"bytes"
 	"github.com/stretchr/testify/assert"
 	"reflect"
 	"testing"
@@ -181,3 +183,27 @@ func TestMrtBgp4mpMessage(t *testing.T) {
 	}
 	assert.Equal(t, reflect.DeepEqual(m1, m2), true)
 }
+
+func TestMrtSplit(t *testing.T) {
+	var b bytes.Buffer
+	numwrite, numread := 10, 0
+	for i := 0; i < numwrite; i++ {
+		msg := NewBGPKeepAliveMessage()
+		m1 := NewBGP4MPMessage(65000, 65001, 1, "192.168.0.1", "192.168.0.2", false, msg)
+		mm, _ := NewMRTMessage(1234, BGP4MP, MESSAGE, m1)
+		b1, err := mm.Serialize()
+		if err != nil {
+			t.Fatal(err)
+		}
+		b.Write(b1)
+	}
+	t.Logf("wrote %d serialized MRT keepalives in the buffer", numwrite)
+	r := bytes.NewReader(b.Bytes())
+	scanner := bufio.NewScanner(r)
+	scanner.Split(SplitMrt)
+	for scanner.Scan() {
+		numread += 1
+	}
+	t.Logf("scanner scanned %d serialized keepalives from the buffer", numread)
+	assert.Equal(t, numwrite, numread)
+}