Fix in-air restarts, protect against an external MAVLink sender exploiting the restart mechanism

3 files changed, 31 insertions, 4 deletions

diff --git a/src/drivers/px4io/px4io.cpp b/src/drivers/px4io/px4io.cpp
index 8458c2fdb..aec6dd3b7 100644
--- a/src/drivers/px4io/px4io.cpp
+++ b/src/drivers/px4io/px4io.cpp
@@ -683,6 +683,25 @@ PX4IO::init()
 
 		/* send command to arm system via command API */
 		vehicle_command_s cmd;
+		/* send this to itself */
+		param_t sys_id_param = param_find("MAV_SYS_ID");
+		param_t comp_id_param = param_find("MAV_COMP_ID");
+
+		int32_t sys_id;
+		int32_t comp_id;
+
+		if (param_get(sys_id_param, &sys_id)) {
+			errx(1, "PRM SYSID");
+		}
+
+		if (param_get(comp_id_param, &comp_id)) {
+			errx(1, "PRM CMPID");
+		}
+
+		cmd.target_system = sys_id;
+		cmd.target_component = comp_id;
+		cmd.source_system = sys_id;
+		cmd.source_component = comp_id;
 		/* request arming */
 		cmd.param1 = 1.0f;
 		cmd.param2 = 0;
@@ -692,10 +711,7 @@ PX4IO::init()
 		cmd.param6 = 0;
 		cmd.param7 = 0;
 		cmd.command = VEHICLE_CMD_COMPONENT_ARM_DISARM;
-		// cmd.target_system = status.system_id;
-		// cmd.target_component = status.component_id;
-		// cmd.source_system = status.system_id;
-		// cmd.source_component = status.component_id;
+
 		/* ask to confirm command */
 		cmd.confirmation =  1;
 
diff --git a/src/modules/commander/commander.cpp b/src/modules/commander/commander.cpp
index 53ed34f46..141b371b3 100644
--- a/src/modules/commander/commander.cpp
+++ b/src/modules/commander/commander.cpp
@@ -484,6 +484,11 @@ bool handle_command(struct vehicle_status_s *status, const struct safety_s *safe
             if (cmd->param1 != 0.0f && (fabsf(cmd->param1 - 1.0f) > 2.0f * FLT_EPSILON)) {
 				mavlink_log_info(mavlink_fd, "Unsupported ARM_DISARM parameter: %.6f", cmd->param1);
             } else {
+
+		// Flick to inair restore first if this comes from an onboard system
+		if (cmd->source_system == status->system_id && cmd->source_component == status->component_id) {
+			status->arming_state = ARMING_STATE_IN_AIR_RESTORE;
+		}
                 transition_result_t arming_res = arm_disarm(cmd->param1 != 0.0f, mavlink_fd, "arm/disarm component command");
                 if (arming_res == TRANSITION_DENIED) {
                     mavlink_log_critical(mavlink_fd, "#audio: REJECTING component arm cmd");
diff --git a/src/modules/mavlink/mavlink_receiver.cpp b/src/modules/mavlink/mavlink_receiver.cpp
index 7c93c1c00..64fc41838 100644
--- a/src/modules/mavlink/mavlink_receiver.cpp
+++ b/src/modules/mavlink/mavlink_receiver.cpp
@@ -217,6 +217,12 @@ MavlinkReceiver::handle_message_command_long(mavlink_message_t *msg)
 			_mavlink->_task_should_exit = true;
 
 		} else {
+
+			if (msg->sysid == mavlink_system.sysid && msg->compid == mavlink_system.compid) {
+				warnx("ignoring CMD spoofed with same SYS/COMP ID");
+				return;
+			}
+
 			struct vehicle_command_s vcmd;
 			memset(&vcmd, 0, sizeof(vcmd));