[PATCH] selftests/acct: share netlink helpers

[Yiyang Chen <cyyzero16@gmail.com>]

Extract the duplicated generic netlink boilerplate (netlink_open,
send_request, get_family_id, and NLA walker macros) from
cgroupstats.c and taskstats_fill_stats_tgid.c into a shared
netlink_helper.{h,c}.

Suggested-by: Andrew Morton <akpm@linux-foundation.org>
Signed-off-by: Yiyang Chen <cyyzero16@gmail.com>
---
 tools/testing/selftests/acct/.gitignore       |   1 +
 tools/testing/selftests/acct/Makefile         |  10 ++
 tools/testing/selftests/acct/cgroupstats.c    | 133 +----------------
 tools/testing/selftests/acct/netlink_helper.c | 116 +++++++++++++++
 tools/testing/selftests/acct/netlink_helper.h |  44 ++++++
 .../acct/taskstats_fill_stats_tgid.c          | 134 ++----------------
 6 files changed, 184 insertions(+), 254 deletions(-)
 create mode 100644 tools/testing/selftests/acct/netlink_helper.c
 create mode 100644 tools/testing/selftests/acct/netlink_helper.h

diff --git a/tools/testing/selftests/acct/.gitignore b/tools/testing/selftests/acct/.gitignore
index 9e9c61c5bfd6..fe0896f54e15 100644
--- a/tools/testing/selftests/acct/.gitignore
+++ b/tools/testing/selftests/acct/.gitignore
@@ -1,4 +1,5 @@
 acct_syscall
 taskstats_fill_stats_tgid
+cgroupstats
 config
 process_log
diff --git a/tools/testing/selftests/acct/Makefile b/tools/testing/selftests/acct/Makefile
index db88d65f5581..93a11a28a636 100644
--- a/tools/testing/selftests/acct/Makefile
+++ b/tools/testing/selftests/acct/Makefile
@@ -3,7 +3,17 @@ TEST_GEN_PROGS := acct_syscall
 TEST_GEN_PROGS += taskstats_fill_stats_tgid
 TEST_GEN_PROGS += cgroupstats
 
+NETLINK_HELPER_PROGS := cgroupstats taskstats_fill_stats_tgid
+
 CFLAGS += -Wall
 LDLIBS += -lpthread
 
 include ../lib.mk
+
+$(NETLINK_HELPER_PROGS): %: %.c netlink_helper.c netlink_helper.h
+	$(call msg,CC,,$@)
+	$(Q)$(LINK.c) $< netlink_helper.c $(LDLIBS) -o $@
+
+$(addprefix $(OUTPUT)/,$(NETLINK_HELPER_PROGS)): $(OUTPUT)/%: %.c netlink_helper.c netlink_helper.h
+	$(call msg,CC,,$@)
+	$(Q)$(LINK.c) $< netlink_helper.c $(LDLIBS) -o $@
diff --git a/tools/testing/selftests/acct/cgroupstats.c b/tools/testing/selftests/acct/cgroupstats.c
index e2836383ed50..0b421a4ca72b 100644
--- a/tools/testing/selftests/acct/cgroupstats.c
+++ b/tools/testing/selftests/acct/cgroupstats.c
@@ -6,139 +6,19 @@
 #include <linux/cgroupstats.h>
 #include <linux/genetlink.h>
 #include <linux/netlink.h>
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <sys/mount.h>
 #include <sys/socket.h>
 #include <sys/types.h>
-#include <sys/mount.h>
-#include <time.h>
 #include <unistd.h>
 
+#include "netlink_helper.h"
 #include "kselftest.h"
 
-#ifndef NLA_ALIGN
-#define NLA_ALIGNTO 4
-#define NLA_ALIGN(len) (((len) + NLA_ALIGNTO - 1) & ~(NLA_ALIGNTO - 1))
-#define NLA_HDRLEN ((int)NLA_ALIGN(sizeof(struct nlattr)))
-#endif
-
-#define RECV_TIMEOUT_SEC 2
-
-static void *nla_data(const struct nlattr *na)
-{
-	return (void *)((char *)na + NLA_HDRLEN);
-}
-
-static int netlink_open(void)
-{
-	struct timeval tv = { .tv_sec = RECV_TIMEOUT_SEC };
-	struct sockaddr_nl addr = {
-		.nl_family = AF_NETLINK,
-		.nl_pid = getpid(),
-	};
-	int fd;
-
-	fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
-	if (fd < 0)
-		return -errno;
-
-	/*
-	 * Ensure that a missing kernel reply fails the individual test
-	 * case instead of hanging the whole test binary.
-	 */
-	if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) {
-		int err = -errno;
-
-		close(fd);
-		return err;
-	}
-
-	if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
-		int err = -errno;
-
-		close(fd);
-		return err;
-	}
-
-	return fd;
-}
-
-static int send_request(int fd, void *buf, size_t len)
-{
-	struct sockaddr_nl addr = {
-		.nl_family = AF_NETLINK,
-	};
-
-	if (sendto(fd, buf, len, 0, (struct sockaddr *)&addr, sizeof(addr)) < 0)
-		return -errno;
-
-	return 0;
-}
-
-static int get_family_id(int fd, const char *name)
-{
-	struct {
-		struct nlmsghdr nlh;
-		struct genlmsghdr genl;
-		char buf[256];
-	} req = { 0 };
-	char resp[8192];
-	struct nlmsghdr *nlh;
-	struct genlmsghdr *genl;
-	struct nlattr *na;
-	int len;
-	int rem;
-	int ret;
-
-	req.nlh.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
-	req.nlh.nlmsg_type = GENL_ID_CTRL;
-	req.nlh.nlmsg_flags = NLM_F_REQUEST;
-	req.nlh.nlmsg_seq = 1;
-	req.nlh.nlmsg_pid = getpid();
-
-	req.genl.cmd = CTRL_CMD_GETFAMILY;
-	req.genl.version = 1;
-
-	na = (struct nlattr *)((char *)&req + NLMSG_ALIGN(req.nlh.nlmsg_len));
-	na->nla_type = CTRL_ATTR_FAMILY_NAME;
-	na->nla_len = NLA_HDRLEN + strlen(name) + 1;
-	memcpy(nla_data(na), name, strlen(name) + 1);
-	req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + NLA_ALIGN(na->nla_len);
-
-	ret = send_request(fd, &req, req.nlh.nlmsg_len);
-	if (ret)
-		return ret;
-
-	len = recv(fd, resp, sizeof(resp), 0);
-	if (len < 0)
-		return -errno;
-
-	for (nlh = (struct nlmsghdr *)resp; NLMSG_OK(nlh, len);
-	     nlh = NLMSG_NEXT(nlh, len)) {
-		if (nlh->nlmsg_type == NLMSG_ERROR) {
-			struct nlmsgerr *err = NLMSG_DATA(nlh);
-
-			return err->error ? err->error : -ENOENT;
-		}
-
-		genl = (struct genlmsghdr *)NLMSG_DATA(nlh);
-		rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN;
-		na = (struct nlattr *)((char *)genl + GENL_HDRLEN);
-		while (rem >= (int)sizeof(*na) &&
-		       na->nla_len >= (int)sizeof(*na) &&
-		       na->nla_len <= rem) {
-			if (na->nla_type == CTRL_ATTR_FAMILY_ID)
-				return *(uint16_t *)nla_data(na);
-			rem -= NLA_ALIGN(na->nla_len);
-			na = (struct nlattr *)((char *)na + NLA_ALIGN(na->nla_len));
-		}
-	}
-
-	return -ENOENT;
-}
-
 static int send_cgroupstats_cmd(int fd, int family_id, uint32_t cgroup_fd,
 				int flags)
 {
@@ -203,15 +83,12 @@ static int recv_cgroupstats_response(int fd, struct cgroupstats *stats)
 
 		rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN;
 		na = (struct nlattr *)((char *)genl + GENL_HDRLEN);
-		while (rem >= (int)sizeof(*na) &&
-		       na->nla_len >= (int)sizeof(*na) &&
-		       na->nla_len <= rem) {
+		while (nla_ok(na, rem)) {
 			if (na->nla_type == CGROUPSTATS_TYPE_CGROUP_STATS) {
 				memcpy(stats, nla_data(na), sizeof(*stats));
 				return 0;
 			}
-			rem -= NLA_ALIGN(na->nla_len);
-			na = (struct nlattr *)((char *)na + NLA_ALIGN(na->nla_len));
+			na = nla_next(na, &rem);
 		}
 	}
 
diff --git a/tools/testing/selftests/acct/netlink_helper.c b/tools/testing/selftests/acct/netlink_helper.c
new file mode 100644
index 000000000000..3ed834f0e770
--- /dev/null
+++ b/tools/testing/selftests/acct/netlink_helper.c
@@ -0,0 +1,116 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <errno.h>
+#include <stdint.h>
+#include <string.h>
+#include <sys/socket.h>
+#include <sys/time.h>
+#include <unistd.h>
+#include <linux/genetlink.h>
+
+#include "netlink_helper.h"
+
+int netlink_open(void)
+{
+	struct timeval tv = { .tv_sec = ACCT_RCV_TIMEOUT_SEC };
+	struct sockaddr_nl addr = {
+		.nl_family = AF_NETLINK,
+		.nl_pid = getpid(),
+	};
+	int fd;
+
+	fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
+	if (fd < 0)
+		return -errno;
+
+	if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) {
+		int err = -errno;
+
+		close(fd);
+		return err;
+	}
+
+	if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
+		int err = -errno;
+
+		close(fd);
+		return err;
+	}
+
+	return fd;
+}
+
+int send_request(int fd, void *buf, size_t len)
+{
+	struct sockaddr_nl addr = {
+		.nl_family = AF_NETLINK,
+	};
+
+	if (sendto(fd, buf, len, 0, (struct sockaddr *)&addr, sizeof(addr)) < 0)
+		return -errno;
+
+	return 0;
+}
+
+/*
+ * Resolve the generic netlink family ID for @name.
+ * Returns the family ID (>= 0) on success, negative errno on failure.
+ */
+int get_family_id(int fd, const char *name)
+{
+	struct {
+		struct nlmsghdr nlh;
+		struct genlmsghdr genl;
+		char buf[256];
+	} req = { 0 };
+	char resp[8192];
+	struct nlmsghdr *nlh;
+	struct genlmsghdr *genl;
+	struct nlattr *na;
+	int len;
+	int rem;
+	int ret;
+
+	req.nlh.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
+	req.nlh.nlmsg_type = GENL_ID_CTRL;
+	req.nlh.nlmsg_flags = NLM_F_REQUEST;
+	req.nlh.nlmsg_seq = 1;
+	req.nlh.nlmsg_pid = getpid();
+
+	req.genl.cmd = CTRL_CMD_GETFAMILY;
+	req.genl.version = 1;
+
+	na = (struct nlattr *)((char *)&req + NLMSG_ALIGN(req.nlh.nlmsg_len));
+	na->nla_type = CTRL_ATTR_FAMILY_NAME;
+	na->nla_len = NLA_HDRLEN + strlen(name) + 1;
+	memcpy(nla_data(na), name, strlen(name) + 1);
+	req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + NLA_ALIGN(na->nla_len);
+
+	ret = send_request(fd, &req, req.nlh.nlmsg_len);
+	if (ret)
+		return ret;
+
+	len = recv(fd, resp, sizeof(resp), 0);
+	if (len < 0)
+		return -errno;
+
+	for (nlh = (struct nlmsghdr *)resp; NLMSG_OK(nlh, len);
+	     nlh = NLMSG_NEXT(nlh, len)) {
+		if (nlh->nlmsg_type == NLMSG_ERROR) {
+			struct nlmsgerr *err = NLMSG_DATA(nlh);
+
+			return err->error ? err->error : -ENOENT;
+		}
+
+		genl = (struct genlmsghdr *)NLMSG_DATA(nlh);
+		rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN;
+		na = (struct nlattr *)((char *)genl + GENL_HDRLEN);
+		while (nla_ok(na, rem)) {
+			if (na->nla_type == CTRL_ATTR_FAMILY_ID)
+				return *(uint16_t *)nla_data(na);
+			na = nla_next(na, &rem);
+		}
+	}
+
+	return -ENOENT;
+}
diff --git a/tools/testing/selftests/acct/netlink_helper.h b/tools/testing/selftests/acct/netlink_helper.h
new file mode 100644
index 000000000000..0320729c4c06
--- /dev/null
+++ b/tools/testing/selftests/acct/netlink_helper.h
@@ -0,0 +1,44 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Shared generic netlink helpers for the acct selftests.
+ */
+#ifndef ACSELFTESTS_ACCT_NETLINK_HELPER_H
+#define ACSELFTESTS_ACCT_NETLINK_HELPER_H
+
+#include <stdbool.h>
+#include <linux/netlink.h>
+
+#ifndef NLA_ALIGNTO
+#define NLA_ALIGNTO 4
+#define NLA_ALIGN(len) (((len) + NLA_ALIGNTO - 1) & ~(NLA_ALIGNTO - 1))
+#define NLA_HDRLEN ((int)NLA_ALIGN(sizeof(struct nlattr)))
+#endif
+
+/* Fail an individual test case instead of hanging the whole binary. */
+#define ACCT_RCV_TIMEOUT_SEC 2
+
+static inline void *nla_data(const struct nlattr *na)
+{
+	return (void *)((char *)na + NLA_HDRLEN);
+}
+
+static inline bool nla_ok(const struct nlattr *na, int remaining)
+{
+	return remaining >= (int)sizeof(*na) &&
+	       na->nla_len >= sizeof(*na) &&
+	       na->nla_len <= remaining;
+}
+
+static inline struct nlattr *nla_next(const struct nlattr *na, int *remaining)
+{
+	int aligned_len = NLA_ALIGN(na->nla_len);
+
+	*remaining -= aligned_len;
+	return (struct nlattr *)((char *)na + aligned_len);
+}
+
+int netlink_open(void);
+int send_request(int fd, void *buf, size_t len);
+int get_family_id(int fd, const char *name);
+
+#endif /* ACSELFTESTS_ACCT_NETLINK_HELPER_H */
diff --git a/tools/testing/selftests/acct/taskstats_fill_stats_tgid.c b/tools/testing/selftests/acct/taskstats_fill_stats_tgid.c
index d6cab4ae26f2..9a4c1554dee3 100644
--- a/tools/testing/selftests/acct/taskstats_fill_stats_tgid.c
+++ b/tools/testing/selftests/acct/taskstats_fill_stats_tgid.c
@@ -16,14 +16,9 @@
 #include <time.h>
 #include <unistd.h>
 
+#include "netlink_helper.h"
 #include "kselftest.h"
 
-#ifndef NLA_ALIGN
-#define NLA_ALIGNTO 4
-#define NLA_ALIGN(len) (((len) + NLA_ALIGNTO - 1) & ~(NLA_ALIGNTO - 1))
-#define NLA_HDRLEN ((int)NLA_ALIGN(sizeof(struct nlattr)))
-#endif
-
 #define BUSY_NS (200ULL * 1000 * 1000)
 
 struct worker_ctx {
@@ -35,26 +30,6 @@ struct worker_ctx {
 
 static unsigned long busy_sink;
 
-static void *taskstats_nla_data(const struct nlattr *na)
-{
-	return (void *)((char *)na + NLA_HDRLEN);
-}
-
-static bool taskstats_nla_ok(const struct nlattr *na, int remaining)
-{
-	return remaining >= (int)sizeof(*na) &&
-	       na->nla_len >= sizeof(*na) &&
-	       na->nla_len <= remaining;
-}
-
-static struct nlattr *taskstats_nla_next(const struct nlattr *na, int *remaining)
-{
-	int aligned_len = NLA_ALIGN(na->nla_len);
-
-	*remaining -= aligned_len;
-	return (struct nlattr *)((char *)na + aligned_len);
-}
-
 static uint64_t timespec_diff_ns(const struct timespec *start,
 				 const struct timespec *end)
 {
@@ -84,99 +59,6 @@ static void burn_cpu_for_ns(uint64_t runtime_ns)
 	busy_sink = acc;
 }
 
-static int netlink_open(void)
-{
-	struct sockaddr_nl addr = {
-		.nl_family = AF_NETLINK,
-		.nl_pid = getpid(),
-	};
-	int fd;
-
-	fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
-	if (fd < 0)
-		return -errno;
-
-	if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
-		int err = -errno;
-
-		close(fd);
-		return err;
-	}
-
-	return fd;
-}
-
-static int send_request(int fd, void *buf, size_t len)
-{
-	struct sockaddr_nl addr = {
-		.nl_family = AF_NETLINK,
-	};
-
-	if (sendto(fd, buf, len, 0, (struct sockaddr *)&addr, sizeof(addr)) < 0)
-		return -errno;
-
-	return 0;
-}
-
-static int get_family_id(int fd, const char *name)
-{
-	struct {
-		struct nlmsghdr nlh;
-		struct genlmsghdr genl;
-		char buf[256];
-	} req = { 0 };
-	char resp[8192];
-	struct nlmsghdr *nlh;
-	struct genlmsghdr *genl;
-	struct nlattr *na;
-	int len;
-	int rem;
-	int ret;
-
-	req.nlh.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
-	req.nlh.nlmsg_type = GENL_ID_CTRL;
-	req.nlh.nlmsg_flags = NLM_F_REQUEST;
-	req.nlh.nlmsg_seq = 1;
-	req.nlh.nlmsg_pid = getpid();
-
-	req.genl.cmd = CTRL_CMD_GETFAMILY;
-	req.genl.version = 1;
-
-	na = (struct nlattr *)((char *)&req + NLMSG_ALIGN(req.nlh.nlmsg_len));
-	na->nla_type = CTRL_ATTR_FAMILY_NAME;
-	na->nla_len = NLA_HDRLEN + strlen(name) + 1;
-	memcpy(taskstats_nla_data(na), name, strlen(name) + 1);
-	req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + NLA_ALIGN(na->nla_len);
-
-	ret = send_request(fd, &req, req.nlh.nlmsg_len);
-	if (ret)
-		return ret;
-
-	len = recv(fd, resp, sizeof(resp), 0);
-	if (len < 0)
-		return -errno;
-
-	for (nlh = (struct nlmsghdr *)resp; NLMSG_OK(nlh, len);
-	     nlh = NLMSG_NEXT(nlh, len)) {
-		if (nlh->nlmsg_type == NLMSG_ERROR) {
-			struct nlmsgerr *err = NLMSG_DATA(nlh);
-
-			return err->error ? err->error : -ENOENT;
-		}
-
-		genl = (struct genlmsghdr *)NLMSG_DATA(nlh);
-		rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN;
-		na = (struct nlattr *)((char *)genl + GENL_HDRLEN);
-		while (taskstats_nla_ok(na, rem)) {
-			if (na->nla_type == CTRL_ATTR_FAMILY_ID)
-				return *(uint16_t *)taskstats_nla_data(na);
-			na = taskstats_nla_next(na, &rem);
-		}
-	}
-
-	return -ENOENT;
-}
-
 static int get_taskstats(int fd, int family_id, uint16_t attr_type, uint32_t id,
 			 struct taskstats *stats)
 {
@@ -209,7 +91,7 @@ static int get_taskstats(int fd, int family_id, uint16_t attr_type, uint32_t id,
 	na = (struct nlattr *)((char *)&req + NLMSG_ALIGN(req.nlh.nlmsg_len));
 	na->nla_type = attr_type;
 	na->nla_len = NLA_HDRLEN + sizeof(id);
-	memcpy(taskstats_nla_data(na), &id, sizeof(id));
+	memcpy(nla_data(na), &id, sizeof(id));
 	req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + NLA_ALIGN(na->nla_len);
 
 	ret = send_request(fd, &req, req.nlh.nlmsg_len);
@@ -231,21 +113,21 @@ static int get_taskstats(int fd, int family_id, uint16_t attr_type, uint32_t id,
 		genl = (struct genlmsghdr *)NLMSG_DATA(nlh);
 		rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN;
 		na = (struct nlattr *)((char *)genl + GENL_HDRLEN);
-		while (taskstats_nla_ok(na, rem)) {
+		while (nla_ok(na, rem)) {
 			if (na->nla_type == TASKSTATS_TYPE_AGGR_PID ||
 			    na->nla_type == TASKSTATS_TYPE_AGGR_TGID) {
-				nested = (struct nlattr *)taskstats_nla_data(na);
+				nested = (struct nlattr *)nla_data(na);
 				nrem = na->nla_len - NLA_HDRLEN;
-				while (taskstats_nla_ok(nested, nrem)) {
+				while (nla_ok(nested, nrem)) {
 					if (nested->nla_type == TASKSTATS_TYPE_STATS) {
-						memcpy(stats, taskstats_nla_data(nested),
+						memcpy(stats, nla_data(nested),
 						       sizeof(*stats));
 						return 0;
 					}
-					nested = taskstats_nla_next(nested, &nrem);
+					nested = nla_next(nested, &nrem);
 				}
 			}
-			na = taskstats_nla_next(na, &rem);
+			na = nla_next(na, &rem);
 		}
 	}
 
-- 
2.43.0