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  */
__kunit_add_resource(struct kunit * test,kunit_resource_init_t init,kunit_resource_free_t free,struct kunit_resource * res,void * data)19 int __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 }
46 EXPORT_SYMBOL_GPL(__kunit_add_resource);
47 
kunit_remove_resource(struct kunit * test,struct kunit_resource * res)48 void 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 }
61 EXPORT_SYMBOL_GPL(kunit_remove_resource);
62 
kunit_destroy_resource(struct kunit * test,kunit_resource_match_t match,void * match_data)63 int 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 }
79 EXPORT_SYMBOL_GPL(kunit_destroy_resource);
80 
81 struct kunit_action_ctx {
82 	struct kunit_resource res;
83 	kunit_action_t *func;
84 	void *ctx;
85 };
86 
__kunit_action_free(struct kunit_resource * res)87 static 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 
kunit_add_action(struct kunit * test,void (* action)(void *),void * ctx)95 int 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 }
114 EXPORT_SYMBOL_GPL(kunit_add_action);
115 
kunit_add_action_or_reset(struct kunit * test,void (* action)(void *),void * ctx)116 int 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 }
125 EXPORT_SYMBOL_GPL(kunit_add_action_or_reset);
126 
__kunit_action_match(struct kunit * test,struct kunit_resource * res,void * match_data)127 static 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 
kunit_remove_action(struct kunit * test,kunit_action_t * action,void * ctx)141 void 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 }
159 EXPORT_SYMBOL_GPL(kunit_remove_action);
160 
kunit_release_action(struct kunit * test,kunit_action_t * action,void * ctx)161 void 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 }
178 EXPORT_SYMBOL_GPL(kunit_release_action);
179