sd-netlink: refcount multicast groups

Track the number of matches installed for a given multicast group, and leave the
group once no matches depend on it.

In order to handle passed-in sockets that are already members of multicast groups
we initialize the refcount based on the membership once we take over the socket.
This way we will leave the socket in the state we found it once we finish with
it.

On kernels that do not fully support reading out the multicast group membership
we fall back to never leaving any groups (as before).

1 files changed, 154 insertions, 1 deletions

diff --git a/src/libsystemd/sd-netlink/netlink-socket.c b/src/libsystemd/sd-netlink/netlink-socket.c
index 84ff7c38c9..e1b14c3ed2 100644
--- a/src/libsystemd/sd-netlink/netlink-socket.c
+++ b/src/libsystemd/sd-netlink/netlink-socket.c
@@ -44,6 +44,65 @@ int socket_open(int family) {
         return fd;
 }
 
+static int broadcast_groups_get(sd_netlink *nl) {
+        _cleanup_free_ uint32_t *groups = NULL;
+        socklen_t len = 0, old_len;
+        unsigned i, j;
+        int r;
+
+        assert(nl);
+        assert(nl->fd > 0);
+
+        r = getsockopt(nl->fd, SOL_NETLINK, NETLINK_LIST_MEMBERSHIPS, NULL, &len);
+        if (r < 0) {
+                if (errno == ENOPROTOOPT) {
+                        nl->broadcast_group_dont_leave = true;
+                        return 0;
+                } else
+                        return -errno;
+        }
+
+        if (len == 0)
+                return 0;
+
+        groups = new0(uint32_t, len);
+        if (!groups)
+                return -ENOMEM;
+
+        old_len = len;
+
+        r = getsockopt(nl->fd, SOL_NETLINK, NETLINK_LIST_MEMBERSHIPS, groups, &len);
+        if (r < 0)
+                return -errno;
+
+        if (old_len != len)
+                return -EIO;
+
+        r = hashmap_ensure_allocated(&nl->broadcast_group_refs, NULL);
+        if (r < 0)
+                return r;
+
+        for (i = 0; i < len; i++) {
+                for (j = 0; j < sizeof(uint32_t) * 8; j ++) {
+                        uint32_t offset;
+                        unsigned group;
+
+                        offset = 1U << j;
+
+                        if (!(groups[i] & offset))
+                                continue;
+
+                        group = i * sizeof(uint32_t) * 8 + j + 1;
+
+                        r = hashmap_put(nl->broadcast_group_refs, UINT_TO_PTR(group), UINT_TO_PTR(1));
+                        if (r < 0)
+                                return r;
+                }
+        }
+
+        return 0;
+}
+
 int socket_bind(sd_netlink *nl) {
         socklen_t addrlen;
         int r, one = 1;
@@ -63,11 +122,32 @@ int socket_bind(sd_netlink *nl) {
         if (r < 0)
                 return -errno;
 
+        r = broadcast_groups_get(nl);
+        if (r < 0)
+                return r;
+
         return 0;
 }
 
+static unsigned broadcast_group_get_ref(sd_netlink *nl, unsigned group) {
+        assert(nl);
+
+        return PTR_TO_UINT(hashmap_get(nl->broadcast_group_refs, UINT_TO_PTR(group)));
+}
 
-int socket_join_broadcast_group(sd_netlink *nl, unsigned group) {
+static int broadcast_group_set_ref(sd_netlink *nl, unsigned group, unsigned n_ref) {
+        int r;
+
+        assert(nl);
+
+        r = hashmap_replace(nl->broadcast_group_refs, UINT_TO_PTR(group), UINT_TO_PTR(n_ref));
+        if (r < 0)
+                return r;
+
+        return 0;
+}
+
+static int broadcast_group_join(sd_netlink *nl, unsigned group) {
         int r;
 
         assert(nl);
@@ -81,6 +161,79 @@ int socket_join_broadcast_group(sd_netlink *nl, unsigned group) {
         return 0;
 }
 
+int socket_broadcast_group_ref(sd_netlink *nl, unsigned group) {
+        unsigned n_ref;
+        int r;
+
+        assert(nl);
+
+        n_ref = broadcast_group_get_ref(nl, group);
+
+        n_ref ++;
+
+        r = hashmap_ensure_allocated(&nl->broadcast_group_refs, NULL);
+        if (r < 0)
+                return r;
+
+        r = broadcast_group_set_ref(nl, group, n_ref);
+        if (r < 0)
+                return r;
+
+        if (n_ref > 1)
+                /* not yet in the group */
+                return 0;
+
+        r = broadcast_group_join(nl, group);
+        if (r < 0)
+                return r;
+
+        return 0;
+}
+
+static int broadcast_group_leave(sd_netlink *nl, unsigned group) {
+        int r;
+
+        assert(nl);
+        assert(nl->fd >= 0);
+        assert(group > 0);
+
+        if (nl->broadcast_group_dont_leave)
+                return 0;
+
+        r = setsockopt(nl->fd, SOL_NETLINK, NETLINK_DROP_MEMBERSHIP, &group, sizeof(group));
+        if (r < 0)
+                return -errno;
+
+        return 0;
+}
+
+int socket_broadcast_group_unref(sd_netlink *nl, unsigned group) {
+        unsigned n_ref;
+        int r;
+
+        assert(nl);
+
+        n_ref = broadcast_group_get_ref(nl, group);
+
+        assert(n_ref > 0);
+
+        n_ref --;
+
+        r = broadcast_group_set_ref(nl, group, n_ref);
+        if (r < 0)
+                return r;
+
+        if (n_ref > 0)
+                /* still refs left */
+                return 0;
+
+        r = broadcast_group_leave(nl, group);
+        if (r < 0)
+                return r;
+
+        return 0;
+}
+
 /* returns the number of bytes sent, or a negative error code */
 int socket_write_message(sd_netlink *nl, sd_netlink_message *m) {
         union {