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
19struct kunit_resource;
20
21typedef int (*kunit_resource_init_t)(struct kunit_resource *, void *);
22typedef 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, &params);
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 */
83struct 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 */
100static 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 */
109static 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 */
136static 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 */
153int __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 */
169static 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
179static inline struct kunit_resource *
180kunit_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 */
191static 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 */
236static inline struct kunit_resource *
237kunit_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(sizeof(*res), 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 */
280static 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(sizeof(*res), 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
299typedef 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 */
309static 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 */
322static inline struct kunit_resource *
323kunit_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 */
350static inline struct kunit_resource *
351kunit_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 */
367int kunit_destroy_resource(struct kunit *test,
368 kunit_resource_match_t match,
369 void *match_data);
370
371static 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 */
388void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
389
390/* A 'deferred action' function to be used with kunit_add_action. */
391typedef 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 */
434int 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 */
459int 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 */
476void 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 */
500void kunit_release_action(struct kunit *test,
501 kunit_action_t *action,
502 void *ctx);
503#endif /* _KUNIT_RESOURCE_H */
504