upgrade.c

// Upgrade is an in-place firmware upgrader for tiny85 chips
// The 'bootloaderAddress' variable in bootloader.h, and the bootloaderData
// progmem array with the bootloader data, are generated by the makefile.
//
// Upgrade will firstly rewrite the interrupt vector table to disable the bootloader,
// rewriting it to just run the upgrade app. Next it erases and writes each page of the
// bootloader in sequence, erasing over any remaining pages leaving them set to 0xFFFF
// Finally upgrader erases it's interrupt table again and fills it with RJMPs to
// bootloaderAddress, effectively bridging the interrupts in to the new bootloader's
// interrupts table.
//
// While upgrade has been written with attiny85 and micronucleus in mind, it should
// work with other bootloaders and other chips with flash self program but no hardware
// bootloader protection, where the bootloader exists at the end of flash
//
// Be very careful to not power down the AVR while upgrader is running.
// If you connect a piezo between pb0 and pb1 you'll hear a bleep when the update
// is complete. You can also connect an LED with pb1 positive and pb0 or gnd negative and
// it will blink

#include "utils.h"
#include <avr/io.h>
#include <avr/interrupt.h>
#include <avr/pgmspace.h>
#include <avr/wdt.h>
#include <avr/boot.h>

// this file is generated by make script
#include "bootloader.h"

void secure_interrupt_vector_table(void);
void write_new_bootloader(void);
void forward_interrupt_vector_table(void);
void beep(void);
void reboot(void);

void load_table(uint16_t address, uint16_t words[SPM_PAGESIZE / 2]);
void erase_page(uint16_t address);
void write_page(uint16_t address, uint16_t words[SPM_PAGESIZE / 2]);

#define TINYVECTOR_RESET_OFFSET     4 // the exact value does not matter since we erase the whole page

int main(void) {
  pinsOff(0xFF); // pull down all pins
  outputs(0xFF); // all to ground - force usb disconnect
  delay(250); // milliseconds
  inputs(0xFF); // let them float
  delay(250);
  cli();

  secure_interrupt_vector_table(); // reset our vector table to it's original state
  write_new_bootloader();
  forward_interrupt_vector_table();

  beep();

  /*
   * Do not directly reboot, because the new bootloader may not wait if e.g. if the entry condition is ENTRY_EXT_RESET.
   * In this case the bootloader jumps to the non existing program and this in turn destroys the bootloader.
   * Line 286 of the bootloader to check for existing program is:
   *   if (bootLoaderStartCondition() || (pgm_read_byte(BOOTLOADER_ADDRESS - TINYVECTOR_RESET_OFFSET + 1) == 0xff)) {
   *
   * So erase the page, the new bootloader checks for an existing "program" and then reboot.
   */
  erase_page(BOOTLOADER_ADDRESS - TINYVECTOR_RESET_OFFSET + 1);
  reboot();

  return 0;
}


// erase first page, removing any interrupt table hooks the bootloader added when
// upgrade was uploaded
void secure_interrupt_vector_table( void ) {
  uint16_t table[ SPM_PAGESIZE / 2 ];
  load_table( 0, table );

  // wipe out any interrupt hooks the bootloader rewrote
  int iter = 0;
  while ( iter < SPM_PAGESIZE / 2 ) {
    table[ iter ] = 0xFFFF;
    iter++;
  }

  erase_page( 0 );
  write_page( 0, table );
}


// erase bootloader's section and write over it with new bootloader code
void write_new_bootloader( void ) {
  uint16_t outgoing_page[ SPM_PAGESIZE / 2 ]; // 64 bytes = 32 words
  int page_addr = 0;
  while ( page_addr < bootloader_size ) {
    // read in one page's worth of data from progmem
    int word_addr = 0;
    while ( word_addr < SPM_PAGESIZE ) {
      int subaddress = ( (int) bootloader ) + page_addr + word_addr;
      if ( subaddress < ( (int) bootloader_end ) ) { // valid boot code
        outgoing_page[ word_addr / 2 ] = pgm_read_word( subaddress );
      } else { // fill last words of last page
        outgoing_page[ word_addr / 2 ] = 0xFFFF;
      }
      word_addr += 2;
    }
    // erase page in destination
    erase_page( bootloader_address + page_addr );
    // write updated page
    write_page( bootloader_address + page_addr, outgoing_page );
    page_addr += SPM_PAGESIZE;
  }
}


#if 0
// erase bootloader's section and write over it with new bootloader code
void write_new_bootloader(void) {
  uint16_t outgoing_page[SPM_PAGESIZE / 2];
  int iter = 0;
  while (iter < sizeof(bootloader_data)) {

    // read in one page's worth of data from progmem
    int word_addr = 0;
    while (word_addr < SPM_PAGESIZE) {
      int subaddress = ((int) bootloader_data) + iter + word_addr;
      if (subaddress >= ((int) bootloader_data) + sizeof(bootloader_data)) {
        outgoing_page[word_addr / 2] = 0xFFFF;
      } else {
        outgoing_page[word_addr / 2] = pgm_read_word(subaddress);
      }

      word_addr += 2;
    }

    // erase page in destination
    erase_page(bootloader_address + iter);
    // write updated page
    write_page(bootloader_address + iter, outgoing_page);

    iter += 64;
  }
}
#endif


// write in forwarding interrupt vector table
void forward_interrupt_vector_table( void ) {
  uint16_t vector_table[ SPM_PAGESIZE / 2 ];

  int iter = 0;
  while ( iter < SPM_PAGESIZE / 2 ) {
    // rjmp to bootloader_address's interrupt table
    vector_table[ iter ] = 0xC000 + ( bootloader_address / 2 ) - 1;
    iter++;
  }

  erase_page( 0 );
  write_page( 0, vector_table );
}


void load_table( uint16_t address, uint16_t words[ SPM_PAGESIZE / 2 ] ) {
  uint16_t subaddress = 0;
  address -= address % SPM_PAGESIZE; // round down to nearest page start

  while (subaddress < SPM_PAGESIZE) {
    words[ subaddress / 2 ] = pgm_read_word( address + subaddress );
    subaddress += 2;
  }
}


void erase_page( uint16_t address ) {
  boot_page_erase( address - ( address % SPM_PAGESIZE ) );
  boot_spm_busy_wait();
}


void write_page( uint16_t address, uint16_t words[ SPM_PAGESIZE / 2 ] ) {
  // fill buffer
  uint16_t iter = 0;
  while ( iter < SPM_PAGESIZE / 2 ) {
    boot_page_fill( address + ( iter * 2 ), words[ iter ] );
    iter++;
  }

  boot_page_write( address );
  boot_spm_busy_wait(); // Wait until the memory is written.
}


// beep for half a second
void beep( void ) {
  outputs( pin(0) | pin(1) );
  pinOff( 1 );

  byte i = 0;
  while ( i < 250 ) {
    delay( 1 );
    pinOn( 1 );
    delay( 1 );
    pinOff( 1 );
    i++;
  }
}


void reboot(void) {
  void (*ptrToFunction)(); // pointer to a function
  ptrToFunction = 0x0000;
  (*ptrToFunction)(); // reset!
}


////////////// Add padding to start of program so no program code could reasonably be erased while program is running
// this never needs to be called - avr-gcc stuff happening: http://www.nongnu.org/avr-libc/user-manual/mem_sections.html
void FakeISR (void) __attribute__ ((naked)) __attribute__ ((section (".init0")));
void FakeISR (void) {
  // 16 nops to pad out first section of program
  asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop");
  asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop"); asm("nop");
}