[PATCH block/for-linus] blkcg: Fix ->pd_alloc_fn() being called with the wrong blkcg on policy activation

[Tejun Heo]

cf09a8ee19ad ("blkcg: pass @q and @blkcg into
blkcg_pol_alloc_pd_fn()") added @blkcg to ->pd_alloc_fn(); however,
blkcg_activate_policy() ends up using pd's allocated for the root
blkcg for all preallocations, so ->pd_init_fn() for non-root blkcgs
can be passed in pd's which are allocated for the root blkcg.

For blk-iocost, this means that ->pd_init_fn() can write beyond the
end of the allocated object as it determines the length of the flex
array at the end based on the blkcg's nesting level.

Signed-off-by: Tejun Heo <t...@kernel.org>
Fixes: cf09a8ee19ad ("blkcg: pass @q and @blkcg into blkcg_pol_alloc_pd_fn()")
---
 block/blk-cgroup.c |   43 +++++++++++++++++++++++++++++--------------
 1 file changed, 29 insertions(+), 14 deletions(-)

--- a/block/blk-cgroup.c
+++ b/block/blk-cgroup.c
@@ -1362,7 +1362,7 @@ int blkcg_activate_policy(struct request
                          const struct blkcg_policy *pol)
 {
        struct blkg_policy_data *pd_prealloc = NULL;
-       struct blkcg_gq *blkg;
+       struct blkcg_gq *blkg, *pinned_blkg = NULL;
        int ret;
 
        if (blkcg_policy_enabled(q, pol))
@@ -1370,15 +1370,7 @@ int blkcg_activate_policy(struct request
 
        if (queue_is_mq(q))
                blk_mq_freeze_queue(q);
-pd_prealloc:
-       if (!pd_prealloc) {
-               pd_prealloc = pol->pd_alloc_fn(GFP_KERNEL, q, &blkcg_root);
-               if (!pd_prealloc) {
-                       ret = -ENOMEM;
-                       goto out_bypass_end;
-               }
-       }
-
+retry:
        spin_lock_irq(&q->queue_lock);
 
        /* blkg_list is pushed at the head, reverse walk to init parents first 
*/
@@ -1388,12 +1380,31 @@ pd_prealloc:
                if (blkg->pd[pol->plid])
                        continue;
 
-               pd = pol->pd_alloc_fn(GFP_NOWAIT | __GFP_NOWARN, q, 
&blkcg_root);
-               if (!pd)
-                       swap(pd, pd_prealloc);
+               if (blkg == pinned_blkg) {
+                       pd = pd_prealloc;
+                       pd_prealloc = NULL;
+               } else {
+                       pd = pol->pd_alloc_fn(GFP_NOWAIT | __GFP_NOWARN, q,
+                                             blkg->blkcg);
+               }
+
                if (!pd) {
+                       if (pinned_blkg)
+                               blkg_put(pinned_blkg);
+                       blkg_get(blkg);
+                       pinned_blkg = blkg;
+
                        spin_unlock_irq(&q->queue_lock);
-                       goto pd_prealloc;
+
+                       kfree(pd_prealloc);
+                       pd_prealloc = pol->pd_alloc_fn(GFP_KERNEL, q,
+                                                      blkg->blkcg);
+                       if (pd_prealloc) {
+                               goto retry;
+                       } else {
+                               ret = -ENOMEM;
+                               goto out_bypass_end;
+                       }
                }
 
                blkg->pd[pol->plid] = pd;
@@ -1403,6 +1414,10 @@ pd_prealloc:
                        pol->pd_init_fn(pd);
        }
 
+       if (pinned_blkg)
+               blkg_put(pinned_blkg);
+       kfree(pd_prealloc);
+
        __set_bit(pol->plid, q->blkcg_pols);
        ret = 0;