1 files changed, 44 insertions, 13 deletions

diff --git a/drivers/char/diag/diag_mux.c b/drivers/char/diag/diag_mux.c
index 6586f5e0cf86..55c5de1ea9fc 100644
--- a/drivers/char/diag/diag_mux.c
+++ b/drivers/char/diag/diag_mux.c
@@ -1,4 +1,4 @@
-/* Copyright (c) 2014-2016, The Linux Foundation. All rights reserved.
+/* Copyright (c) 2014-2017, The Linux Foundation. All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify
  * it under the terms of the GNU General Public License version 2 and
@@ -133,21 +133,43 @@ int diag_mux_queue_read(int proc)
 int diag_mux_write(int proc, unsigned char *buf, int len, int ctx)
 {
 	struct diag_logger_t *logger = NULL;
-	int peripheral;
+	int peripheral, upd;
 
 	if (proc < 0 || proc >= NUM_MUX_PROC)
 		return -EINVAL;
 	if (!diag_mux)
 		return -EIO;
 
-	peripheral = GET_BUF_PERIPHERAL(ctx);
-	if (peripheral > NUM_PERIPHERALS)
-		return -EINVAL;
-
-	if (MD_PERIPHERAL_MASK(peripheral) & diag_mux->mux_mask)
-		logger = diag_mux->md_ptr;
-	else
-		logger = diag_mux->usb_ptr;
+	upd = GET_PD_CTXT(ctx);
+	if (upd) {
+		switch (upd) {
+		case DIAG_ID_MPSS:
+			upd = PERIPHERAL_MODEM;
+			break;
+		case UPD_WLAN:
+			break;
+		default:
+			pr_err("diag: invalid pd ctxt= %d\n", upd);
+			return -EINVAL;
+		}
+		if (((MD_PERIPHERAL_MASK(upd)) &
+			(diag_mux->mux_mask)) &&
+			driver->md_session_map[upd])
+			logger = diag_mux->md_ptr;
+		else
+			logger = diag_mux->usb_ptr;
+	} else {
+
+		peripheral = GET_BUF_PERIPHERAL(ctx);
+		if (peripheral > NUM_PERIPHERALS)
+			return -EINVAL;
+
+		if (MD_PERIPHERAL_MASK(peripheral) &
+			diag_mux->mux_mask)
+			logger = diag_mux->md_ptr;
+		else
+			logger = diag_mux->usb_ptr;
+	}
 
 	if (logger && logger->log_ops && logger->log_ops->write)
 		return logger->log_ops->write(proc, buf, len, ctx);
@@ -159,9 +181,17 @@ int diag_mux_close_peripheral(int proc, uint8_t peripheral)
 	struct diag_logger_t *logger = NULL;
 	if (proc < 0 || proc >= NUM_MUX_PROC)
 		return -EINVAL;
+
 	/* Peripheral should account for Apps data as well */
-	if (peripheral > NUM_PERIPHERALS)
-		return -EINVAL;
+	if (peripheral > NUM_PERIPHERALS) {
+		if (driver->num_pd_session) {
+			if (peripheral > NUM_MD_SESSIONS)
+				return -EINVAL;
+		} else {
+			return -EINVAL;
+		}
+	}
+
 	if (!diag_mux)
 		return -EIO;
 
@@ -182,7 +212,8 @@ int diag_mux_switch_logging(int *req_mode, int *peripheral_mask)
 	if (!req_mode)
 		return -EINVAL;
 
-	if (*peripheral_mask <= 0 || *peripheral_mask > DIAG_CON_ALL) {
+	if (*peripheral_mask <= 0 ||
+		(*peripheral_mask > (DIAG_CON_ALL | DIAG_CON_UPD_ALL))) {
 		pr_err("diag: mask %d in %s\n", *peripheral_mask, __func__);
 		return -EINVAL;
 	}