From de30757d2181e78eaa11957e31ef420e833abc90 Mon Sep 17 00:00:00 2001
From: Marcin Niestroj <m.niestroj@emb.dev>
Date: Tue, 20 Feb 2024 09:18:14 +0100
Subject: [PATCH] drivers: wifi: esp_at: escape SSID and PSK

According to [1], SSID and PSK need to be escaped:

  Escape character syntax is needed if SSID or password contains any
  special characters, such as , or " or \

Implement character escaping to fix connection attempt to WiFi APs
containing special characters as part of SSID, like "My_2,4GHz_AP".
Increase "connect command" buffer length to handle worst-case scenario of
all the SSID and PSK characters being special characters.

[1] https://docs.espressif.com/projects/esp-at/en/release-v2.4.0.0/esp32/AT_Command_Set/Wi-Fi_AT_Commands.html#id6

Signed-off-by: Marcin Niestroj <m.niestroj@emb.dev>
---
 drivers/wifi/esp_at/esp.c | 101 ++++++++++++++++++++++++++++++++------
 drivers/wifi/esp_at/esp.h |   2 +-
 2 files changed, 88 insertions(+), 15 deletions(-)

diff --git a/drivers/wifi/esp_at/esp.c b/drivers/wifi/esp_at/esp.c
index bfc8b448356..6b63d3b3b15 100644
--- a/drivers/wifi/esp_at/esp.c
+++ b/drivers/wifi/esp_at/esp.c
@@ -996,11 +996,71 @@ out:
 	esp_flags_clear(dev, EDF_STA_CONNECTING);
 }
 
+static int esp_conn_cmd_append(struct esp_data *data, size_t *off,
+			       const char *chunk, size_t chunk_len)
+{
+	char *str_end = &data->conn_cmd[sizeof(data->conn_cmd)];
+	char *str = &data->conn_cmd[*off];
+	const char *chunk_end = chunk + chunk_len;
+
+	for (; chunk < chunk_end; chunk++) {
+		if (str_end - str < 1) {
+			return -ENOSPC;
+		}
+
+		*str = *chunk;
+		str++;
+	}
+
+	*off = str - data->conn_cmd;
+
+	return 0;
+}
+
+#define esp_conn_cmd_append_literal(data, off, chunk)			\
+	esp_conn_cmd_append(data, off, chunk, sizeof(chunk) - 1)
+
+static int esp_conn_cmd_escape_and_append(struct esp_data *data, size_t *off,
+					  const char *chunk, size_t chunk_len)
+{
+	char *str_end = &data->conn_cmd[sizeof(data->conn_cmd)];
+	char *str = &data->conn_cmd[*off];
+	const char *chunk_end = chunk + chunk_len;
+
+	for (; chunk < chunk_end; chunk++) {
+		switch (*chunk) {
+		case ',':
+		case '\\':
+		case '"':
+			if (str_end - str < 2) {
+				return -ENOSPC;
+			}
+
+			*str = '\\';
+			str++;
+
+			break;
+		}
+
+		if (str_end - str < 1) {
+			return -ENOSPC;
+		}
+
+		*str = *chunk;
+		str++;
+	}
+
+	*off = str - data->conn_cmd;
+
+	return 0;
+}
+
 static int esp_mgmt_connect(const struct device *dev,
 			    struct wifi_connect_req_params *params)
 {
 	struct esp_data *data = dev->data;
-	int len;
+	size_t off = 0;
+	int err;
 
 	if (!net_if_is_carrier_ok(data->net_iface) ||
 	    !net_if_is_admin_up(data->net_iface)) {
@@ -1013,21 +1073,34 @@ static int esp_mgmt_connect(const struct device *dev,
 
 	esp_flags_set(data, EDF_STA_CONNECTING);
 
-	len = snprintk(data->conn_cmd, sizeof(data->conn_cmd),
-		       "AT+"_CWJAP"=\"");
-	memcpy(&data->conn_cmd[len], params->ssid, params->ssid_length);
-	len += params->ssid_length;
-
-	len += snprintk(&data->conn_cmd[len],
-				sizeof(data->conn_cmd) - len, "\",\"");
-
-	if (params->security == WIFI_SECURITY_TYPE_PSK) {
-		memcpy(&data->conn_cmd[len], params->psk, params->psk_length);
-		len += params->psk_length;
+	err = esp_conn_cmd_append_literal(data, &off, "AT+"_CWJAP"=\"");
+	if (err) {
+		return err;
 	}
 
-	len += snprintk(&data->conn_cmd[len], sizeof(data->conn_cmd) - len,
-			"\"");
+	err = esp_conn_cmd_escape_and_append(data, &off,
+					     params->ssid, params->ssid_length);
+	if (err) {
+		return err;
+	}
+
+	err = esp_conn_cmd_append_literal(data, &off, "\",\"");
+	if (err) {
+		return err;
+	}
+
+	if (params->security == WIFI_SECURITY_TYPE_PSK) {
+		err = esp_conn_cmd_escape_and_append(data, &off,
+						     params->psk, params->psk_length);
+		if (err) {
+			return err;
+		}
+	}
+
+	err = esp_conn_cmd_append_literal(data, &off, "\"");
+	if (err) {
+		return err;
+	}
 
 	k_work_submit_to_queue(&data->workq, &data->connect_work);
 
diff --git a/drivers/wifi/esp_at/esp.h b/drivers/wifi/esp_at/esp.h
index bda67cb7ebf..c775bf1f14a 100644
--- a/drivers/wifi/esp_at/esp.h
+++ b/drivers/wifi/esp_at/esp.h
@@ -77,7 +77,7 @@ extern "C" {
 	STRINGIFY(_UART_BAUD)",8,1,0,"_FLOW_CONTROL
 
 #define CONN_CMD_MAX_LEN (sizeof("AT+"_CWJAP"=\"\",\"\"") + \
-			  WIFI_SSID_MAX_LEN + WIFI_PSK_MAX_LEN)
+			  WIFI_SSID_MAX_LEN * 2 + WIFI_PSK_MAX_LEN * 2)
 
 #if defined(CONFIG_WIFI_ESP_AT_DNS_USE)
 #define ESP_MAX_DNS	MIN(3, CONFIG_DNS_RESOLVER_MAX_SERVERS)