[PATCH v4] virtio_balloon: prime stats vq after virtio_device_ready()

[PATCH v4] virtio_balloon: prime stats vq after virtio_device_ready()

[Michael S. Tsirkin]

The virtio spec requires the driver not to kick the device before
DRIVER_OK is set.  init_vqs() primes the stats virtqueue with a buffer
and kicks the device before virtio_device_ready() is called in
virtballoon_probe(), violating this requirement.

Further, if the device responds to the early kick by processing the
buffer before DRIVER_OK, stats_request() fires and queues
update_balloon_stats_work.  Should probe then fail and free vb, the work
runs against freed memory.

To fix, move buffer setup to after DRIVER_OK. Be careful to
disable update_balloon_stats_work while this is going on,
to make sure it does not race with the setup.

Testing: tested that stats still work after the change.

Fixes: 9564e138b1f6 ("virtio: Add memory statistics reporting to the balloon driver (V4)")
Reported-by: Sashiko:gemini-3.1-pro-preview
Cc: David Hildenbrand <david@kernel.org>
Assisted-by: Claude:claude-sonnet-4-6
Signed-off-by: Michael S. Tsirkin <mst@redhat.com>

---
v4: extract stats_update_and_send() helper; call get_buf() before
    updating stats buffer in stats_handle_request().
    also ignore add bug failures (never trigger anyway) (reported by sashiko)
v3: add disable_work comment and BUG_ON -> WARN_ON_ONCE (David Hildenbrand)
v2: check that work enable/disable is balanced; explain how add_outbuf
    never fails in probe/restore
 drivers/virtio/virtio_balloon.c | 67 +++++++++++++++++++--------------
 1 file changed, 39 insertions(+), 28 deletions(-)

diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
index 088b3a0e6ce6..4c4b9bea4356 100644
--- a/drivers/virtio/virtio_balloon.c
+++ b/drivers/virtio/virtio_balloon.c
@@ -417,6 +417,17 @@ static unsigned int update_balloon_stats(struct virtio_balloon *vb)
 	return idx;
 }
 
+static void stats_update_and_send(struct virtio_balloon *vb)
+{
+	struct scatterlist sg;
+	unsigned int num_stats;
+
+	num_stats = update_balloon_stats(vb);
+	sg_init_one(&sg, vb->stats, sizeof(vb->stats[0]) * num_stats);
+	virtqueue_add_outbuf(vb->stats_vq, &sg, 1, vb, GFP_KERNEL);
+	virtqueue_kick(vb->stats_vq);
+}
+
 /*
  * While most virtqueues communicate guest-initiated requests to the hypervisor,
  * the stats queue operates in reverse.  The driver initializes the virtqueue
@@ -440,18 +451,11 @@ static void stats_request(struct virtqueue *vq)
 
 static void stats_handle_request(struct virtio_balloon *vb)
 {
-	struct virtqueue *vq;
-	struct scatterlist sg;
-	unsigned int len, num_stats;
+	unsigned int len;
 
-	num_stats = update_balloon_stats(vb);
-
-	vq = vb->stats_vq;
-	if (!virtqueue_get_buf(vq, &len))
+	if (!virtqueue_get_buf(vb->stats_vq, &len))
 		return;
-	sg_init_one(&sg, vb->stats, sizeof(vb->stats[0]) * num_stats);
-	virtqueue_add_outbuf(vq, &sg, 1, vb, GFP_KERNEL);
-	virtqueue_kick(vq);
+	stats_update_and_send(vb);
 }
 
 static inline s64 towards_target(struct virtio_balloon *vb)
@@ -611,25 +615,9 @@ static int init_vqs(struct virtio_balloon *vb)
 	vb->inflate_vq = vqs[VIRTIO_BALLOON_VQ_INFLATE];
 	vb->deflate_vq = vqs[VIRTIO_BALLOON_VQ_DEFLATE];
 	if (virtio_has_feature(vb->vdev, VIRTIO_BALLOON_F_STATS_VQ)) {
-		struct scatterlist sg;
-		unsigned int num_stats;
 		vb->stats_vq = vqs[VIRTIO_BALLOON_VQ_STATS];
-
-		/*
-		 * Prime this virtqueue with one buffer so the hypervisor can
-		 * use it to signal us later (it can't be broken yet!).
-		 */
-		num_stats = update_balloon_stats(vb);
-
-		sg_init_one(&sg, vb->stats, sizeof(vb->stats[0]) * num_stats);
-		err = virtqueue_add_outbuf(vb->stats_vq, &sg, 1, vb,
-					   GFP_KERNEL);
-		if (err) {
-			dev_warn(&vb->vdev->dev, "%s: add stat_vq failed\n",
-				 __func__);
-			return err;
-		}
-		virtqueue_kick(vb->stats_vq);
+		/* Prevent update_balloon_stats_work from accessing the stats vq. */
+		disable_work(&vb->update_balloon_stats_work);
 	}
 
 	if (virtio_has_feature(vb->vdev, VIRTIO_BALLOON_F_FREE_PAGE_HINT))
@@ -916,6 +904,25 @@ static int virtio_balloon_register_shrinker(struct virtio_balloon *vb)
 	return 0;
 }
 
+static void setup_vqs(struct virtio_balloon *vb)
+{
+	bool ret;
+
+	if (!virtio_has_feature(vb->vdev, VIRTIO_BALLOON_F_STATS_VQ))
+		return;
+
+	/*
+	 * Prime this virtqueue with one buffer so the hypervisor can
+	 * use it to signal us later (it can't be broken yet!).
+	 */
+	stats_update_and_send(vb);
+
+	ret = enable_and_queue_work(system_freezable_wq,
+				    &vb->update_balloon_stats_work);
+	/* Make sure we balanced enable/disable, or we won't report stats. */
+	WARN_ON_ONCE(!ret);
+}
+
 static int virtballoon_probe(struct virtio_device *vdev)
 {
 	struct virtio_balloon *vb;
@@ -1059,6 +1066,8 @@ static int virtballoon_probe(struct virtio_device *vdev)
 
 	virtio_device_ready(vdev);
 
+	setup_vqs(vb);
+
 	if (towards_target(vb))
 		virtballoon_changed(vdev);
 	return 0;
@@ -1148,6 +1157,8 @@ static int virtballoon_restore(struct virtio_device *vdev)
 
 	virtio_device_ready(vdev);
 
+	setup_vqs(vb);
+
 	if (towards_target(vb))
 		virtballoon_changed(vdev);
 	update_balloon_size(vb);
-- 
MST

Re: [PATCH v4] virtio_balloon: prime stats vq after virtio_device_ready()

[David Hildenbrand (Arm)]

On 7/7/26 23:00, Michael S. Tsirkin wrote:
> The virtio spec requires the driver not to kick the device before
> DRIVER_OK is set.  init_vqs() primes the stats virtqueue with a buffer
> and kicks the device before virtio_device_ready() is called in
> virtballoon_probe(), violating this requirement.
> 
> Further, if the device responds to the early kick by processing the
> buffer before DRIVER_OK, stats_request() fires and queues
> update_balloon_stats_work.  Should probe then fail and free vb, the work
> runs against freed memory.
> 
> To fix, move buffer setup to after DRIVER_OK. Be careful to
> disable update_balloon_stats_work while this is going on,
> to make sure it does not race with the setup.
> 
> Testing: tested that stats still work after the change.
> 
> Fixes: 9564e138b1f6 ("virtio: Add memory statistics reporting to the balloon driver (V4)")
> Reported-by: Sashiko:gemini-3.1-pro-preview
> Cc: David Hildenbrand <david@kernel.org>
> Assisted-by: Claude:claude-sonnet-4-6
> Signed-off-by: Michael S. Tsirkin <mst@redhat.com>
> 
> ---
> v4: extract stats_update_and_send() helper; call get_buf() before
>     updating stats buffer in stats_handle_request().
>     also ignore add bug failures (never trigger anyway) (reported by sashiko)
> v3: add disable_work comment and BUG_ON -> WARN_ON_ONCE (David Hildenbrand)
> v2: check that work enable/disable is balanced; explain how add_outbuf
>     never fails in probe/restore
>  drivers/virtio/virtio_balloon.c | 67 +++++++++++++++++++--------------
>  1 file changed, 39 insertions(+), 28 deletions(-)
> 
> diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
> index 088b3a0e6ce6..4c4b9bea4356 100644
> --- a/drivers/virtio/virtio_balloon.c
> +++ b/drivers/virtio/virtio_balloon.c
> @@ -417,6 +417,17 @@ static unsigned int update_balloon_stats(struct virtio_balloon *vb)
>  	return idx;
>  }
>  
> +static void stats_update_and_send(struct virtio_balloon *vb)
> +{
> +	struct scatterlist sg;
> +	unsigned int num_stats;

Could do

	const unsigned int num_stats = update_balloon_stats(vb);


LGTM.

-- 
Cheers,

David

Re: [PATCH v4] virtio_balloon: prime stats vq after virtio_device_ready()

[Michael S. Tsirkin]

On Wed, Jul 08, 2026 at 10:13:58AM +0200, David Hildenbrand (Arm) wrote:
> On 7/7/26 23:00, Michael S. Tsirkin wrote:
> > The virtio spec requires the driver not to kick the device before
> > DRIVER_OK is set.  init_vqs() primes the stats virtqueue with a buffer
> > and kicks the device before virtio_device_ready() is called in
> > virtballoon_probe(), violating this requirement.
> > 
> > Further, if the device responds to the early kick by processing the
> > buffer before DRIVER_OK, stats_request() fires and queues
> > update_balloon_stats_work.  Should probe then fail and free vb, the work
> > runs against freed memory.
> > 
> > To fix, move buffer setup to after DRIVER_OK. Be careful to
> > disable update_balloon_stats_work while this is going on,
> > to make sure it does not race with the setup.
> > 
> > Testing: tested that stats still work after the change.
> > 
> > Fixes: 9564e138b1f6 ("virtio: Add memory statistics reporting to the balloon driver (V4)")
> > Reported-by: Sashiko:gemini-3.1-pro-preview
> > Cc: David Hildenbrand <david@kernel.org>
> > Assisted-by: Claude:claude-sonnet-4-6
> > Signed-off-by: Michael S. Tsirkin <mst@redhat.com>
> > 
> > ---
> > v4: extract stats_update_and_send() helper; call get_buf() before
> >     updating stats buffer in stats_handle_request().
> >     also ignore add bug failures (never trigger anyway) (reported by sashiko)
> > v3: add disable_work comment and BUG_ON -> WARN_ON_ONCE (David Hildenbrand)
> > v2: check that work enable/disable is balanced; explain how add_outbuf
> >     never fails in probe/restore
> >  drivers/virtio/virtio_balloon.c | 67 +++++++++++++++++++--------------
> >  1 file changed, 39 insertions(+), 28 deletions(-)
> > 
> > diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
> > index 088b3a0e6ce6..4c4b9bea4356 100644
> > --- a/drivers/virtio/virtio_balloon.c
> > +++ b/drivers/virtio/virtio_balloon.c
> > @@ -417,6 +417,17 @@ static unsigned int update_balloon_stats(struct virtio_balloon *vb)
> >  	return idx;
> >  }
> >  
> > +static void stats_update_and_send(struct virtio_balloon *vb)
> > +{
> > +	struct scatterlist sg;
> > +	unsigned int num_stats;
> 
> Could do
> 
> 	const unsigned int num_stats = update_balloon_stats(vb);
> 


Indeed, a bit nicer. Thanks.

> LGTM.
> 
> -- 
> Cheers,
> 
> David