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 
__kunit_static_stub_resource_free(struct kunit_resource * res)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. */
__kunit_static_stub_resource_match(struct kunit * test,struct kunit_resource * res,void * match_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. */
__kunit_get_static_stub_address_impl(struct kunit * test,void * real_fn_addr)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 
kunit_deactivate_static_stub(struct kunit * test,void * real_fn_addr)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  */
__kunit_activate_static_stub(struct kunit * test,void * real_fn_addr,void * replacement_addr)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