Merge tag 'tee-bus-callback-for-6.20' of git://git.kernel.org/pub/scm/linux/kernel/git/jenswi/linux-tee into soc/drivers

TEE bus callback for 6.20

- Move from generic device_driver to TEE bus-specific callbacks
- Add module_tee_client_driver() and registration helpers to reduce
  boilerplate
- Convert several client drivers (TPM, KEYS, firmware, EFI, hwrng,
  and RTC)
- Update documentation and fix kernel-doc warnings

* tag 'tee-bus-callback-for-6.20' of git://git.kernel.org/pub/scm/linux/kernel/git/jenswi/linux-tee:
  tpm/tpm_ftpm_tee: Fix kdoc after function renames
  tpm/tpm_ftpm_tee: Make use of tee bus methods
  tpm/tpm_ftpm_tee: Make use of tee specific driver registration
  KEYS: trusted: Make use of tee bus methods
  KEYS: trusted: Migrate to use tee specific driver registration function
  firmware: tee_bnxt: Make use of tee bus methods
  firmware: tee_bnxt: Make use of module_tee_client_driver()
  firmware: arm_scmi: Make use of tee bus methods
  firmware: arm_scmi: optee: Make use of module_tee_client_driver()
  efi: stmm: Make use of tee bus methods
  efi: stmm: Make use of module_tee_client_driver()
  hwrng: optee - Make use of tee bus methods
  hwrng: optee - Make use of module_tee_client_driver()
  rtc: optee: Make use of tee bus methods
  rtc: optee: Migrate to use tee specific driver registration function
  tee: Adapt documentation to cover recent additions
  tee: Add probe, remove and shutdown bus callbacks to tee_client_driver
  tee: Add some helpers to reduce boilerplate for tee client drivers

Signed-off-by: Arnd Bergmann <arnd@arndb.de>

10 files changed, 166 insertions, 140 deletions

diff --git a/Documentation/driver-api/tee.rst b/Documentation/driver-api/tee.rst
index 5eaeb8103988..4d58ac0712c1 100644
--- a/Documentation/driver-api/tee.rst
+++ b/Documentation/driver-api/tee.rst
@@ -43,24 +43,12 @@ snippet would look like::
 	MODULE_DEVICE_TABLE(tee, client_id_table);
 
 	static struct tee_client_driver client_driver = {
+		.probe		= client_probe,
+		.remove		= client_remove,
 		.id_table	= client_id_table,
 		.driver		= {
 			.name		= DRIVER_NAME,
-			.bus		= &tee_bus_type,
-			.probe		= client_probe,
-			.remove		= client_remove,
 		},
 	};
 
-	static int __init client_init(void)
-	{
-		return driver_register(&client_driver.driver);
-	}
-
-	static void __exit client_exit(void)
-	{
-		driver_unregister(&client_driver.driver);
-	}
-
-	module_init(client_init);
-	module_exit(client_exit);
+	module_tee_client_driver(client_driver);
diff --git a/drivers/char/hw_random/optee-rng.c b/drivers/char/hw_random/optee-rng.c
index 96b5d546d136..5a3fa0b38497 100644
--- a/drivers/char/hw_random/optee-rng.c
+++ b/drivers/char/hw_random/optee-rng.c
@@ -211,9 +211,9 @@ static int optee_ctx_match(struct tee_ioctl_version_data *ver, const void *data)
 		return 0;
 }
 
-static int optee_rng_probe(struct device *dev)
+static int optee_rng_probe(struct tee_client_device *rng_device)
 {
-	struct tee_client_device *rng_device = to_tee_client_device(dev);
+	struct device *dev = &rng_device->dev;
 	int ret = 0, err = -ENODEV;
 	struct tee_ioctl_open_session_arg sess_arg;
 
@@ -261,12 +261,10 @@ out_ctx:
 	return err;
 }
 
-static int optee_rng_remove(struct device *dev)
+static void optee_rng_remove(struct tee_client_device *tee_dev)
 {
 	tee_client_close_session(pvt_data.ctx, pvt_data.session_id);
 	tee_client_close_context(pvt_data.ctx);
-
-	return 0;
 }
 
 static const struct tee_client_device_id optee_rng_id_table[] = {
@@ -278,27 +276,15 @@ static const struct tee_client_device_id optee_rng_id_table[] = {
 MODULE_DEVICE_TABLE(tee, optee_rng_id_table);
 
 static struct tee_client_driver optee_rng_driver = {
+	.probe		= optee_rng_probe,
+	.remove		= optee_rng_remove,
 	.id_table	= optee_rng_id_table,
 	.driver		= {
 		.name		= DRIVER_NAME,
-		.bus		= &tee_bus_type,
-		.probe		= optee_rng_probe,
-		.remove		= optee_rng_remove,
 	},
 };
 
-static int __init optee_rng_mod_init(void)
-{
-	return driver_register(&optee_rng_driver.driver);
-}
-
-static void __exit optee_rng_mod_exit(void)
-{
-	driver_unregister(&optee_rng_driver.driver);
-}
-
-module_init(optee_rng_mod_init);
-module_exit(optee_rng_mod_exit);
+module_tee_client_driver(optee_rng_driver);
 
 MODULE_LICENSE("GPL v2");
 MODULE_AUTHOR("Sumit Garg <sumit.garg@linaro.org>");
diff --git a/drivers/char/tpm/tpm_ftpm_tee.c b/drivers/char/tpm/tpm_ftpm_tee.c
index 4e63c30aeaf1..b82490439633 100644
--- a/drivers/char/tpm/tpm_ftpm_tee.c
+++ b/drivers/char/tpm/tpm_ftpm_tee.c
@@ -163,13 +163,13 @@ static int ftpm_tee_match(struct tee_ioctl_version_data *ver, const void *data)
 }
 
 /**
- * ftpm_tee_probe() - initialize the fTPM
+ * ftpm_tee_probe_generic() - initialize the fTPM
  * @dev: the device description.
  *
  * Return:
  *	On success, 0. On failure, -errno.
  */
-static int ftpm_tee_probe(struct device *dev)
+static int ftpm_tee_probe_generic(struct device *dev)
 {
 	int rc;
 	struct tpm_chip *chip;
@@ -251,21 +251,28 @@ out_tee_session:
 	return rc;
 }
 
+static int ftpm_tee_probe(struct tee_client_device *tcdev)
+{
+	struct device *dev = &tcdev->dev;
+
+	return ftpm_tee_probe_generic(dev);
+}
+
 static int ftpm_plat_tee_probe(struct platform_device *pdev)
 {
 	struct device *dev = &pdev->dev;
 
-	return ftpm_tee_probe(dev);
+	return ftpm_tee_probe_generic(dev);
 }
 
 /**
- * ftpm_tee_remove() - remove the TPM device
+ * ftpm_tee_remove_generic() - remove the TPM device
  * @dev: the device description.
  *
  * Return:
  *	0 always.
  */
-static int ftpm_tee_remove(struct device *dev)
+static void ftpm_tee_remove_generic(struct device *dev)
 {
 	struct ftpm_tee_private *pvt_data = dev_get_drvdata(dev);
 
@@ -285,15 +292,20 @@ static int ftpm_tee_remove(struct device *dev)
 	tee_client_close_context(pvt_data->ctx);
 
 	/* memory allocated with devm_kzalloc() is freed automatically */
+}
 
-	return 0;
+static void ftpm_tee_remove(struct tee_client_device *tcdev)
+{
+	struct device *dev = &tcdev->dev;
+
+	ftpm_tee_remove_generic(dev);
 }
 
 static void ftpm_plat_tee_remove(struct platform_device *pdev)
 {
 	struct device *dev = &pdev->dev;
 
-	ftpm_tee_remove(dev);
+	ftpm_tee_remove_generic(dev);
 }
 
 /**
@@ -335,12 +347,11 @@ static const struct tee_client_device_id optee_ftpm_id_table[] = {
 MODULE_DEVICE_TABLE(tee, optee_ftpm_id_table);
 
 static struct tee_client_driver ftpm_tee_driver = {
+	.probe		= ftpm_tee_probe,
+	.remove		= ftpm_tee_remove,
 	.id_table	= optee_ftpm_id_table,
 	.driver		= {
 		.name		= "optee-ftpm",
-		.bus		= &tee_bus_type,
-		.probe		= ftpm_tee_probe,
-		.remove		= ftpm_tee_remove,
 	},
 };
 
@@ -352,7 +363,7 @@ static int __init ftpm_mod_init(void)
 	if (rc)
 		return rc;
 
-	rc = driver_register(&ftpm_tee_driver.driver);
+	rc = tee_client_driver_register(&ftpm_tee_driver);
 	if (rc) {
 		platform_driver_unregister(&ftpm_tee_plat_driver);
 		return rc;
@@ -364,7 +375,7 @@ static int __init ftpm_mod_init(void)
 static void __exit ftpm_mod_exit(void)
 {
 	platform_driver_unregister(&ftpm_tee_plat_driver);
-	driver_unregister(&ftpm_tee_driver.driver);
+	tee_client_driver_unregister(&ftpm_tee_driver);
 }
 
 module_init(ftpm_mod_init);
diff --git a/drivers/firmware/arm_scmi/transports/optee.c b/drivers/firmware/arm_scmi/transports/optee.c
index dc0f46340153..07ae18d5279d 100644
--- a/drivers/firmware/arm_scmi/transports/optee.c
+++ b/drivers/firmware/arm_scmi/transports/optee.c
@@ -529,8 +529,9 @@ static const struct of_device_id scmi_of_match[] = {
 DEFINE_SCMI_TRANSPORT_DRIVER(scmi_optee, scmi_optee_driver, scmi_optee_desc,
 			     scmi_of_match, core);
 
-static int scmi_optee_service_probe(struct device *dev)
+static int scmi_optee_service_probe(struct tee_client_device *scmi_pta)
 {
+	struct device *dev = &scmi_pta->dev;
 	struct scmi_optee_agent *agent;
 	struct tee_context *tee_ctx;
 	int ret;
@@ -578,24 +579,22 @@ err:
 	return ret;
 }
 
-static int scmi_optee_service_remove(struct device *dev)
+static void scmi_optee_service_remove(struct tee_client_device *scmi_pta)
 {
 	struct scmi_optee_agent *agent = scmi_optee_private;
 
 	if (!scmi_optee_private)
-		return -EINVAL;
+		return;
 
 	platform_driver_unregister(&scmi_optee_driver);
 
 	if (!list_empty(&scmi_optee_private->channel_list))
-		return -EBUSY;
+		return;
 
 	/* Ensure cleared reference is visible before resources are released */
 	smp_store_mb(scmi_optee_private, NULL);
 
 	tee_client_close_context(agent->tee_ctx);
-
-	return 0;
 }
 
 static const struct tee_client_device_id scmi_optee_service_id[] = {
@@ -609,26 +608,15 @@ static const struct tee_client_device_id scmi_optee_service_id[] = {
 MODULE_DEVICE_TABLE(tee, scmi_optee_service_id);
 
 static struct tee_client_driver scmi_optee_service_driver = {
-	.id_table	= scmi_optee_service_id,
-	.driver		= {
+	.probe = scmi_optee_service_probe,
+	.remove = scmi_optee_service_remove,
+	.id_table = scmi_optee_service_id,
+	.driver = {
 		.name = "scmi-optee",
-		.bus = &tee_bus_type,
-		.probe = scmi_optee_service_probe,
-		.remove = scmi_optee_service_remove,
 	},
 };
 
-static int __init scmi_transport_optee_init(void)
-{
-	return driver_register(&scmi_optee_service_driver.driver);
-}
-module_init(scmi_transport_optee_init);
-
-static void __exit scmi_transport_optee_exit(void)
-{
-	driver_unregister(&scmi_optee_service_driver.driver);
-}
-module_exit(scmi_transport_optee_exit);
+module_tee_client_driver(scmi_optee_service_driver);
 
 MODULE_AUTHOR("Etienne Carriere <etienne.carriere@foss.st.com>");
 MODULE_DESCRIPTION("SCMI OPTEE Transport driver");
diff --git a/drivers/firmware/broadcom/tee_bnxt_fw.c b/drivers/firmware/broadcom/tee_bnxt_fw.c
index 40e3183a3d11..a706c84eb2b6 100644
--- a/drivers/firmware/broadcom/tee_bnxt_fw.c
+++ b/drivers/firmware/broadcom/tee_bnxt_fw.c
@@ -181,9 +181,9 @@ static int optee_ctx_match(struct tee_ioctl_version_data *ver, const void *data)
 	return (ver->impl_id == TEE_IMPL_ID_OPTEE);
 }
 
-static int tee_bnxt_fw_probe(struct device *dev)
+static int tee_bnxt_fw_probe(struct tee_client_device *bnxt_device)
 {
-	struct tee_client_device *bnxt_device = to_tee_client_device(dev);
+	struct device *dev = &bnxt_device->dev;
 	int ret, err = -ENODEV;
 	struct tee_ioctl_open_session_arg sess_arg;
 	struct tee_shm *fw_shm_pool;
@@ -231,17 +231,15 @@ out_ctx:
 	return err;
 }
 
-static int tee_bnxt_fw_remove(struct device *dev)
+static void tee_bnxt_fw_remove(struct tee_client_device *bnxt_device)
 {
 	tee_shm_free(pvt_data.fw_shm_pool);
 	tee_client_close_session(pvt_data.ctx, pvt_data.session_id);
 	tee_client_close_context(pvt_data.ctx);
 	pvt_data.ctx = NULL;
-
-	return 0;
 }
 
-static void tee_bnxt_fw_shutdown(struct device *dev)
+static void tee_bnxt_fw_shutdown(struct tee_client_device *bnxt_device)
 {
 	tee_shm_free(pvt_data.fw_shm_pool);
 	tee_client_close_session(pvt_data.ctx, pvt_data.session_id);
@@ -258,28 +256,16 @@ static const struct tee_client_device_id tee_bnxt_fw_id_table[] = {
 MODULE_DEVICE_TABLE(tee, tee_bnxt_fw_id_table);
 
 static struct tee_client_driver tee_bnxt_fw_driver = {
+	.probe		= tee_bnxt_fw_probe,
+	.remove		= tee_bnxt_fw_remove,
+	.shutdown	= tee_bnxt_fw_shutdown,
 	.id_table	= tee_bnxt_fw_id_table,
 	.driver		= {
 		.name		= KBUILD_MODNAME,
-		.bus		= &tee_bus_type,
-		.probe		= tee_bnxt_fw_probe,
-		.remove		= tee_bnxt_fw_remove,
-		.shutdown	= tee_bnxt_fw_shutdown,
 	},
 };
 
-static int __init tee_bnxt_fw_mod_init(void)
-{
-	return driver_register(&tee_bnxt_fw_driver.driver);
-}
-
-static void __exit tee_bnxt_fw_mod_exit(void)
-{
-	driver_unregister(&tee_bnxt_fw_driver.driver);
-}
-
-module_init(tee_bnxt_fw_mod_init);
-module_exit(tee_bnxt_fw_mod_exit);
+module_tee_client_driver(tee_bnxt_fw_driver);
 
 MODULE_AUTHOR("Vikas Gupta <vikas.gupta@broadcom.com>");
 MODULE_DESCRIPTION("Broadcom bnxt firmware manager");
diff --git a/drivers/firmware/efi/stmm/tee_stmm_efi.c b/drivers/firmware/efi/stmm/tee_stmm_efi.c
index 65c0fe1ba275..7b04dd649629 100644
--- a/drivers/firmware/efi/stmm/tee_stmm_efi.c
+++ b/drivers/firmware/efi/stmm/tee_stmm_efi.c
@@ -520,8 +520,9 @@ static void tee_stmm_restore_efivars_generic_ops(void)
 	efivars_generic_ops_register();
 }
 
-static int tee_stmm_efi_probe(struct device *dev)
+static int tee_stmm_efi_probe(struct tee_client_device *tee_dev)
 {
+	struct device *dev = &tee_dev->dev;
 	struct tee_ioctl_open_session_arg sess_arg;
 	efi_status_t ret;
 	int rc;
@@ -571,37 +572,23 @@ static int tee_stmm_efi_probe(struct device *dev)
 	return 0;
 }
 
-static int tee_stmm_efi_remove(struct device *dev)
+static void tee_stmm_efi_remove(struct tee_client_device *dev)
 {
 	tee_stmm_restore_efivars_generic_ops();
-
-	return 0;
 }
 
 MODULE_DEVICE_TABLE(tee, tee_stmm_efi_id_table);
 
 static struct tee_client_driver tee_stmm_efi_driver = {
 	.id_table	= tee_stmm_efi_id_table,
+	.probe		= tee_stmm_efi_probe,
+	.remove		= tee_stmm_efi_remove,
 	.driver		= {
 		.name		= "tee-stmm-efi",
-		.bus		= &tee_bus_type,
-		.probe		= tee_stmm_efi_probe,
-		.remove		= tee_stmm_efi_remove,
 	},
 };
 
-static int __init tee_stmm_efi_mod_init(void)
-{
-	return driver_register(&tee_stmm_efi_driver.driver);
-}
-
-static void __exit tee_stmm_efi_mod_exit(void)
-{
-	driver_unregister(&tee_stmm_efi_driver.driver);
-}
-
-module_init(tee_stmm_efi_mod_init);
-module_exit(tee_stmm_efi_mod_exit);
+module_tee_client_driver(tee_stmm_efi_driver);
 
 MODULE_LICENSE("GPL");
 MODULE_AUTHOR("Ilias Apalodimas <ilias.apalodimas@linaro.org>");
diff --git a/drivers/rtc/rtc-optee.c b/drivers/rtc/rtc-optee.c
index 184c6d142801..eefde789d194 100644
--- a/drivers/rtc/rtc-optee.c
+++ b/drivers/rtc/rtc-optee.c
@@ -547,9 +547,9 @@ static int optee_ctx_match(struct tee_ioctl_version_data *ver, const void *data)
 		return 0;
 }
 
-static int optee_rtc_probe(struct device *dev)
+static int optee_rtc_probe(struct tee_client_device *rtc_device)
 {
-	struct tee_client_device *rtc_device = to_tee_client_device(dev);
+	struct device *dev = &rtc_device->dev;
 	struct tee_ioctl_open_session_arg sess2_arg = {0};
 	struct tee_ioctl_open_session_arg sess_arg = {0};
 	struct optee_rtc *priv;
@@ -682,8 +682,9 @@ out_ctx:
 	return err;
 }
 
-static int optee_rtc_remove(struct device *dev)
+static void optee_rtc_remove(struct tee_client_device *rtc_device)
 {
+	struct device *dev = &rtc_device->dev;
 	struct optee_rtc *priv = dev_get_drvdata(dev);
 
 	if (priv->features & TA_RTC_FEATURE_ALARM) {
@@ -696,8 +697,6 @@ static int optee_rtc_remove(struct device *dev)
 	tee_shm_free(priv->shm);
 	tee_client_close_session(priv->ctx, priv->session_id);
 	tee_client_close_context(priv->ctx);
-
-	return 0;
 }
 
 static int optee_rtc_suspend(struct device *dev)
@@ -724,27 +723,15 @@ MODULE_DEVICE_TABLE(tee, optee_rtc_id_table);
 
 static struct tee_client_driver optee_rtc_driver = {
 	.id_table	= optee_rtc_id_table,
+	.probe		= optee_rtc_probe,
+	.remove		= optee_rtc_remove,
 	.driver		= {
 		.name		= "optee_rtc",
-		.bus		= &tee_bus_type,
-		.probe		= optee_rtc_probe,
-		.remove		= optee_rtc_remove,
 		.pm		= pm_sleep_ptr(&optee_rtc_pm_ops),
 	},
 };
 
-static int __init optee_rtc_mod_init(void)
-{
-	return driver_register(&optee_rtc_driver.driver);
-}
-
-static void __exit optee_rtc_mod_exit(void)
-{
-	driver_unregister(&optee_rtc_driver.driver);
-}
-
-module_init(optee_rtc_mod_init);
-module_exit(optee_rtc_mod_exit);
+module_tee_client_driver(optee_rtc_driver);
 
 MODULE_LICENSE("GPL v2");
 MODULE_AUTHOR("Clément Léger <clement.leger@bootlin.com>");
diff --git a/drivers/tee/tee_core.c b/drivers/tee/tee_core.c
index d65d47cc154e..a015730664c7 100644
--- a/drivers/tee/tee_core.c
+++ b/drivers/tee/tee_core.c
@@ -1398,13 +1398,97 @@ static int tee_client_device_uevent(const struct device *dev,
 	return add_uevent_var(env, "MODALIAS=tee:%pUb", dev_id);
 }
 
+static int tee_client_device_probe(struct device *dev)
+{
+	struct tee_client_device *tcdev = to_tee_client_device(dev);
+	struct tee_client_driver *drv = to_tee_client_driver(dev->driver);
+
+	if (drv->probe)
+		return drv->probe(tcdev);
+	else
+		return 0;
+}
+
+static void tee_client_device_remove(struct device *dev)
+{
+	struct tee_client_device *tcdev = to_tee_client_device(dev);
+	struct tee_client_driver *drv = to_tee_client_driver(dev->driver);
+
+	if (drv->remove)
+		drv->remove(tcdev);
+}
+
+static void tee_client_device_shutdown(struct device *dev)
+{
+	struct tee_client_device *tcdev = to_tee_client_device(dev);
+	struct tee_client_driver *drv = to_tee_client_driver(dev->driver);
+
+	if (dev->driver && drv->shutdown)
+		drv->shutdown(tcdev);
+}
+
 const struct bus_type tee_bus_type = {
 	.name		= "tee",
 	.match		= tee_client_device_match,
 	.uevent		= tee_client_device_uevent,
+	.probe		= tee_client_device_probe,
+	.remove		= tee_client_device_remove,
+	.shutdown	= tee_client_device_shutdown,
 };
 EXPORT_SYMBOL_GPL(tee_bus_type);
 
+static int tee_client_device_probe_legacy(struct tee_client_device *tcdev)
+{
+	struct device *dev = &tcdev->dev;
+	struct device_driver *driver = dev->driver;
+
+	return driver->probe(dev);
+}
+
+static void tee_client_device_remove_legacy(struct tee_client_device *tcdev)
+{
+	struct device *dev = &tcdev->dev;
+	struct device_driver *driver = dev->driver;
+
+	driver->remove(dev);
+}
+
+static void tee_client_device_shutdown_legacy(struct tee_client_device *tcdev)
+{
+	struct device *dev = &tcdev->dev;
+	struct device_driver *driver = dev->driver;
+
+	driver->shutdown(dev);
+}
+
+int __tee_client_driver_register(struct tee_client_driver *tee_driver,
+				 struct module *owner)
+{
+	tee_driver->driver.owner = owner;
+	tee_driver->driver.bus = &tee_bus_type;
+
+	/*
+	 * Drivers that have callbacks set for tee_driver->driver need updating
+	 * to use the callbacks in tee_driver instead. driver_register() warns
+	 * about that, so no need to warn here, too.
+	 */
+	if (!tee_driver->probe && tee_driver->driver.probe)
+		tee_driver->probe = tee_client_device_probe_legacy;
+	if (!tee_driver->remove && tee_driver->driver.remove)
+		tee_driver->remove = tee_client_device_remove_legacy;
+	if (!tee_driver->shutdown && tee_driver->driver.probe)
+		tee_driver->shutdown = tee_client_device_shutdown_legacy;
+
+	return driver_register(&tee_driver->driver);
+}
+EXPORT_SYMBOL_GPL(__tee_client_driver_register);
+
+void tee_client_driver_unregister(struct tee_client_driver *tee_driver)
+{
+	driver_unregister(&tee_driver->driver);
+}
+EXPORT_SYMBOL_GPL(tee_client_driver_unregister);
+
 static int __init tee_init(void)
 {
 	int rc;
diff --git a/include/linux/tee_drv.h b/include/linux/tee_drv.h
index 88a6f9697c89..e561a26f537a 100644
--- a/include/linux/tee_drv.h
+++ b/include/linux/tee_drv.h
@@ -315,6 +315,9 @@ struct tee_client_device {
  * @driver:		driver structure
  */
 struct tee_client_driver {
+	int (*probe)(struct tee_client_device *);
+	void (*remove)(struct tee_client_device *);
+	void (*shutdown)(struct tee_client_device *);
 	const struct tee_client_device_id *id_table;
 	struct device_driver driver;
 };
@@ -322,4 +325,13 @@ struct tee_client_driver {
 #define to_tee_client_driver(d) \
 		container_of_const(d, struct tee_client_driver, driver)
 
+#define tee_client_driver_register(drv) \
+        __tee_client_driver_register(drv, THIS_MODULE)
+int __tee_client_driver_register(struct tee_client_driver *, struct module *);
+void tee_client_driver_unregister(struct tee_client_driver *);
+
+#define module_tee_client_driver(__tee_client_driver) \
+	module_driver(__tee_client_driver, tee_client_driver_register, \
+		      tee_client_driver_unregister)
+
 #endif /*__TEE_DRV_H*/
diff --git a/security/keys/trusted-keys/trusted_tee.c b/security/keys/trusted-keys/trusted_tee.c
index aa3d477de6db..6e465c8bef5e 100644
--- a/security/keys/trusted-keys/trusted_tee.c
+++ b/security/keys/trusted-keys/trusted_tee.c
@@ -202,9 +202,9 @@ static int optee_ctx_match(struct tee_ioctl_version_data *ver, const void *data)
 		return 0;
 }
 
-static int trusted_key_probe(struct device *dev)
+static int trusted_key_probe(struct tee_client_device *rng_device)
 {
-	struct tee_client_device *rng_device = to_tee_client_device(dev);
+	struct device *dev = &rng_device->dev;
 	int ret;
 	struct tee_ioctl_open_session_arg sess_arg;
 
@@ -244,13 +244,11 @@ out_ctx:
 	return ret;
 }
 
-static int trusted_key_remove(struct device *dev)
+static void trusted_key_remove(struct tee_client_device *dev)
 {
 	unregister_key_type(&key_type_trusted);
 	tee_client_close_session(pvt_data.ctx, pvt_data.session_id);
 	tee_client_close_context(pvt_data.ctx);
-
-	return 0;
 }
 
 static const struct tee_client_device_id trusted_key_id_table[] = {
@@ -261,23 +259,22 @@ static const struct tee_client_device_id trusted_key_id_table[] = {
 MODULE_DEVICE_TABLE(tee, trusted_key_id_table);
 
 static struct tee_client_driver trusted_key_driver = {
+	.probe		= trusted_key_probe,
+	.remove		= trusted_key_remove,
 	.id_table	= trusted_key_id_table,
 	.driver		= {
 		.name		= DRIVER_NAME,
-		.bus		= &tee_bus_type,
-		.probe		= trusted_key_probe,
-		.remove		= trusted_key_remove,
 	},
 };
 
 static int trusted_tee_init(void)
 {
-	return driver_register(&trusted_key_driver.driver);
+	return tee_client_driver_register(&trusted_key_driver);
 }
 
 static void trusted_tee_exit(void)
 {
-	driver_unregister(&trusted_key_driver.driver);
+	tee_client_driver_unregister(&trusted_key_driver);
 }
 
 struct trusted_key_ops trusted_key_tee_ops = {