diff --git a/target/linux/ipq807x/base-files/etc/init.d/bootcount b/target/linux/ipq807x/base-files/etc/init.d/bootcount
new file mode 100755
index 000000000..280c109cf
--- /dev/null
+++ b/target/linux/ipq807x/base-files/etc/init.d/bootcount
@@ -0,0 +1,14 @@
+#!/bin/sh /etc/rc.common
+
+START=99
+
+boot() {
+	case $(board_name) in
+	xiaomi,ax3600)
+		# OTA handling should not be used. Reset it just in case.
+		fw_setenv flag_ota_reboot 0
+		# Not strictly needed but useful to handle partition crash condition
+		fw_setenv flag_boot_success 1
+	;;
+	esac
+}
\ No newline at end of file
diff --git a/target/linux/ipq807x/base-files/lib/upgrade/platform.sh b/target/linux/ipq807x/base-files/lib/upgrade/platform.sh
new file mode 100755
index 000000000..be37dcc3d
--- /dev/null
+++ b/target/linux/ipq807x/base-files/lib/upgrade/platform.sh
@@ -0,0 +1,43 @@
+PART_NAME=firmware
+REQUIRE_IMAGE_METADATA=1
+
+RAMFS_COPY_BIN='fw_printenv fw_setenv'
+RAMFS_COPY_DATA='/etc/fw_env.config /var/lock/fw_printenv.lock'
+
+platform_check_image() {
+	return 0;
+}
+
+platform_do_upgrade() {
+	case "$(board_name)" in
+	xiaomi,ax3600)
+		part_num="$(fw_printenv -n flag_boot_rootfs)"
+		if [ "$part_num" -eq "0" ]; then
+			CI_UBIPART="rootfs_1"
+			target_num=1
+			# Reset fail flag for the current partition
+			# With both partition set to fail, the partition 2 (bit 1)
+			# is loaded
+			fw_setenv flag_try_sys2_failed 0
+		else
+			CI_UBIPART="rootfs"
+			target_num=0
+			# Reset fail flag for the current partition
+			# or uboot will skip the loading of this partition
+			fw_setenv flag_try_sys1_failed 0
+		fi
+
+		# Tell uboot to switch partition
+		fw_setenv flag_boot_rootfs $target_num
+		fw_setenv flag_last_success $target_num
+
+		# Reset success flag
+		fw_setenv flag_boot_success 0
+
+		nand_do_upgrade "$1"
+		;;
+	*)
+		default_do_upgrade "$1"
+		;;
+	esac
+}
\ No newline at end of file