From 0d9bfa4d18da72597950d8a903b56fb1d6e9c830 Mon Sep 17 00:00:00 2001
From: turekt <32360115+turekt@users.noreply.github.com>
Date: Sat, 23 Sep 2023 17:08:35 +0200
Subject: [PATCH] Fix overflow in Flush by using receiveAckAware and handling
 the overrun flag (#237)

Fixes https://github.com/google/nftables/issues/235
Added support for messages having overrun flag
Changed `conn.Receive` call to `receiveAckAware` in `Flush`
---
 conn.go | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/conn.go b/conn.go
index 711d7f6..5c8f51d 100644
--- a/conn.go
+++ b/conn.go
@@ -154,14 +154,33 @@ func receiveAckAware(nlconn *netlink.Conn, sentMsgFlags netlink.HeaderFlags) ([]
 		return reply, nil
 	}
 
-	// Dump flag is not set, we expect an ack
+	if len(reply) != 0 {
+		last := reply[len(reply)-1]
+		for re := last.Header.Type; (re&netlink.Overrun) == netlink.Overrun && (re&netlink.Done) != netlink.Done; re = last.Header.Type {
+			// we are not finished, the message is overrun
+			r, err := nlconn.Receive()
+			if err != nil {
+				return nil, err
+			}
+			reply = append(reply, r...)
+			last = reply[len(reply)-1]
+		}
+
+		if last.Header.Type == netlink.Error && binaryutil.BigEndian.Uint32(last.Data[:4]) == 0 {
+			// we have already collected an ack
+			return reply, nil
+		}
+	}
+
+	// Now we expect an ack
 	ack, err := nlconn.Receive()
 	if err != nil {
 		return nil, err
 	}
 
 	if len(ack) == 0 {
-		return nil, errors.New("received an empty ack")
+		// received an empty ack?
+		return reply, nil
 	}
 
 	msg := ack[0]
@@ -232,10 +251,7 @@ func (cc *Conn) Flush() error {
 
 	// Fetch the requested acknowledgement for each message we sent.
 	for _, msg := range cc.messages {
-		if msg.Header.Flags&netlink.Acknowledge == 0 {
-			continue // message did not request an acknowledgement
-		}
-		if _, err := conn.Receive(); err != nil {
+		if _, err := receiveAckAware(conn, msg.Header.Flags); err != nil {
 			return fmt.Errorf("conn.Receive: %w", err)
 		}
 	}