aoc: add protect on service read/write timeout

Prevent NULL pointer during aoc SSR
Block free flow when service write read active

Bug: 209403953
Signed-off-by: yixuanjiang <yixuanjiang@google.com>
Change-Id: Ib0a0f4082018fc237cdd789c7dfba3cee03d760b

1 files changed, 77 insertions, 41 deletions

diff --git a/aoc.c b/aoc.c
index 627f68e..7ef6725 100644
--- a/aoc.c
+++ b/aoc.c
@@ -116,7 +116,7 @@ struct aoc_prvdata {
 	void *map_handler_ctx;
 
 	struct delayed_work monitor_work;
-	bool aoc_process_active;
+	atomic_t aoc_process_active;
 
 	struct device *dev;
 	struct iommu_domain *domain;
@@ -1039,7 +1039,6 @@ EXPORT_SYMBOL_GPL(aoc_service_read);
 ssize_t aoc_service_read_timeout(struct aoc_service_dev *dev, uint8_t *buffer,
 				 size_t count, long timeout)
 {
-	const struct device *parent;
 	struct aoc_prvdata *prvdata;
 	aoc_service *service;
 
@@ -1053,22 +1052,32 @@ ssize_t aoc_service_read_timeout(struct aoc_service_dev *dev, uint8_t *buffer,
 	if (dev->dead)
 		return -ENODEV;
 
-	if (aoc_state != AOC_STATE_ONLINE)
-		return -EBUSY;
+	if (!aoc_platform_device)
+		return -ENODEV;
 
-	parent = dev->dev.parent;
-	prvdata = dev_get_drvdata(parent);
+	prvdata = platform_get_drvdata(aoc_platform_device);
+	if (!prvdata)
+		return -ENODEV;
+
+	atomic_inc(&prvdata->aoc_process_active);
+	if (aoc_state != AOC_STATE_ONLINE || work_busy(&prvdata->watchdog_work)) {
+		ret = -EBUSY;
+		goto err;
+	}
 
 	service_number = dev->service_index;
 	service = service_at_index(prvdata, dev->service_index);
 
 	if (!aoc_is_valid_dram_address(prvdata, service)) {
 		WARN_ONCE(1, "aoc service %d has invalid DRAM region", service_number);
-		return -ENODEV;
+		ret = -ENODEV;
+		goto err;
 	}
 
-	if (aoc_service_message_slots(service, AOC_UP) == 0)
-		return -EBADF;
+	if (aoc_service_message_slots(service, AOC_UP) == 0) {
+		ret = -EBADF;
+		goto err;
+	}
 
 	if (!aoc_service_can_read_message(service, AOC_UP)) {
 		set_bit(service_number, &read_blocked_mask);
@@ -1080,28 +1089,37 @@ ssize_t aoc_service_read_timeout(struct aoc_service_dev *dev, uint8_t *buffer,
 		clear_bit(service_number, &read_blocked_mask);
 	}
 
-	if (dev->dead)
-		return -ENODEV;
-
-	if (aoc_state != AOC_STATE_ONLINE)
-		return -ENODEV;
+	if (dev->dead || (aoc_state != AOC_STATE_ONLINE)) {
+		ret = -ENODEV;
+		goto err;
+	}
 
 	if (ret < 0)
-		return ret;
+		goto err;
 
 	/* AoC timed out */
-	if (ret == 0)
-		return -ETIMEDOUT;
+	if (ret == 0) {
+		ret = -ETIMEDOUT;
+		goto err;
+	}
 
 	if (!aoc_service_is_ring(service) &&
 	    count < aoc_service_current_message_size(service, prvdata->ipc_base,
-						     AOC_UP))
-		return -EFBIG;
+						     AOC_UP)) {
+		ret = -EFBIG;
+		goto err;
+	}
 
 	msg_size = count;
 	aoc_service_read_message(service, prvdata->ipc_base, AOC_UP, buffer,
 				 &msg_size);
 
+err:
+	atomic_dec(&prvdata->aoc_process_active);
+
+	if (ret < 0)
+		return ret;
+
 	return msg_size;
 }
 EXPORT_SYMBOL_GPL(aoc_service_read_timeout);
@@ -1182,7 +1200,6 @@ EXPORT_SYMBOL_GPL(aoc_service_write);
 ssize_t aoc_service_write_timeout(struct aoc_service_dev *dev, const uint8_t *buffer,
 				  size_t count, long timeout)
 {
-	const struct device *parent;
 	struct aoc_prvdata *prvdata;
 
 	aoc_service *service;
@@ -1196,25 +1213,37 @@ ssize_t aoc_service_write_timeout(struct aoc_service_dev *dev, const uint8_t *bu
 	if (dev->dead)
 		return -ENODEV;
 
-	if (aoc_state != AOC_STATE_ONLINE)
+	if (!aoc_platform_device)
 		return -ENODEV;
 
-	parent = dev->dev.parent;
-	prvdata = dev_get_drvdata(parent);
+	prvdata = platform_get_drvdata(aoc_platform_device);
+	if (!prvdata)
+		return -ENODEV;
+
+	atomic_inc(&prvdata->aoc_process_active);
+	if (aoc_state != AOC_STATE_ONLINE || work_busy(&prvdata->watchdog_work)) {
+		ret = -EBUSY;
+		goto err;
+	}
 
 	service_number = dev->service_index;
 	service = service_at_index(prvdata, service_number);
 
 	if (!aoc_is_valid_dram_address(prvdata, service)) {
 		WARN_ONCE(1, "aoc service %d has invalid DRAM region", service_number);
-		return -ENODEV;
+		ret = -ENODEV;
+		goto err;
 	}
 
-	if (aoc_service_message_slots(service, AOC_DOWN) == 0)
-		return -EBADF;
+	if (aoc_service_message_slots(service, AOC_DOWN) == 0) {
+		ret = -EBADF;
+		goto err;
+	}
 
-	if (count > aoc_service_message_size(service, AOC_DOWN))
-		return -EFBIG;
+	if (count > aoc_service_message_size(service, AOC_DOWN)) {
+		ret = -EFBIG;
+		goto err;
+	}
 
 	if (!aoc_service_can_write_message(service, AOC_DOWN)) {
 		set_bit(service_number, &write_blocked_mask);
@@ -1226,17 +1255,18 @@ ssize_t aoc_service_write_timeout(struct aoc_service_dev *dev, const uint8_t *bu
 		clear_bit(service_number, &write_blocked_mask);
 	}
 
-	if (dev->dead)
-		return -ENODEV;
-
-	if (aoc_state != AOC_STATE_ONLINE)
-		return -ENODEV;
+	if (dev->dead || (aoc_state != AOC_STATE_ONLINE)) {
+		ret = -ENODEV;
+		goto err;
+	}
 
 	if (ret < 0)
-		return ret;
+		goto err;
 
-	if (ret == 0)
-		return -ETIMEDOUT;
+	if (ret == 0) {
+		ret = -ETIMEDOUT;
+		goto err;
+	}
 
 	ret = aoc_service_write_message(service, prvdata->ipc_base, AOC_DOWN,
 					buffer, count);
@@ -1244,6 +1274,12 @@ ssize_t aoc_service_write_timeout(struct aoc_service_dev *dev, const uint8_t *bu
 	if (!aoc_service_is_ring(service) || aoc_ring_is_push(service))
 		signal_aoc(prvdata->mbox_channels[interrupt].channel);
 
+err:
+	atomic_dec(&prvdata->aoc_process_active);
+
+	if (ret < 0)
+		return ret;
+
 	return count;
 }
 EXPORT_SYMBOL_GPL(aoc_service_write_timeout);
@@ -2190,8 +2226,8 @@ static void aoc_take_offline(struct aoc_prvdata *prvdata)
 		pr_notice("taking aoc offline\n");
 		aoc_state = AOC_STATE_OFFLINE;
 
-		/* wait until aoc_process_services finish */
-		while (prvdata->aoc_process_active);
+		/* wait until aoc_process or service write/read finish */
+		while (!!atomic_read(&prvdata->aoc_process_active));
 
 		bus_for_each_dev(&aoc_bus_type, NULL, NULL, aoc_remove_device);
 
@@ -2227,11 +2263,11 @@ static void aoc_process_services(struct aoc_prvdata *prvdata, int offset)
 	int services;
 	int i;
 
+	atomic_inc(&prvdata->aoc_process_active);
+
 	if (aoc_state != AOC_STATE_ONLINE || work_busy(&prvdata->watchdog_work))
 		goto exit;
 
-	prvdata->aoc_process_active = true;
-
 	services = aoc_num_services();
 	for (i = 0; i < services; i++) {
 		service_dev = service_dev_at_index(prvdata, i);
@@ -2252,7 +2288,7 @@ static void aoc_process_services(struct aoc_prvdata *prvdata, int offset)
 		}
 	}
 exit:
-	prvdata->aoc_process_active = false;
+	atomic_dec(&prvdata->aoc_process_active);
 }
 
 void aoc_set_map_handler(struct aoc_service_dev *dev, aoc_map_handler handler,