1// SPDX-License-Identifier: GPL-2.0
  2/*
  3 * KUnit function redirection (static stubbing) API.
  4 *
  5 * Copyright (C) 2022, Google LLC.
  6 * Author: David Gow <davidgow@google.com>
  7 */
  8
  9#include <kunit/test.h>
 10#include <kunit/static_stub.h>
 11#include "hooks-impl.h"
 12
 13
 14/* Context for a static stub. This is stored in the resource data. */
 15struct kunit_static_stub_ctx {
 16	void *real_fn_addr;
 17	void *replacement_addr;
 18};
 19
 20static void __kunit_static_stub_resource_free(struct kunit_resource *res)
 21{
 22	kfree(res->data);
 23}
 24
 25/* Matching function for kunit_find_resource(). match_data is real_fn_addr. */
 26static bool __kunit_static_stub_resource_match(struct kunit *test,
 27						struct kunit_resource *res,
 28						void *match_real_fn_addr)
 29{
 30	/* This pointer is only valid if res is a static stub resource. */
 31	struct kunit_static_stub_ctx *ctx = res->data;
 32
 33	/* Make sure the resource is a static stub resource. */
 34	if (res->free != &__kunit_static_stub_resource_free)
 35		return false;
 36
 37	return ctx->real_fn_addr == match_real_fn_addr;
 38}
 39
 40/* Hook to return the address of the replacement function. */
 41void *__kunit_get_static_stub_address_impl(struct kunit *test, void *real_fn_addr)
 42{
 43	struct kunit_resource *res;
 44	struct kunit_static_stub_ctx *ctx;
 45	void *replacement_addr;
 46
 47	res = kunit_find_resource(test,
 48				  __kunit_static_stub_resource_match,
 49				  real_fn_addr);
 50
 51	if (!res)
 52		return NULL;
 53
 54	ctx = res->data;
 55	replacement_addr = ctx->replacement_addr;
 56	kunit_put_resource(res);
 57	return replacement_addr;
 58}
 59
 60void kunit_deactivate_static_stub(struct kunit *test, void *real_fn_addr)
 61{
 62	struct kunit_resource *res;
 63
 64	KUNIT_ASSERT_PTR_NE_MSG(test, real_fn_addr, NULL,
 65				"Tried to deactivate a NULL stub.");
 66
 67	/* Look up the existing stub for this function. */
 68	res = kunit_find_resource(test,
 69				  __kunit_static_stub_resource_match,
 70				  real_fn_addr);
 71
 72	/* Error out if the stub doesn't exist. */
 73	KUNIT_ASSERT_PTR_NE_MSG(test, res, NULL,
 74				"Tried to deactivate a nonexistent stub.");
 75
 76	/* Free the stub. We 'put' twice, as we got a reference
 77	 * from kunit_find_resource()
 78	 */
 79	kunit_remove_resource(test, res);
 80	kunit_put_resource(res);
 81}
 82EXPORT_SYMBOL_GPL(kunit_deactivate_static_stub);
 83
 84/* Helper function for kunit_activate_static_stub(). The macro does
 85 * typechecking, so use it instead.
 86 */
 87void __kunit_activate_static_stub(struct kunit *test,
 88				  void *real_fn_addr,
 89				  void *replacement_addr)
 90{
 91	struct kunit_static_stub_ctx *ctx;
 92	struct kunit_resource *res;
 93
 94	KUNIT_ASSERT_PTR_NE_MSG(test, real_fn_addr, NULL,
 95				"Tried to activate a stub for function NULL");
 96
 97	/* If the replacement address is NULL, deactivate the stub. */
 98	if (!replacement_addr) {
 99		kunit_deactivate_static_stub(test, replacement_addr);
100		return;
101	}
102
103	/* Look up any existing stubs for this function, and replace them. */
104	res = kunit_find_resource(test,
105				  __kunit_static_stub_resource_match,
106				  real_fn_addr);
107	if (res) {
108		ctx = res->data;
109		ctx->replacement_addr = replacement_addr;
110
111		/* We got an extra reference from find_resource(), so put it. */
112		kunit_put_resource(res);
113	} else {
114		ctx = kmalloc(sizeof(*ctx), GFP_KERNEL);
115		KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx);
116		ctx->real_fn_addr = real_fn_addr;
117		ctx->replacement_addr = replacement_addr;
118		res = kunit_alloc_resource(test, NULL,
119				     &__kunit_static_stub_resource_free,
120				     GFP_KERNEL, ctx);
121	}
122}
123EXPORT_SYMBOL_GPL(__kunit_activate_static_stub);