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(objp: 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 | match: __kunit_static_stub_resource_match, |
49 | match_data: 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 | match: __kunit_static_stub_resource_match, |
70 | match_data: 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 | match: __kunit_static_stub_resource_match, |
106 | match_data: 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(size: 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 | free: &__kunit_static_stub_resource_free, |
120 | GFP_KERNEL, context: ctx); |
121 | } |
122 | } |
123 | EXPORT_SYMBOL_GPL(__kunit_activate_static_stub); |
124 | |