[PATCH RFC 6/8] vsock: allow datagrams to use directional transports

@@ -465,6 +465,7 @@ static struct virtio_transport vhost_transport = {
 		.dgram_dequeue            = virtio_transport_dgram_dequeue,
 		.dgram_bind               = virtio_transport_dgram_bind,
 		.dgram_allow              = vhost_transport_dgram_allow,
+		.dgram_transport_match    = virtio_transport_dgram_transport_match,
 
 		.stream_enqueue           = virtio_transport_stream_enqueue,
 		.stream_dequeue           = virtio_transport_stream_dequeue,

@@ -14,6 +14,7 @@ struct virtio_vsock_skb_cb {
 	bool tap_delivered;
 	/* Datagram payload is still charged to bytes_unsent. */
 	bool tx_accounted;
+	bool to_host;
 	u32 offset;
 };
 

@@ -173,6 +174,7 @@ struct virtio_vsock_sock {
 
 struct virtio_vsock_pkt_info {
 	u32 remote_cid, remote_port;
+	const struct virtio_transport *t_ops;
 	struct vsock_sock *vsk;
 	struct msghdr *msg;
 	struct net *net;

@@ -263,6 +265,7 @@ bool virtio_transport_stream_allow(struct vsock_sock *vsk, u32 cid, u32 port);
 int virtio_transport_dgram_bind(struct vsock_sock *vsk,
 				struct sockaddr_vm *addr);
 bool virtio_transport_dgram_allow(struct vsock_sock *vsk, u32 cid, u32 port);
+bool virtio_transport_dgram_transport_match(const struct vsock_transport *t);
 
 int virtio_transport_connect(struct vsock_sock *vsk);
 

@@ -276,6 +279,7 @@ virtio_transport_stream_enqueue(struct vsock_sock *vsk,
 				size_t len);
 int
 virtio_transport_dgram_enqueue(struct vsock_sock *vsk,
+			       const struct vsock_transport *transport,
 			       struct sockaddr_vm *remote_addr,
 			       struct msghdr *msg,
 			       size_t len);

@@ -121,9 +121,13 @@ struct vsock_transport {
 	int (*dgram_bind)(struct vsock_sock *, struct sockaddr_vm *);
 	int (*dgram_dequeue)(struct vsock_sock *vsk, struct msghdr *msg,
 			     size_t len, int flags);
-	int (*dgram_enqueue)(struct vsock_sock *, struct sockaddr_vm *,
-			     struct msghdr *, size_t len);
+	int (*dgram_enqueue)(struct vsock_sock *vsk,
+			     const struct vsock_transport *transport,
+			     struct sockaddr_vm *remote_addr,
+			     struct msghdr *msg, size_t len);
 	bool (*dgram_allow)(struct vsock_sock *vsk, u32 cid, u32 port);
+	/* Optional: another transport can use this socket's datagram state. */
+	bool (*dgram_transport_match)(const struct vsock_transport *transport);
 	/* Optional payload length override, called with the receive queue locked. */
 	size_t (*dgram_payload_len)(const struct sk_buff *skb);
 

@@ -240,6 +244,9 @@ void vsock_remove_sock(struct vsock_sock *vsk);
 void vsock_for_each_connected_socket(struct vsock_transport *transport,
 				     void (*fn)(struct sock *sk));
 int vsock_assign_transport(struct vsock_sock *vsk, struct vsock_sock *psk);
+const struct vsock_transport *
+vsock_dgram_get_transport(struct vsock_sock *vsk,
+			  struct sockaddr_vm *remote_addr);
 bool vsock_find_cid(unsigned int cid);
 void vsock_linger(struct sock *sk);
 

@@ -136,6 +136,7 @@
 #include <linux/types.h>
 #include <linux/bitops.h>
 #include <linux/cred.h>
+#include <linux/err.h>
 #include <linux/errqueue.h>
 #include <linux/init.h>
 #include <linux/io.h>

@@ -599,6 +600,7 @@ static void vsock_deassign_transport(struct vsock_sock *vsk)
 int vsock_assign_transport(struct vsock_sock *vsk, struct vsock_sock *psk)
 {
 	const struct vsock_transport *new_transport;
+	const struct vsock_transport *fallback_transport = NULL;
 	struct sock *sk = sk_vsock(vsk);
 	unsigned int remote_cid = vsk->remote_addr.svm_cid;
 	__u8 remote_flags;

@@ -622,7 +624,12 @@ int vsock_assign_transport(struct vsock_sock *vsk, struct vsock_sock *psk)
 
 	switch (sk->sk_type) {
 	case SOCK_DGRAM:
-		new_transport = transport_dgram;
+		if (transport_dgram)
+			new_transport = transport_dgram;
+		else if (transport_g2h)
+			new_transport = transport_g2h;
+		else
+			new_transport = transport_h2g;
 		break;
 	case SOCK_STREAM:
 	case SOCK_SEQPACKET:

@@ -661,6 +668,11 @@ int vsock_assign_transport(struct vsock_sock *vsk, struct vsock_sock *psk)
 		ret = -ENODEV;
 		goto err;
 	}
+	if (sk->sk_type == SOCK_DGRAM && !transport_dgram &&
+	    new_transport == transport_g2h && transport_h2g &&
+	    transport_h2g != new_transport &&
+	    try_module_get(transport_h2g->module))
+		fallback_transport = transport_h2g;
 
 	/* It's safe to release the mutex after a successful try_module_get().
 	 * Whichever transport `new_transport` points at, it won't go away until

@@ -697,6 +709,14 @@ int vsock_assign_transport(struct vsock_sock *vsk, struct vsock_sock *psk)
 	}
 
 	ret = new_transport->init(vsk, psk);
+	if (ret == -ESOCKTNOSUPPORT && fallback_transport) {
+		module_put(new_transport->module);
+		new_transport = fallback_transport;
+		fallback_transport = NULL;
+		ret = new_transport->init(vsk, psk);
+	}
+	if (fallback_transport)
+		module_put(fallback_transport->module);
 	if (ret) {
 		module_put(new_transport->module);
 		return ret;

@@ -711,6 +731,98 @@ int vsock_assign_transport(struct vsock_sock *vsk, struct vsock_sock *psk)
 }
 EXPORT_SYMBOL_GPL(vsock_assign_transport);
 
+/**
+ * vsock_dgram_get_transport - select a transport for a datagram destination
+ * @vsk: socket whose datagram state will be used
+ * @remote_addr: destination address
+ *
+ * The caller must hold the socket lock and release the returned transport's
+ * module reference after use. Dedicated datagram transports retain their
+ * original selection; compatible transports share the socket's existing state.
+ * Parent fallback sets VMADDR_FLAG_TO_HOST to preserve that route on connect.
+ *
+ * Return: a referenced transport or an ERR_PTR on failure.
+ */
+const struct vsock_transport *
+vsock_dgram_get_transport(struct vsock_sock *vsk,
+			  struct sockaddr_vm *remote_addr)
+{
+	const struct vsock_transport *owner = vsk->transport;
+	const struct vsock_transport *transport;
+	struct sock *sk = sk_vsock(vsk);
+	u32 cid = remote_addr->svm_cid;
+	u32 local_cid;
+	bool to_host = false;
+	int err = -ENODEV;
+
+	if (!owner)
+		return ERR_PTR(-ENODEV);
+
+	mutex_lock(&vsock_register_mutex);
+	if (!owner->dgram_transport_match) {
+		transport = owner;
+	} else if (vsock_use_local_transport(cid)) {
+		transport = transport_local;
+	} else if (cid <= VMADDR_CID_HOST ||
+		   (remote_addr->svm_flags & VMADDR_FLAG_TO_HOST)) {
+		transport = transport_g2h;
+	} else if (transport_h2g &&
+		   (!transport_h2g->has_remote_cid ||
+		    transport_h2g->has_remote_cid(vsk, cid))) {
+		transport = transport_h2g;
+	} else if (sock_net(sk)->vsock.g2h_fallback && transport_g2h &&
+		   transport_g2h->has_remote_cid &&
+		   transport_g2h->has_remote_cid(vsk, cid)) {
+		transport = transport_g2h;
+		to_host = true;
+	} else {
+		transport = transport_h2g;
+	}
+
+	if (!transport)
+		goto out;
+	if (owner->dgram_transport_match &&
+	    !owner->dgram_transport_match(transport)) {
+		err = -ESOCKTNOSUPPORT;
+		goto out;
+	}
+	/* Loopback source addresses must route replies back to loopback. */
+	if (owner->dgram_transport_match && transport == transport_local &&
+	    vsk->local_addr.svm_cid != VMADDR_CID_ANY &&
+	    !vsock_use_local_transport(vsk->local_addr.svm_cid)) {
+		err = -EADDRNOTAVAIL;
+		goto out;
+	}
+	if (!try_module_get(transport->module))
+		goto out;
+	mutex_unlock(&vsock_register_mutex);
+
+	if (owner->dgram_transport_match) {
+		local_cid = transport->get_local_cid();
+		if (vsk->local_addr.svm_cid != VMADDR_CID_ANY &&
+		    local_cid != VMADDR_CID_LOCAL &&
+		    vsk->local_addr.svm_cid != local_cid) {
+			err = -EADDRNOTAVAIL;
+			goto put_transport;
+		}
+	}
+	if (!transport->dgram_allow(vsk, cid, remote_addr->svm_port)) {
+		err = -EINVAL;
+		goto put_transport;
+	}
+	if (to_host)
+		remote_addr->svm_flags |= VMADDR_FLAG_TO_HOST;
+	return transport;
+
+put_transport:
+	module_put(transport->module);
+	return ERR_PTR(err);
+out:
+	mutex_unlock(&vsock_register_mutex);
+	return ERR_PTR(err);
+}
+EXPORT_SYMBOL_GPL(vsock_dgram_get_transport);
+
 /*
  * Provide safe access to static transport_{h2g,g2h,dgram,local} callbacks.
  * Otherwise we may race with module removal. Do not use on `vsk->transport`.

@@ -1040,7 +1152,7 @@ static void __vsock_release(struct sock *sk, int level)
 
 	if (vsk->transport)
 		vsk->transport->release(vsk);
-	else if (sock_type_connectible(sk->sk_type))
+	else
 		vsock_remove_sock(vsk);
 
 	sock_orphan(sk);

@@ -1483,7 +1595,8 @@ static int vsock_dgram_sendmsg(struct socket *sock, struct msghdr *msg,
 		 */
 
 		if (remote_addr->svm_cid == VMADDR_CID_ANY)
-			remote_addr->svm_cid = transport->get_local_cid();
+			remote_addr->svm_cid = transport->dgram_transport_match ?
+				VMADDR_CID_LOCAL : transport->get_local_cid();
 
 		if (!vsock_addr_bound(remote_addr)) {
 			err = -EINVAL;

@@ -1493,7 +1606,8 @@ static int vsock_dgram_sendmsg(struct socket *sock, struct msghdr *msg,
 		remote_addr = &vsk->remote_addr;
 
 		if (remote_addr->svm_cid == VMADDR_CID_ANY)
-			remote_addr->svm_cid = transport->get_local_cid();
+			remote_addr->svm_cid = transport->dgram_transport_match ?
+				VMADDR_CID_LOCAL : transport->get_local_cid();
 
 		/* XXX Should connect() or this function ensure remote_addr is
 		 * bound?

@@ -1507,13 +1621,14 @@ static int vsock_dgram_sendmsg(struct socket *sock, struct msghdr *msg,
 		goto out;
 	}
 
-	if (!transport->dgram_allow(vsk, remote_addr->svm_cid,
-				    remote_addr->svm_port)) {
-		err = -EINVAL;
+	transport = vsock_dgram_get_transport(vsk, remote_addr);
+	if (IS_ERR(transport)) {
+		err = PTR_ERR(transport);
 		goto out;
 	}
 
-	err = transport->dgram_enqueue(vsk, remote_addr, msg, len);
+	err = transport->dgram_enqueue(vsk, transport, remote_addr, msg, len);
+	module_put(transport->module);
 
 out:
 	release_sock(sk);

@@ -1527,6 +1642,7 @@ static int vsock_dgram_connect(struct socket *sock,
 	struct sock *sk;
 	struct vsock_sock *vsk;
 	struct sockaddr_vm *remote_addr;
+	const struct vsock_transport *transport;
 
 	sk = sock->sk;
 	vsk = vsock_sk(sk);

@@ -1548,11 +1664,16 @@ static int vsock_dgram_connect(struct socket *sock,
 	if (err)
 		goto out;
 
-	if (!vsk->transport->dgram_allow(vsk, remote_addr->svm_cid,
-					 remote_addr->svm_port)) {
-		err = -EINVAL;
+	if (remote_addr->svm_cid == VMADDR_CID_ANY &&
+	    vsk->transport->dgram_transport_match)
+		remote_addr->svm_cid = VMADDR_CID_LOCAL;
+
+	transport = vsock_dgram_get_transport(vsk, remote_addr);
+	if (IS_ERR(transport)) {
+		err = PTR_ERR(transport);
 		goto out;
 	}
+	module_put(transport->module);
 
 	memcpy(&vsk->remote_addr, remote_addr, sizeof(vsk->remote_addr));
 	sock->state = SS_CONNECTED;

@@ -443,6 +443,9 @@ static int hvs_sock_init(struct vsock_sock *vsk, struct vsock_sock *psk)
 	struct hvsock *hvs;
 	struct sock *sk = sk_vsock(vsk);
 
+	if (sk->sk_type == SOCK_DGRAM)
+		return -ESOCKTNOSUPPORT;
+
 	hvs = kzalloc_obj(*hvs);
 	if (!hvs)
 		return -ENOMEM;

@@ -563,6 +566,7 @@ static int hvs_dgram_dequeue(struct vsock_sock *vsk, struct msghdr *msg,
 }
 
 static int hvs_dgram_enqueue(struct vsock_sock *vsk,
+			     const struct vsock_transport *transport,
 			     struct sockaddr_vm *remote, struct msghdr *msg,
 			     size_t dgram_len)
 {

@@ -585,6 +585,7 @@ static struct virtio_transport virtio_transport = {
 		.dgram_dequeue            = virtio_transport_dgram_dequeue,
 		.dgram_enqueue            = virtio_transport_dgram_enqueue,
 		.dgram_allow              = virtio_transport_dgram_allow_check,
+		.dgram_transport_match    = virtio_transport_dgram_transport_match,
 
 		.stream_dequeue           = virtio_transport_stream_dequeue,
 		.stream_enqueue           = virtio_transport_stream_enqueue,

@@ -322,13 +322,18 @@ static int virtio_transport_send_pkt_info(struct vsock_sock *vsk,
 
 	info->type = virtio_transport_get_type(sk_vsock(vsk));
 
-	t_ops = virtio_transport_get_ops(vsk);
+	t_ops = info->t_ops ?: virtio_transport_get_ops(vsk);
 	if (unlikely(!t_ops))
 		return -EFAULT;
 
 	src_cid = t_ops->transport.get_local_cid();
+	/* Preserve explicit datagram bindings so loopback replies can find them. */
+	if (info->type == VIRTIO_VSOCK_TYPE_DGRAM &&
+	    src_cid == VMADDR_CID_LOCAL &&
+	    vsk->local_addr.svm_cid != VMADDR_CID_ANY)
+		src_cid = vsk->local_addr.svm_cid;
 	src_port = vsk->local_addr.svm_port;
-	if (!info->remote_cid) {
+	if (info->type != VIRTIO_VSOCK_TYPE_DGRAM && !info->remote_cid) {
 		dst_cid	= vsk->remote_addr.svm_cid;
 		dst_port = vsk->remote_addr.svm_port;
 	} else {

@@ -1066,6 +1071,8 @@ virtio_transport_dgram_dequeue(struct vsock_sock *vsk,
 		vsock_addr_init(vm_addr, le64_to_cpu(hdr->src_cid),
 				le32_to_cpu(hdr->src_port));
 		msg->msg_namelen = sizeof(*vm_addr);
+		if (VIRTIO_VSOCK_SKB_CB(skb)->to_host)
+			vm_addr->svm_flags = VMADDR_FLAG_TO_HOST;
 	}
 
 	if (flags & MSG_TRUNC)

@@ -1288,6 +1295,13 @@ bool virtio_transport_dgram_allow(struct vsock_sock *vsk, u32 cid, u32 port)
 }
 EXPORT_SYMBOL_GPL(virtio_transport_dgram_allow);
 
+bool virtio_transport_dgram_transport_match(const struct vsock_transport *t)
+{
+	return t && t->dgram_transport_match ==
+		virtio_transport_dgram_transport_match;
+}
+EXPORT_SYMBOL_GPL(virtio_transport_dgram_transport_match);
+
 int virtio_transport_connect(struct vsock_sock *vsk)
 {
 	struct virtio_vsock_pkt_info info = {

@@ -1318,12 +1332,14 @@ EXPORT_SYMBOL_GPL(virtio_transport_shutdown);
 
 int
 virtio_transport_dgram_enqueue(struct vsock_sock *vsk,
+			       const struct vsock_transport *transport,
 			       struct sockaddr_vm *remote_addr,
 			       struct msghdr *msg,
 			       size_t dgram_len)
 {
 	struct virtio_vsock_pkt_info info = {
 		.op = VIRTIO_VSOCK_OP_RW,
+		.t_ops = container_of(transport, struct virtio_transport, transport),
 		.msg = msg,
 		.pkt_len = dgram_len,
 		.vsk = vsk,

@@ -1861,6 +1877,33 @@ static bool virtio_transport_valid_type(u16 type)
 	       (type == VIRTIO_VSOCK_TYPE_DGRAM);
 }
 
+static bool
+virtio_transport_dgram_source_allow(struct vsock_sock *vsk,
+				    const struct vsock_transport *transport,
+				    const struct sockaddr_vm *src)
+{
+	const struct vsock_transport *peer_transport;
+	struct sockaddr_vm peer = vsk->remote_addr;
+	bool allow;
+
+	if (!virtio_transport_dgram_transport_match(vsk->transport) ||
+	    !virtio_transport_dgram_transport_match(transport))
+		return false;
+	/* Shutdown preserves the peer; only disconnect removes it. */
+	if (sk_vsock(vsk)->sk_socket->state == SS_UNCONNECTED)
+		return true;
+
+	peer_transport = vsock_dgram_get_transport(vsk, &peer);
+	if (IS_ERR(peer_transport))
+		return false;
+	allow = peer_transport == transport &&
+		src->svm_port == vsk->remote_addr.svm_port &&
+		(src->svm_cid == vsk->remote_addr.svm_cid ||
+		 transport->get_local_cid() == VMADDR_CID_LOCAL);
+	module_put(peer_transport->module);
+	return allow;
+}
+
 /* We are under the virtio-vsock's vsock->rx_lock or vhost-vsock's vq->mutex
  * lock.
  */

@@ -1911,14 +1954,14 @@ void virtio_transport_recv_pkt(struct virtio_transport *t,
 		vsk = vsock_sk(sk);
 		lock_sock(sk);
 		if (sock_flag(sk, SOCK_DEAD) ||
-		    vsk->transport != &t->transport ||
-		    (sk->sk_socket->state != SS_UNCONNECTED &&
-		     !vsock_check_source(vsk, &t->transport, &src))) {
+		    !virtio_transport_dgram_source_allow(vsk, &t->transport, &src)) {
 			release_sock(sk);
 			sock_put(sk);
 			goto free_pkt;
 		}
 
+		VIRTIO_VSOCK_SKB_CB(skb)->to_host =
+			t->transport.get_local_cid() > VMADDR_CID_HOST;
 		if (sock_queue_rcv_skb(sk, skb))
 			kfree_skb(skb);
 

@@ -1723,6 +1723,7 @@ static int vmci_transport_dgram_bind(struct vsock_sock *vsk,
 
 static int vmci_transport_dgram_enqueue(
 	struct vsock_sock *vsk,
+	const struct vsock_transport *transport,
 	struct sockaddr_vm *remote_addr,
 	struct msghdr *msg,
 	size_t len)

@@ -77,6 +77,7 @@ static struct virtio_transport loopback_transport = {
 		.dgram_dequeue            = virtio_transport_dgram_dequeue,
 		.dgram_enqueue            = virtio_transport_dgram_enqueue,
 		.dgram_allow              = virtio_transport_dgram_allow,
+		.dgram_transport_match    = virtio_transport_dgram_transport_match,
 
 		.stream_dequeue           = virtio_transport_stream_dequeue,
 		.stream_enqueue           = virtio_transport_stream_enqueue,