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); 124