From 849abae41de5062d454bd8c68326ade67f7e183b Mon Sep 17 00:00:00 2001
From: ok-home <ok-home@ya.ru>
Date: Mon, 25 Sep 2023 10:42:33 +0700
Subject: [PATCH] tmp rollback

---
 private_include/ota_ws_private.h |  6 ++-
 source/ota_ws.html               | 77 ++++++++++++++++++++++----------
 source/ota_ws_esp.c              | 32 +++++++++++++
 source/ota_ws_http.c             | 33 +++++++++++---
 4 files changed, 117 insertions(+), 31 deletions(-)

diff --git a/private_include/ota_ws_private.h b/private_include/ota_ws_private.h
index 110b90b..a391bc6 100644
--- a/private_include/ota_ws_private.h
+++ b/private_include/ota_ws_private.h
@@ -11,10 +11,14 @@
 #define OTA_END "otaEnd"
 #define OTA_ERROR "otaError"
 #define OTA_CANCEL "otaCancel"
+#define OTA_CHECK_ROLLBACK "otaCheckRollback"
+#define OTA_PROCESS_ROLLBACK "otaProcessRollback"
 
 #define OTA_CHUNK_SIZE 4096*2
 
 esp_err_t start_ota_ws(void);
 esp_err_t write_ota_ws(int data_read, uint8_t *ota_write_data);
 esp_err_t end_ota_ws(void);
-esp_err_t abort_ota_ws(void);
\ No newline at end of file
+esp_err_t abort_ota_ws(void);
+bool check_ota_ws_rollback_enable(void);
+esp_err_t rollback_ota_ws(bool rollback);
\ No newline at end of file
diff --git a/source/ota_ws.html b/source/ota_ws.html
index 44858ef..29fc417 100644
--- a/source/ota_ws.html
+++ b/source/ota_ws.html
@@ -69,28 +69,38 @@
     <div class="column">
         <button class="btn" id="goHome">На главную</button>
     </div>
-    <div class="cl1" style="display:none">
-        <label class="cl01" for="otaFile">Ota File Name</label>
-        <input class="cl02" type="file" id="otaFile" placeholder="select file" onchange="readOtaFile(this)">
-    </div>
-    <div class="column" style="display:block" id="otaFileSelectVisible">
-        <button class="btn" id="otaFileSelect" onclick="document.getElementById('otaFile').click()">File Select</button>
-    </div>
-
-    <div class="column" style="display:none" id="otaStartVisible">
-        <button class="btn" id="otaStartCancel">Ota Start</button>
-    </div>
-    <div class="column" style="display:none" id="otaReStartVisible">
-        <button class="btn" id="otaReStart">Ota ReStart</button>
-    </div>
-    <div id="otaProgressVisible" style="display:none">
-        <div class="cl1">
-            <progress class="cl02" id="otaPogress" max=100 value=0>
+    <div id="rollback" style="display:none">
+        <div class="column" >
+            <button class="btn" id="otaVerifyApp">Ota APP veryfyed</button>
+        </div>
+        <div class="column">
+            <button class="btn" id="otaRollback">Ota APP rollback and restart</button>
         </div>
     </div>
+    <div id="update" style="display:block">
+        <div class="cl1" style="display:none">
+            <label class="cl01" for="otaFile">Ota File Name</label>
+            <input class="cl02" type="file" id="otaFile" placeholder="select file" onchange="readOtaFile(this)">
+        </div>
+        <div class="column" style="display:block" id="otaFileSelectVisible">
+            <button class="btn" id="otaFileSelect" onclick="document.getElementById('otaFile').click()">File
+                Select</button>
+        </div>
 
+        <div class="column" style="display:none" id="otaStartVisible">
+            <button class="btn" id="otaStartCancel">Ota Start</button>
+        </div>
+        <div class="column" style="display:none" id="otaReStartVisible">
+            <button class="btn" id="otaReStart">Ota ReStart</button>
+        </div>
+        <div id="otaProgressVisible" style="display:none">
+            <div class="cl1">
+                <progress class="cl02" id="otaPogress" max=100 value=0>
+            </div>
+        </div>
+    </div>
     <script>
-        
+
         let otaData;
         let otaSetChunkSize = 0;
         let otaStartsegment = 0;
@@ -105,8 +115,8 @@
 
             reader.onload = function () {
                 otaData = new Uint8Array(reader.result);
-//                console.log(reader.result);
-//                console.log(otaData.length);
+                //                console.log(reader.result);
+                //                console.log(otaData.length);
                 document.getElementById("otaStartVisible").style.display = "block";
                 document.getElementById("otaProgressVisible").style.display = "none";
                 document.getElementById("otaReStartVisible").style.display = "none";
@@ -123,7 +133,7 @@
 
                 socket.send(JSON.stringify({ name: "otaSize", value: otaData.length }));
                 //console.log(JSON.stringify({ name: "otaSize", value: otaData.length }));
-                
+
                 otaStarted = 1;
                 this.innerHTML = "Click to Cancel";
                 document.getElementById("otaFileSelect").disabled = true;
@@ -156,8 +166,8 @@
             receiveWsData(JSON.stringify({ name: "otaSetChunkSize", value: 1024 }));
             while (otaStartsegment + otaSetChunkSize <= otaData.length && otaStarted == 1) {
                 await sleep(1000);
-                if(otaStarted == 1){
-                receiveWsData(JSON.stringify({ name: "otaGetChunk", value: otaStartsegment }));
+                if (otaStarted == 1) {
+                    receiveWsData(JSON.stringify({ name: "otaGetChunk", value: otaStartsegment }));
                 }
                 otaStartsegment += otaSetChunkSize;
             }
@@ -205,6 +215,11 @@
                         document.getElementById("otaReStartVisible").style.display = "block";
                         document.getElementById("otaReStart").innerHTML = "Transfer Cancel " + obj.value;
                         document.getElementById("otaReStart").disabled = true;
+                        break;
+                    case "otaCheckRollback":
+                        document.getElementById("rollback").style.display = "block";
+                        document.getElementById("update").style.display = "none";
+                        break;
                 }
             }
             catch
@@ -214,7 +229,21 @@
         };
     </script>
 
-    <script>  // Прием, обработка и отправка данных в WS
+    <script>  // rollback
+        document.getElementById("otaVerifyApp").addEventListener("click", function (e) {
+        socket.send(JSON.stringify({ name: "otaProcessRollback", value: "false" }));
+        document.getElementById("rollback").style.display = "none";
+        document.getElementById("update").style.display = "block";
+            console.log(JSON.stringify({ name: "otaProcessRollback", value: "false" }));
+        });
+
+        document.getElementById("otaRollback").addEventListener("click", function (e) {
+        socket.send(JSON.stringify({ name: "otaProcessRollback", value: "true" }));
+        document.getElementById("rollback").style.display = "none";
+        document.getElementById("update").style.display = "block";
+            console.log(JSON.stringify({ name: "otaProcessRollback", value: "true" }));
+        });
+
     </script>
 
     <script> // основной старт скрипта, открыть сокет
diff --git a/source/ota_ws_esp.c b/source/ota_ws_esp.c
index 421970d..1a196c6 100644
--- a/source/ota_ws_esp.c
+++ b/source/ota_ws_esp.c
@@ -103,4 +103,36 @@ esp_err_t end_ota_ws(void)
 esp_err_t abort_ota_ws(void)
 {
     return esp_ota_abort(update_handle);
+}
+// false - rollback disable
+// true - rollback enable
+bool check_ota_ws_rollback_enable(void)
+{
+#ifdef CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE
+    esp_ota_img_states_t ota_state_running_part;
+    const esp_partition_t *running = esp_ota_get_running_partition();
+    if (esp_ota_get_state_partition(running, &ota_state_running_part) == ESP_OK) {
+        if (ota_state_running_part == ESP_OTA_IMG_PENDING_VERIFY) {
+            ESP_LOGI(TAG, "Running app has ESP_OTA_IMG_PENDING_VERIFY state");
+            return true;
+        }
+    }
+#endif
+    return false;    
+}
+// rollback == true - rollback
+// rollback == false - app valid? no rollback
+esp_err_t rollback_ota_ws(bool rollback)
+{
+#ifdef CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE
+    if(rollback == false)
+    {
+        return esp_ota_mark_app_valid_cancel_rollback(); // app valid
+    }
+    else
+    {
+        return esp_ota_mark_app_invalid_rollback_and_reboot(); // app rolback & reboot
+    }
+#endif
+    return ESP_FAIL;
 }
\ No newline at end of file
diff --git a/source/ota_ws_http.c b/source/ota_ws_http.c
index 6ee3808..3210ce8 100644
--- a/source/ota_ws_http.c
+++ b/source/ota_ws_http.c
@@ -67,18 +67,24 @@ static esp_err_t send_json_string(char *str, httpd_req_t *req)
 }
 static esp_err_t ota_ws_handler(httpd_req_t *req)
 {
-
-    if (req->method == HTTP_GET)
-    {
-        ESP_LOGI(TAG, "Handshake done, the new connection was opened");
-        return ESP_OK;
-    }
     char json_key[64] = {0};
     char json_value[64] = {0};
     char json_str[128] = {0};
 
     httpd_ws_frame_t ws_pkt;
     uint8_t *buf = NULL;
+
+
+    if (req->method == HTTP_GET)
+    {
+        ESP_LOGI(TAG, "Handshake done, the new connection was opened");
+        if(check_ota_ws_rollback_enable())
+        {
+        snprintf(json_str, sizeof(json_str), "{\"name\":\"%s\",\"value\":\"%s\" }", OTA_CHECK_ROLLBACK, "true");
+        send_json_string(json_str, req);
+        }
+        return ESP_OK;
+    }
     memset(&ws_pkt, 0, sizeof(httpd_ws_frame_t));
     /* Set max_len = 0 to get the frame len */
     esp_err_t ret = httpd_ws_recv_frame(req, &ws_pkt, 0);
@@ -147,10 +153,25 @@ static esp_err_t ota_ws_handler(httpd_req_t *req)
             ret = ESP_OK;
             goto _recv_ret;
         }
+        if (strncmp(json_key, OTA_PROCESS_ROLLBACK, sizeof(OTA_PROCESS_ROLLBACK)) == 0) // process rollback ?
+        {
+            ESP_LOGI(TAG,"rollback command %s %s",json_key,json_value)
+            if(srtncmp(json_value,"true",sizeof("true")) == 0)
+            {
+                ret = rollback_ota_ws(true); // rollback
+            }
+            else
+            {
+                ret = rollback_ota_ws(false);  // app veryfied
+            }
+            goto _recv_ret;
+        }
+
         if (strncmp(json_key, OTA_RESTART_ESP, sizeof(OTA_RESTART_ESP)) == 0) // cancel ota
         {
             esp_restart();
         }
+        
     }
     else if (ws_pkt.type == HTTPD_WS_TYPE_BINARY && ota_started)
     {