Replace NL_KEEP code with proper message reference counting

Adds reference counting to netlink messages so callbacks
can hold on to a message without using the broken keep
message flag.

5 files changed, 35 insertions, 50 deletions

diff --git a/include/netlink-types.h b/include/netlink-types.h
index be5fc4e..238b131 100644
--- a/include/netlink-types.h
+++ b/include/netlink-types.h
@@ -129,6 +129,7 @@ struct nl_msg
 	struct ucred		nm_creds;
 	struct nlmsghdr *	nm_nlh;
 	size_t			nm_size;
+	int			nm_refcnt;
 };
 
 struct rtnl_link_map
diff --git a/include/netlink/handlers.h b/include/netlink/handlers.h
index 38c9ba0..266dc44 100644
--- a/include/netlink/handlers.h
+++ b/include/netlink/handlers.h
@@ -65,24 +65,6 @@ enum nl_cb_action {
 };
 
 /**
- * Callback action modifiers
- * @ingroup cb
- *
- * These should be ORed to the callback actions defined by enum
- * nl_cb_action.
- */
-enum nl_cb_action_mods {
-	/** Callee keeps the message, don't free */
-	NL_KEEP_MSG = 0x1000,
-#define NL_KEEP_MSG NL_KEEP_MSG	/* for config testing */
-};
-
-
-/* backwards compatibility */
-#define NL_PROCEED NL_OK
-#define NL_EXIT NL_STOP
-
-/**
  * Callback kinds
  * @ingroup cb
  */
diff --git a/include/netlink/msg.h b/include/netlink/msg.h
index 1cb1305..e331f42 100644
--- a/include/netlink/msg.h
+++ b/include/netlink/msg.h
@@ -81,6 +81,7 @@ extern int		  nlmsg_expand(struct nl_msg *, size_t);
 extern struct nlmsghdr *  nlmsg_put(struct nl_msg *, uint32_t, uint32_t,
 				    int, int, int);
 extern struct nlmsghdr *  nlmsg_hdr(struct nl_msg *);
+extern void		  nlmsg_get(struct nl_msg *);
 extern void		  nlmsg_free(struct nl_msg *);
 
 /* attribute modification */
diff --git a/lib/msg.c b/lib/msg.c
index ad6f5a1..d08d057 100644
--- a/lib/msg.c
+++ b/lib/msg.c
@@ -372,6 +372,8 @@ static struct nl_msg *__nlmsg_alloc(size_t len)
 	if (!nm)
 		goto errout;
 
+	nm->nm_refcnt = 1;
+
 	nm->nm_nlh = malloc(len);
 	if (!nm->nm_nlh)
 		goto errout;
@@ -645,21 +647,39 @@ struct nlmsghdr *nlmsg_hdr(struct nl_msg *n)
 }
 
 /**
- * Free a netlink message
- * @arg n		netlink message
- *
- * Destroys a netlink message and frees up all used memory.
+ * Acquire a reference on a netlink message
+ * @arg msg		message to acquire reference from
+ */
+void nlmsg_get(struct nl_msg *msg)
+{
+	msg->nm_refcnt++;
+	NL_DBG(4, "New reference to message %p, total %d\n",
+	       msg, msg->nm_refcnt);
+}
+
+/**
+ * Release a reference from an netlink message
+ * @arg msg		message to release reference from
  *
- * @pre The message must be unused.
+ * Frees memory after the last reference has been released.
  */
-void nlmsg_free(struct nl_msg *n)
+void nlmsg_free(struct nl_msg *msg)
 {
-	if (!n)
+	if (!msg)
 		return;
 
-	free(n->nm_nlh);
-	free(n);
-	NL_DBG(2, "msg %p: Freed\n", n);
+	msg->nm_refcnt--;
+	NL_DBG(4, "Returned message reference %p, %d remaining\n",
+	       msg, msg->nm_refcnt);
+
+	if (msg->nm_refcnt < 0)
+		BUG();
+
+	if (msg->nm_refcnt <= 0) {
+		free(msg->nm_nlh);
+		free(msg);
+		NL_DBG(2, "msg %p: Freed\n", msg);
+	}
 }
 
 /** @} */
diff --git a/lib/nl.c b/lib/nl.c
index 715dfeb..80a920a 100644
--- a/lib/nl.c
+++ b/lib/nl.c
@@ -551,9 +551,6 @@ abort:
 #define NL_CB_CALL(cb, type, msg) \
 do { \
 	err = nl_cb_call(cb, type, msg); \
-	if (free_msg && (err & NL_KEEP_MSG))	 \
-		free_msg = 0;			 \
-	err &= ~NL_KEEP_MSG;			 \
 	switch (err) { \
 	case NL_OK: \
 		err = 0; \
@@ -567,23 +564,10 @@ do { \
 	} \
 } while (0)
 
-/*
- * NOTE: on handling freeing of the message data
- *
- * By default, the message data is freed after handling is done. In
- * order to allow a callback using it after exiting the message
- * handling loop, it can return NL_KEEP_MSG ORed to it's return code.
- * 
- * Once the freeing of the message is disabled, it cannot be activated
- * again; this way, if a callback decides to switch it off because it
- * will keep the allocated data, another one cannot activate it, have
- * it freed and cause a race condition with later access to that (now
- * freed) data.
- */ 
 static int recvmsgs(struct nl_sock *sk, struct nl_cb *cb)
 {
 	int n, err = 0, multipart = 0;
-	unsigned char *buf = NULL, free_msg = 1;
+	unsigned char *buf = NULL;
 	struct nlmsghdr *hdr;
 	struct sockaddr_nl nla = {0};
 	struct nl_msg *msg = NULL;
@@ -605,9 +589,7 @@ continue_reading:
 	while (nlmsg_ok(hdr, n)) {
 		NL_DBG(3, "recgmsgs(%p): Processing valid message...\n", sk);
 
-		if (free_msg)
-			nlmsg_free(msg);
-		free_msg = 1;	/* By default, we free the message data */
+		nlmsg_free(msg);
 		msg = nlmsg_convert(hdr);
 		if (!msg) {
 			err = -NLE_NOMEM;
@@ -741,8 +723,7 @@ skip:
 		hdr = nlmsg_next(hdr, &n);
 	}
 	
-	if (free_msg)
-		nlmsg_free(msg);
+	nlmsg_free(msg);
 	free(buf);
 	free(creds);
 	buf = NULL;