tcp_transport: fix possible buffer overflow in ws transport connect

closes IDF-692

diff --git a/components/tcp_transport/transport_ws.c b/components/tcp_transport/transport_ws.c
index c8f3da31c66b2241486ff5db82dace301f304aaa..08c669d08e3d401d802ff4c66f3f83f4307f93e7 100644 (file)
--- a/components/tcp_transport/transport_ws.c
+++ b/components/tcp_transport/transport_ws.c
@@ -103,14 +103,26 @@ static int ws_connect(esp_transport_handle_t t, const char *host, int port, int
                          ws->path,
                          host, port,
                          client_key);
-    if (ws->sub_protocol) {
-        len += snprintf(ws->buffer + len, DEFAULT_WS_BUFFER - len, "Sec-WebSocket-Protocol: %s\r\n", ws->sub_protocol);
-    }
-    len += snprintf(ws->buffer + len, DEFAULT_WS_BUFFER - len, "\r\n");
     if (len <= 0 || len >= DEFAULT_WS_BUFFER) {
         ESP_LOGE(TAG, "Error in request generation, %d", len);
         return -1;
     }
+    if (ws->sub_protocol) {
+        int r = snprintf(ws->buffer + len, DEFAULT_WS_BUFFER - len, "Sec-WebSocket-Protocol: %s\r\n", ws->sub_protocol);
+        len += r;
+        if (r <= 0 || len >= DEFAULT_WS_BUFFER) {
+            ESP_LOGE(TAG, "Error in request generation"
+                          "(snprintf of subprotocol returned %d, desired request len: %d, buffer size: %d", r, len, DEFAULT_WS_BUFFER);
+            return -1;
+        }
+    }
+    int r = snprintf(ws->buffer + len, DEFAULT_WS_BUFFER - len, "\r\n");
+    len += r;
+    if (r <= 0 || len >= DEFAULT_WS_BUFFER) {
+        ESP_LOGE(TAG, "Error in request generation"
+                       "(snprintf of header terminal returned %d, desired request len: %d, buffer size: %d", r, len, DEFAULT_WS_BUFFER);
+        return -1;
+    }
     ESP_LOGD(TAG, "Write upgrate request\r\n%s", ws->buffer);
     if (esp_transport_write(ws->parent, ws->buffer, len, timeout_ms) <= 0) {
         ESP_LOGE(TAG, "Error write Upgrade header %s", ws->buffer);