nltest: split multi-part messages into smaller batches

nltest/nltest.go

@@ -185,13 +185,13 @@ func (c *socket) Receive() ([]netlink.Message, error) {
 		}
 	}
 
-	// When a multi-part message is detected, return all messages except for the
-	// final "multi-part done", so that a second call to Receive from netlink.Conn
-	// will drain that message.
+	// When a multi-part message is detected, return messages in batches, so that
+	// multiple calls to Receive from netlink.Conn are needed to receive all
+	// messages.
 	if multi {
-		last := c.msgs[len(c.msgs)-1]
-		ret := c.msgs[:len(c.msgs)-1]
-		c.msgs = []netlink.Message{last}
+		batchSize := (len(c.msgs) + 1) / 2
+		ret := c.msgs[:batchSize]
+		c.msgs = c.msgs[batchSize:]
 
 		return ret, c.err
 	}

nltest/nltest_test.go

@@ -101,6 +101,84 @@ func TestConnReceiveError(t *testing.T) {
 	}
 }
 
+func TestConnReceiveMultipart(t *testing.T) {
+	msgs := []netlink.Message{
+		{
+			Data: []byte{0x00, 0x00, 0x00, 0x01},
+			Header: netlink.Header{
+				Flags: netlink.Multi,
+			},
+		},
+		{
+			Data: []byte{0x00, 0x00, 0x00, 0x02},
+			Header: netlink.Header{
+				Flags: netlink.Multi,
+			},
+		},
+		{
+			Data: []byte{0x00, 0x00, 0x00, 0x03},
+			Header: netlink.Header{
+				Flags: netlink.Multi,
+			},
+		},
+		{
+			Data: []byte{0x00, 0x00, 0x00, 0x04},
+			Header: netlink.Header{
+				Flags: netlink.Multi,
+			},
+		},
+		{
+			Header: netlink.Header{
+				Type:  netlink.Done,
+				Flags: netlink.Multi,
+			},
+		},
+	}
+
+	responded := false
+	c := nltest.Dial(func(_ []netlink.Message) ([]netlink.Message, error) {
+		// This is necessary so that nltest does not assume that the subsequent call
+		// to Receive is a multicast response. This would cause it to rerun this
+		// callback and return the same messages again, instead of simulating no
+		// more messages coming.
+		// TODO: Is there a better way to handle mutlicast responses from nltest?
+		if !responded {
+			responded = true
+			return nltest.Multipart(msgs)
+		}
+		return nil, io.EOF
+	})
+	defer c.Close()
+
+	// Send an empty request to trigger the multipart response.
+	if _, err := c.Send(netlink.Message{}); err != nil {
+		t.Fatalf("failed to send request: %v", err)
+	}
+
+	got, err := c.Receive()
+	if err != nil {
+		t.Fatalf("failed to receive messages: %v", err)
+	}
+
+	// Expect all messages but the one with the Done type.
+	want := msgs[:len(msgs)-1]
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Fatalf("unexpected multipart messages (-want +got):\n%s", diff)
+	}
+
+	// Any subsequent call to Receive should return no messages, since they're
+	// all drained from the previous call.
+	got, err = c.Receive()
+	if err != nil {
+		t.Fatalf("failed to receive messages: %v", err)
+	}
+
+	want = nil
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Fatalf("unexpected messages after multipart response (-want +got):\n%s", diff)
+	}
+}
+
 func TestConnExecuteOK(t *testing.T) {
 	req := netlink.Message{
 		Header: netlink.Header{