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);