From d62fd4899a8e61fd2a27a9b1f3b5d5705b0a8e9c Mon Sep 17 00:00:00 2001
From: Angus Gratton <gus@projectgus.com>
Date: Mon, 23 May 2016 12:23:04 +1000
Subject: [PATCH] ota_basic example cleanup

---
 examples/ota_basic/ota_basic.c | 115 ++++++++++++++++++++-------------
 1 file changed, 69 insertions(+), 46 deletions(-)

diff --git a/examples/ota_basic/ota_basic.c b/examples/ota_basic/ota_basic.c
index 9ab3f20..61b50b1 100644
--- a/examples/ota_basic/ota_basic.c
+++ b/examples/ota_basic/ota_basic.c
@@ -1,6 +1,8 @@
 /* A very simple OTA example
  *
- * Binds a TCP socket, reads an image from it over TFTP and then flashes live.
+ * Tries to run both a TFTP client and a TFTP server simultaneously, either will accept a TTP firmware and update it.
+ *
+ * Not a realistic OTA setup, this needs adapting (choose either client or server) before you'd want to use it.
  *
  * For more information about esp-open-rtos OTA see https://github.com/SuperHouse/esp-open-rtos/wiki/OTA-Update-Configuration
  *
@@ -18,6 +20,7 @@
 #include "ota-tftp.h"
 #include "rboot-api.h"
 
+/* TFTP client will request this image filenames from this server */
 #define TFTP_IMAGE_SERVER "192.168.1.23"
 #define TFTP_IMAGE_FILENAME1 "firmware1.bin"
 #define TFTP_IMAGE_FILENAME2 "firmware2.bin"
@@ -25,6 +28,68 @@
 /* Output of the command 'sha256sum firmware1.bin' */
 static const char *FIRMWARE1_SHA256 = "88199daff8b9e76975f685ec7f95bc1df3c61bd942a33a54a40707d2a41e5488";
 
+/* Example function to TFTP download a firmware file and verify its SHA256 before
+   booting into it.
+*/
+static void tftpclient_download_and_verify_file1(int slot, rboot_config *conf)
+{
+    printf("Downloading %s to slot %d...\n", TFTP_IMAGE_FILENAME1, slot);
+    int res = ota_tftp_download(TFTP_IMAGE_SERVER, TFTP_PORT+1, TFTP_IMAGE_FILENAME1, 1000, slot);
+    printf("ota_tftp_download %s result %d\n", TFTP_IMAGE_FILENAME1, res);
+
+    if (res != 0) {
+        return;
+    }
+
+    printf("Looks valid, calculating SHA256...\n");
+    uint32_t length;
+    bool valid = rboot_verify_image(conf->roms[slot], &length, NULL);
+    static mbedtls_sha256_context ctx;
+    mbedtls_sha256_init(&ctx);
+    mbedtls_sha256_starts(&ctx, 0);
+    valid = valid && rboot_digest_image(conf->roms[slot], length, (rboot_digest_update_fn)mbedtls_sha256_update, &ctx);
+    static uint8_t hash_result[32];
+    mbedtls_sha256_finish(&ctx, hash_result);
+    mbedtls_sha256_free(&ctx);
+
+    if(!valid)
+    {
+        printf("Not valid after all :(\n");
+        return;
+    }
+
+    printf("Image SHA256 = ");
+    valid = true;
+    for(int i = 0; i < sizeof(hash_result); i++) {
+        char hexbuf[3];
+        snprintf(hexbuf, 3, "%02x", hash_result[i]);
+        printf(hexbuf);
+        if(strncmp(hexbuf, FIRMWARE1_SHA256+i*2, 2))
+            valid = false;
+    }
+    printf("\n");
+
+    if(!valid) {
+        printf("Downloaded image SHA256 didn't match expected '%s'\n", FIRMWARE1_SHA256);
+        return;
+    }
+
+    printf("SHA256 Matches. Rebooting into slot %d...\n", slot);
+    rboot_set_current_rom(slot);
+    sdk_system_restart();
+}
+
+/* Much simpler function that just downloads a file via TFTP into an rboot slot.
+
+   (
+ */
+static void tftpclient_download_file2(int slot)
+{
+    printf("Downloading %s to slot %d...\n", TFTP_IMAGE_FILENAME2, slot);
+    int res = ota_tftp_download(TFTP_IMAGE_SERVER, TFTP_PORT+1, TFTP_IMAGE_FILENAME2, 1000, slot);
+    printf("ota_tftp_download %s result %d\n", TFTP_IMAGE_FILENAME2, res);
+}
+
 void tftp_client_task(void *pvParameters)
 {
     printf("TFTP client task starting...\n");
@@ -43,53 +108,10 @@ void tftp_client_task(void *pvParameters)
        Note: example will reboot into FILENAME1 if it is successfully downloaded, but FILENAME2 is ignored.
     */
     while(1) {
-        printf("Downloading %s to slot %d...\n", TFTP_IMAGE_FILENAME1, slot);
-        int res = ota_tftp_download(TFTP_IMAGE_SERVER, TFTP_PORT, TFTP_IMAGE_FILENAME1, 1000, slot);
-        printf("ota_tftp_download %s result %d\n", TFTP_IMAGE_FILENAME1, res);
-        if(res == 0) {
-            printf("Looks valid, calculating SHA256...\n");
-            uint32_t length;
-            bool valid = rboot_verify_image(conf.roms[slot], &length, NULL);
-            static mbedtls_sha256_context ctx;
-            mbedtls_sha256_init(&ctx);
-            mbedtls_sha256_starts(&ctx, 0);
-            valid = valid && rboot_digest_image(conf.roms[slot], length, (rboot_digest_update_fn)mbedtls_sha256_update, &ctx);
-            static uint8_t hash_result[32];
-            mbedtls_sha256_finish(&ctx, hash_result);
-            mbedtls_sha256_free(&ctx);
-
-            if(!valid)
-            {
-                printf("Not valid after all :(\n");
-            }
-            else
-            {
-                printf("Image SHA256 = ");
-                bool valid = true;
-                for(int i = 0; i < sizeof(hash_result); i++) {
-                    char hexbuf[3];
-                    snprintf(hexbuf, 3, "%02x", hash_result[i]);
-                    printf(hexbuf);
-                    if(strncmp(hexbuf, FIRMWARE1_SHA256+i*2, 2))
-                        valid = false;
-
-                }
-                printf("\n");
-                if (valid) {
-                    printf("SHA256 Matches. Rebooting into slot %d...\n", slot);
-                    rboot_set_current_rom(slot);
-                    sdk_system_restart();
-                }
-                else {
-                    printf("Downloaded image SHA256 didn't match expected '%s'\n", FIRMWARE1_SHA256);
-                }
-            }
-        }
+        tftpclient_download_and_verify_file1(slot, &conf);
         vTaskDelay(5000 / portTICK_RATE_MS);
 
-        printf("Downloading %s to slot %d...\n", TFTP_IMAGE_FILENAME2, slot);
-        res = ota_tftp_download(TFTP_IMAGE_SERVER, TFTP_PORT, TFTP_IMAGE_FILENAME2, 1000, slot);
-        printf("ota_tftp_download %s result %d\n", TFTP_IMAGE_FILENAME2, res);
+        tftpclient_download_file2(slot);
         vTaskDelay(5000 / portTICK_RATE_MS);
     }
 }
@@ -114,6 +136,7 @@ void user_init(void)
     sdk_wifi_set_opmode(STATION_MODE);
     sdk_wifi_station_set_config(&config);
 
+    printf("Starting TFTP server...");
     ota_tftp_init_server(TFTP_PORT);
     xTaskCreate(&tftp_client_task, (signed char *)"tftp_client", 2048, NULL, 2, NULL);
 }