optee: fix tee out of memory failure seen during kexec reboot

The following out of memory errors are seen on kexec reboot
from the optee core.

[    0.368428] tee_bnxt_fw optee-clnt0: tee_shm_alloc failed
[    0.368461] tee_bnxt_fw: probe of optee-clnt0 failed with error -22

tee_shm_release() is not invoked on dma shm buffer.

Implement .shutdown() method to handle the release of the buffers
correctly.

More info:
OP-TEE/optee_os#3637

Signed-off-by: Allen Pais <apais@linux.microsoft.com>

drivers/tee/optee/core.c

@@ -582,36 +582,64 @@ static optee_invoke_fn *get_invoke_func(struct device *dev)
 	return ERR_PTR(-EINVAL);
 }
 
-static int optee_remove(struct platform_device *pdev)
+static int __optee_shutoff(struct platform_device *pdev, bool shutdown)
 {
 	struct optee *optee = platform_get_drvdata(pdev);
 
-	/*
-	 * Ask OP-TEE to free all cached shared memory objects to decrease
-	 * reference counters and also avoid wild pointers in secure world
-	 * into the old shared memory range.
-	 */
-	optee_disable_shm_cache(optee);
+	if (shutdown) {
+		optee_disable_shm_cache(optee);
+	} else {
+		/*
+		 * Ask OP-TEE to free all cached shared memory
+		 * objects to decrease reference counters and
+		 * also avoid wild pointers in secure world
+		 * into the old shared memory range.
+		 */
+		optee_disable_shm_cache(optee);
 
-	/*
-	 * The two devices have to be unregistered before we can free the
-	 * other resources.
-	 */
-	tee_device_unregister(optee->supp_teedev);
-	tee_device_unregister(optee->teedev);
+		/*
+		 * The two devices have to be unregistered before
+		 * we can free the other resources.
+		 */
+		tee_device_unregister(optee->supp_teedev);
+		tee_device_unregister(optee->teedev);
 
-	tee_shm_pool_free(optee->pool);
-	if (optee->memremaped_shm)
-		memunmap(optee->memremaped_shm);
-	optee_wait_queue_exit(&optee->wait_queue);
-	optee_supp_uninit(&optee->supp);
-	mutex_destroy(&optee->call_queue.mutex);
+		tee_shm_pool_free(optee->pool);
+		if (optee->memremaped_shm)
+			memunmap(optee->memremaped_shm);
+		optee_wait_queue_exit(&optee->wait_queue);
+		optee_supp_uninit(&optee->supp);
+		mutex_destroy(&optee->call_queue.mutex);
 
-	kfree(optee);
+		kfree(optee);
+	}
 
 	return 0;
 }
 
+/* optee_remove - Device Removal Routine
+ * @pdev: platform device information struct
+ *
+ * optee_remove is called by platform subsystem to alter the driver
+ * that it should release the device
+ */
+static int optee_remove(struct platform_device *pdev)
+{
+	return __optee_shutoff(pdev, false);
+}
+
+/* optee_shutdown - Device Removal Routine
+ * @pdev: platform device information struct
+ *
+ * platform_shutdown is called by the platform subsystem to alter
+ * the driver that a shutdown/reboot(or kexec) is happening and
+ * device must be disabled.
+ */
+static void optee_shutdown(struct platform_device *pdev)
+{
+	__optee_shutoff(pdev, true);
+}
+
 static int optee_probe(struct platform_device *pdev)
 {
 	optee_invoke_fn *invoke_fn;
@@ -738,6 +766,7 @@ MODULE_DEVICE_TABLE(of, optee_dt_match);
 static struct platform_driver optee_driver = {
 	.probe  = optee_probe,
 	.remove = optee_remove,
+	.shutdown = optee_shutdown,
 	.driver = {
 		.name = "optee",
 		.of_match_table = optee_dt_match,