Line data    Source code

       1             : // SPDX-License-Identifier: GPL-2.0
       2             : /*
       3             :  * KUnit function redirection (static stubbing) API.
       4             :  *
       5             :  * Copyright (C) 2022, Google LLC.
       6             :  * Author: David Gow <davidgow@google.com>
       7             :  */
       8             : 
       9             : #include <kunit/test.h>
      10             : #include <kunit/static_stub.h>
      11             : #include "hooks-impl.h"
      12             : 
      13             : 
      14             : /* Context for a static stub. This is stored in the resource data. */
      15             : struct kunit_static_stub_ctx {
      16             :         void *real_fn_addr;
      17             :         void *replacement_addr;
      18             : };
      19             : 
      20           0 : static void __kunit_static_stub_resource_free(struct kunit_resource *res)
      21             : {
      22           0 :         kfree(res->data);
      23           0 : }
      24             : 
      25             : /* Matching function for kunit_find_resource(). match_data is real_fn_addr. */
      26           0 : static bool __kunit_static_stub_resource_match(struct kunit *test,
      27             :                                                 struct kunit_resource *res,
      28             :                                                 void *match_real_fn_addr)
      29             : {
      30             :         /* This pointer is only valid if res is a static stub resource. */
      31           0 :         struct kunit_static_stub_ctx *ctx = res->data;
      32             : 
      33             :         /* Make sure the resource is a static stub resource. */
      34           0 :         if (res->free != &__kunit_static_stub_resource_free)
      35             :                 return false;
      36             : 
      37           0 :         return ctx->real_fn_addr == match_real_fn_addr;
      38             : }
      39             : 
      40             : /* Hook to return the address of the replacement function. */
      41           0 : void *__kunit_get_static_stub_address_impl(struct kunit *test, void *real_fn_addr)
      42             : {
      43             :         struct kunit_resource *res;
      44             :         struct kunit_static_stub_ctx *ctx;
      45             :         void *replacement_addr;
      46             : 
      47           0 :         res = kunit_find_resource(test,
      48             :                                   __kunit_static_stub_resource_match,
      49             :                                   real_fn_addr);
      50             : 
      51           0 :         if (!res)
      52             :                 return NULL;
      53             : 
      54           0 :         ctx = res->data;
      55           0 :         replacement_addr = ctx->replacement_addr;
      56           0 :         kunit_put_resource(res);
      57           0 :         return replacement_addr;
      58             : }
      59             : 
      60           0 : void kunit_deactivate_static_stub(struct kunit *test, void *real_fn_addr)
      61             : {
      62             :         struct kunit_resource *res;
      63             : 
      64           0 :         KUNIT_ASSERT_PTR_NE_MSG(test, real_fn_addr, NULL,
      65             :                                 "Tried to deactivate a NULL stub.");
      66             : 
      67             :         /* Look up the existing stub for this function. */
      68           0 :         res = kunit_find_resource(test,
      69             :                                   __kunit_static_stub_resource_match,
      70             :                                   real_fn_addr);
      71             : 
      72             :         /* Error out if the stub doesn't exist. */
      73           0 :         KUNIT_ASSERT_PTR_NE_MSG(test, res, NULL,
      74             :                                 "Tried to deactivate a nonexistent stub.");
      75             : 
      76             :         /* Free the stub. We 'put' twice, as we got a reference
      77             :          * from kunit_find_resource()
      78             :          */
      79           0 :         kunit_remove_resource(test, res);
      80           0 :         kunit_put_resource(res);
      81           0 : }
      82             : EXPORT_SYMBOL_GPL(kunit_deactivate_static_stub);
      83             : 
      84             : /* Helper function for kunit_activate_static_stub(). The macro does
      85             :  * typechecking, so use it instead.
      86             :  */
      87           0 : void __kunit_activate_static_stub(struct kunit *test,
      88             :                                   void *real_fn_addr,
      89             :                                   void *replacement_addr)
      90             : {
      91             :         struct kunit_static_stub_ctx *ctx;
      92             :         struct kunit_resource *res;
      93             : 
      94           0 :         KUNIT_ASSERT_PTR_NE_MSG(test, real_fn_addr, NULL,
      95             :                                 "Tried to activate a stub for function NULL");
      96             : 
      97             :         /* If the replacement address is NULL, deactivate the stub. */
      98           0 :         if (!replacement_addr) {
      99           0 :                 kunit_deactivate_static_stub(test, replacement_addr);
     100           0 :                 return;
     101             :         }
     102             : 
     103             :         /* Look up any existing stubs for this function, and replace them. */
     104           0 :         res = kunit_find_resource(test,
     105             :                                   __kunit_static_stub_resource_match,
     106             :                                   real_fn_addr);
     107           0 :         if (res) {
     108           0 :                 ctx = res->data;
     109           0 :                 ctx->replacement_addr = replacement_addr;
     110             : 
     111             :                 /* We got an extra reference from find_resource(), so put it. */
     112             :                 kunit_put_resource(res);
     113             :         } else {
     114           0 :                 ctx = kmalloc(sizeof(*ctx), GFP_KERNEL);
     115           0 :                 KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx);
     116           0 :                 ctx->real_fn_addr = real_fn_addr;
     117           0 :                 ctx->replacement_addr = replacement_addr;
     118           0 :                 res = kunit_alloc_resource(test, NULL,
     119             :                                      &__kunit_static_stub_resource_free,
     120             :                                      GFP_KERNEL, ctx);
     121             :         }
     122             : }
     123             : EXPORT_SYMBOL_GPL(__kunit_activate_static_stub);