lib/lufa/Bootloaders/DFU/BootloaderDFU.c

   1 /*
   2              LUFA Library
   3      Copyright (C) Dean Camera, 2017.
   4
   5   dean [at] fourwalledcubicle [dot] com
   6            www.lufa-lib.org
   7 */
   8
   9 /*
  10   Copyright 2017  Dean Camera (dean [at] fourwalledcubicle [dot] com)
  11
  12   Permission to use, copy, modify, distribute, and sell this
  13   software and its documentation for any purpose is hereby granted
  14   without fee, provided that the above copyright notice appear in
  15   all copies and that both that the copyright notice and this
  16   permission notice and warranty disclaimer appear in supporting
  17   documentation, and that the name of the author not be used in
  18   advertising or publicity pertaining to distribution of the
  19   software without specific, written prior permission.
  20
  21   The author disclaims all warranties with regard to this
  22   software, including all implied warranties of merchantability
  23   and fitness.  In no event shall the author be liable for any
  24   special, indirect or consequential damages or any damages
  25   whatsoever resulting from loss of use, data or profits, whether
  26   in an action of contract, negligence or other tortious action,
  27   arising out of or in connection with the use or performance of
  28   this software.
  29 */
  30
  31 /** \file
  32  *
  33  *  Main source file for the DFU class bootloader. This file contains the complete bootloader logic.
  34  */
  35
  36 #define  INCLUDE_FROM_BOOTLOADER_C
  37 #include "BootloaderDFU.h"
  38
  39 /** Flag to indicate if the bootloader is currently running in secure mode, disallowing memory operations
  40  *  other than erase. This is initially set to the value set by SECURE_MODE, and cleared by the bootloader
  41  *  once a memory erase has completed in a bootloader session.
  42  */
  43 static bool IsSecure = SECURE_MODE;
  44
  45 /** Flag to indicate if the bootloader should be running, or should exit and allow the application code to run
  46  *  via a soft reset. When cleared, the bootloader will abort, the USB interface will shut down and the application
  47  *  jumped to via an indirect jump to location 0x0000 (or other location specified by the host).
  48  */
  49 static bool RunBootloader = true;
  50
  51 /** Flag to indicate if the bootloader is waiting to exit. When the host requests the bootloader to exit and
  52  *  jump to the application address it specifies, it sends two sequential commands which must be properly
  53  *  acknowledged. Upon reception of the first the RunBootloader flag is cleared and the WaitForExit flag is set,
  54  *  causing the bootloader to wait for the final exit command before shutting down.
  55  */
  56 static bool WaitForExit = false;
  57
  58 /** Current DFU state machine state, one of the values in the DFU_State_t enum. */
  59 static uint8_t DFU_State = dfuIDLE;
  60
  61 /** Status code of the last executed DFU command. This is set to one of the values in the DFU_Status_t enum after
  62  *  each operation, and returned to the host when a Get Status DFU request is issued.
  63  */
  64 static uint8_t DFU_Status = OK;
  65
  66 /** Data containing the DFU command sent from the host. */
  67 static DFU_Command_t SentCommand;
  68
  69 /** Response to the last issued Read Data DFU command. Unlike other DFU commands, the read command
  70  *  requires a single byte response from the bootloader containing the read data when the next DFU_UPLOAD command
  71  *  is issued by the host.
  72  */
  73 static uint8_t ResponseByte;
  74
  75 /** Pointer to the start of the user application. By default this is 0x0000 (the reset vector), however the host
  76  *  may specify an alternate address when issuing the application soft-start command.
  77  */
  78 static AppPtr_t AppStartPtr = (AppPtr_t)0x0000;
  79
  80 /** 64-bit flash page number. This is concatenated with the current 16-bit address on USB AVRs containing more than
  81  *  64KB of flash memory.
  82  */
  83 static uint8_t Flash64KBPage = 0;
  84
  85 /** Memory start address, indicating the current address in the memory being addressed (either FLASH or EEPROM
  86  *  depending on the issued command from the host).
  87  */
  88 static uint16_t StartAddr = 0x0000;
  89
  90 /** Memory end address, indicating the end address to read from/write to in the memory being addressed (either FLASH
  91  *  of EEPROM depending on the issued command from the host).
  92  */
  93 static uint16_t EndAddr = 0x0000;
  94
  95 /** Magic lock for forced application start. If the HWBE fuse is programmed and BOOTRST is unprogrammed, the bootloader
  96  *  will start if the /HWB line of the AVR is held low and the system is reset. However, if the /HWB line is still held
  97  *  low when the application attempts to start via a watchdog reset, the bootloader will re-start. If set to the value
  98  *  \ref MAGIC_BOOT_KEY the special init function \ref Application_Jump_Check() will force the application to start.
  99  */
 100 uint16_t MagicBootKey ATTR_NO_INIT;
 101
 102
 103 /** Special startup routine to check if the bootloader was started via a watchdog reset, and if the magic application
 104  *  start key has been loaded into \ref MagicBootKey. If the bootloader started via the watchdog and the key is valid,
 105  *  this will force the user application to start via a software jump.
 106  */
 107 void Application_Jump_Check(void)
 108 {
 109         bool JumpToApplication = false;
 110
 111         #if (BOARD == BOARD_LEONARDO)
 112                 /* Enable pull-up on the IO13 pin so we can use it to select the mode */
 113                 PORTC |= (1 << 7);
 114                 Delay_MS(10);
 115
 116                 /* If IO13 is not jumpered to ground, start the user application instead */
 117                 JumpToApplication = ((PINC & (1 << 7)) != 0);
 118
 119                 /* Disable pull-up after the check has completed */
 120                 PORTC &= ~(1 << 7);
 121         #elif ((BOARD == BOARD_XPLAIN) || (BOARD == BOARD_XPLAIN_REV1))
 122                 /* Disable JTAG debugging */
 123                 JTAG_DISABLE();
 124
 125                 /* Enable pull-up on the JTAG TCK pin so we can use it to select the mode */
 126                 PORTF |= (1 << 4);
 127                 Delay_MS(10);
 128
 129                 /* If the TCK pin is not jumpered to ground, start the user application instead */
 130                 JumpToApplication = ((PINF & (1 << 4)) != 0);
 131
 132                 /* Re-enable JTAG debugging */
 133                 JTAG_ENABLE();
 134         #else
 135                 /* Check if the device's BOOTRST fuse is set */
 136                 if (boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS) & FUSE_BOOTRST)
 137                 {
 138                         /* If the reset source was not an external reset or the key is correct, clear it and jump to the application */
 139                         //if (!(MCUSR & (1 << EXTRF)) || (MagicBootKey == MAGIC_BOOT_KEY))
 140                         //  JumpToApplication = true;
 141
 142                         /* Clear reset source */
 143                         MCUSR &= ~(1 << EXTRF);
 144                 }
 145                 else
 146                 {
 147                         /* If the reset source was the bootloader and the key is correct, clear it and jump to the application;
 148                          * this can happen in the HWBE fuse is set, and the HBE pin is low during the watchdog reset */
 149                         //if ((MCUSR & (1 << WDRF)) && (MagicBootKey == MAGIC_BOOT_KEY))
 150                         //      JumpToApplication = true;
 151
 152                         /* Clear reset source */
 153                         MCUSR &= ~(1 << WDRF);
 154                 }
 155         #endif
 156
 157         /* Don't run the user application if the reset vector is blank (no app loaded) */
 158         bool ApplicationValid = (pgm_read_word_near(0) != 0xFFFF);
 159
 160         /* If a request has been made to jump to the user application, honor it */
 161         if (JumpToApplication && ApplicationValid)
 162         {
 163                 /* Turn off the watchdog */
 164                 MCUSR &= ~(1 << WDRF);
 165                 wdt_disable();
 166
 167                 /* Clear the boot key and jump to the user application */
 168                 MagicBootKey = 0;
 169
 170                 // cppcheck-suppress constStatement
 171                 ((void (*)(void))0x0000)();
 172         }
 173 }
 174
 175 /** Main program entry point. This routine configures the hardware required by the bootloader, then continuously
 176  *  runs the bootloader processing routine until instructed to soft-exit, or hard-reset via the watchdog to start
 177  *  the loaded application code.
 178  */
 179 int main(void)
 180 {
 181         /* Configure hardware required by the bootloader */
 182         SetupHardware();
 183
 184         /* Turn on first LED on the board to indicate that the bootloader has started */
 185         LEDs_SetAllLEDs(LEDS_LED1 | LEDS_LED2);
 186
 187         /* Enable global interrupts so that the USB stack can function */
 188         GlobalInterruptEnable();
 189
 190
 191         #if (BOARD == BOARD_QMK)
 192                 uint16_t keypress = 0;
 193         #endif
 194
 195         /* Run the USB management task while the bootloader is supposed to be running */
 196         while (RunBootloader || WaitForExit) {
 197           USB_USBTask();
 198           #if (BOARD == BOARD_QMK)
 199                 bool pressed = (PIN(QMK_ESC_INPUT) & NUM(QMK_ESC_INPUT));
 200                 if ((DFU_State == dfuIDLE) && (keypress > 5000) && pressed) {
 201                         break;
 202                 }
 203                 if (pressed) {
 204                         keypress++;
 205                 } else {
 206                         keypress = 0;
 207                 }
 208
 209           #endif
 210         }
 211
 212         /* Reset configured hardware back to their original states for the user application */
 213         ResetHardware();
 214
 215         /* Start the user application */
 216         AppStartPtr();
 217 }
 218
 219 /** Configures all hardware required for the bootloader. */
 220 static void SetupHardware(void)
 221 {
 222         /* Disable watchdog if enabled by bootloader/fuses */
 223         MCUSR &= ~(1 << WDRF);
 224         wdt_disable();
 225
 226         /* Disable clock division */
 227         clock_prescale_set(clock_div_1);
 228
 229         /* Relocate the interrupt vector table to the bootloader section */
 230         MCUCR = (1 << IVCE);
 231         MCUCR = (1 << IVSEL);
 232
 233         #if (BOARD == BOARD_QMK)
 234                 // output setup
 235                 DDR(QMK_ESC_OUTPUT) |= NUM(QMK_ESC_OUTPUT);
 236                 PORT(QMK_ESC_OUTPUT) |= NUM(QMK_ESC_OUTPUT);
 237
 238                 // input setup
 239                 DDR(QMK_ESC_INPUT) |= NUM(QMK_ESC_INPUT);
 240         #endif
 241
 242         /* Initialize the USB and other board hardware drivers */
 243         USB_Init();
 244         LEDs_Init();
 245
 246         /* Bootloader active LED toggle timer initialization */
 247         TIMSK1 = (1 << TOIE1);
 248         TCCR1B = ((1 << CS11) | (1 << CS10));
 249
 250 }
 251
 252 /** Resets all configured hardware required for the bootloader back to their original states. */
 253 static void ResetHardware(void)
 254 {
 255         /* Shut down the USB and other board hardware drivers */
 256         USB_Disable();
 257         LEDs_Disable();
 258
 259         /* Disable Bootloader active LED toggle timer */
 260         TIMSK1 = 0;
 261         TCCR1B = 0;
 262
 263         /* Relocate the interrupt vector table back to the application section */
 264         MCUCR = (1 << IVCE);
 265         MCUCR = 0;
 266
 267         #if (BOARD == BOARD_QMK)
 268                 DDR(QMK_ESC_OUTPUT) = PORT(QMK_ESC_OUTPUT) = DDR(QMK_ESC_INPUT) = PORT(QMK_ESC_INPUT) = 0;
 269         #endif
 270 }
 271
 272 /** ISR to periodically toggle the LEDs on the board to indicate that the bootloader is active. */
 273 ISR(TIMER1_OVF_vect, ISR_BLOCK)
 274 {
 275         LEDs_ToggleLEDs(LEDS_LED1 | LEDS_LED2);
 276 }
 277
 278 /** Event handler for the USB_ControlRequest event. This is used to catch and process control requests sent to
 279  *  the device from the USB host before passing along unhandled control requests to the library for processing
 280  *  internally.
 281  */
 282 void EVENT_USB_Device_ControlRequest(void)
 283 {
 284         /* Ignore any requests that aren't directed to the DFU interface */
 285         if ((USB_ControlRequest.bmRequestType & (CONTROL_REQTYPE_TYPE | CONTROL_REQTYPE_RECIPIENT)) !=
 286             (REQTYPE_CLASS | REQREC_INTERFACE))
 287         {
 288                 return;
 289         }
 290
 291         /* Activity - toggle indicator LEDs */
 292         LEDs_ToggleLEDs(LEDS_LED1 | LEDS_LED2);
 293
 294         /* Get the size of the command and data from the wLength value */
 295         SentCommand.DataSize = USB_ControlRequest.wLength;
 296
 297         switch (USB_ControlRequest.bRequest)
 298         {
 299                 case DFU_REQ_DNLOAD:
 300                         Endpoint_ClearSETUP();
 301
 302                         /* Check if bootloader is waiting to terminate */
 303                         if (WaitForExit)
 304                         {
 305                                 /* Bootloader is terminating - process last received command */
 306                                 ProcessBootloaderCommand();
 307
 308                                 /* Indicate that the last command has now been processed - free to exit bootloader */
 309                                 WaitForExit = false;
 310                         }
 311
 312                         /* If the request has a data stage, load it into the command struct */
 313                         if (SentCommand.DataSize)
 314                         {
 315                                 while (!(Endpoint_IsOUTReceived()))
 316                                 {
 317                                         if (USB_DeviceState == DEVICE_STATE_Unattached)
 318                                           return;
 319                                 }
 320
 321                                 /* First byte of the data stage is the DNLOAD request's command */
 322                                 SentCommand.Command = Endpoint_Read_8();
 323
 324                                 /* One byte of the data stage is the command, so subtract it from the total data bytes */
 325                                 SentCommand.DataSize--;
 326
 327                                 /* Load in the rest of the data stage as command parameters */
 328                                 for (uint8_t DataByte = 0; (DataByte < sizeof(SentCommand.Data)) &&
 329                                      Endpoint_BytesInEndpoint(); DataByte++)
 330                                 {
 331                                         SentCommand.Data[DataByte] = Endpoint_Read_8();
 332                                         SentCommand.DataSize--;
 333                                 }
 334
 335                                 /* Process the command */
 336                                 ProcessBootloaderCommand();
 337                         }
 338
 339                         /* Check if currently downloading firmware */
 340                         if (DFU_State == dfuDNLOAD_IDLE)
 341                         {
 342                                 if (!(SentCommand.DataSize))
 343                                 {
 344                                         DFU_State = dfuIDLE;
 345                                 }
 346                                 else
 347                                 {
 348                                         /* Throw away the filler bytes before the start of the firmware */
 349                                         DiscardFillerBytes(DFU_FILLER_BYTES_SIZE);
 350
 351                                         /* Throw away the packet alignment filler bytes before the start of the firmware */
 352                                         DiscardFillerBytes(StartAddr % FIXED_CONTROL_ENDPOINT_SIZE);
 353
 354                                         /* Calculate the number of bytes remaining to be written */
 355                                         uint16_t BytesRemaining = ((EndAddr - StartAddr) + 1);
 356
 357                                         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))        // Write flash
 358                                         {
 359                                                 /* Calculate the number of words to be written from the number of bytes to be written */
 360                                                 uint16_t WordsRemaining = (BytesRemaining >> 1);
 361
 362                                                 union
 363                                                 {
 364                                                         uint16_t Words[2];
 365                                                         uint32_t Long;
 366                                                 } CurrFlashAddress                 = {.Words = {StartAddr, Flash64KBPage}};
 367
 368                                                 uint32_t CurrFlashPageStartAddress = CurrFlashAddress.Long;
 369                                                 uint8_t  WordsInFlashPage          = 0;
 370
 371                                                 while (WordsRemaining--)
 372                                                 {
 373                                                         /* Check if endpoint is empty - if so clear it and wait until ready for next packet */
 374                                                         if (!(Endpoint_BytesInEndpoint()))
 375                                                         {
 376                                                                 Endpoint_ClearOUT();
 377
 378                                                                 while (!(Endpoint_IsOUTReceived()))
 379                                                                 {
 380                                                                         if (USB_DeviceState == DEVICE_STATE_Unattached)
 381                                                                           return;
 382                                                                 }
 383                                                         }
 384
 385                                                         /* Write the next word into the current flash page */
 386                                                         boot_page_fill(CurrFlashAddress.Long, Endpoint_Read_16_LE());
 387
 388                                                         /* Adjust counters */
 389                                                         WordsInFlashPage      += 1;
 390                                                         CurrFlashAddress.Long += 2;
 391
 392                                                         /* See if an entire page has been written to the flash page buffer */
 393                                                         if ((WordsInFlashPage == (SPM_PAGESIZE >> 1)) || !(WordsRemaining))
 394                                                         {
 395                                                                 /* Commit the flash page to memory */
 396                                                                 boot_page_write(CurrFlashPageStartAddress);
 397                                                                 boot_spm_busy_wait();
 398
 399                                                                 /* Check if programming incomplete */
 400                                                                 if (WordsRemaining)
 401                                                                 {
 402                                                                         CurrFlashPageStartAddress = CurrFlashAddress.Long;
 403                                                                         WordsInFlashPage          = 0;
 404
 405                                                                         /* Erase next page's temp buffer */
 406                                                                         boot_page_erase(CurrFlashAddress.Long);
 407                                                                         boot_spm_busy_wait();
 408                                                                 }
 409                                                         }
 410                                                 }
 411
 412                                                 /* Once programming complete, start address equals the end address */
 413                                                 StartAddr = EndAddr;
 414
 415                                                 /* Re-enable the RWW section of flash */
 416                                                 boot_rww_enable();
 417                                         }
 418                                         else                                                   // Write EEPROM
 419                                         {
 420                                                 while (BytesRemaining--)
 421                                                 {
 422                                                         /* Check if endpoint is empty - if so clear it and wait until ready for next packet */
 423                                                         if (!(Endpoint_BytesInEndpoint()))
 424                                                         {
 425                                                                 Endpoint_ClearOUT();
 426
 427                                                                 while (!(Endpoint_IsOUTReceived()))
 428                                                                 {
 429                                                                         if (USB_DeviceState == DEVICE_STATE_Unattached)
 430                                                                           return;
 431                                                                 }
 432                                                         }
 433
 434                                                         /* Read the byte from the USB interface and write to to the EEPROM */
 435                                                         eeprom_update_byte((uint8_t*)StartAddr, Endpoint_Read_8());
 436
 437                                                         /* Adjust counters */
 438                                                         StartAddr++;
 439                                                 }
 440                                         }
 441
 442                                         /* Throw away the currently unused DFU file suffix */
 443                                         DiscardFillerBytes(DFU_FILE_SUFFIX_SIZE);
 444                                 }
 445                         }
 446
 447                         Endpoint_ClearOUT();
 448
 449                         Endpoint_ClearStatusStage();
 450
 451                         break;
 452                 case DFU_REQ_UPLOAD:
 453                         Endpoint_ClearSETUP();
 454
 455                         while (!(Endpoint_IsINReady()))
 456                         {
 457                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 458                                   return;
 459                         }
 460
 461                         if (DFU_State != dfuUPLOAD_IDLE)
 462                         {
 463                                 if ((DFU_State == dfuERROR) && IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))       // Blank Check
 464                                 {
 465                                         /* Blank checking is performed in the DFU_DNLOAD request - if we get here we've told the host
 466                                            that the memory isn't blank, and the host is requesting the first non-blank address */
 467                                         Endpoint_Write_16_LE(StartAddr);
 468                                 }
 469                                 else
 470                                 {
 471                                         /* Idle state upload - send response to last issued command */
 472                                         Endpoint_Write_8(ResponseByte);
 473                                 }
 474                         }
 475                         else
 476                         {
 477                                 /* Determine the number of bytes remaining in the current block */
 478                                 uint16_t BytesRemaining = ((EndAddr - StartAddr) + 1);
 479
 480                                 if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))            // Read FLASH
 481                                 {
 482                                         /* Calculate the number of words to be written from the number of bytes to be written */
 483                                         uint16_t WordsRemaining = (BytesRemaining >> 1);
 484
 485                                         union
 486                                         {
 487                                                 uint16_t Words[2];
 488                                                 uint32_t Long;
 489                                         } CurrFlashAddress = {.Words = {StartAddr, Flash64KBPage}};
 490
 491                                         while (WordsRemaining--)
 492                                         {
 493                                                 /* Check if endpoint is full - if so clear it and wait until ready for next packet */
 494                                                 if (Endpoint_BytesInEndpoint() == FIXED_CONTROL_ENDPOINT_SIZE)
 495                                                 {
 496                                                         Endpoint_ClearIN();
 497
 498                                                         while (!(Endpoint_IsINReady()))
 499                                                         {
 500                                                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 501                                                                   return;
 502                                                         }
 503                                                 }
 504
 505                                                 /* Read the flash word and send it via USB to the host */
 506                                                 #if (FLASHEND > 0xFFFF)
 507                                                         Endpoint_Write_16_LE(pgm_read_word_far(CurrFlashAddress.Long));
 508                                                 #else
 509                                                         Endpoint_Write_16_LE(pgm_read_word(CurrFlashAddress.Long));
 510                                                 #endif
 511
 512                                                 /* Adjust counters */
 513                                                 CurrFlashAddress.Long += 2;
 514                                         }
 515
 516                                         /* Once reading is complete, start address equals the end address */
 517                                         StartAddr = EndAddr;
 518                                 }
 519                                 else if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x02))       // Read EEPROM
 520                                 {
 521                                         while (BytesRemaining--)
 522                                         {
 523                                                 /* Check if endpoint is full - if so clear it and wait until ready for next packet */
 524                                                 if (Endpoint_BytesInEndpoint() == FIXED_CONTROL_ENDPOINT_SIZE)
 525                                                 {
 526                                                         Endpoint_ClearIN();
 527
 528                                                         while (!(Endpoint_IsINReady()))
 529                                                         {
 530                                                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 531                                                                   return;
 532                                                         }
 533                                                 }
 534
 535                                                 /* Read the EEPROM byte and send it via USB to the host */
 536                                                 Endpoint_Write_8(eeprom_read_byte((uint8_t*)StartAddr));
 537
 538                                                 /* Adjust counters */
 539                                                 StartAddr++;
 540                                         }
 541                                 }
 542
 543                                 /* Return to idle state */
 544                                 DFU_State = dfuIDLE;
 545                         }
 546
 547                         Endpoint_ClearIN();
 548
 549                         Endpoint_ClearStatusStage();
 550                         break;
 551                 case DFU_REQ_GETSTATUS:
 552                         Endpoint_ClearSETUP();
 553
 554                         while (!(Endpoint_IsINReady()))
 555                         {
 556                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 557                                   return;
 558                         }
 559
 560                         /* Write 8-bit status value */
 561                         Endpoint_Write_8(DFU_Status);
 562
 563                         /* Write 24-bit poll timeout value */
 564                         Endpoint_Write_8(0);
 565                         Endpoint_Write_16_LE(0);
 566
 567                         /* Write 8-bit state value */
 568                         Endpoint_Write_8(DFU_State);
 569
 570                         /* Write 8-bit state string ID number */
 571                         Endpoint_Write_8(0);
 572
 573                         Endpoint_ClearIN();
 574
 575                         Endpoint_ClearStatusStage();
 576                         break;
 577                 case DFU_REQ_CLRSTATUS:
 578                         Endpoint_ClearSETUP();
 579
 580                         /* Reset the status value variable to the default OK status */
 581                         DFU_Status = OK;
 582
 583                         Endpoint_ClearStatusStage();
 584                         break;
 585                 case DFU_REQ_GETSTATE:
 586                         Endpoint_ClearSETUP();
 587
 588                         while (!(Endpoint_IsINReady()))
 589                         {
 590                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 591                                   return;
 592                         }
 593
 594                         /* Write the current device state to the endpoint */
 595                         Endpoint_Write_8(DFU_State);
 596
 597                         Endpoint_ClearIN();
 598
 599                         Endpoint_ClearStatusStage();
 600                         break;
 601                 case DFU_REQ_ABORT:
 602                         Endpoint_ClearSETUP();
 603
 604                         /* Reset the current state variable to the default idle state */
 605                         DFU_State = dfuIDLE;
 606
 607                         Endpoint_ClearStatusStage();
 608                         break;
 609         }
 610 }
 611
 612 /** Routine to discard the specified number of bytes from the control endpoint stream. This is used to
 613  *  discard unused bytes in the stream from the host, including the memory program block suffix.
 614  *
 615  *  \param[in] NumberOfBytes  Number of bytes to discard from the host from the control endpoint
 616  */
 617 static void DiscardFillerBytes(uint8_t NumberOfBytes)
 618 {
 619         while (NumberOfBytes--)
 620         {
 621                 if (!(Endpoint_BytesInEndpoint()))
 622                 {
 623                         Endpoint_ClearOUT();
 624
 625                         /* Wait until next data packet received */
 626                         while (!(Endpoint_IsOUTReceived()))
 627                         {
 628                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 629                                   return;
 630                         }
 631                 }
 632                 else
 633                 {
 634                         Endpoint_Discard_8();
 635                 }
 636         }
 637 }
 638
 639 /** Routine to process an issued command from the host, via a DFU_DNLOAD request wrapper. This routine ensures
 640  *  that the command is allowed based on the current secure mode flag value, and passes the command off to the
 641  *  appropriate handler function.
 642  */
 643 static void ProcessBootloaderCommand(void)
 644 {
 645         /* Check if device is in secure mode */
 646         if (IsSecure)
 647         {
 648                 /* Don't process command unless it is a READ or chip erase command */
 649                 if (!(((SentCommand.Command == COMMAND_WRITE)             &&
 650                         IS_TWOBYTE_COMMAND(SentCommand.Data, 0x00, 0xFF)) ||
 651                            (SentCommand.Command == COMMAND_READ)))
 652                 {
 653                         /* Set the state and status variables to indicate the error */
 654                         DFU_State  = dfuERROR;
 655                         DFU_Status = errWRITE;
 656
 657                         /* Stall command */
 658                         Endpoint_StallTransaction();
 659
 660                         /* Don't process the command */
 661                         return;
 662                 }
 663         }
 664
 665         /* Dispatch the required command processing routine based on the command type */
 666         switch (SentCommand.Command)
 667         {
 668                 case COMMAND_PROG_START:
 669                         ProcessMemProgCommand();
 670                         break;
 671                 case COMMAND_DISP_DATA:
 672                         ProcessMemReadCommand();
 673                         break;
 674                 case COMMAND_WRITE:
 675                         ProcessWriteCommand();
 676                         break;
 677                 case COMMAND_READ:
 678                         ProcessReadCommand();
 679                         break;
 680                 case COMMAND_CHANGE_BASE_ADDR:
 681                         if (IS_TWOBYTE_COMMAND(SentCommand.Data, 0x03, 0x00))              // Set 64KB flash page command
 682                           Flash64KBPage = SentCommand.Data[2];
 683
 684                         break;
 685         }
 686 }
 687
 688 /** Routine to concatenate the given pair of 16-bit memory start and end addresses from the host, and store them
 689  *  in the StartAddr and EndAddr global variables.
 690  */
 691 static void LoadStartEndAddresses(void)
 692 {
 693         union
 694         {
 695                 uint8_t  Bytes[2];
 696                 uint16_t Word;
 697         } Address[2] = {{.Bytes = {SentCommand.Data[2], SentCommand.Data[1]}},
 698                         {.Bytes = {SentCommand.Data[4], SentCommand.Data[3]}}};
 699
 700         /* Load in the start and ending read addresses from the sent data packet */
 701         StartAddr = Address[0].Word;
 702         EndAddr   = Address[1].Word;
 703 }
 704
 705 /** Handler for a Memory Program command issued by the host. This routine handles the preparations needed
 706  *  to write subsequent data from the host into the specified memory.
 707  */
 708 static void ProcessMemProgCommand(void)
 709 {
 710         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00) ||                          // Write FLASH command
 711             IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))                            // Write EEPROM command
 712         {
 713                 /* Load in the start and ending read addresses */
 714                 LoadStartEndAddresses();
 715
 716                 /* If FLASH is being written to, we need to pre-erase the first page to write to */
 717                 if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))
 718                 {
 719                         union
 720                         {
 721                                 uint16_t Words[2];
 722                                 uint32_t Long;
 723                         } CurrFlashAddress = {.Words = {StartAddr, Flash64KBPage}};
 724
 725                         /* Erase the current page's temp buffer */
 726                         boot_page_erase(CurrFlashAddress.Long);
 727                         boot_spm_busy_wait();
 728                 }
 729
 730                 /* Set the state so that the next DNLOAD requests reads in the firmware */
 731                 DFU_State = dfuDNLOAD_IDLE;
 732         }
 733 }
 734
 735 /** Handler for a Memory Read command issued by the host. This routine handles the preparations needed
 736  *  to read subsequent data from the specified memory out to the host, as well as implementing the memory
 737  *  blank check command.
 738  */
 739 static void ProcessMemReadCommand(void)
 740 {
 741         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00) ||                          // Read FLASH command
 742         IS_ONEBYTE_COMMAND(SentCommand.Data, 0x02))                            // Read EEPROM command
 743         {
 744                 /* Load in the start and ending read addresses */
 745                 LoadStartEndAddresses();
 746
 747                 /* Set the state so that the next UPLOAD requests read out the firmware */
 748                 DFU_State = dfuUPLOAD_IDLE;
 749         }
 750         else if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))                       // Blank check FLASH command
 751         {
 752                 uint32_t CurrFlashAddress = 0;
 753
 754                 while (CurrFlashAddress < (uint32_t)BOOT_START_ADDR)
 755                 {
 756                         /* Check if the current byte is not blank */
 757                         #if (FLASHEND > 0xFFFF)
 758                         if (pgm_read_byte_far(CurrFlashAddress) != 0xFF)
 759                         #else
 760                         if (pgm_read_byte(CurrFlashAddress) != 0xFF)
 761                         #endif
 762                         {
 763                                 /* Save the location of the first non-blank byte for response back to the host */
 764                                 Flash64KBPage = (CurrFlashAddress >> 16);
 765                                 StartAddr     = CurrFlashAddress;
 766
 767                                 /* Set state and status variables to the appropriate error values */
 768                                 DFU_State  = dfuERROR;
 769                                 DFU_Status = errCHECK_ERASED;
 770
 771                                 break;
 772                         }
 773
 774                         CurrFlashAddress++;
 775                 }
 776         }
 777 }
 778
 779 /** Handler for a Data Write command issued by the host. This routine handles non-programming commands such as
 780  *  bootloader exit (both via software jumps and hardware watchdog resets) and flash memory erasure.
 781  */
 782 static void ProcessWriteCommand(void)
 783 {
 784         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x03))                            // Start application
 785         {
 786                 /* Indicate that the bootloader is terminating */
 787                 WaitForExit = true;
 788
 789                 /* Check if data supplied for the Start Program command - no data executes the program */
 790                 if (SentCommand.DataSize)
 791                 {
 792                         if (SentCommand.Data[1] == 0x01)                                   // Start via jump
 793                         {
 794                                 union
 795                                 {
 796                                         uint8_t  Bytes[2];
 797                                         AppPtr_t FuncPtr;
 798                                 } Address = {.Bytes = {SentCommand.Data[4], SentCommand.Data[3]}};
 799
 800                                 /* Load in the jump address into the application start address pointer */
 801                                 AppStartPtr = Address.FuncPtr;
 802                         }
 803                 }
 804                 else
 805                 {
 806                         if (SentCommand.Data[1] == 0x00)                                   // Start via watchdog
 807                         {
 808                                 /* Unlock the forced application start mode of the bootloader if it is restarted */
 809                                 MagicBootKey = MAGIC_BOOT_KEY;
 810
 811                                 /* Start the watchdog to reset the AVR once the communications are finalized */
 812                                 wdt_enable(WDTO_250MS);
 813                         }
 814                         else                                                               // Start via jump
 815                         {
 816                                 /* Set the flag to terminate the bootloader at next opportunity if a valid application has been loaded */
 817                                 if (pgm_read_word_near(0) == 0xFFFF)
 818                                   RunBootloader = false;
 819                         }
 820                 }
 821         }
 822         else if (IS_TWOBYTE_COMMAND(SentCommand.Data, 0x00, 0xFF))                 // Erase flash
 823         {
 824                 uint32_t CurrFlashAddress = 0;
 825
 826                 /* Clear the application section of flash */
 827                 while (CurrFlashAddress < (uint32_t)BOOT_START_ADDR)
 828                 {
 829                         boot_page_erase(CurrFlashAddress);
 830                         boot_spm_busy_wait();
 831                         boot_page_write(CurrFlashAddress);
 832                         boot_spm_busy_wait();
 833
 834                         CurrFlashAddress += SPM_PAGESIZE;
 835                 }
 836
 837                 /* Re-enable the RWW section of flash as writing to the flash locks it out */
 838                 boot_rww_enable();
 839
 840                 /* Memory has been erased, reset the security bit so that programming/reading is allowed */
 841                 IsSecure = false;
 842         }
 843 }
 844
 845 /** Handler for a Data Read command issued by the host. This routine handles bootloader information retrieval
 846  *  commands such as device signature and bootloader version retrieval.
 847  */
 848 static void ProcessReadCommand(void)
 849 {
 850         const uint8_t BootloaderInfo[3] = {BOOTLOADER_VERSION, BOOTLOADER_ID_BYTE1, BOOTLOADER_ID_BYTE2};
 851         const uint8_t SignatureInfo[4]  = {0x58, AVR_SIGNATURE_1, AVR_SIGNATURE_2, AVR_SIGNATURE_3};
 852
 853         uint8_t DataIndexToRead    = SentCommand.Data[1];
 854         bool    ReadAddressInvalid = false;
 855
 856         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))                        // Read bootloader info
 857         {
 858                 if (DataIndexToRead < 3)
 859                   ResponseByte = BootloaderInfo[DataIndexToRead];
 860                 else
 861                   ReadAddressInvalid = true;
 862         }
 863         else if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))                    // Read signature byte
 864         {
 865                 switch (DataIndexToRead)
 866                 {
 867                         case 0x30:
 868                                 ResponseByte = SignatureInfo[0];
 869                                 break;
 870                         case 0x31:
 871                                 ResponseByte = SignatureInfo[1];
 872                                 break;
 873                         case 0x60:
 874                                 ResponseByte = SignatureInfo[2];
 875                                 break;
 876                         case 0x61:
 877                                 ResponseByte = SignatureInfo[3];
 878                                 break;
 879                         default:
 880                                 ReadAddressInvalid = true;
 881                                 break;
 882                 }
 883         }
 884
 885         if (ReadAddressInvalid)
 886         {
 887                 /* Set the state and status variables to indicate the error */
 888                 DFU_State  = dfuERROR;
 889                 DFU_Status = errADDRESS;
 890         }
 891 }