1 | /* SPDX-License-Identifier: GPL-2.0 */ |
2 | /* |
3 | * KUnit resource API for test managed resources (allocations, etc.). |
4 | * |
5 | * Copyright (C) 2022, Google LLC. |
6 | * Author: Daniel Latypov <dlatypov@google.com> |
7 | */ |
8 | |
9 | #ifndef _KUNIT_RESOURCE_H |
10 | #define _KUNIT_RESOURCE_H |
11 | |
12 | #include <kunit/test.h> |
13 | |
14 | #include <linux/kref.h> |
15 | #include <linux/list.h> |
16 | #include <linux/slab.h> |
17 | #include <linux/spinlock.h> |
18 | |
19 | struct kunit_resource; |
20 | |
21 | typedef int (*kunit_resource_init_t)(struct kunit_resource *, void *); |
22 | typedef void (*kunit_resource_free_t)(struct kunit_resource *); |
23 | |
24 | /** |
25 | * struct kunit_resource - represents a *test managed resource* |
26 | * @data: for the user to store arbitrary data. |
27 | * @name: optional name |
28 | * @free: a user supplied function to free the resource. |
29 | * |
30 | * Represents a *test managed resource*, a resource which will automatically be |
31 | * cleaned up at the end of a test case. This cleanup is performed by the 'free' |
32 | * function. The struct kunit_resource itself is freed automatically with |
33 | * kfree() if it was allocated by KUnit (e.g., by kunit_alloc_resource()), but |
34 | * must be freed by the user otherwise. |
35 | * |
36 | * Resources are reference counted so if a resource is retrieved via |
37 | * kunit_alloc_and_get_resource() or kunit_find_resource(), we need |
38 | * to call kunit_put_resource() to reduce the resource reference count |
39 | * when finished with it. Note that kunit_alloc_resource() does not require a |
40 | * kunit_resource_put() because it does not retrieve the resource itself. |
41 | * |
42 | * Example: |
43 | * |
44 | * .. code-block:: c |
45 | * |
46 | * struct kunit_kmalloc_params { |
47 | * size_t size; |
48 | * gfp_t gfp; |
49 | * }; |
50 | * |
51 | * static int kunit_kmalloc_init(struct kunit_resource *res, void *context) |
52 | * { |
53 | * struct kunit_kmalloc_params *params = context; |
54 | * res->data = kmalloc(params->size, params->gfp); |
55 | * |
56 | * if (!res->data) |
57 | * return -ENOMEM; |
58 | * |
59 | * return 0; |
60 | * } |
61 | * |
62 | * static void kunit_kmalloc_free(struct kunit_resource *res) |
63 | * { |
64 | * kfree(res->data); |
65 | * } |
66 | * |
67 | * void *kunit_kmalloc(struct kunit *test, size_t size, gfp_t gfp) |
68 | * { |
69 | * struct kunit_kmalloc_params params; |
70 | * |
71 | * params.size = size; |
72 | * params.gfp = gfp; |
73 | * |
74 | * return kunit_alloc_resource(test, kunit_kmalloc_init, |
75 | * kunit_kmalloc_free, gfp, ¶ms); |
76 | * } |
77 | * |
78 | * Resources can also be named, with lookup/removal done on a name |
79 | * basis also. kunit_add_named_resource(), kunit_find_named_resource() |
80 | * and kunit_destroy_named_resource(). Resource names must be |
81 | * unique within the test instance. |
82 | */ |
83 | struct kunit_resource { |
84 | void *data; |
85 | const char *name; |
86 | kunit_resource_free_t free; |
87 | |
88 | /* private: internal use only. */ |
89 | struct kref refcount; |
90 | struct list_head node; |
91 | bool should_kfree; |
92 | }; |
93 | |
94 | /** |
95 | * kunit_get_resource() - Hold resource for use. Should not need to be used |
96 | * by most users as we automatically get resources |
97 | * retrieved by kunit_find_resource*(). |
98 | * @res: resource |
99 | */ |
100 | static inline void kunit_get_resource(struct kunit_resource *res) |
101 | { |
102 | kref_get(kref: &res->refcount); |
103 | } |
104 | |
105 | /* |
106 | * Called when refcount reaches zero via kunit_put_resource(); |
107 | * should not be called directly. |
108 | */ |
109 | static inline void kunit_release_resource(struct kref *kref) |
110 | { |
111 | struct kunit_resource *res = container_of(kref, struct kunit_resource, |
112 | refcount); |
113 | |
114 | if (res->free) |
115 | res->free(res); |
116 | |
117 | /* 'res' is valid here, as if should_kfree is set, res->free may not free |
118 | * 'res' itself, just res->data |
119 | */ |
120 | if (res->should_kfree) |
121 | kfree(objp: res); |
122 | } |
123 | |
124 | /** |
125 | * kunit_put_resource() - When caller is done with retrieved resource, |
126 | * kunit_put_resource() should be called to drop |
127 | * reference count. The resource list maintains |
128 | * a reference count on resources, so if no users |
129 | * are utilizing a resource and it is removed from |
130 | * the resource list, it will be freed via the |
131 | * associated free function (if any). Only |
132 | * needs to be used if we alloc_and_get() or |
133 | * find() resource. |
134 | * @res: resource |
135 | */ |
136 | static inline void kunit_put_resource(struct kunit_resource *res) |
137 | { |
138 | kref_put(kref: &res->refcount, release: kunit_release_resource); |
139 | } |
140 | |
141 | /** |
142 | * __kunit_add_resource() - Internal helper to add a resource. |
143 | * |
144 | * res->should_kfree is not initialised. |
145 | * @test: The test context object. |
146 | * @init: a user-supplied function to initialize the result (if needed). If |
147 | * none is supplied, the resource data value is simply set to @data. |
148 | * If an init function is supplied, @data is passed to it instead. |
149 | * @free: a user-supplied function to free the resource (if needed). |
150 | * @res: The resource. |
151 | * @data: value to pass to init function or set in resource data field. |
152 | */ |
153 | int __kunit_add_resource(struct kunit *test, |
154 | kunit_resource_init_t init, |
155 | kunit_resource_free_t free, |
156 | struct kunit_resource *res, |
157 | void *data); |
158 | |
159 | /** |
160 | * kunit_add_resource() - Add a *test managed resource*. |
161 | * @test: The test context object. |
162 | * @init: a user-supplied function to initialize the result (if needed). If |
163 | * none is supplied, the resource data value is simply set to @data. |
164 | * If an init function is supplied, @data is passed to it instead. |
165 | * @free: a user-supplied function to free the resource (if needed). |
166 | * @res: The resource. |
167 | * @data: value to pass to init function or set in resource data field. |
168 | */ |
169 | static inline int kunit_add_resource(struct kunit *test, |
170 | kunit_resource_init_t init, |
171 | kunit_resource_free_t free, |
172 | struct kunit_resource *res, |
173 | void *data) |
174 | { |
175 | res->should_kfree = false; |
176 | return __kunit_add_resource(test, init, free, res, data); |
177 | } |
178 | |
179 | static inline struct kunit_resource * |
180 | kunit_find_named_resource(struct kunit *test, const char *name); |
181 | |
182 | /** |
183 | * kunit_add_named_resource() - Add a named *test managed resource*. |
184 | * @test: The test context object. |
185 | * @init: a user-supplied function to initialize the resource data, if needed. |
186 | * @free: a user-supplied function to free the resource data, if needed. |
187 | * @res: The resource. |
188 | * @name: name to be set for resource. |
189 | * @data: value to pass to init function or set in resource data field. |
190 | */ |
191 | static inline int kunit_add_named_resource(struct kunit *test, |
192 | kunit_resource_init_t init, |
193 | kunit_resource_free_t free, |
194 | struct kunit_resource *res, |
195 | const char *name, |
196 | void *data) |
197 | { |
198 | struct kunit_resource *existing; |
199 | |
200 | if (!name) |
201 | return -EINVAL; |
202 | |
203 | existing = kunit_find_named_resource(test, name); |
204 | if (existing) { |
205 | kunit_put_resource(res: existing); |
206 | return -EEXIST; |
207 | } |
208 | |
209 | res->name = name; |
210 | res->should_kfree = false; |
211 | |
212 | return __kunit_add_resource(test, init, free, res, data); |
213 | } |
214 | |
215 | /** |
216 | * kunit_alloc_and_get_resource() - Allocates and returns a *test managed resource*. |
217 | * @test: The test context object. |
218 | * @init: a user supplied function to initialize the resource. |
219 | * @free: a user supplied function to free the resource (if needed). |
220 | * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL |
221 | * @context: for the user to pass in arbitrary data to the init function. |
222 | * |
223 | * Allocates a *test managed resource*, a resource which will automatically be |
224 | * cleaned up at the end of a test case. See &struct kunit_resource for an |
225 | * example. |
226 | * |
227 | * This is effectively identical to kunit_alloc_resource, but returns the |
228 | * struct kunit_resource pointer, not just the 'data' pointer. It therefore |
229 | * also increments the resource's refcount, so kunit_put_resource() should be |
230 | * called when you've finished with it. |
231 | * |
232 | * Note: KUnit needs to allocate memory for a kunit_resource object. You must |
233 | * specify an @internal_gfp that is compatible with the use context of your |
234 | * resource. |
235 | */ |
236 | static inline struct kunit_resource * |
237 | kunit_alloc_and_get_resource(struct kunit *test, |
238 | kunit_resource_init_t init, |
239 | kunit_resource_free_t free, |
240 | gfp_t internal_gfp, |
241 | void *context) |
242 | { |
243 | struct kunit_resource *res; |
244 | int ret; |
245 | |
246 | res = kzalloc(size: sizeof(*res), flags: internal_gfp); |
247 | if (!res) |
248 | return NULL; |
249 | |
250 | res->should_kfree = true; |
251 | |
252 | ret = __kunit_add_resource(test, init, free, res, data: context); |
253 | if (!ret) { |
254 | /* |
255 | * bump refcount for get; kunit_resource_put() should be called |
256 | * when done. |
257 | */ |
258 | kunit_get_resource(res); |
259 | return res; |
260 | } |
261 | return NULL; |
262 | } |
263 | |
264 | /** |
265 | * kunit_alloc_resource() - Allocates a *test managed resource*. |
266 | * @test: The test context object. |
267 | * @init: a user supplied function to initialize the resource. |
268 | * @free: a user supplied function to free the resource (if needed). |
269 | * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL |
270 | * @context: for the user to pass in arbitrary data to the init function. |
271 | * |
272 | * Allocates a *test managed resource*, a resource which will automatically be |
273 | * cleaned up at the end of a test case. See &struct kunit_resource for an |
274 | * example. |
275 | * |
276 | * Note: KUnit needs to allocate memory for a kunit_resource object. You must |
277 | * specify an @internal_gfp that is compatible with the use context of your |
278 | * resource. |
279 | */ |
280 | static inline void *kunit_alloc_resource(struct kunit *test, |
281 | kunit_resource_init_t init, |
282 | kunit_resource_free_t free, |
283 | gfp_t internal_gfp, |
284 | void *context) |
285 | { |
286 | struct kunit_resource *res; |
287 | |
288 | res = kzalloc(size: sizeof(*res), flags: internal_gfp); |
289 | if (!res) |
290 | return NULL; |
291 | |
292 | res->should_kfree = true; |
293 | if (!__kunit_add_resource(test, init, free, res, data: context)) |
294 | return res->data; |
295 | |
296 | return NULL; |
297 | } |
298 | |
299 | typedef bool (*kunit_resource_match_t)(struct kunit *test, |
300 | struct kunit_resource *res, |
301 | void *match_data); |
302 | |
303 | /** |
304 | * kunit_resource_name_match() - Match a resource with the same name. |
305 | * @test: Test case to which the resource belongs. |
306 | * @res: The resource. |
307 | * @match_name: The name to match against. |
308 | */ |
309 | static inline bool kunit_resource_name_match(struct kunit *test, |
310 | struct kunit_resource *res, |
311 | void *match_name) |
312 | { |
313 | return res->name && strcmp(res->name, match_name) == 0; |
314 | } |
315 | |
316 | /** |
317 | * kunit_find_resource() - Find a resource using match function/data. |
318 | * @test: Test case to which the resource belongs. |
319 | * @match: match function to be applied to resources/match data. |
320 | * @match_data: data to be used in matching. |
321 | */ |
322 | static inline struct kunit_resource * |
323 | kunit_find_resource(struct kunit *test, |
324 | kunit_resource_match_t match, |
325 | void *match_data) |
326 | { |
327 | struct kunit_resource *res, *found = NULL; |
328 | unsigned long flags; |
329 | |
330 | spin_lock_irqsave(&test->lock, flags); |
331 | |
332 | list_for_each_entry_reverse(res, &test->resources, node) { |
333 | if (match(test, res, (void *)match_data)) { |
334 | found = res; |
335 | kunit_get_resource(res: found); |
336 | break; |
337 | } |
338 | } |
339 | |
340 | spin_unlock_irqrestore(lock: &test->lock, flags); |
341 | |
342 | return found; |
343 | } |
344 | |
345 | /** |
346 | * kunit_find_named_resource() - Find a resource using match name. |
347 | * @test: Test case to which the resource belongs. |
348 | * @name: match name. |
349 | */ |
350 | static inline struct kunit_resource * |
351 | kunit_find_named_resource(struct kunit *test, |
352 | const char *name) |
353 | { |
354 | return kunit_find_resource(test, match: kunit_resource_name_match, |
355 | match_data: (void *)name); |
356 | } |
357 | |
358 | /** |
359 | * kunit_destroy_resource() - Find a kunit_resource and destroy it. |
360 | * @test: Test case to which the resource belongs. |
361 | * @match: Match function. Returns whether a given resource matches @match_data. |
362 | * @match_data: Data passed into @match. |
363 | * |
364 | * RETURNS: |
365 | * 0 if kunit_resource is found and freed, -ENOENT if not found. |
366 | */ |
367 | int kunit_destroy_resource(struct kunit *test, |
368 | kunit_resource_match_t match, |
369 | void *match_data); |
370 | |
371 | static inline int kunit_destroy_named_resource(struct kunit *test, |
372 | const char *name) |
373 | { |
374 | return kunit_destroy_resource(test, match: kunit_resource_name_match, |
375 | match_data: (void *)name); |
376 | } |
377 | |
378 | /** |
379 | * kunit_remove_resource() - remove resource from resource list associated with |
380 | * test. |
381 | * @test: The test context object. |
382 | * @res: The resource to be removed. |
383 | * |
384 | * Note that the resource will not be immediately freed since it is likely |
385 | * the caller has a reference to it via alloc_and_get() or find(); |
386 | * in this case a final call to kunit_put_resource() is required. |
387 | */ |
388 | void kunit_remove_resource(struct kunit *test, struct kunit_resource *res); |
389 | |
390 | /* A 'deferred action' function to be used with kunit_add_action. */ |
391 | typedef void (kunit_action_t)(void *); |
392 | |
393 | /** |
394 | * KUNIT_DEFINE_ACTION_WRAPPER() - Wrap a function for use as a deferred action. |
395 | * |
396 | * @wrapper: The name of the new wrapper function define. |
397 | * @orig: The original function to wrap. |
398 | * @arg_type: The type of the argument accepted by @orig. |
399 | * |
400 | * Defines a wrapper for a function which accepts a single, pointer-sized |
401 | * argument. This wrapper can then be passed to kunit_add_action() and |
402 | * similar. This should be used in preference to casting a function |
403 | * directly to kunit_action_t, as casting function pointers will break |
404 | * control flow integrity (CFI), leading to crashes. |
405 | */ |
406 | #define KUNIT_DEFINE_ACTION_WRAPPER(wrapper, orig, arg_type) \ |
407 | static void wrapper(void *in) \ |
408 | { \ |
409 | arg_type arg = (arg_type)in; \ |
410 | orig(arg); \ |
411 | } |
412 | |
413 | |
414 | /** |
415 | * kunit_add_action() - Call a function when the test ends. |
416 | * @test: Test case to associate the action with. |
417 | * @action: The function to run on test exit |
418 | * @ctx: Data passed into @func |
419 | * |
420 | * Defer the execution of a function until the test exits, either normally or |
421 | * due to a failure. @ctx is passed as additional context. All functions |
422 | * registered with kunit_add_action() will execute in the opposite order to that |
423 | * they were registered in. |
424 | * |
425 | * This is useful for cleaning up allocated memory and resources, as these |
426 | * functions are called even if the test aborts early due to, e.g., a failed |
427 | * assertion. |
428 | * |
429 | * See also: devm_add_action() for the devres equivalent. |
430 | * |
431 | * Returns: |
432 | * 0 on success, an error if the action could not be deferred. |
433 | */ |
434 | int kunit_add_action(struct kunit *test, kunit_action_t *action, void *ctx); |
435 | |
436 | /** |
437 | * kunit_add_action_or_reset() - Call a function when the test ends. |
438 | * @test: Test case to associate the action with. |
439 | * @action: The function to run on test exit |
440 | * @ctx: Data passed into @func |
441 | * |
442 | * Defer the execution of a function until the test exits, either normally or |
443 | * due to a failure. @ctx is passed as additional context. All functions |
444 | * registered with kunit_add_action() will execute in the opposite order to that |
445 | * they were registered in. |
446 | * |
447 | * This is useful for cleaning up allocated memory and resources, as these |
448 | * functions are called even if the test aborts early due to, e.g., a failed |
449 | * assertion. |
450 | * |
451 | * If the action cannot be created (e.g., due to the system being out of memory), |
452 | * then action(ctx) will be called immediately, and an error will be returned. |
453 | * |
454 | * See also: devm_add_action_or_reset() for the devres equivalent. |
455 | * |
456 | * Returns: |
457 | * 0 on success, an error if the action could not be deferred. |
458 | */ |
459 | int kunit_add_action_or_reset(struct kunit *test, kunit_action_t *action, |
460 | void *ctx); |
461 | |
462 | /** |
463 | * kunit_remove_action() - Cancel a matching deferred action. |
464 | * @test: Test case the action is associated with. |
465 | * @action: The deferred function to cancel. |
466 | * @ctx: The context passed to the deferred function to trigger. |
467 | * |
468 | * Prevent an action deferred via kunit_add_action() from executing when the |
469 | * test terminates. |
470 | * |
471 | * If the function/context pair was deferred multiple times, only the most |
472 | * recent one will be cancelled. |
473 | * |
474 | * See also: devm_remove_action() for the devres equivalent. |
475 | */ |
476 | void kunit_remove_action(struct kunit *test, |
477 | kunit_action_t *action, |
478 | void *ctx); |
479 | |
480 | /** |
481 | * kunit_release_action() - Run a matching action call immediately. |
482 | * @test: Test case the action is associated with. |
483 | * @action: The deferred function to trigger. |
484 | * @ctx: The context passed to the deferred function to trigger. |
485 | * |
486 | * Execute a function deferred via kunit_add_action()) immediately, rather than |
487 | * when the test ends. |
488 | * |
489 | * If the function/context pair was deferred multiple times, it will only be |
490 | * executed once here. The most recent deferral will no longer execute when |
491 | * the test ends. |
492 | * |
493 | * kunit_release_action(test, func, ctx); |
494 | * is equivalent to |
495 | * func(ctx); |
496 | * kunit_remove_action(test, func, ctx); |
497 | * |
498 | * See also: devm_release_action() for the devres equivalent. |
499 | */ |
500 | void kunit_release_action(struct kunit *test, |
501 | kunit_action_t *action, |
502 | void *ctx); |
503 | #endif /* _KUNIT_RESOURCE_H */ |
504 | |