From 8d315111e7dd8d21ad6721160661d3590516b243 Mon Sep 17 00:00:00 2001
From: nekomona <nekomona@163.com>
Date: Wed, 12 Jul 2023 19:28:21 +0800
Subject: [PATCH] Add IP check for remote command

---
 src/debug.h                   |  1 +
 src/network/connection.h      |  2 ++
 src/network/remotecmd.cpp     | 28 ++++++++++++++++++++++++----
 src/network/wifihandler.cpp   |  1 +
 src/serial/serialcommands.cpp | 19 +++++++++++++++++++
 5 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/src/debug.h b/src/debug.h
index 82e5bbd97..91ea7b4a2 100644
--- a/src/debug.h
+++ b/src/debug.h
@@ -48,6 +48,7 @@
 #define serialDebug false // Set to true to get Serial output for debugging
 #define serialBaudRate 115200
 #define USE_REMOTE_COMMAND true
+#define ALLOW_REMOTE_WIFI_PROV false
 #define LED_INTERVAL_STANDBY 10000
 #define PRINT_STATE_EVERY_MS 60000
 
diff --git a/src/network/connection.h b/src/network/connection.h
index 3ca015626..2f12cfb3e 100644
--- a/src/network/connection.h
+++ b/src/network/connection.h
@@ -103,6 +103,8 @@ class Connection {
 	);
 #endif
 
+friend RemoteCmd;
+
 private:
 	void updateSensorState(Sensor* const sensor1, Sensor* const sensor2);
 
diff --git a/src/network/remotecmd.cpp b/src/network/remotecmd.cpp
index 442c4cb75..eb38c9108 100644
--- a/src/network/remotecmd.cpp
+++ b/src/network/remotecmd.cpp
@@ -1,5 +1,7 @@
 #include "remotecmd.h"
 
+#include "GlobalVars.h"
+
 namespace SlimeVR {
 namespace Network {
 
@@ -17,10 +19,28 @@ void RemoteCmd::update() {
 			r_Logger.info("Remote command multi-connection dropped");
 		} else {
 			rcmdClient = rcmdServer.accept();
-			r_Logger.info(
-				"Remote command from %s connected",
-				rcmdClient.remoteIP().toString().c_str()
-			);
+			if (networkConnection.isConnected()) {
+				// Only accept if rcmdClient have the same remote IP as udpmanager
+				if (rcmdClient.remoteIP() = networkConnection.m_ServerHost) {
+					rcmdClient.stop();
+				}
+			}
+#if !ALLOW_REMOTE_WIFI_PROV
+			else {
+				rcmdClient.stop();
+			}
+#endif
+			if (rcmdClient.connected()) {
+				r_Logger.info(
+					"Remote command from %s connected",
+					rcmdClient.remoteIP().toString().c_str()
+				);
+			} else {
+				r_Logger.info(
+					"Remote command from %s dropped",
+					rcmdClient.remoteIP().toString().c_str()
+				);
+			}
 		}
 	}
 }
diff --git a/src/network/wifihandler.cpp b/src/network/wifihandler.cpp
index e701a6afc..62422d1c8 100644
--- a/src/network/wifihandler.cpp
+++ b/src/network/wifihandler.cpp
@@ -121,6 +121,7 @@ void WiFiNetwork::setUp() {
 void onConnected() {
     WiFiNetwork::stopProvisioning();
     statusManager.setStatus(SlimeVR::Status::WIFI_CONNECTING, false);
+	networkRemoteCmd.reset();
     isWifiConnected = true;
     hadWifi = true;
     wifiHandlerLogger.info("Connected successfully to SSID '%s', ip address %s", WiFi.SSID().c_str(), WiFi.localIP().toString().c_str());
diff --git a/src/serial/serialcommands.cpp b/src/serial/serialcommands.cpp
index af8dc9435..649e82d0f 100644
--- a/src/serial/serialcommands.cpp
+++ b/src/serial/serialcommands.cpp
@@ -38,6 +38,7 @@ namespace SerialCommands {
     CmdCallback<6> cmdCallbacks;
     CmdParser cmdParser;
     CmdBuffer<64> cmdBuffer;
+	bool cmdFromRemote = false;
 
     void cmdSet(CmdParser * parser) {
         if(parser->getParamCount() != 1 && parser->equalCmdParam(1, "WIFI")  ) {
@@ -84,6 +85,10 @@ namespace SerialCommands {
     }
 
     void cmdGet(CmdParser * parser) {
+#if USE_REMOTE_COMMAND && ALLOW_REMOTE_WIFI_PROV
+		if (cmdFromRemote && !networkConnection.isConnected()) return;
+#endif
+
         if (parser->getParamCount() < 2) {
             return;
         }
@@ -141,11 +146,19 @@ namespace SerialCommands {
     }
 
     void cmdReboot(CmdParser * parser) {
+#if USE_REMOTE_COMMAND && ALLOW_REMOTE_WIFI_PROV
+		if (cmdFromRemote && !networkConnection.isConnected()) return;
+#endif
+
         logger.info("REBOOT");
         ESP.restart();
     }
 
     void cmdFactoryReset(CmdParser * parser) {
+#if USE_REMOTE_COMMAND && ALLOW_REMOTE_WIFI_PROV
+		if (cmdFromRemote && !networkConnection.isConnected()) return;
+#endif
+
         logger.info("FACTORY RESET");
 
         configuration.reset();
@@ -171,6 +184,10 @@ namespace SerialCommands {
     }
 
     void cmdTemperatureCalibration(CmdParser* parser) {
+#if USE_REMOTE_COMMAND && ALLOW_REMOTE_WIFI_PROV
+		if (cmdFromRemote && !networkConnection.isConnected()) return;
+#endif
+
         if (parser->getParamCount() > 1) {
             if (parser->equalCmdParam(1, "PRINT")) {
                 sensorManager.getFirst()->printTemperatureCalibrationState();
@@ -212,9 +229,11 @@ namespace SerialCommands {
         #if USE_REMOTE_COMMAND
         if (networkRemoteCmd.isConnected()) {
 			Stream & networkStream = networkRemoteCmd.getStream();
+			cmdFromRemote = true;
 			while (networkStream.available()) {
 				cmdCallbacks.updateCmdProcessing(&cmdParser, &cmdBuffer, &networkStream);
 			}
+			cmdFromRemote = false;
 		}
         #endif
     }