vhost-net: Free ubuf when vhost_dev_set_owner fails

Signed-off-by: Asias He <asias@redhat.com>
Signed-off-by: Michael S. Tsirkin <mst@redhat.com>

diff --git a/drivers/vhost/net.c b/drivers/vhost/net.c
index a3645bd163d8ed0de10af23dbb72d30aab516fd3..354665a3ecd808b7333fcb51bcebcf99d33245c0 100644 (file)
--- a/drivers/vhost/net.c
+++ b/drivers/vhost/net.c
@@ -146,6 +146,19 @@ void vhost_ubuf_put_and_wait(struct vhost_ubuf_ref *ubufs)
        kfree(ubufs);
 }
 
+static void vhost_net_clear_ubuf_info(struct vhost_net *n)
+{
+
+       bool zcopy;
+       int i;
+
+       for (i = 0; i < n->dev.nvqs; ++i) {
+               zcopy = vhost_zcopy_mask & (0x1 << i);
+               if (zcopy)
+                       kfree(n->vqs[i].ubuf_info);
+       }
+}
+
 int vhost_net_set_ubuf_info(struct vhost_net *n)
 {
        bool zcopy;
@@ -1027,6 +1040,23 @@ static int vhost_net_set_features(struct vhost_net *n, u64 features)
        return 0;
 }
 
+static long vhost_net_set_owner(struct vhost_net *n)
+{
+       int r;
+
+       mutex_lock(&n->dev.mutex);
+       r = vhost_net_set_ubuf_info(n);
+       if (r)
+               goto out;
+       r = vhost_dev_set_owner(&n->dev);
+       if (r)
+               vhost_net_clear_ubuf_info(n);
+       vhost_net_flush(n);
+out:
+       mutex_unlock(&n->dev.mutex);
+       return r;
+}
+
 static long vhost_net_ioctl(struct file *f, unsigned int ioctl,
                            unsigned long arg)
 {
@@ -1055,19 +1085,15 @@ static long vhost_net_ioctl(struct file *f, unsigned int ioctl,
                return vhost_net_set_features(n, features);
        case VHOST_RESET_OWNER:
                return vhost_net_reset_owner(n);
+       case VHOST_SET_OWNER:
+               return vhost_net_set_owner(n);
        default:
                mutex_lock(&n->dev.mutex);
-               if (ioctl == VHOST_SET_OWNER) {
-                       r = vhost_net_set_ubuf_info(n);
-                       if (r)
-                               goto out;
-               }
                r = vhost_dev_ioctl(&n->dev, ioctl, argp);
                if (r == -ENOIOCTLCMD)
                        r = vhost_vring_ioctl(&n->dev, ioctl, argp);
                else
                        vhost_net_flush(n);
-out:
                mutex_unlock(&n->dev.mutex);
                return r;
        }