lib/kunit/static_stub.c

   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 static void __kunit_static_stub_resource_free(struct kunit_resource *res)
  21 {
  22         kfree(res->data);
  23 }
  24
  25 /* Matching function for kunit_find_resource(). match_data is real_fn_addr. */
  26 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         struct kunit_static_stub_ctx *ctx = res->data;
  32
  33         /* Make sure the resource is a static stub resource. */
  34         if (res->free != &__kunit_static_stub_resource_free)
  35                 return false;
  36
  37         return ctx->real_fn_addr == match_real_fn_addr;
  38 }
  39
  40 /* Hook to return the address of the replacement function. */
  41 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         res = kunit_find_resource(test,
  48                                   __kunit_static_stub_resource_match,
  49                                   real_fn_addr);
  50
  51         if (!res)
  52                 return NULL;
  53
  54         ctx = res->data;
  55         replacement_addr = ctx->replacement_addr;
  56         kunit_put_resource(res);
  57         return replacement_addr;
  58 }
  59
  60 void kunit_deactivate_static_stub(struct kunit *test, void *real_fn_addr)
  61 {
  62         struct kunit_resource *res;
  63
  64         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         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         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         kunit_remove_resource(test, res);
  80         kunit_put_resource(res);
  81 }
  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 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         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         if (!replacement_addr) {
  99                 kunit_deactivate_static_stub(test, replacement_addr);
 100                 return;
 101         }
 102
 103         /* Look up any existing stubs for this function, and replace them. */
 104         res = kunit_find_resource(test,
 105                                   __kunit_static_stub_resource_match,
 106                                   real_fn_addr);
 107         if (res) {
 108                 ctx = res->data;
 109                 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                 ctx = kmalloc(sizeof(*ctx), GFP_KERNEL);
 115                 KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx);
 116                 ctx->real_fn_addr = real_fn_addr;
 117                 ctx->replacement_addr = replacement_addr;
 118                 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);