Tooltron

#include <avr/io.h>

#include <avr/boot.h>

#include <avr/pgmspace.h>

#include "tooltron.h"

#include "uart.h"

#define TRUE        0

#define FALSE        1

#define ADDR    18

// Error thresholds

#define MAX_TIMEOUT 60000   // Seconds to wait before exiting bootloader mode

#define MAX_RETRIES 5       // Number of times to retry before giving up

// Memory locations

#define MAIN_ADDR 0x000     // Where the user code starts

#define BOOT_START 0x400    // Where the bootloader starts

//Status LED

#define LED_DDR  DDRB

#define LED_PORT PORTB

#define LED      PORTB1

//Function prototypes

void (*main_start)(void) = BOOT_START/2 - 1;

// Packet handler states

typedef enum {

    sd,

    src,

    dest,

    comd,

    read,

    cs,

    ack

} state_t;

typedef union {

    uint8_t bytes[2];

    int16_t sword;

} rjump_t;

void init_uart(uint16_t baud) {

        // Set baud rate

        UBRRH = (uint8_t)(baud>>8);

        UBRRL = (uint8_t)baud;

        // Enable RX/TX

        UCSRB = _BV(RXEN) | _BV(TXEN);

        // Enable the TXEN pin as output

        DDRD |= TX_EN;

    uart_toggle_transmit(UART_TX_OFF);

}

int8_t uart_get_byte(uint8_t *output_byte) {

    if (UCSRA & _BV(RXC)) {

        *output_byte = UDR;

        return 0;

    } else {

        return -1;

    }

}

void uart_send_byte(uint8_t data) {

        //Waits until current transmit is done

    while (!(UCSRA & _BV(UDRE)));

    // Enable writes and send

        uart_toggle_transmit(UART_TX_ON);

    UDR = data;

    // Waits until the transmit is done

    while(!(UCSRA & _BV(TXC)));

    uart_toggle_transmit(UART_TX_OFF);

    UCSRA |= _BV(TXC);

        return;

}

void uart_toggle_transmit(uint8_t state) {

        if (state == UART_TX_ON) {

                PORTD |= TX_EN;

        } else {

                PORTD &= ~TX_EN;

        }

}

char parse_packet(uint8_t *mbuf) {

  uint8_t r;        // Byte from the network

  uint8_t crc;      // Running checksum of the packet

  uint8_t cmd;      // The command received

  uint8_t pos;      // Position in the message buffer

  uint8_t lim;      // Max number of bytes to read into the message buf

  state_t state;    // State machine

  uint16_t count;

  r = 0;

  crc = 0;

  cmd = 0;

  pos = 0;

  lim = 0;

  state = sd;

  count = 0;

  while (1) {

    // Wait for the next byte

    while ((uart_get_byte(&r)) < 0) {

        if (count >= MAX_TIMEOUT) {

            return TT_BAD;

        } else {

            count++;

        }

    }

    switch (state) {

        case sd:

            if (r == DELIM) {

                state = src;

            }

            break;

        case src:

            if (r == DELIM) {

                state = src;

            } else {

                crc = r;

                state = dest;

            }

            break;

        case dest:

            if (r == DELIM) {

                state = src;

            } else if (r == ADDR) {

                crc ^= r;

                state = comd;

            } else {

                state = sd;

            }

            break;

        case comd:

            cmd = r;

            crc ^= r;

            if (r == DELIM) {

                state = src;

            } else if (r == TT_PROGM) {

                lim = PROGM_PACKET_SIZE;

                state = read;

            } else if (r == TT_PROGD) {

                lim = PROGD_PACKET_SIZE;

                state = read;

            } else {

                state = cs;

            }

            break;

        case read:

            mbuf[pos] = r;

            crc ^= r;

            pos++;

            if (pos == lim) {

                state = cs;

            }

            break;

        case cs:

            if (r == crc) {

                return cmd;

            } else {

                return TT_BAD;

            }

            break;

        default:

            return TT_BAD;

    }

  }

}

void send_packet(uint8_t cmd) {

    uart_send_byte(DELIM);

    uart_send_byte(ADDR);

    uart_send_byte(SERVER);

    uart_send_byte(cmd);

    uart_send_byte(ACK_CRC ^ cmd);

}

// SPM_PAGESIZE is set to 32 bytes

void onboard_program_write(uint16_t page, uint8_t *buf) {

  uint16_t i;

  boot_page_erase (page);

  boot_spm_busy_wait ();      // Wait until the memory is erased.

  for (i=0; i < SPM_PAGESIZE; i+=2){

    // Set up little-endian word.

    boot_page_fill (page + i, buf[i] | (buf[i+1] <<8));

  }

  boot_page_write (page);     // Store buffer in flash page.

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

}

int main(void) {

  uint8_t mbuf[PROGD_PACKET_SIZE];

  rjump_t jbuf;

  uint16_t caddr = MAIN_ADDR;

  uint8_t iteration;

  uint8_t resp;

  uint16_t prog_len;

  uint8_t i;

  uint8_t retries;

  iteration = 0;

  retries = 0;

  init_uart(51); //MAGIC NUMBER??

  //set LED pin as output

  LED_DDR |= 0x07;

  PORTB = 0x07;

  //Start bootloading process

  send_packet(TT_BOOT);

  resp = parse_packet(mbuf);

  // Enter programming mode 

  if (resp == TT_PROGM) {

    prog_len = mbuf[0];

    prog_len |= mbuf[1] << 8;

    // This will insert a NOP into the user code jump in case

    // the programming fails

    for (i = 0; i < PROGD_PACKET_SIZE; i++) {

      mbuf[i]= 0;

    }

    onboard_program_write(BOOT_START - SPM_PAGESIZE, mbuf);

  // Run user code

  } else {

      main_start();

  }

  send_packet(TT_ACK);

  while(1) {

      resp = parse_packet(mbuf);

      if (resp == TT_PROGD) {

          // We need to muck with the reset vector jump in the first page

          if (iteration == 0) {

              // Store the jump to user code

              jbuf.bytes[0] = mbuf[0];

              jbuf.bytes[1] = mbuf[1];

              // Rewrite the user code jump to be correct since we are

              // using relative jumps (rjmp)

              jbuf.sword &= 0x0FFF;

              jbuf.sword -= (BOOT_START >> 1) - 1;

              jbuf.sword &= 0x0FFF;

              jbuf.sword |= 0xC000;

              // Rewrite the reset vector to jump to the bootloader

              mbuf[0] = (BOOT_START/2 - 1) & 0xFF;

              mbuf[1] = 0xC0 | (((BOOT_START/2 - 1) >> 8) & 0x0F);

              iteration = 1;

          }

          // Write the page to the flash

          onboard_program_write(caddr, mbuf);

          caddr += PROGD_PACKET_SIZE;

          retries = 0;

      } else {

          send_packet(TT_NACK);

          retries++;

          // If we failed too many times, reset. The main_start should

          // no-op back to the start of the bootloader

          if (retries > MAX_RETRIES) {

              main_start();

          }

      }

      send_packet(TT_ACK);

      // Once we write the last packet we must override the jump to

      // user code to point to the correct address

      if (prog_len <= PROGD_PACKET_SIZE) {

          for (i = 0; i < PROGD_PACKET_SIZE; i++) {

              mbuf[i]= 0;

          }

          mbuf[PROGD_PACKET_SIZE-2] = jbuf.bytes[0];

          mbuf[PROGD_PACKET_SIZE-1] = jbuf.bytes[1];

          onboard_program_write(BOOT_START - SPM_PAGESIZE, mbuf);

          main_start();

      } else { 

          prog_len -= PROGD_PACKET_SIZE;

      }

  }

  // Should never get here

  return -1;

}