Bootloaders/DFU/BootloaderDFU.c

   1 /*
   2              LUFA Library
   3      Copyright (C) Dean Camera, 2021.
   4
   5   dean [at] fourwalledcubicle [dot] com
   6            www.lufa-lib.org
   7 */
   8
   9 /*
  10   Copyright 2021  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                 {
 136                         /* If the reset source was the bootloader and the key is correct, clear it and jump to the application;
 137                          * this can happen in the HWBE fuse is set, and the HBE pin is low during the watchdog reset */
 138                         if ((MCUSR & (1 << WDRF)) && (MagicBootKey == MAGIC_BOOT_KEY))
 139                                 JumpToApplication = true;
 140
 141                         /* Clear reset source */
 142                         MCUSR &= ~(1 << WDRF);
 143                 }
 144         #endif
 145
 146         /* Don't run the user application if the reset vector is blank (no app loaded) */
 147         bool ApplicationValid = (pgm_read_word_near(0) != 0xFFFF);
 148
 149         /* If a request has been made to jump to the user application, honor it */
 150         if (JumpToApplication && ApplicationValid)
 151         {
 152                 /* Turn off the watchdog */
 153                 MCUSR &= ~(1 << WDRF);
 154                 wdt_disable();
 155
 156                 /* Clear the boot key and jump to the user application */
 157                 MagicBootKey = 0;
 158
 159                 // cppcheck-suppress constStatement
 160                 ((void (*)(void))0x0000)();
 161         }
 162 }
 163
 164
 165 static volatile bool stayinbootloader;
 166
 167 /** Main program entry point. This routine configures the hardware required by the bootloader, then continuously
 168  *  runs the bootloader processing routine until instructed to soft-exit, or hard-reset via the watchdog to start
 169  *  the loaded application code.
 170  */
 171 int main(void)
 172 {
 173         /* Configure hardware required by the bootloader */
 174         SetupHardware();
 175
 176         /* Turn on first LED on the board to indicate that the bootloader has started */
 177         //LEDs_SetAllLEDs(LEDS_LED1);
 178
 179         /* Enable global interrupts so that the USB stack can function */
 180         GlobalInterruptEnable();
 181
 182         /* Run the USB management task while the bootloader is supposed to be running */
 183         /*if bit_is_clear(PINB,PB5)     // PB5 is Digital 9 on Arduino Pro Micro
 184         {
 185                 loop_until_bit_is_set(PINB,PB5);
 186
 187                 while ((RunBootloader || WaitForExit) && bit_is_set(PINB,PB5))
 188                         USB_USBTask();
 189
 190                 loop_until_bit_is_clear(PINB,PB5);
 191         }*/
 192
 193         stayinbootloader = false;
 194
 195         uint16_t i = 0;
 196         while (RunBootloader || WaitForExit)
 197         {
 198                 USB_USBTask();
 199
 200                 if (!stayinbootloader)
 201                 {
 202                         _delay_ms(1);
 203                         if (i++ > 5000)
 204                         {
 205                                 break;
 206                         }
 207                 }
 208                 else
 209                 {
 210                         i = 0;
 211                 }
 212         }
 213
 214         /* Wait a short time to end all USB transactions and then disconnect */
 215         _delay_us(1000);
 216
 217         /* Reset configured hardware back to their original states for the user application */
 218         ResetHardware();
 219
 220         /* Start the user application */
 221         AppStartPtr();
 222 }
 223
 224 /** Configures all hardware required for the bootloader. */
 225 static void SetupHardware(void)
 226 {
 227         /* Disable watchdog if enabled by bootloader/fuses */
 228         MCUSR &= ~(1 << WDRF);
 229         wdt_disable();
 230
 231         /* Disable clock division */
 232         clock_prescale_set(clock_div_1);
 233
 234         /* Relocate the interrupt vector table to the bootloader section */
 235         MCUCR = (1 << IVCE);
 236         MCUCR = (1 << IVSEL);
 237
 238         /* Initialize the USB and other board hardware drivers */
 239         USB_Init();
 240         //LEDs_Init();
 241         DDRB = 1;
 242         PORTB = _BV(PB5);
 243         DDRD = 0b00100000;
 244         PORTD = 0;
 245
 246         /* Bootloader active LED toggle timer initialization */
 247         TIMSK1 = (1 << TOIE1);
 248         TCCR1B = ((1 << CS11) | (1 << CS10));
 249 }
 250
 251 /** Resets all configured hardware required for the bootloader back to their original states. */
 252 static void ResetHardware(void)
 253 {
 254         /* Shut down the USB and other board hardware drivers */
 255         USB_Disable();
 256         //LEDs_Disable();
 257         DDRB = 0;
 258         PORTB = 0;
 259         DDRD = 0;
 260         PORTD = 0;
 261
 262         /* Disable Bootloader active LED toggle timer */
 263         TIMSK1 = 0;
 264         TCCR1B = 0;
 265
 266         /* Relocate the interrupt vector table back to the application section */
 267         MCUCR = (1 << IVCE);
 268         MCUCR = 0;
 269 }
 270
 271 /** ISR to periodically toggle the LEDs on the board to indicate that the bootloader is active. */
 272 ISR(TIMER1_OVF_vect, ISR_BLOCK)
 273 {
 274         //LEDs_ToggleLEDs(LEDS_LED1 | LEDS_LED2);
 275         PORTB &= ~_BV(PB0);
 276         _delay_ms(5);
 277         PORTB |= _BV(PB0);
 278 }
 279
 280 /** Event handler for the USB_ControlRequest event. This is used to catch and process control requests sent to
 281  *  the device from the USB host before passing along unhandled control requests to the library for processing
 282  *  internally.
 283  */
 284 void EVENT_USB_Device_ControlRequest(void)
 285 {
 286         /* Ignore any requests that aren't directed to the DFU interface */
 287         if ((USB_ControlRequest.bmRequestType & (CONTROL_REQTYPE_TYPE | CONTROL_REQTYPE_RECIPIENT)) !=
 288             (REQTYPE_CLASS | REQREC_INTERFACE))
 289         {
 290                 return;
 291         }
 292
 293 stayinbootloader = true;
 294
 295         /* Activity - toggle indicator LEDs */
 296         //LEDs_ToggleLEDs(LEDS_LED1 | LEDS_LED2);
 297         PORTB &= ~_BV(PB0);
 298         _delay_ms(5);
 299         PORTB |= _BV(PB0);
 300
 301         /* Get the size of the command and data from the wLength value */
 302         SentCommand.DataSize = USB_ControlRequest.wLength;
 303
 304         switch (USB_ControlRequest.bRequest)
 305         {
 306                 case DFU_REQ_DNLOAD:
 307                         Endpoint_ClearSETUP();
 308
 309                         /* Check if bootloader is waiting to terminate */
 310                         if (WaitForExit)
 311                         {
 312                                 /* Bootloader is terminating - process last received command */
 313                                 ProcessBootloaderCommand();
 314
 315                                 /* Indicate that the last command has now been processed - free to exit bootloader */
 316                                 WaitForExit = false;
 317                         }
 318
 319                         /* If the request has a data stage, load it into the command struct */
 320                         if (SentCommand.DataSize)
 321                         {
 322                                 while (!(Endpoint_IsOUTReceived()))
 323                                 {
 324                                         if (USB_DeviceState == DEVICE_STATE_Unattached)
 325                                           return;
 326                                 }
 327
 328                                 /* First byte of the data stage is the DNLOAD request's command */
 329                                 SentCommand.Command = Endpoint_Read_8();
 330
 331                                 /* One byte of the data stage is the command, so subtract it from the total data bytes */
 332                                 SentCommand.DataSize--;
 333
 334                                 /* Load in the rest of the data stage as command parameters */
 335                                 for (uint8_t DataByte = 0; (DataByte < sizeof(SentCommand.Data)) &&
 336                                      Endpoint_BytesInEndpoint(); DataByte++)
 337                                 {
 338                                         SentCommand.Data[DataByte] = Endpoint_Read_8();
 339                                         SentCommand.DataSize--;
 340                                 }
 341
 342                                 /* Process the command */
 343                                 ProcessBootloaderCommand();
 344                         }
 345
 346                         /* Check if currently downloading firmware */
 347                         if (DFU_State == dfuDNLOAD_IDLE)
 348                         {
 349                                 if (!(SentCommand.DataSize))
 350                                 {
 351                                         DFU_State = dfuIDLE;
 352                                 }
 353                                 else
 354                                 {
 355                                         /* Throw away the filler bytes before the start of the firmware */
 356                                         DiscardFillerBytes(DFU_FILLER_BYTES_SIZE);
 357
 358                                         /* Throw away the packet alignment filler bytes before the start of the firmware */
 359                                         DiscardFillerBytes(StartAddr % FIXED_CONTROL_ENDPOINT_SIZE);
 360
 361                                         /* Calculate the number of bytes remaining to be written */
 362                                         uint16_t BytesRemaining = ((EndAddr - StartAddr) + 1);
 363
 364                                         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))        // Write flash
 365                                         {
 366                                                 /* Calculate the number of words to be written from the number of bytes to be written */
 367                                                 uint16_t WordsRemaining = (BytesRemaining >> 1);
 368
 369                                                 union
 370                                                 {
 371                                                         uint16_t Words[2];
 372                                                         uint32_t Long;
 373                                                 } CurrFlashAddress                 = {.Words = {StartAddr, Flash64KBPage}};
 374
 375                                                 uint32_t CurrFlashPageStartAddress = CurrFlashAddress.Long;
 376                                                 uint8_t  WordsInFlashPage          = 0;
 377
 378                                                 while (WordsRemaining--)
 379                                                 {
 380                                                         /* Check if endpoint is empty - if so clear it and wait until ready for next packet */
 381                                                         if (!(Endpoint_BytesInEndpoint()))
 382                                                         {
 383                                                                 Endpoint_ClearOUT();
 384
 385                                                                 while (!(Endpoint_IsOUTReceived()))
 386                                                                 {
 387                                                                         if (USB_DeviceState == DEVICE_STATE_Unattached)
 388                                                                           return;
 389                                                                 }
 390                                                         }
 391
 392                                                         /* Write the next word into the current flash page */
 393                                                         BootloaderAPI_FillWord(CurrFlashAddress.Long, Endpoint_Read_16_LE());
 394
 395                                                         /* Adjust counters */
 396                                                         WordsInFlashPage      += 1;
 397                                                         CurrFlashAddress.Long += 2;
 398
 399                                                         /* See if an entire page has been written to the flash page buffer */
 400                                                         if ((WordsInFlashPage == (SPM_PAGESIZE >> 1)) || !(WordsRemaining))
 401                                                         {
 402                                                                 /* Commit the flash page to memory */
 403                                                                 BootloaderAPI_WritePage(CurrFlashPageStartAddress);
 404
 405                                                                 /* Check if programming incomplete */
 406                                                                 if (WordsRemaining)
 407                                                                 {
 408                                                                         CurrFlashPageStartAddress = CurrFlashAddress.Long;
 409                                                                         WordsInFlashPage          = 0;
 410
 411                                                                         /* Erase next page's temp buffer */
 412                                                                         BootloaderAPI_ErasePage(CurrFlashAddress.Long);
 413                                                                 }
 414                                                         }
 415                                                 }
 416
 417                                                 /* Once programming complete, start address equals the end address */
 418                                                 StartAddr = EndAddr;
 419                                         }
 420                                         else                                                   // Write EEPROM
 421                                         {
 422                                                 while (BytesRemaining--)
 423                                                 {
 424                                                         /* Check if endpoint is empty - if so clear it and wait until ready for next packet */
 425                                                         if (!(Endpoint_BytesInEndpoint()))
 426                                                         {
 427                                                                 Endpoint_ClearOUT();
 428
 429                                                                 while (!(Endpoint_IsOUTReceived()))
 430                                                                 {
 431                                                                         if (USB_DeviceState == DEVICE_STATE_Unattached)
 432                                                                           return;
 433                                                                 }
 434                                                         }
 435
 436                                                         /* Read the byte from the USB interface and write to to the EEPROM */
 437                                                         eeprom_update_byte((uint8_t*)StartAddr, Endpoint_Read_8());
 438
 439                                                         /* Adjust counters */
 440                                                         StartAddr++;
 441                                                 }
 442                                         }
 443
 444                                         /* Throw away the currently unused DFU file suffix */
 445                                         DiscardFillerBytes(DFU_FILE_SUFFIX_SIZE);
 446                                 }
 447                         }
 448
 449                         Endpoint_ClearOUT();
 450
 451                         Endpoint_ClearStatusStage();
 452
 453                         break;
 454                 case DFU_REQ_UPLOAD:
 455                         Endpoint_ClearSETUP();
 456
 457                         while (!(Endpoint_IsINReady()))
 458                         {
 459                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 460                                   return;
 461                         }
 462
 463                         if (DFU_State != dfuUPLOAD_IDLE)
 464                         {
 465                                 if ((DFU_State == dfuERROR) && IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))       // Blank Check
 466                                 {
 467                                         /* Blank checking is performed in the DFU_DNLOAD request - if we get here we've told the host
 468                                            that the memory isn't blank, and the host is requesting the first non-blank address */
 469                                         Endpoint_Write_16_LE(StartAddr);
 470                                 }
 471                                 else
 472                                 {
 473                                         /* Idle state upload - send response to last issued command */
 474                                         Endpoint_Write_8(ResponseByte);
 475                                 }
 476                         }
 477                         else
 478                         {
 479                                 /* Determine the number of bytes remaining in the current block */
 480                                 uint16_t BytesRemaining = ((EndAddr - StartAddr) + 1);
 481
 482                                 if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))            // Read FLASH
 483                                 {
 484                                         /* Calculate the number of words to be written from the number of bytes to be written */
 485                                         uint16_t WordsRemaining = (BytesRemaining >> 1);
 486
 487                                         union
 488                                         {
 489                                                 uint16_t Words[2];
 490                                                 uint32_t Long;
 491                                         } CurrFlashAddress = {.Words = {StartAddr, Flash64KBPage}};
 492
 493                                         while (WordsRemaining--)
 494                                         {
 495                                                 /* Check if endpoint is full - if so clear it and wait until ready for next packet */
 496                                                 if (Endpoint_BytesInEndpoint() == FIXED_CONTROL_ENDPOINT_SIZE)
 497                                                 {
 498                                                         Endpoint_ClearIN();
 499
 500                                                         while (!(Endpoint_IsINReady()))
 501                                                         {
 502                                                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 503                                                                   return;
 504                                                         }
 505                                                 }
 506
 507                                                 /* Read the flash word and send it via USB to the host */
 508                                                 #if (FLASHEND > 0xFFFF)
 509                                                         Endpoint_Write_16_LE(pgm_read_word_far(CurrFlashAddress.Long));
 510                                                 #else
 511                                                         Endpoint_Write_16_LE(pgm_read_word(CurrFlashAddress.Long));
 512                                                 #endif
 513
 514                                                 /* Adjust counters */
 515                                                 CurrFlashAddress.Long += 2;
 516                                         }
 517
 518                                         /* Once reading is complete, start address equals the end address */
 519                                         StartAddr = EndAddr;
 520                                 }
 521                                 else if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x02))       // Read EEPROM
 522                                 {
 523                                         while (BytesRemaining--)
 524                                         {
 525                                                 /* Check if endpoint is full - if so clear it and wait until ready for next packet */
 526                                                 if (Endpoint_BytesInEndpoint() == FIXED_CONTROL_ENDPOINT_SIZE)
 527                                                 {
 528                                                         Endpoint_ClearIN();
 529
 530                                                         while (!(Endpoint_IsINReady()))
 531                                                         {
 532                                                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 533                                                                   return;
 534                                                         }
 535                                                 }
 536
 537                                                 /* Read the EEPROM byte and send it via USB to the host */
 538                                                 Endpoint_Write_8(eeprom_read_byte((uint8_t*)StartAddr));
 539
 540                                                 /* Adjust counters */
 541                                                 StartAddr++;
 542                                         }
 543                                 }
 544
 545                                 /* Return to idle state */
 546                                 DFU_State = dfuIDLE;
 547                         }
 548
 549                         Endpoint_ClearIN();
 550
 551                         Endpoint_ClearStatusStage();
 552                         break;
 553                 case DFU_REQ_GETSTATUS:
 554                         Endpoint_ClearSETUP();
 555
 556                         while (!(Endpoint_IsINReady()))
 557                         {
 558                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 559                                   return;
 560                         }
 561
 562                         /* Write 8-bit status value */
 563                         Endpoint_Write_8(DFU_Status);
 564
 565                         /* Write 24-bit poll timeout value */
 566                         Endpoint_Write_8(0);
 567                         Endpoint_Write_16_LE(0);
 568
 569                         /* Write 8-bit state value */
 570                         Endpoint_Write_8(DFU_State);
 571
 572                         /* Write 8-bit state string ID number */
 573                         Endpoint_Write_8(0);
 574
 575                         Endpoint_ClearIN();
 576
 577                         Endpoint_ClearStatusStage();
 578                         break;
 579                 case DFU_REQ_CLRSTATUS:
 580                         Endpoint_ClearSETUP();
 581
 582                         /* Reset the status value variable to the default OK status */
 583                         DFU_Status = OK;
 584
 585                         Endpoint_ClearStatusStage();
 586                         break;
 587                 case DFU_REQ_GETSTATE:
 588                         Endpoint_ClearSETUP();
 589
 590                         while (!(Endpoint_IsINReady()))
 591                         {
 592                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 593                                   return;
 594                         }
 595
 596                         /* Write the current device state to the endpoint */
 597                         Endpoint_Write_8(DFU_State);
 598
 599                         Endpoint_ClearIN();
 600
 601                         Endpoint_ClearStatusStage();
 602                         break;
 603                 case DFU_REQ_ABORT:
 604                         Endpoint_ClearSETUP();
 605
 606                         /* Reset the current state variable to the default idle state */
 607                         DFU_State = dfuIDLE;
 608
 609                         Endpoint_ClearStatusStage();
 610                         break;
 611         }
 612 }
 613
 614 /** Routine to discard the specified number of bytes from the control endpoint stream. This is used to
 615  *  discard unused bytes in the stream from the host, including the memory program block suffix.
 616  *
 617  *  \param[in] NumberOfBytes  Number of bytes to discard from the host from the control endpoint
 618  */
 619 static void DiscardFillerBytes(uint8_t NumberOfBytes)
 620 {
 621         while (NumberOfBytes--)
 622         {
 623                 if (!(Endpoint_BytesInEndpoint()))
 624                 {
 625                         Endpoint_ClearOUT();
 626
 627                         /* Wait until next data packet received */
 628                         while (!(Endpoint_IsOUTReceived()))
 629                         {
 630                                 if (USB_DeviceState == DEVICE_STATE_Unattached)
 631                                   return;
 632                         }
 633                 }
 634                 else
 635                 {
 636                         Endpoint_Discard_8();
 637                 }
 638         }
 639 }
 640
 641 /** Routine to process an issued command from the host, via a DFU_DNLOAD request wrapper. This routine ensures
 642  *  that the command is allowed based on the current secure mode flag value, and passes the command off to the
 643  *  appropriate handler function.
 644  */
 645 static void ProcessBootloaderCommand(void)
 646 {
 647         /* Check if device is in secure mode */
 648         if (IsSecure)
 649         {
 650                 /* Don't process command unless it is a READ or chip erase command */
 651                 if (!(((SentCommand.Command == COMMAND_WRITE)             &&
 652                         IS_TWOBYTE_COMMAND(SentCommand.Data, 0x00, 0xFF)) ||
 653                            (SentCommand.Command == COMMAND_READ)))
 654                 {
 655                         /* Set the state and status variables to indicate the error */
 656                         DFU_State  = dfuERROR;
 657                         DFU_Status = errWRITE;
 658
 659                         /* Stall command */
 660                         Endpoint_StallTransaction();
 661
 662                         /* Don't process the command */
 663                         return;
 664                 }
 665         }
 666
 667         /* Dispatch the required command processing routine based on the command type */
 668         switch (SentCommand.Command)
 669         {
 670                 case COMMAND_PROG_START:
 671                         ProcessMemProgCommand();
 672                         break;
 673                 case COMMAND_DISP_DATA:
 674                         ProcessMemReadCommand();
 675                         break;
 676                 case COMMAND_WRITE:
 677                         ProcessWriteCommand();
 678                         break;
 679                 case COMMAND_READ:
 680                         ProcessReadCommand();
 681                         break;
 682                 case COMMAND_CHANGE_BASE_ADDR:
 683                         if (IS_TWOBYTE_COMMAND(SentCommand.Data, 0x03, 0x00))              // Set 64KB flash page command
 684                           Flash64KBPage = SentCommand.Data[2];
 685
 686                         break;
 687         }
 688 }
 689
 690 /** Routine to concatenate the given pair of 16-bit memory start and end addresses from the host, and store them
 691  *  in the StartAddr and EndAddr global variables.
 692  */
 693 static void LoadStartEndAddresses(void)
 694 {
 695         union
 696         {
 697                 uint8_t  Bytes[2];
 698                 uint16_t Word;
 699         } Address[2] = {{.Bytes = {SentCommand.Data[2], SentCommand.Data[1]}},
 700                         {.Bytes = {SentCommand.Data[4], SentCommand.Data[3]}}};
 701
 702         /* Load in the start and ending read addresses from the sent data packet */
 703         StartAddr = Address[0].Word;
 704         EndAddr   = Address[1].Word;
 705 }
 706
 707 /** Handler for a Memory Program command issued by the host. This routine handles the preparations needed
 708  *  to write subsequent data from the host into the specified memory.
 709  */
 710 static void ProcessMemProgCommand(void)
 711 {
 712         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00) ||                          // Write FLASH command
 713             IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))                            // Write EEPROM command
 714         {
 715                 /* Load in the start and ending read addresses */
 716                 LoadStartEndAddresses();
 717
 718                 /* If FLASH is being written to, we need to pre-erase the first page to write to */
 719                 if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))
 720                 {
 721                         union
 722                         {
 723                                 uint16_t Words[2];
 724                                 uint32_t Long;
 725                         } CurrFlashAddress = {.Words = {StartAddr, Flash64KBPage}};
 726
 727                         /* Erase the current page's temp buffer */
 728                         BootloaderAPI_ErasePage(CurrFlashAddress.Long);
 729                 }
 730
 731                 /* Set the state so that the next DNLOAD requests reads in the firmware */
 732                 DFU_State = dfuDNLOAD_IDLE;
 733         }
 734 }
 735
 736 /** Handler for a Memory Read command issued by the host. This routine handles the preparations needed
 737  *  to read subsequent data from the specified memory out to the host, as well as implementing the memory
 738  *  blank check command.
 739  */
 740 static void ProcessMemReadCommand(void)
 741 {
 742         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00) ||                          // Read FLASH command
 743         IS_ONEBYTE_COMMAND(SentCommand.Data, 0x02))                            // Read EEPROM command
 744         {
 745                 /* Load in the start and ending read addresses */
 746                 LoadStartEndAddresses();
 747
 748                 /* Set the state so that the next UPLOAD requests read out the firmware */
 749                 DFU_State = dfuUPLOAD_IDLE;
 750         }
 751         else if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))                       // Blank check FLASH command
 752         {
 753                 uint32_t CurrFlashAddress = 0;
 754
 755                 while (CurrFlashAddress < (uint32_t)BOOT_START_ADDR)
 756                 {
 757                         /* Check if the current byte is not blank */
 758                         #if (FLASHEND > 0xFFFF)
 759                         if (pgm_read_byte_far(CurrFlashAddress) != 0xFF)
 760                         #else
 761                         if (pgm_read_byte(CurrFlashAddress) != 0xFF)
 762                         #endif
 763                         {
 764                                 /* Save the location of the first non-blank byte for response back to the host */
 765                                 Flash64KBPage = (CurrFlashAddress >> 16);
 766                                 StartAddr     = CurrFlashAddress;
 767
 768                                 /* Set state and status variables to the appropriate error values */
 769                                 DFU_State  = dfuERROR;
 770                                 DFU_Status = errCHECK_ERASED;
 771
 772                                 break;
 773                         }
 774
 775                         CurrFlashAddress++;
 776                 }
 777         }
 778 }
 779
 780 /** Handler for a Data Write command issued by the host. This routine handles non-programming commands such as
 781  *  bootloader exit (both via software jumps and hardware watchdog resets) and flash memory erasure.
 782  */
 783 static void ProcessWriteCommand(void)
 784 {
 785         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x03))                            // Start application
 786         {
 787                 /* Indicate that the bootloader is terminating */
 788                 WaitForExit = true;
 789
 790                 /* Check if data supplied for the Start Program command - no data executes the program */
 791                 if (SentCommand.DataSize)
 792                 {
 793                         if (SentCommand.Data[1] == 0x01)                                   // Start via jump
 794                         {
 795                                 union
 796                                 {
 797                                         uint8_t  Bytes[2];
 798                                         AppPtr_t FuncPtr;
 799                                 } Address = {.Bytes = {SentCommand.Data[4], SentCommand.Data[3]}};
 800
 801                                 /* Load in the jump address into the application start address pointer */
 802                                 AppStartPtr = Address.FuncPtr;
 803                         }
 804                 }
 805                 else
 806                 {
 807                         if (SentCommand.Data[1] == 0x00)                                   // Start via watchdog
 808                         {
 809                                 /* Unlock the forced application start mode of the bootloader if it is restarted */
 810                                 MagicBootKey = MAGIC_BOOT_KEY;
 811
 812                                 /* Start the watchdog to reset the AVR once the communications are finalized */
 813                                 wdt_enable(WDTO_250MS);
 814                         }
 815                         else                                                               // Start via jump
 816                         {
 817                                 /* Set the flag to terminate the bootloader at next opportunity if a valid application has been loaded */
 818                                 if (pgm_read_word_near(0) == 0xFFFF)
 819                                   RunBootloader = false;
 820                         }
 821                 }
 822         }
 823         else if (IS_TWOBYTE_COMMAND(SentCommand.Data, 0x00, 0xFF))                 // Erase flash
 824         {
 825                 /* Clear the application section of flash */
 826                 for (uint32_t CurrFlashAddress = 0; CurrFlashAddress < (uint32_t)BOOT_START_ADDR; CurrFlashAddress += SPM_PAGESIZE)
 827                         BootloaderAPI_ErasePage(CurrFlashAddress);
 828
 829                 /* Memory has been erased, reset the security bit so that programming/reading is allowed */
 830                 IsSecure = false;
 831         }
 832 }
 833
 834 /** Handler for a Data Read command issued by the host. This routine handles bootloader information retrieval
 835  *  commands such as device signature and bootloader version retrieval.
 836  */
 837 static void ProcessReadCommand(void)
 838 {
 839         const uint8_t BootloaderInfo[3] = {BOOTLOADER_VERSION, BOOTLOADER_ID_BYTE1, BOOTLOADER_ID_BYTE2};
 840         const uint8_t SignatureInfo[4]  = {0x58, AVR_SIGNATURE_1, AVR_SIGNATURE_2, AVR_SIGNATURE_3};
 841
 842         uint8_t DataIndexToRead    = SentCommand.Data[1];
 843         bool    ReadAddressInvalid = false;
 844
 845         if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x00))                        // Read bootloader info
 846         {
 847                 if (DataIndexToRead < 3)
 848                   ResponseByte = BootloaderInfo[DataIndexToRead];
 849                 else
 850                   ReadAddressInvalid = true;
 851         }
 852         else if (IS_ONEBYTE_COMMAND(SentCommand.Data, 0x01))                    // Read signature byte
 853         {
 854                 switch (DataIndexToRead)
 855                 {
 856                         case 0x30:
 857                                 ResponseByte = SignatureInfo[0];
 858                                 break;
 859                         case 0x31:
 860                                 ResponseByte = SignatureInfo[1];
 861                                 break;
 862                         case 0x60:
 863                                 ResponseByte = SignatureInfo[2];
 864                                 break;
 865                         case 0x61:
 866                                 ResponseByte = SignatureInfo[3];
 867                                 break;
 868                         default:
 869                                 ReadAddressInvalid = true;
 870                                 break;
 871                 }
 872         }
 873
 874         if (ReadAddressInvalid)
 875         {
 876                 /* Set the state and status variables to indicate the error */
 877                 DFU_State  = dfuERROR;
 878                 DFU_Status = errADDRESS;
 879         }
 880 }