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#include <kunit/resource.h>
10#include <kunit/test.h>
11#include <linux/kref.h>
12
13/*
14 * Used for static resources and when a kunit_resource * has been created by
15 * kunit_alloc_resource().  When an init function is supplied, @data is passed
16 * into the init function; otherwise, we simply set the resource data field to
17 * the data value passed in. Doesn't initialize res->should_kfree.
18 */
19int __kunit_add_resource(struct kunit *test,
20			 kunit_resource_init_t init,
21			 kunit_resource_free_t free,
22			 struct kunit_resource *res,
23			 void *data)
24{
25	int ret = 0;
26	unsigned long flags;
27
28	res->free = free;
29	kref_init(&res->refcount);
30
31	if (init) {
32		ret = init(res, data);
33		if (ret)
34			return ret;
35	} else {
36		res->data = data;
37	}
38
39	spin_lock_irqsave(&test->lock, flags);
40	list_add_tail(&res->node, &test->resources);
41	/* refcount for list is established by kref_init() */
42	spin_unlock_irqrestore(&test->lock, flags);
43
44	return ret;
45}
46EXPORT_SYMBOL_GPL(__kunit_add_resource);
47
48void kunit_remove_resource(struct kunit *test, struct kunit_resource *res)
49{
50	unsigned long flags;
51	bool was_linked;
52
53	spin_lock_irqsave(&test->lock, flags);
54	was_linked = !list_empty(&res->node);
55	list_del_init(&res->node);
56	spin_unlock_irqrestore(&test->lock, flags);
57
58	if (was_linked)
59		kunit_put_resource(res);
60}
61EXPORT_SYMBOL_GPL(kunit_remove_resource);
62
63int kunit_destroy_resource(struct kunit *test, kunit_resource_match_t match,
64			   void *match_data)
65{
66	struct kunit_resource *res = kunit_find_resource(test, match,
67							 match_data);
68
69	if (!res)
70		return -ENOENT;
71
72	kunit_remove_resource(test, res);
73
74	/* We have a reference also via _find(); drop it. */
75	kunit_put_resource(res);
76
77	return 0;
78}
79EXPORT_SYMBOL_GPL(kunit_destroy_resource);
80
81struct kunit_action_ctx {
82	struct kunit_resource res;
83	kunit_action_t *func;
84	void *ctx;
85};
86
87static void __kunit_action_free(struct kunit_resource *res)
88{
89	struct kunit_action_ctx *action_ctx = container_of(res, struct kunit_action_ctx, res);
90
91	action_ctx->func(action_ctx->ctx);
92}
93
94
95int kunit_add_action(struct kunit *test, void (*action)(void *), void *ctx)
96{
97	struct kunit_action_ctx *action_ctx;
98
99	KUNIT_ASSERT_NOT_NULL_MSG(test, action, "Tried to action a NULL function!");
100
101	action_ctx = kzalloc(sizeof(*action_ctx), GFP_KERNEL);
102	if (!action_ctx)
103		return -ENOMEM;
104
105	action_ctx->func = action;
106	action_ctx->ctx = ctx;
107
108	action_ctx->res.should_kfree = true;
109	/* As init is NULL, this cannot fail. */
110	__kunit_add_resource(test, NULL, __kunit_action_free, &action_ctx->res, action_ctx);
111
112	return 0;
113}
114EXPORT_SYMBOL_GPL(kunit_add_action);
115
116int kunit_add_action_or_reset(struct kunit *test, void (*action)(void *),
117			      void *ctx)
118{
119	int res = kunit_add_action(test, action, ctx);
120
121	if (res)
122		action(ctx);
123	return res;
124}
125EXPORT_SYMBOL_GPL(kunit_add_action_or_reset);
126
127static bool __kunit_action_match(struct kunit *test,
128				struct kunit_resource *res, void *match_data)
129{
130	struct kunit_action_ctx *match_ctx = (struct kunit_action_ctx *)match_data;
131	struct kunit_action_ctx *res_ctx = container_of(res, struct kunit_action_ctx, res);
132
133	/* Make sure this is a free function. */
134	if (res->free != __kunit_action_free)
135		return false;
136
137	/* Both the function and context data should match. */
138	return (match_ctx->func == res_ctx->func) && (match_ctx->ctx == res_ctx->ctx);
139}
140
141void kunit_remove_action(struct kunit *test,
142			kunit_action_t *action,
143			void *ctx)
144{
145	struct kunit_action_ctx match_ctx;
146	struct kunit_resource *res;
147
148	match_ctx.func = action;
149	match_ctx.ctx = ctx;
150
151	res = kunit_find_resource(test, __kunit_action_match, &match_ctx);
152	if (res) {
153		/* Remove the free function so we don't run the action. */
154		res->free = NULL;
155		kunit_remove_resource(test, res);
156		kunit_put_resource(res);
157	}
158}
159EXPORT_SYMBOL_GPL(kunit_remove_action);
160
161void kunit_release_action(struct kunit *test,
162			 kunit_action_t *action,
163			 void *ctx)
164{
165	struct kunit_action_ctx match_ctx;
166	struct kunit_resource *res;
167
168	match_ctx.func = action;
169	match_ctx.ctx = ctx;
170
171	res = kunit_find_resource(test, __kunit_action_match, &match_ctx);
172	if (res) {
173		kunit_remove_resource(test, res);
174		/* We have to put() this here, else free won't be called. */
175		kunit_put_resource(res);
176	}
177}
178EXPORT_SYMBOL_GPL(kunit_release_action);
179