[PATCH net-next] net: skbuff: replace ubuf_info ops ptr with type id

@@ -665,12 +665,10 @@ static ssize_t tap_get_user(struct tap_queue *q, void *msg_control,
 		skb_set_network_header(skb, depth);
 
 	/* copy skb_ubuf_info for callback when skb has no error */
-	if (zerocopy) {
+	if (zerocopy)
 		skb_zcopy_init(skb, msg_control);
-	} else if (msg_control) {
-		struct ubuf_info *uarg = msg_control;
-		uarg->ops->complete(NULL, uarg, false);
-	}
+	else if (msg_control)
+		net_zcopy_complete(NULL, msg_control, false);
 
 	dev_queue_xmit(skb);
 	rcu_read_unlock();

@@ -1954,12 +1954,10 @@ static ssize_t tun_get_user(struct tun_struct *tun, struct tun_file *tfile,
 		skb->protocol = eth_type_trans(skb, tun->dev);
 
 	/* copy skb_ubuf_info for callback when skb has no error */
-	if (zerocopy) {
+	if (zerocopy)
 		skb_zcopy_init(skb, msg_control);
-	} else if (msg_control) {
-		struct ubuf_info *uarg = msg_control;
-		uarg->ops->complete(NULL, uarg, false);
-	}
+	else if (msg_control)
+		net_zcopy_complete(NULL, msg_control, false);
 
 	skb_reset_network_header(skb);
 	skb_probe_transport_header(skb);

@@ -390,9 +390,6 @@ bool xenvif_rx_queue_tail(struct xenvif_queue *queue, struct sk_buff *skb);
 
 void xenvif_carrier_on(struct xenvif *vif);
 
-/* Callbacks from stack when TX packet can be released */
-extern const struct ubuf_info_ops xenvif_ubuf_ops;
-
 static inline pending_ring_idx_t nr_pending_reqs(struct xenvif_queue *queue)
 {
 	return MAX_PENDING_REQS -

@@ -593,7 +593,7 @@ int xenvif_init_queue(struct xenvif_queue *queue)
 
 	for (i = 0; i < MAX_PENDING_REQS; i++) {
 		queue->pending_tx_info[i].callback_struct = (struct ubuf_info_msgzc)
-			{ { .ops = &xenvif_ubuf_ops },
+			{ { .type = UBUF_TYPE_XENVIF },
 			  { { .ctx = NULL,
 			      .desc = i } } };
 		queue->grant_tx_handle[i] = NETBACK_INVALID_HANDLE;

@@ -1158,7 +1158,7 @@ static int xenvif_handle_frag_list(struct xenvif_queue *queue, struct sk_buff *s
 	uarg = skb_shinfo(skb)->destructor_arg;
 	/* increase inflight counter to offset decrement in callback */
 	atomic_inc(&queue->inflight_packets);
-	uarg->ops->complete(NULL, uarg, true);
+	net_zcopy_complete(NULL, uarg, true);
 	skb_shinfo(skb)->destructor_arg = NULL;
 
 	/* Fill the skb with the new (local) frags. */

@@ -1315,10 +1315,6 @@ static void xenvif_zerocopy_callback(struct sk_buff *skb,
 	xenvif_skb_zerocopy_complete(queue);
 }
 
-const struct ubuf_info_ops xenvif_ubuf_ops = {
-	.complete = xenvif_zerocopy_callback,
-};
-
 static inline void xenvif_tx_dealloc_action(struct xenvif_queue *queue)
 {
 	struct gnttab_unmap_grant_ref *gop;

@@ -1760,6 +1756,8 @@ static int __init netback_init(void)
 		fatal_skb_slots = XEN_NETBK_LEGACY_SLOTS_MAX;
 	}
 
+	net_zcopy_register(UBUF_TYPE_XENVIF, xenvif_zerocopy_callback);
+
 	rc = xenvif_xenbus_init();
 	if (rc)
 		goto failed_init;

@@ -1771,6 +1769,7 @@ static int __init netback_init(void)
 	return 0;
 
 failed_init:
+	net_zcopy_register(UBUF_TYPE_XENVIF, NULL);
 	return rc;
 }
 

@@ -1782,6 +1781,7 @@ static void __exit netback_fini(void)
 	debugfs_remove_recursive(xen_netback_dbg_root);
 #endif /* CONFIG_DEBUG_FS */
 	xenvif_xenbus_fini();
+	net_zcopy_register(UBUF_TYPE_XENVIF, NULL);
 }
 module_exit(netback_fini);
 

@@ -431,10 +431,6 @@ static void vhost_zerocopy_complete(struct sk_buff *skb,
 	rcu_read_unlock_bh();
 }
 
-static const struct ubuf_info_ops vhost_ubuf_ops = {
-	.complete = vhost_zerocopy_complete,
-};
-
 static inline unsigned long busy_clock(void)
 {
 	return local_clock() >> 10;

@@ -926,7 +922,7 @@ static void handle_tx_zerocopy(struct vhost_net *net, struct socket *sock)
 			vq->heads[nvq->upend_idx].len = VHOST_DMA_IN_PROGRESS;
 			ubuf->ctx = nvq->ubufs;
 			ubuf->desc = nvq->upend_idx;
-			ubuf->ubuf.ops = &vhost_ubuf_ops;
+			ubuf->ubuf.type = UBUF_TYPE_VHOST;
 			ubuf->ubuf.flags = SKBFL_ZEROCOPY_FRAG;
 			refcount_set(&ubuf->ubuf.refcnt, 1);
 			msg.msg_control = &ctl;

@@ -1895,15 +1891,25 @@ static struct miscdevice vhost_net_misc = {
 
 static int __init vhost_net_init(void)
 {
+	int rc;
+
 	if (experimental_zcopytx)
 		vhost_net_enable_zcopy(VHOST_NET_VQ_TX);
-	return misc_register(&vhost_net_misc);
+
+	net_zcopy_register(UBUF_TYPE_VHOST, vhost_zerocopy_complete);
+
+	rc = misc_register(&vhost_net_misc);
+	if (rc)
+		net_zcopy_register(UBUF_TYPE_VHOST, NULL);
+
+	return rc;
 }
 module_init(vhost_net_init);
 
 static void __exit vhost_net_exit(void)
 {
 	misc_deregister(&vhost_net_misc);
+	net_zcopy_register(UBUF_TYPE_VHOST, NULL);
 }
 module_exit(vhost_net_exit);
 

@@ -2,10 +2,19 @@
 #ifndef _LINUX_IO_URING_NET_H
 #define _LINUX_IO_URING_NET_H
 
+#include <linux/errno.h>
+#include <linux/types.h>
+#include <net/net_debug.h>
+
 struct io_uring_cmd;
+struct sk_buff;
+struct ubuf_info;
 
 #if defined(CONFIG_IO_URING)
 int io_uring_cmd_sock(struct io_uring_cmd *cmd, unsigned int issue_flags);
+void io_tx_ubuf_complete(struct sk_buff *skb, struct ubuf_info *uarg,
+			 bool success);
+int io_link_skb(struct sk_buff *skb, struct ubuf_info *uarg);
 
 #else
 static inline int io_uring_cmd_sock(struct io_uring_cmd *cmd,

@@ -13,6 +22,17 @@ static inline int io_uring_cmd_sock(struct io_uring_cmd *cmd,
 {
 	return -EOPNOTSUPP;
 }
+
+static inline void io_tx_ubuf_complete(struct sk_buff *skb,
+				       struct ubuf_info *uarg, bool success)
+{
+	DEBUG_NET_WARN_ON_ONCE(1);
+}
+
+static inline int io_link_skb(struct sk_buff *skb, struct ubuf_info *uarg)
+{
+	return -EOPNOTSUPP;
+}
 #endif
 
 #endif

@@ -534,13 +534,18 @@ enum {
 #define SKBFL_ALL_ZEROCOPY	(SKBFL_ZEROCOPY_FRAG | SKBFL_PURE_ZEROCOPY | \
 				 SKBFL_DONT_ORPHAN | SKBFL_MANAGED_FRAG_REFS)
 
-struct ubuf_info_ops {
-	void (*complete)(struct sk_buff *, struct ubuf_info *,
-			 bool zerocopy_success);
-	/* has to be compatible with skb_zcopy_set() */
-	int (*link_skb)(struct sk_buff *skb, struct ubuf_info *uarg);
+enum ubuf_info_type {
+	UBUF_TYPE_NONE = 0,
+	UBUF_TYPE_MSG_ZEROCOPY,
+	UBUF_TYPE_IO_URING,
+	UBUF_TYPE_TPACKET,
+	UBUF_TYPE_VHOST,
+	UBUF_TYPE_XENVIF,
 };
 
+typedef void (*ubuf_complete_fn)(struct sk_buff *skb, struct ubuf_info *uarg,
+				 bool zerocopy_success);
+
 /*
  * The callback notifies userspace to release buffers when skb DMA is done in
  * lower device, the skb last reference should be 0 when calling this.

@@ -550,8 +555,8 @@ struct ubuf_info_ops {
  * The desc field is used to track userspace buffer index.
  */
 struct ubuf_info {
-	const struct ubuf_info_ops *ops;
 	refcount_t refcnt;
+	u8 type;
 	u8 flags;
 };
 

@@ -1757,7 +1762,9 @@ static inline void skb_set_end_offset(struct sk_buff *skb, unsigned int offset)
 }
 #endif
 
-extern const struct ubuf_info_ops msg_zerocopy_ubuf_ops;
+void net_zcopy_register(enum ubuf_info_type type, ubuf_complete_fn fn);
+void net_zcopy_complete(struct sk_buff *skb, struct ubuf_info *uarg,
+			bool zerocopy_success);
 
 struct ubuf_info *msg_zerocopy_realloc(struct sock *sk, size_t size,
 				       struct ubuf_info *uarg, bool devmem);

@@ -1843,13 +1850,13 @@ static inline void skb_zcopy_set(struct sk_buff *skb, struct ubuf_info *uarg,
 static inline void net_zcopy_put(struct ubuf_info *uarg)
 {
 	if (uarg)
-		uarg->ops->complete(NULL, uarg, true);
+		net_zcopy_complete(NULL, uarg, true);
 }
 
 static inline void net_zcopy_put_abort(struct ubuf_info *uarg, bool have_uref)
 {
 	if (uarg) {
-		if (uarg->ops == &msg_zerocopy_ubuf_ops)
+		if (uarg->type == UBUF_TYPE_MSG_ZEROCOPY)
 			msg_zerocopy_put_abort(uarg, have_uref);
 		else if (have_uref)
 			net_zcopy_put(uarg);

@@ -1862,7 +1869,7 @@ static inline void skb_zcopy_clear(struct sk_buff *skb, bool zerocopy_success)
 	struct ubuf_info *uarg = skb_zcopy(skb);
 
 	if (uarg) {
-		uarg->ops->complete(skb, uarg, zerocopy_success);
+		net_zcopy_complete(skb, uarg, zerocopy_success);
 
 		skb_shinfo(skb)->flags &= ~SKBFL_ALL_ZEROCOPY;
 	}

@@ -10,8 +10,6 @@
 #include "notif.h"
 #include "rsrc.h"
 
-static const struct ubuf_info_ops io_ubuf_ops;
-
 static void io_notif_tw_complete(struct io_tw_req tw_req, io_tw_token_t tw)
 {
 	struct io_kiocb *notif = tw_req.req;

@@ -68,7 +66,7 @@ void io_tx_ubuf_complete(struct sk_buff *skb, struct ubuf_info *uarg,
 	__io_req_task_work_add(notif, tw_flags);
 }
 
-static int io_link_skb(struct sk_buff *skb, struct ubuf_info *uarg)
+int io_link_skb(struct sk_buff *skb, struct ubuf_info *uarg)
 {
 	struct io_notif_data *nd, *prev_nd;
 	struct io_kiocb *prev_notif, *notif;

@@ -89,7 +87,7 @@ static int io_link_skb(struct sk_buff *skb, struct ubuf_info *uarg)
 	if (unlikely(nd->head != nd || nd->next))
 		return -EEXIST;
 	/* don't mix zc providers */
-	if (unlikely(prev_uarg->ops != &io_ubuf_ops))
+	if (unlikely(prev_uarg->type != UBUF_TYPE_IO_URING))
 		return -EEXIST;
 
 	prev_nd = container_of(prev_uarg, struct io_notif_data, uarg);

@@ -107,11 +105,6 @@ static int io_link_skb(struct sk_buff *skb, struct ubuf_info *uarg)
 	return 0;
 }
 
-static const struct ubuf_info_ops io_ubuf_ops = {
-	.complete = io_tx_ubuf_complete,
-	.link_skb = io_link_skb,
-};
-
 struct io_kiocb *io_alloc_notif(struct io_ring_ctx *ctx)
 	__must_hold(&ctx->uring_lock)
 {

@@ -137,7 +130,7 @@ struct io_kiocb *io_alloc_notif(struct io_ring_ctx *ctx)
 	nd->head = nd;
 
 	nd->uarg.flags = IO_NOTIF_UBUF_FLAGS;
-	nd->uarg.ops = &io_ubuf_ops;
+	nd->uarg.type = UBUF_TYPE_IO_URING;
 	refcount_set(&nd->uarg.refcnt, 1);
 	return notif;
 }

@@ -4,6 +4,7 @@
 #include <linux/uio.h>
 #include <net/sock.h>
 #include <linux/nospec.h>
+#include <linux/io_uring/net.h>
 
 #include "rsrc.h"
 

@@ -24,8 +25,6 @@ struct io_notif_data {
 };
 
 struct io_kiocb *io_alloc_notif(struct io_ring_ctx *ctx);
-void io_tx_ubuf_complete(struct sk_buff *skb, struct ubuf_info *uarg,
-			 bool success);
 
 static inline struct io_notif_data *io_notif_to_data(struct io_kiocb *notif)
 {

@@ -92,6 +92,8 @@
 #include <linux/capability.h>
 #include <linux/user_namespace.h>
 #include <linux/indirect_call_wrapper.h>
+#include <linux/io_uring/net.h>
+#include <linux/static_call.h>
 #include <linux/textsearch.h>
 
 #include "dev.h"

@@ -1737,7 +1739,7 @@ static struct ubuf_info *msg_zerocopy_alloc(struct sock *sk, size_t size,
 		return NULL;
 	}
 
-	uarg->ubuf.ops = &msg_zerocopy_ubuf_ops;
+	uarg->ubuf.type = UBUF_TYPE_MSG_ZEROCOPY;
 	uarg->id = ((u32)atomic_inc_return(&sk->sk_zckey)) - 1;
 	uarg->len = 1;
 	uarg->bytelen = size;

@@ -1763,7 +1765,7 @@ struct ubuf_info *msg_zerocopy_realloc(struct sock *sk, size_t size,
 		u32 bytelen, next;
 
 		/* there might be non MSG_ZEROCOPY users */
-		if (uarg->ops != &msg_zerocopy_ubuf_ops)
+		if (uarg->type != UBUF_TYPE_MSG_ZEROCOPY)
 			return NULL;
 
 		/* realloc only when socket is locked (TCP, UDP cork),

@@ -1898,10 +1900,72 @@ void msg_zerocopy_put_abort(struct ubuf_info *uarg, bool have_uref)
 }
 EXPORT_SYMBOL_GPL(msg_zerocopy_put_abort);
 
-const struct ubuf_info_ops msg_zerocopy_ubuf_ops = {
-	.complete = msg_zerocopy_complete,
-};
-EXPORT_SYMBOL_GPL(msg_zerocopy_ubuf_ops);
+static void net_zcopy_unregistered(struct sk_buff *skb, struct ubuf_info *uarg,
+				   bool zerocopy_success)
+{
+	DEBUG_NET_WARN_ON_ONCE(1);
+}
+
+DEFINE_STATIC_CALL(ubuf_complete_tpacket, net_zcopy_unregistered);
+DEFINE_STATIC_CALL(ubuf_complete_vhost, net_zcopy_unregistered);
+DEFINE_STATIC_CALL(ubuf_complete_xenvif, net_zcopy_unregistered);
+
+void net_zcopy_register(enum ubuf_info_type type, ubuf_complete_fn fn)
+{
+	ubuf_complete_fn cur;
+	bool reg = fn;
+
+	if (!fn)
+		fn = net_zcopy_unregistered;
+
+	switch (type) {
+	case UBUF_TYPE_TPACKET:
+		cur = static_call_query(ubuf_complete_tpacket);
+		static_call_update(ubuf_complete_tpacket, fn);
+		break;
+	case UBUF_TYPE_VHOST:
+		cur = static_call_query(ubuf_complete_vhost);
+		static_call_update(ubuf_complete_vhost, fn);
+		break;
+	case UBUF_TYPE_XENVIF:
+		cur = static_call_query(ubuf_complete_xenvif);
+		static_call_update(ubuf_complete_xenvif, fn);
+		break;
+	default:
+		WARN_ON_ONCE(1);
+		return;
+	}
+
+	WARN_ON_ONCE(reg ? cur != net_zcopy_unregistered :
+			   cur == net_zcopy_unregistered);
+}
+EXPORT_SYMBOL_GPL(net_zcopy_register);
+
+void net_zcopy_complete(struct sk_buff *skb, struct ubuf_info *uarg,
+			bool zerocopy_success)
+{
+	switch (uarg->type) {
+	case UBUF_TYPE_MSG_ZEROCOPY:
+		msg_zerocopy_complete(skb, uarg, zerocopy_success);
+		break;
+	case UBUF_TYPE_IO_URING:
+		io_tx_ubuf_complete(skb, uarg, zerocopy_success);
+		break;
+	case UBUF_TYPE_TPACKET:
+		static_call(ubuf_complete_tpacket)(skb, uarg, zerocopy_success);
+		break;
+	case UBUF_TYPE_VHOST:
+		static_call(ubuf_complete_vhost)(skb, uarg, zerocopy_success);
+		break;
+	case UBUF_TYPE_XENVIF:
+		static_call(ubuf_complete_xenvif)(skb, uarg, zerocopy_success);
+		break;
+	default:
+		DEBUG_NET_WARN_ON_ONCE(1);
+		break;
+	}
+}
+EXPORT_SYMBOL_GPL(net_zcopy_complete);
 
 int skb_zerocopy_iter_stream(struct sock *sk, struct sk_buff *skb,
 			     struct msghdr *msg, int len,

@@ -1910,8 +1974,8 @@ int skb_zerocopy_iter_stream(struct sock *sk, struct sk_buff *skb,
 {
 	int err, orig_len = skb->len;
 
-	if (uarg->ops->link_skb) {
-		err = uarg->ops->link_skb(skb, uarg);
+	if (uarg->type == UBUF_TYPE_IO_URING) {
+		err = io_link_skb(skb, uarg);
 		if (err)
 			return err;
 	} else {

@@ -2606,10 +2606,6 @@ static void tpacket_ubuf_complete(struct sk_buff *skb, struct ubuf_info *uarg,
 	sk_free(&po->sk);
 }
 
-static const struct ubuf_info_ops tpacket_ubuf_ops = {
-	.complete = tpacket_ubuf_complete,
-};
-
 static int tpacket_fill_skb(struct packet_sock *po, struct sk_buff *skb,
 		struct net_device *dev, void *data, int tp_len,
 		__be16 proto, unsigned char *addr, int hlen, int copylen,

@@ -2942,7 +2938,7 @@ static int tpacket_snd(struct packet_sock *po, struct msghdr *msg)
 		}
 		uarg->po = po;
 		uarg->ph = ph;
-		uarg->ubuf.ops = &tpacket_ubuf_ops;
+		uarg->ubuf.type = UBUF_TYPE_TPACKET;
 		uarg->ubuf.flags = SKBFL_ZEROCOPY_FRAG;
 		refcount_set(&uarg->ubuf.refcnt, 1);
 

@@ -4931,6 +4927,7 @@ static void __exit packet_exit(void)
 	proto_unregister(&packet_proto);
 	unregister_netdevice_notifier(&packet_netdev_notifier);
 	unregister_pernet_subsys(&packet_net_ops);
+	net_zcopy_register(UBUF_TYPE_TPACKET, NULL);
 }
 
 static int __init packet_init(void)

@@ -4946,6 +4943,8 @@ static int __init packet_init(void)
 	rc = proto_register(&packet_proto, 0);
 	if (rc)
 		goto out_notifier;
+
+	net_zcopy_register(UBUF_TYPE_TPACKET, tpacket_ubuf_complete);
 	rc = sock_register(&packet_family_ops);
 	if (rc)
 		goto out_proto;

@@ -4953,6 +4952,7 @@ static int __init packet_init(void)
 	return 0;
 
 out_proto:
+	net_zcopy_register(UBUF_TYPE_TPACKET, NULL);
 	proto_unregister(&packet_proto);
 out_notifier:
 	unregister_netdevice_notifier(&packet_netdev_notifier);