umcumgr/src/image.c

//
// Created by nock on 26/03/19.
//

#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <netinet/in.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <math.h>

#include "base64.h"
#include "crc.h"
#include "fixkey-msg.h"
#include "fixkey-payload.h"
#include "gpio.h"
#include "image.h"
#include "nmgr.h"
#include "timer.h"
#include "serial_util.h"
#include <tinycbor/cbor.h>

#ifdef DO_UPLOAD_SHA
#include "sha256.h"
#endif

#define IMAGE_CHUNK_SIZE 124 //216 // 124 for identical output to newtmgr
#define MACRO_CHUNK_SIZE 302 //432 // 302 ""

/*
 * Larger values above greatly increase speed, but seem to break older
 * versions of mcuboot (like the one shipped with inital EM firmware)
 */

#define CRC16_INITIAL_CRC       0       /* what to seed crc16 with */
#define CRC_CITT_POLYMINAL 0x1021

static int serial_fd = -1;

void image_help(void){
    fprintf(stderr, "umcumgr image [version|upload]\n");
    exit(1);
}

static char *hex(unsigned char in) {
    static char out[3] = {0};
    static char lookup[] = {
        '0', '1', '2', '3', '4', '5', '6', '7',
        '8', '9', 'a', 'b', 'c', 'd', 'e', 'f'
    };
    out[0] = lookup[in >> 4];
    out[1] = lookup[in & 0xf];
    out[2] = 0;
    return out;
}

static void read_response(uint8_t *dest, size_t *len) {
    static uint8_t buf[BOOT_SERIAL_IN_MAX*2] = {0};
    static uint8_t *p = buf;
    timer_start(3);
    while (true) {
        int bytes_read = read(serial_fd, p, BOOT_SERIAL_IN_MAX);
        if (bytes_read < 0) {
            if (errno == EAGAIN) {
                if (timer_expired()) {
                    fprintf(stderr, "Timeout waiting for bootloader\n");
                    exit(12);
                }
                usleep(250000);
                continue;
            }
            fprintf(stderr, "Failed to read from serial port: %s\n", strerror(errno));
            exit(2);
        }
        p += bytes_read;
        uint8_t *nl_p = memchr(buf, '\n', p-buf);
        if (nl_p) {
            *len = nl_p - buf;
            memcpy(dest, buf, *len);
            uint16_t new_buf_len = p - nl_p;
            if (new_buf_len) {
                p = buf;
                return;
            }
            memcpy(buf, nl_p+1, new_buf_len);
            p = buf+new_buf_len;
            break;
        }
    }
}

static void hex_print(char *in, size_t len) {
    char tmp[len*2+1];
    tmp[len] = 0;
    for (int i = 0; i < len; i++) {
        strcat(tmp, hex(in[i]));
    }
    printf("%s\n", tmp);
}

static void printlen(void *src, size_t len){
    uint8_t *p = src;
    for (size_t i = 0; i < len; i++) {
        //fprintf(stderr, "printlen: %u %s %c\n", i, hex(p[i]), p[i]);
        putchar(p[i]);
    }
}

static void assert_pkt_start(uint8_t *buf){
    assert(buf[0] == SHELL_NLIP_PKT_START1);
    assert(buf[1] == SHELL_NLIP_PKT_START2);
}

static void cbor_require_type(CborValue *it, CborType type) {
    CborType actual_type = cbor_value_get_type(it);
    if (actual_type != type) {
        fprintf(stderr, "Cbor deserialization error. Expected type %d, received %d\n", type, actual_type);
        exit(1);
    }
}

static char *get_version(void) {
    int rc = write(serial_fd, IMAGE_LIST_CMD, sizeof(IMAGE_LIST_CMD) - 1);
    if (rc < sizeof(IMAGE_LIST_CMD) - 1) {
        fprintf(stderr, "Failed to send command: %s\n", strerror(errno));
        return NULL;
    }
    uint8_t tmp[BOOT_SERIAL_IN_MAX];
    size_t in_len;

    read_response(tmp, &in_len);

    // Process the response
    CborParser parser;
    CborValue it;

    // Validate packet start marker
    if (in_len < 2) {
        fprintf(stderr, "packet too small, bailing");
        exit(5);
    }
    assert_pkt_start(tmp);

    // Decode the payload
    size_t out_len;
    uint8_t *decoded_buf = base64_decode(tmp+2, in_len-2, &out_len);

    if (!decoded_buf) {
        fprintf(stderr, "Failed to decode base64, bailing\n");
        exit(5);
    }

    // Unclear why this would be useful in a line oriented (cannonical mode) program
    uint16_t pkt_len = ntohs(*(uint16_t *)decoded_buf);
    uint16_t crc = crc16(decoded_buf+sizeof(pkt_len), pkt_len, CRC_CITT_POLYMINAL, CRC16_INITIAL_CRC, true);
    if (crc) {
        fprintf(stderr, "CRC failed on response packet");
        exit(6);
    }
    uint8_t *hdr_start = decoded_buf + sizeof(pkt_len);

    // The header seems a bit pointless... at least for serial
    struct nmgr_hdr *hdr = (struct nmgr_hdr *)(decoded_buf + sizeof(pkt_len));

    uint8_t *p = (uint8_t *)hdr + sizeof(struct nmgr_hdr);
    cbor_parser_init(p, out_len, 0, &parser, &it);
    cbor_require_type(&it, CborMapType);
    CborValue map;
    cbor_value_enter_container(&it, &map);
    cbor_require_type(&map, CborTextStringType);
    cbor_value_advance(&map);
    cbor_require_type(&map, CborArrayType);
    CborValue img_array;
    cbor_value_enter_container(&map, &img_array);

    char *version = NULL;
    while (!cbor_value_at_end(&img_array)) {
        assert(cbor_value_get_type(&img_array) == CborMapType);
        CborValue version_tag;
        CborError err = cbor_value_map_find_value(&img_array, "version", &version_tag);
        if (cbor_value_get_type(&version_tag) == CborInvalidType) {
            cbor_value_advance(&img_array);
            continue;
        }
        size_t len;
        cbor_value_dup_text_string(&version_tag, &version, &len, NULL);
        break;
    }
    if (!version) {
        printf("No applicaiton image found\n");
    }
    free(decoded_buf);
    return version;
}

static void wrap_and_send_pkt(uint8_t *data, size_t len) {
    static uint8_t seq = 67;
    uint8_t b64_buf[BOOT_SERIAL_OUT_MAX*2];
    uint16_t tot_len = 2 + len + sizeof(struct nmgr_hdr) + 2; // +2 is len +2 is crc
    uint16_t net_len = htons(tot_len-2); // tot_len doesn't include the len field

    struct nmgr_hdr hdr = {0};
    hdr.nh_op = NMGR_OP_WRITE;
    hdr.nh_group = htons(MGMT_GROUP_ID_IMAGE);
    hdr.nh_id = IMGMGR_NMGR_ID_UPLOAD;
    hdr.nh_seq = seq++;
    hdr.nh_len = htons(len);

    memcpy(b64_buf, &net_len, sizeof(net_len));

#define HDR_OFFSET 2
#define DATA_OFFSET (HDR_OFFSET+sizeof(struct nmgr_hdr))
#define CRC_OFFSET(len) (DATA_OFFSET+len)

    memcpy(b64_buf+HDR_OFFSET, &hdr, sizeof(struct nmgr_hdr));
    memcpy(b64_buf+DATA_OFFSET, data, len);

    uint16_t crc = crc16(b64_buf+HDR_OFFSET, tot_len-4, CRC_CITT_POLYMINAL, CRC16_INITIAL_CRC, true);
    crc = htons(crc);
    memcpy(b64_buf+CRC_OFFSET(len), &crc, sizeof(crc));

    size_t out_len;
    uint8_t *out_buf = base64_encode(b64_buf, tot_len, &out_len);

    uint8_t *pos = out_buf;
    while (pos < out_buf+out_len) {
        uint8_t start_pkt[] = {SHELL_NLIP_PKT_START1, SHELL_NLIP_PKT_START2};
        uint8_t start_data[] = {SHELL_NLIP_DATA_START1, SHELL_NLIP_DATA_START2};

        if (pos == out_buf) {
            write(serial_fd, start_pkt, sizeof(start_pkt));
        } else {
            usleep(20000);
            write(serial_fd, start_data, sizeof(start_data));
        }

        size_t remaining = out_buf + out_len - pos;
        size_t to_write = remaining > IMAGE_CHUNK_SIZE ? IMAGE_CHUNK_SIZE : remaining;
        write(serial_fd, pos, to_write);
        write(serial_fd, "\n", 1);
        pos += to_write;
    }
    free(out_buf);
    fsync(serial_fd);
}

void do_upload(char *filename){
    size_t pos = 0, len;
    uint8_t const *payload;
    if (filename) {
        int fd = open(filename, O_RDONLY);
        if (fd < 0) {
            fprintf(stderr, "Failed to open \"%s\": %s\n", filename, strerror(errno));
            exit(2);
        }

        struct stat st;
        fstat(fd, &st);
        len = st.st_size;
        uint8_t *payload_mut = calloc(1, len);
        if (!payload_mut) {
            fprintf(stderr, "Failed alloc %s:%d", __FILE__, __LINE__);
            exit(9);
        }
        size_t bytes_read = 0;
        while (bytes_read < len) {
            errno = 0;
            int r = read(fd, &payload_mut[bytes_read], len);
            if (r < 0 && errno != EAGAIN) {
                fprintf(stderr, "Failed to read file: %s:%d\n", __FILE__, __LINE__);
                exit(9);
            }
            bytes_read += r;
        }
        payload = payload_mut;
    } else {
        payload = fixkey_payload_bin;
        len = fixkey_payload_len;
    }

    while (pos < len){
        size_t data_chunk_len;
        if (pos == 0) {
            data_chunk_len = MACRO_CHUNK_SIZE;
        } else {
            data_chunk_len = MACRO_CHUNK_SIZE + 45;
        }
        CborEncoder root, map;
        uint8_t cbor_buf[BOOT_SERIAL_OUT_MAX*2];
        cbor_encoder_init(&root, cbor_buf, BOOT_SERIAL_OUT_MAX*2, 0);
#ifdef DO_UPLOAD_SHA
        cbor_encoder_create_map(&root, &map, pos == 0 ? 4 : 2);
#else
        cbor_encoder_create_map(&root, &map, pos == 0 ? 3 : 2);
#endif

        cbor_encode_text_stringz(&map, "data");

        // Most packets will be `data_chunk_len`, but not the last one
        size_t actual_size = len - pos > data_chunk_len ? data_chunk_len : len - pos;

        uint8_t data[data_chunk_len];
/*        int actual_size;
        errno = 0;
        if ((actual_size = read(fd, data, data_chunk_len)) < data_chunk_len) {
            if (errno != 0) {
                fprintf(stderr, "Error reading the firmware file: %s\n", strerror(errno));
            }
        }*/
        cbor_encode_byte_string(&map, &payload[pos], data_chunk_len);

        if (pos == 0) {
            cbor_encode_text_stringz(&map, "len");
            cbor_encode_uint(&map, len);
        }

        cbor_encode_text_stringz(&map, "off");
        cbor_encode_uint(&map, pos);

#ifdef DO_UPLOAD_SHA
        if (pos == 0) {
            cbor_encode_text_stringz(&map, "sha");
            SHA256_CTX sha_ctx;
            sha256_init(&sha_ctx);
            sha256_update(&sha_ctx, payload, len);
            uint8_t sha[32];
            sha256_final(&sha_ctx, (BYTE *)&sha);
            cbor_encode_byte_string(&map, sha, 32);
        }
#endif /* DO_UPLOAD_SHA */

        cbor_encoder_close_container(&root, &map);
        size_t cbor_size = cbor_encoder_get_buffer_size(&root, cbor_buf);

        wrap_and_send_pkt(cbor_buf, cbor_size);

        uint8_t resp_buf[BOOT_SERIAL_IN_MAX*2];
        size_t in_len;
        read_response(resp_buf, &in_len);

        size_t out_len;
        uint8_t *out = base64_decode(resp_buf+2, in_len-2, &out_len);

        uint16_t l;
        memcpy(&l, out, sizeof(l));
        l = ntohs(l);

        if (crc16(out+2, l, CRC_CITT_POLYMINAL, CRC16_INITIAL_CRC, true)) {
            fprintf(stderr, "Bad CRC, continuing\n");
        }

        CborParser parser;
        CborValue it, rc;
        cbor_parser_init(out+sizeof(struct nmgr_hdr)+2, out_len-sizeof(struct nmgr_hdr)-2, 0, &parser, &it);
        cbor_require_type(&it, CborMapType);
        if (cbor_value_map_find_value(&it, "rc", &rc)) {
            fprintf(stderr, "Failed to find an response code %s:%d\n", __FILE__, __LINE__);
        }
        int r;
        cbor_value_get_int(&rc, &r);
        if (r) {
            fprintf(stderr, "Got an error response (%d) to a sent chunk, bailing\n", r);
            exit(3);
        }
        if (cbor_value_map_find_value(&it, "off", &rc)) {
            fprintf(stderr, "Failed to find an offset in the response %s:%d\n", __FILE__, __LINE__);
        } else {
            cbor_value_get_int(&rc, &r);
            uint8_t percent_complete = round(((double)r/len)*100);
            printf("\033[80D%d/%d Written (%d%%)", r, len, percent_complete);
            fflush(stdout);
        }
        pos += actual_size;
        free(out);
    }
    printf("\nSuccess.\n");
    exit(0);
}

void fix_and_upload(char *filename) {
    char *version = get_version();
    if (!strcmp(version, "0.3.2.0") ) {
        // Upload bootloader update payload
        printf("Bootloader update required, uploading...\n");
        do_upload(NULL);

        // Reset into controller
        gpio_ble_req_controller();

        // Update bootloader
        write(serial_fd, fixkey_msg, sizeof(fixkey_msg));
        if (fsync(serial_fd) < 0) {
            fprintf(stderr, "fsync of serial device failed.\n");
        }

        // Reset into new bootloader
        gpio_ble_req_bootloader();
    }

    // Do the update
    printf("Writing %s into slot 1\n", filename);
    do_upload(filename);
}

void image_main(int fd, int argc, char **argv){
    gpio_ble_req_bootloader();
    serial_fd = fd;
    if (argc < 1) {
        image_help();
        return;
    }
    if (!strcmp(argv[0], "version")) {
        char *version = get_version();
        printf("%s\n", version);
        free(version);
    } else if (!strcmp(argv[0], "upload")) {
        if (argc < 2) {
            fprintf(stderr, "umcumgr image upload <binary file>\n");
        }
        fix_and_upload(argv[1]);
    } else {
        image_help();
    }
    gpio_ble_req_controller();
}